initial streamer implementation

utils.py

@@ -1,12 +1,10 @@
 import re
-import time
-import json
 import asyncio
 import bittensor as bt
 from aiohttp import web
-from responses import TextStreamResponse
 from collections import Counter
 from prompting.rewards import DateRewardModel, FloatDiffModel
+from validators.streamer import AsyncResponseDataStreamer
 
 UNSUCCESSFUL_RESPONSE_PATTERNS = [
     "I'm sorry",
@@ -136,46 +134,27 @@ def guess_task_name(challenge: str):
     return "qa"
 
 
+# Simulate the stream synapse for the echo endpoint
+class EchoAsyncIterator:
+    def __init__(self, message: str, k: int, delay: float):
+        self.message = message
+        self.k = k
+        self.delay = delay
+
+    async def __aiter__(self):
+        for _ in range(self.k):
+            for word in self.message.split():
+                yield [word]
+                await asyncio.sleep(self.delay)
+
+
 async def echo_stream(request: web.Request) -> web.StreamResponse:
     request_data = request["data"]
     k = request_data.get("k", 1)
     message = "\n\n".join(request_data["messages"])
 
-    # Create a StreamResponse
-    response = web.StreamResponse(
-        status=200, reason="OK", headers={"Content-Type": "application/json"}
-    )
-    await response.prepare(request)
-
-    completion = ""
-    chunks = []
-    chunks_timings = []
-    start_time = time.time()
-    # Echo the message k times with a timeout between each chunk
-    for _ in range(k):
-        for word in message.split():
-            chunk = f"{word} "
-            await response.write(chunk.encode("utf-8"))
-            completion += chunk
-            await asyncio.sleep(0.3)
-            bt.logging.info(f"Echoed: {chunk}")
-
-            chunks.append(chunk)
-            chunks_timings.append(time.time() - start_time)
-
-    completion = completion.strip()
-
-    # Prepare final JSON chunk
-    response_data = TextStreamResponse(
-        streamed_chunks=chunks,
-        streamed_chunks_timings=chunks_timings,
-        completion=completion,
-        timing=time.time() - start_time,
-    ).to_dict()
-
-    # Send the final JSON as part of the stream
-    await response.write(json.dumps(response_data).encode("utf-8"))
-
-    # Finalize the response
-    await response.write_eof()
-    return response
+
+    echo_iterator = EchoAsyncIterator(message, k, delay=0.3)
+    streamer = AsyncResponseDataStreamer(echo_iterator, selected_uid=0, delay=0.3)
+
+    return await streamer.stream(request)

validators/__init__.py

@@ -1,2 +1,3 @@
 from .base import QueryValidatorParams, ValidatorAPI, MockValidator
 from .sn1_validator_wrapper import S1ValidatorAPI
+from .streamer import AsyncResponseDataStreamer

validators/sn1_validator_wrapper.py

@@ -1,21 +1,13 @@
-import json
-import utils
-import torch
-import traceback
-import time
 import random
 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
-from prompting.dendrite import DendriteResponseEvent
+from prompting.protocol import StreamPromptingSynapse
 from .base import QueryValidatorParams, ValidatorAPI
 from aiohttp.web_response import Response, StreamResponse
-from deprecated import deprecated
 from dataclasses import dataclass
 from typing import List
-from responses import TextStreamResponse
+from .streamer import AsyncResponseDataStreamer
 
 
 @dataclass
@@ -29,154 +21,35 @@ class S1ValidatorAPI(ValidatorAPI):
     def __init__(self):
         self.validator = Validator()
 
-    @deprecated(
-        reason="This function is deprecated. Validators use stream synapse now, use get_stream_response instead."
-    )
-    async def get_response(self, params: QueryValidatorParams) -> Response:
-        try:
-            # Guess the task name of current request
-            task_name = utils.guess_task_name(params.messages[-1])
-
-            # Get the list of uids to query for this step.
-            uids = get_random_uids(
-                self.validator, k=params.k_miners, exclude=params.exclude or []
-            ).tolist()
-            axons = [self.validator.metagraph.axons[uid] for uid in uids]
-
-            # Make calls to the network with the prompt.
-            bt.logging.info(f"Calling dendrite")
-            responses = await self.validator.dendrite(
-                axons=axons,
-                synapse=PromptingSynapse(roles=params.roles, messages=params.messages),
-                timeout=params.timeout,
-            )
-
-            bt.logging.info(f"Creating DendriteResponseEvent:\n {responses}")
-            # Encapsulate the responses in a response event (dataclass)
-            response_event = DendriteResponseEvent(
-                responses, torch.LongTensor(uids), params.timeout
-            )
-
-            # convert dict to json
-            response = response_event.__state_dict__()
-
-            response["completion_is_valid"] = valid = list(
-                map(utils.completion_is_valid, response["completions"])
-            )
-            valid_completions = [
-                response["completions"][i] for i, v in enumerate(valid) if v
-            ]
-
-            response["task_name"] = task_name
-            response["ensemble_result"] = utils.ensemble_result(
-                valid_completions, task_name=task_name, prefer=params.prefer
-            )
-
-            bt.logging.info(f"Response:\n {response}")
-            return Response(
-                status=200,
-                reason="I can't believe it's not butter!",
-                text=json.dumps(response),
-            )
-
-        except Exception:
-            bt.logging.error(
-                f"Encountered in {self.__class__.__name__}:get_response:\n{traceback.format_exc()}"
-            )
-            return Response(status=500, reason="Internal error")
-
-    async def process_response(
-        self, response: StreamResponse, async_generator: Awaitable
-    ) -> ProcessedStreamResponse:
-        """Process a single response asynchronously."""
-        # Initialize chunk with a default value
-        chunk = None
-        # Initialize chunk array to accumulate streamed chunks
-        chunks = []
-        chunks_timings = []
-
-        start_time = time.time()
-        last_sent_index = 0
-        async for chunk in async_generator:
-            if isinstance(chunk, list):
-                # Chunks are currently returned in string arrays, so we need to concatenate them
-                concatenated_chunks = "".join(chunk)
-                new_data = concatenated_chunks[last_sent_index:]
-
-                if new_data:
-                    await response.write(new_data.encode("utf-8"))
-                    bt.logging.info(f"Received new chunk from miner: {chunk}")
-                    last_sent_index += len(new_data)
-                    chunks.extend(chunk)
-                    chunks_timings.append(time.time() - start_time)
-
-        if chunk is not None and isinstance(chunk, StreamPromptingSynapse):
-            # Assuming the last chunk holds the last value yielded which should be a synapse with the completion filled
-            return ProcessedStreamResponse(
-                synapse=chunk,
-                streamed_chunks=chunks,
-                streamed_chunks_timings=chunks_timings,
-            )
-        else:
-            raise ValueError("The last chunks not a StreamPrompting synapse")
 
     async def get_stream_response(self, params: QueryValidatorParams) -> StreamResponse:
-        response = StreamResponse(status=200, reason="OK")
-        response.headers["Content-Type"] = "application/json"
-
-        await response.prepare(params.request)  # Prepare and send the headers
-
-        try:
-            # Guess the task name of current request
-            # task_name = utils.guess_task_name(params.messages[-1])
-
-            # Get the list of uids to query for this step.
-            uids = get_random_uids(
-                self.validator, k=params.k_miners, exclude=params.exclude or []
-            ).tolist()
-            axons = [self.validator.metagraph.axons[uid] for uid in uids]
-
-            # Make calls to the network with the prompt.
-            bt.logging.info(f"Calling dendrite")
-            start_time = time.time()
-
-            streams_responses = await self.validator.dendrite(
-                axons=axons,
-                synapse=StreamPromptingSynapse(
-                    roles=params.roles, messages=params.messages
-                ),
-                timeout=params.timeout,
-                deserialize=False,
-                streaming=True,
-            )
-
-            uid_stream_dict = dict(zip(uids, streams_responses))
-
-            random_uid, random_stream = random.choice(list(uid_stream_dict.items()))
-            processed_response = await self.process_response(response, random_stream)
-
-            # Prepare final JSON chunk
-            response_data = json.dumps(
-                TextStreamResponse(
-                    streamed_chunks=processed_response.streamed_chunks,
-                    streamed_chunks_timings=processed_response.streamed_chunks_timings,
-                    uid=random_uid,
-                    completion=processed_response.synapse.completion,
-                    timing=time.time() - start_time,
-                ).to_dict()
-            )
-
-            # Send the final JSON as part of the stream
-            await response.write(json.dumps(response_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
-
+        # Guess the task name of current request
+        # task_name = utils.guess_task_name(params.messages[-1])
+
+        # Get the list of uids to query for this step.
+        uids = get_random_uids(
+            self.validator, k=params.k_miners, exclude=params.exclude or []
+        ).tolist()
+        axons = [self.validator.metagraph.axons[uid] for uid in uids]
+
+        # Make calls to the network with the prompt.
+        bt.logging.info(f"Calling dendrite")
+
+        streams_responses = await self.validator.dendrite(
+            axons=axons,
+            synapse=StreamPromptingSynapse(
+                roles=params.roles, messages=params.messages
+            ),
+            timeout=params.timeout,
+            deserialize=False,
+            streaming=True,
+        )
+        uid_stream_dict = dict(zip(uids, streams_responses))
+        random_uid, random_stream = random.choice(list(uid_stream_dict.items()))                        
+        
+        # Creates a streamer from the selected stream
+        streamer = AsyncResponseDataStreamer(async_iterator=random_stream, selected_uid=random_uid)        
+        response = await streamer.stream(params.request)                        
         return response
 
     async def query_validator(self, params: QueryValidatorParams) -> Response:

validators/streamer.py

@@ -0,0 +1,121 @@
+import json
+import time
+import traceback
+import bittensor as bt
+from pydantic import BaseModel
+from datetime import datetime
+from typing import AsyncIterator, Optional, List
+from aiohttp import web, web_response
+from prompting.protocol import StreamPromptingSynapse
+
+
+class StreamChunk(BaseModel):
+    delta: str
+    finish_reason: Optional[str]
+    accumulated_chunks: List[str]
+    accumulated_chunks_timings: List[float]
+    timestamp: str
+    sequence_number: int
+    selected_uid: int
+    
+    def encode(self, encoding: str) -> bytes:
+        data = json.dumps(self.dict(), indent=4)
+        return data.encode(encoding)
+
+
+class StreamError(BaseModel):
+    error: str
+    timestamp: str
+    sequence_number: int
+    finish_reason: str = 'error'
+
+    def encode(self, encoding: str) -> bytes:
+        data = json.dumps(self.dict(), indent=4)
+        return data.encode(encoding)
+
+
+class AsyncResponseDataStreamer:
+    def __init__(self, async_iterator: AsyncIterator, selected_uid:int, delay: float = 0.1):
+        self.async_iterator = async_iterator        
+        self.delay = delay
+        self.selected_uid = selected_uid
+        self.accumulated_chunks: List[str] = []
+        self.accumulated_chunks_timings: List[float] = []
+        self.finish_reason: str = None
+        self.sequence_number: int = 0
+
+    async def stream(self, request: web.Request) -> web_response.StreamResponse:    
+        response = web_response.StreamResponse(status=200, reason="OK")
+        response.headers["Content-Type"] = "application/json"
+        await response.prepare(request)  # Prepare and send the headers
+                        
+        try:
+            start_time = time.time()
+            async for chunk in self.async_iterator:
+                 if isinstance(chunk, list):
+                    # Chunks are currently returned in string arrays, so we need to concatenate them
+                    concatenated_chunks = "".join(chunk)
+                    self.accumulated_chunks.append(concatenated_chunks)                    
+                    self.accumulated_chunks_timings.append(time.time() - start_time)
+                    # Gets new response state
+                    self.sequence_number += 1
+                    new_response_state = self._create_chunk_response(concatenated_chunks)                    
+                    # Writes the new response state to the response                                    
+                    await response.write(new_response_state.encode('utf-8'))
+                                
+            if chunk is not None and isinstance(chunk, StreamPromptingSynapse):                                                            
+                self.finish_reason = "completed"
+                self.sequence_number += 1
+                # Assuming the last chunk holds the last value yielded which should be a synapse with the completion filled
+                synapse = chunk 
+                final_state = self._create_chunk_response(synapse.completion)
+                await response.write(final_state.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")
+            error_response = self._create_error_response(str(e))
+            response.write(error_response.encode('utf-8'))
+        finally:
+            await response.write_eof()  # Ensure to close the response properly
+            return response
+
+    def _create_chunk_response(self, chunk: str) -> StreamChunk:
+        """
+        Creates a StreamChunk object with the current state.
+
+        :param chunk: List of strings representing the current chunk.
+        :return: StreamChunk object.
+        """
+        return StreamChunk(
+            delta=chunk,
+            finish_reason=self.finish_reason,
+            accumulated_chunks=self.accumulated_chunks,
+            accumulated_chunks_timings=self.accumulated_chunks_timings,
+            timestamp=self._current_timestamp(),
+            sequence_number=self.sequence_number,
+            selected_uid=self.selected_uid
+        )
+
+    def _create_error_response(self, error_message: str) -> StreamError:
+        """
+        Creates a StreamError object.
+
+        :param error_message: Error message to include in the StreamError.
+        :return: StreamError object.
+        """
+        return StreamError(
+            error=error_message,
+            timestamp=self._current_timestamp(),
+            sequence_number=self.sequence_number
+        )
+
+    def _current_timestamp(self) -> str:
+        """
+        Returns the current timestamp in ISO format.
+
+        :return: Current timestamp as a string.
+        """
+        return datetime.utcnow().isoformat()