forward.py

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