force streaming calls to all endpoints

server.py

@@ -2,21 +2,16 @@ import asyncio
 import utils
 import bittensor as bt
 from aiohttp import web
-from aiohttp.web_response import Response
 from validators import S1ValidatorAPI, QueryValidatorParams, ValidatorAPI
 from middlewares import api_key_middleware, json_parsing_middleware
 
 """
-# test
+# Test chat endpoint with curl
 ```
-curl -X POST http://0.0.0.0:10000/chat/ -H "api_key: hello" -d '{"k": 5, "timeout": 3, "roles": ["user"], "messages": ["hello world"]}'
+curl --no-buffer -X POST http://0.0.0.0:10000/chat/ -H "api_key: hey-michal" -d '{"k": 5, "timeout": 15, "roles": ["user"], "messages": ["on what exact date did the 21st century begin??"]}'
 
-curl -X POST http://0.0.0.0:10000/chat/ -H "api_key: hey-michal" -d '{"k": 5, "timeout": 3, "roles": ["user"], "messages": ["on what exact date did the 21st century begin?"]}'
-
-curl -X POST http://0.0.0.0:10000/chat/ -H "api_key: hey-michal" -d '{"k": 5, "timeout": 15, "roles": ["user"], "messages": ["who you really are?"]}'
-
-# stream
-curl --no-buffer -X POST http://129.146.127.82:10000/echo/ -H "api_key: hey-michal" -d '{"k": 3, "timeout": 0.2, "roles": ["user"], "messages": ["i need to tell you something important but first"]}'
+# echo stream test endpoint
+curl --no-buffer -X POST http://0.0.0.0:10000/echo/ -H "api_key: hey-michal" -d '{"k": 3, "timeout": 0.2, "roles": ["user"], "messages": ["i need to tell you something important but first"]}'
 ```
 
 TROUBLESHOOT
@@ -25,7 +20,15 @@ check if port is open
 sudo ufw allow 10000/tcp
 sudo ufw allow 10000/tcp
 ```
-# run
+
+---
+
+# Run Chattensor
+
+## With vanilla python:
+python server.py --neuron.model_id mock --wallet.name sn1 --wallet.hotkey v1 --netuid 1 --neuron.tasks math --neuron.task_p 1 --neuron.device cpu --subtensor.network local
+
+## With PM2:
 ```
 EXPECTED_ACCESS_KEY="hey-michal" pm2 start app.py --interpreter python3 --name app -- --neuron.model_id mock --wallet.name sn1 --wallet.hotkey v1 --netuid 1 --neuron.tasks math --neuron.task_p 1 --neuron.device cpu
 ```
@@ -38,7 +41,7 @@ add --mock to test the echo stream
 """
 
 
-async def chat(request: web.Request) -> Response:
+async def chat(request: web.Request) -> web.StreamResponse:
     """
     Chat endpoint for the validator.
     """
@@ -51,9 +54,8 @@ async def chat(request: web.Request) -> Response:
     return response
 
 
-async def echo_stream(request, request_data):
-    request_data = request["data"]
-    return await utils.echo_stream(request_data)
+async def echo_stream(request: web.Request) -> web.StreamResponse:    
+    return await utils.echo_stream(request)
 
 
 class ValidatorApplication(web.Application):

utils.py

@@ -1,8 +1,10 @@
 import re
-import bittensor as bt
 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
 
@@ -134,47 +136,45 @@ def guess_task_name(challenge: str):
     return "qa"
 
 
-async def echo_stream(request_data: dict):
-    k = request_data.get("k", 1)
-    exclude = request_data.get("exclude", [])
-    timeout = request_data.get("timeout", 0.2)
+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": "text/plain"}
+        status=200, reason="OK", headers={"Content-Type": "application/json"}
     )
-    await response.prepare()
+    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} "
+            chunk = f"{word} "            
             await response.write(chunk.encode("utf-8"))
             completion += chunk
-            time.sleep(timeout)
+            await asyncio.sleep(.3)
             bt.logging.info(f"Echoed: {chunk}")
-
+            
+            chunks.append(chunk)
+            chunks_timings.append(time.time() - start_time)
+            
     completion = completion.strip()
 
