integration fixes

forward.py

@@ -1,244 +1,245 @@
-# 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
+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

server.py

@@ -37,16 +37,14 @@ add --mock to test the echo stream
 async def chat(request: web.Request) -> Response:
     """
     Chat endpoint for the validator.
-    """    
-    request_data = request['data']
-    params = QueryValidatorParams.from_dict(request_data)    
-    # TODO: SET STREAM AS DEFAULT
-    stream = request_data.get('stream', True)        
+    """        
+    params = QueryValidatorParams.from_request(request)    
+
     
     # Access the validator from the application context
     validator: ValidatorAPI = request.app['validator']
     
-    response = await validator.query_validator(params, stream=stream)
+    response = await validator.query_validator(params)
     return response
 
 

validators/base.py

@@ -1,7 +1,7 @@
 from abc import ABC, abstractmethod
 from typing import List
 from dataclasses import dataclass
-from aiohttp.web_response import Response
+from aiohttp.web import Response, Request
 
 @dataclass
 class QueryValidatorParams:
@@ -11,16 +11,20 @@ class QueryValidatorParams:
     messages: List[str]
     timeout: int
     prefer: str
+    request: Request
     
     @staticmethod
-    def from_dict(data: dict):
+    def from_request(request: Request):
+        data = request['data']
+        
         return QueryValidatorParams(            
             k_miners=data.get('k', 10),
             exclude=data.get('exclude', []),
             roles=data['roles'],
             messages=data['messages'],
             timeout=data.get('timeout', 10),
-            prefer=data.get('prefer', 'longest')
+            prefer=data.get('prefer', 'longest'),
+            request=request
         )
 
 class ValidatorAPI(ABC):

validators/sn1_validator_wrapper.py

@@ -2,7 +2,9 @@ import json
 import utils
 import torch
 import traceback
+import asyncio
 import bittensor as bt
+from typing import Awaitable
 from prompting.validator import Validator
 from prompting.utils.uids import get_random_uids
 from prompting.protocol import PromptingSynapse, StreamPromptingSynapse
@@ -55,11 +57,30 @@ class S1ValidatorAPI(ValidatorAPI):
             return Response(status=500, reason="Internal error")
         
         
+    async def process_response(self, response: StreamResponse, 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}")
+                
+                # TODO: SET PROPER IMPLEMENTATION TO RETURN CHUNK
+                if chunk is not None:
+                    json_data = json.dumps(chunk)
+                    await response.write(json_data.encode('utf-8'))
+                    
+        except Exception as e:
+            bt.logging.error(f'Encountered an error in {self.__class__.__name__}:get_stream_response:\n{traceback.format_exc()}')
+            response.set_status(500, reason="Internal error")
+            await response.write(json.dumps({'error': str(e)}).encode('utf-8'))
+        finally:
+            await response.write_eof()  # Ensure to close the response properly
+        
     async def get_stream_response(self, params:QueryValidatorParams) -> StreamResponse:
         response = StreamResponse(status=200, reason="OK")
         response.headers['Content-Type'] = 'application/json'
 
-        await response.prepare()  # Prepare and send the headers
+        await response.prepare(params.request)  # Prepare and send the headers
         
         try:
             # Guess the task name of current request
@@ -78,14 +99,11 @@ class S1ValidatorAPI(ValidatorAPI):
                 deserialize=False,
                 streaming=True,
             )
-
-            # Asynchronous iteration over streaming responses
-            async for stream_result in streams_responses:
-                if stream_result is not None:
-                    # Convert stream result to JSON and write to the response stream
-                    json_data = json.dumps(stream_result)
-                    await response.write(json_data.encode('utf-8'))
-
+            
+            tasks = [self.process_response(uid, res) for uid, res in dict(zip(uids, streams_responses))]                        
+            results = await asyncio.gather(*tasks, return_exceptions=True)
+            
+            # TODO: Continue implementation, business decision needs to be made on how to handle the results      
         except Exception as e:
             bt.logging.error(f'Encountered an error in {self.__class__.__name__}:get_stream_response:\n{traceback.format_exc()}')
             response.set_status(500, reason="Internal error")
@@ -94,12 +112,15 @@ class S1ValidatorAPI(ValidatorAPI):
             await response.write_eof()  # Ensure to close the response properly
 
         return response
-            
 
-    async def query_validator(self, params:QueryValidatorParams, stream: bool = True) -> Response:        
+
+    async def query_validator(self, params:QueryValidatorParams) -> Response:        
+        # TODO: SET STREAM AS DEFAULT
+        stream = params.request.get('stream', False)        
+        
         if stream:
             return await self.get_stream_response(params)
         else:
             # DEPRECATED
             return await self.get_response(params)
-        
+