Add stream forward

forward.py

@@ -0,0 +1,244 @@
+import time
+import sys
+import asyncio
+import numpy as np
+import bittensor as bt
+import traceback
+from typing import List, Dict, Awaitable
+from prompting.agent import HumanAgent
+from prompting.dendrite import DendriteResponseEvent
+from prompting.conversation import create_task
+from prompting.protocol import StreamPromptingSynapse
+from prompting.rewards import RewardResult
+from prompting.utils.uids import get_random_uids
+from prompting.utils.logging import log_event
+from prompting.utils.misc import async_log, serialize_exception_to_string
+from dataclasses import dataclass
+
+@async_log
+async def generate_reference(agent):    
+    loop = asyncio.get_running_loop()
+    result = await loop.run_in_executor(None, agent.task.generate_reference, agent.llm_pipeline)
+    return result    
+
+@async_log
+async def execute_dendrite_call(dendrite_call):
+    responses = await dendrite_call
+    return responses
+
+@dataclass
+class StreamResult:
+    synapse: StreamPromptingSynapse = None
+    exception: BaseException = None
+    uid: int = None
+
+
+async def process_response(uid: int, async_generator: Awaitable):
+    """Process a single response asynchronously."""
+    try:
+        chunk = None  # Initialize chunk with a default value
+        async for chunk in async_generator:  # most important loop, as this is where we acquire the final synapse.
+            bt.logging.debug(f"\nchunk for uid {uid}: {chunk}")
+
+        if chunk is not None:
+            synapse = chunk  # last object yielded is the synapse itself with completion filled
+
+            # Assuming chunk holds the last value yielded which should be a synapse
+            if isinstance(synapse, StreamPromptingSynapse):
+                return synapse
+
+        bt.logging.debug(
+            f"Synapse is not StreamPromptingSynapse. Miner uid {uid} completion set to '' "
+        )
+    except Exception as e:
+        # bt.logging.error(f"Error in generating reference or handling responses: {e}", exc_info=True)
+        traceback_details = traceback.format_exc()
+        bt.logging.error(
+            f"Error in generating reference or handling responses for uid {uid}: {e}\n{traceback_details}"
+        )
+
+        failed_synapse = StreamPromptingSynapse(
+            roles=["user"], messages=["failure"], completion=""
+        )
+
+        return failed_synapse
+
+
+@async_log
+async def handle_response(responses: Dict[int, Awaitable]) -> List[StreamResult]:
+    """The handle_response function is responsible for creating asyncio tasks around acquiring streamed miner chunks
+    and processing them asynchronously. It then pairs the results with their original UIDs and returns a list of StreamResults.
+
+    Args:
+        responses (Dict[int, Awaitable]): Responses contains awaitables that are used to acquire streamed miner chunks.
+
+    Raises:
+        ValueError
+
+    Returns:
+        List[StreamResult]: DataClass containing the synapse, exception, and uid
+    """
+    tasks_with_uid = [
+        (uid, responses[uid]) for uid, _ in responses.items()
+    ]  # Pair UIDs with their tasks
+
+    # Start tasks, preserving order and their associated UIDs
+    tasks = [process_response(uid, resp) for uid, resp in tasks_with_uid]
+
+    results = await asyncio.gather(*tasks, return_exceptions=True)
+
+    mapped_results = []
+    # Pair each result with its original uid
+    for (uid, _), result in zip(tasks_with_uid, results):
+        # If the result is a StreamPromptingSynapse, the response was successful and the stream result is added without exceptions
+        if isinstance(result, StreamPromptingSynapse):
+            mapped_results.append(StreamResult(synapse=result, uid=uid))
+
+        # If the result is an exception, the response was unsuccessful and the stream result is added with the exception and an empty synapse
+        elif isinstance(result, BaseException):
+            failed_synapse = StreamPromptingSynapse(
+                roles=["user"], messages=["failure"], completion=""
+            )
+            mapped_results.append(
+                StreamResult(synapse=failed_synapse, exception=result, uid=uid)
+            )
+
+        # If the result is neither an error or a StreamSynapse, log the error and raise a ValueError
+        else:
+            bt.logging.error(f"Unexpected result type for UID {uid}: {result}")
+            raise ValueError(f"Unexpected result type for UID {uid}: {result}")
+
+    return mapped_results
+
+
+@async_log
+async def generate_reference(agent: HumanAgent):
+    loop = asyncio.get_running_loop()
+    result = await loop.run_in_executor(
+        None, agent.task.generate_reference, agent.llm_pipeline
+    )
+    return result
+
+
+def log_stream_results(stream_results: List[StreamResult]):
+    failed_responses = [
+        response for response in stream_results if response.exception is not None
+    ]
+    empty_responses = [
+        response
+        for response in stream_results
+        if response.exception is None and response.synapse.completion == ""
+    ]
+    non_empty_responses = [
+        response
+        for response in stream_results
+        if response.exception is None and response.synapse.completion != ""
+    ]
+
+    bt.logging.info(f"Total of non_empty responses: ({len(non_empty_responses)})")
+    bt.logging.info(f"Total of empty responses: ({len(empty_responses)})")
+    bt.logging.info(
+        f"Total of failed responses: ({len(failed_responses)}):\n {failed_responses}"
+    )
+
+    for failed_response in failed_responses:
+        formatted_exception = serialize_exception_to_string(failed_response.exception)
+        bt.logging.error(
+            f"Failed response for uid {failed_response.uid}: {formatted_exception}"
+        )
+
+
+async def run_step(
+    self, agent: HumanAgent, k: int, timeout: float, exclude: list = None
+):
+    """Executes a single step of the agent, which consists of:
+    - Getting a list of uids to query
+    - Querying the network
+    - Rewarding the network
+    - Updating the scores
+    - Logging the event
+
+    Args:
+        agent (HumanAgent): The agent to run the step for.
+        k (int): The number of uids to query.
+        timeout (float): The timeout for the queries.
+        exclude (list, optional): The list of uids to exclude from the query. Defaults to [].
+    """
+
+    bt.logging.debug("run_step", agent.task.name)
+
+    # Record event start time.
+    start_time = time.time()
+    # Get the list of uids to query for this step.
+    uids = get_random_uids(self, k=k, exclude=exclude or []).to(self.device)
+    uids_cpu = uids.cpu().tolist()
+
+    axons = [self.metagraph.axons[uid] for uid in uids]
+
+    # Directly call dendrite and process responses in parallel
+    streams_responses = await self.dendrite(
+        axons=axons,
+        synapse=StreamPromptingSynapse(roles=["user"], messages=[agent.challenge]),
+        timeout=timeout,
+        deserialize=False,
+        streaming=True,
+    )
+
+    # Prepare the task for handling stream responses
+    handle_stream_responses_task = asyncio.create_task(
+        handle_response(responses=dict(zip(uids_cpu, streams_responses)))
+    )
+
+    if not agent.task.static_reference:
+        reference_generation_task = generate_reference(agent)
+        _, stream_results = await asyncio.gather(
+            reference_generation_task, handle_stream_responses_task
+        )
+    else:
+        stream_results = await handle_stream_responses_task
+
+    log_stream_results(stream_results)
+
+    all_synapses_results = [stream_result.synapse for stream_result in stream_results]   
+            
+    # Encapsulate the responses in a response event (dataclass)
+    response_event = DendriteResponseEvent(
+        responses=all_synapses_results, uids=uids, timeout=timeout
+    )
+
+    bt.logging.info(f"Created DendriteResponseEvent:\n {response_event}")
+    # Reward the responses and get the reward result (dataclass)
+    # This contains a list of RewardEvents but can be exported as a dict (column-wise) for logging etc
+    reward_result = RewardResult(
+        self.reward_pipeline,
+        agent=agent,
+        response_event=response_event,
+        device=self.device,
+    )
+    bt.logging.info(f"Created RewardResult:\n {reward_result}")
+
+    # The original idea was that the agent is 'satisfied' when it gets a good enough response (e.g. reward critera is met, such as ROUGE>threshold)
+    agent.update_progress(
+        top_reward=reward_result.rewards.max(),
+        top_response=response_event.completions[reward_result.rewards.argmax()],
+    )
+
+    self.update_scores(reward_result.rewards, uids)
+
+    stream_results_uids = [stream_result.uid for stream_result in stream_results]
+    stream_results_exceptions = [
+        serialize_exception_to_string(stream_result.exception)
+        for stream_result in stream_results
+    ]
+    # Log the step event.
+    event = {
+        "block": self.block,
+        "step_time": time.time() - start_time,
+        "stream_results_uids": stream_results_uids,
+        "stream_results_exceptions": stream_results_exceptions,
+        **agent.__state_dict__(full=self.config.neuron.log_full),
+        **reward_result.__state_dict__(full=self.config.neuron.log_full),
+        **response_event.__state_dict__(),
+    }
+
+    return event