-    # Prepare final JSON chunk
-    json_chunk = json.dumps(
-        {
-            "uids": [0],
-            "completion": completion,
-            "completions": [completion.strip()],
-            "timings": [0],
-            "status_messages": ["Went well!"],
-            "status_codes": [200],
-            "completion_is_valid": [True],
-            "task_name": "echo",
-            "ensemble_result": {},
-        }
-    )
+    # 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(f"\n\nJSON_RESPONSE_BEGIN:\n{json_chunk}".encode("utf-8"))
+    await response.write(json.dumps(response_data).encode("utf-8"))
 
     # Finalize the response
     await response.write_eof()

validators/base.py

@@ -1,7 +1,7 @@
 from abc import ABC, abstractmethod
 from typing import List
 from dataclasses import dataclass
-from aiohttp.web import Response, Request
+from aiohttp.web import Response, Request, StreamResponse
 
 
 @dataclass
@@ -31,10 +31,10 @@ class QueryValidatorParams:
 
 class ValidatorAPI(ABC):
     @abstractmethod
-    async def query_validator(self, params: QueryValidatorParams) -> Response:
+    async def query_validator(self, params: QueryValidatorParams) -> StreamResponse:
         pass
 
 
 class MockValidator(ValidatorAPI):
-    async def query_validator(self, params: QueryValidatorParams) -> Response:
+    async def query_validator(self, params: QueryValidatorParams) -> StreamResponse:
         ...

validators/sn1_validator_wrapper.py

@@ -2,7 +2,7 @@ import json
 import utils
 import torch
 import traceback
-import asyncio
+import time
 import random
 import bittensor as bt
 from typing import Awaitable
@@ -13,6 +13,16 @@ from prompting.dendrite import DendriteResponseEvent
 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
+
+
+@dataclass
+class ProcessedStreamResponse:
+    streamed_chunks: List[str]
+    streamed_chunks_timings: List[float]
+    synapse: StreamPromptingSynapse
 
 
 class S1ValidatorAPI(ValidatorAPI):
@@ -77,13 +87,39 @@ class S1ValidatorAPI(ValidatorAPI):
 
     async def process_response(
         self, response: StreamResponse, async_generator: Awaitable
-    ):
+    ) -> ProcessedStreamResponse:
         """Process a single response asynchronously."""
-        chunk = None  # Initialize chunk with a default value
-        async for chunk in async_generator:
-            if chunk is not None and hasattr(chunk, 'completion'):
-                # Directly write the string encoded as UTF-8 bytes
-                await response.write(chunk.completion.encode('utf-8'))
+        # 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 chunkis not a StreamPrompting synapse")
+
 
 
     async def get_stream_response(self, params: QueryValidatorParams) -> StreamResponse:
@@ -94,7 +130,7 @@ class S1ValidatorAPI(ValidatorAPI):
 
         try:
             # Guess the task name of current request
-            task_name = utils.guess_task_name(params.messages[-1])
+            # task_name = utils.guess_task_name(params.messages[-1])
 
             # Get the list of uids to query for this step.
             uids = get_random_uids(
@@ -103,7 +139,8 @@ class S1ValidatorAPI(ValidatorAPI):
             axons = [self.validator.metagraph.axons[uid] for uid in uids]
 
             # Make calls to the network with the prompt.
-            bt.logging.info(f"Calling dendrite")                         
+            bt.logging.info(f"Calling dendrite")
+            start_time = time.time()
             
             streams_responses = await self.validator.dendrite(
                 axons=axons,
@@ -115,8 +152,22 @@ class S1ValidatorAPI(ValidatorAPI):
                 streaming=True,
             )
             
-            random_stream = random.choice(streams_responses)            
-            await self.process_response(response, random_stream)                            
+            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()}"
@@ -128,12 +179,5 @@ class S1ValidatorAPI(ValidatorAPI):
 
         return response
 
-    async def query_validator(self, params: QueryValidatorParams) -> Response:
-        # TODO: SET STREAM AS DEFAULT
-        stream = params.request.get("stream", True)
-
-        if stream:
-            return await self.get_stream_response(params)
-        else:
-            # DEPRECATED
-            return await self.get_response(params)
+    async def query_validator(self, params: QueryValidatorParams) -> Response:                
+        return await self.get_stream_response(params)