Merge pull request #4 from macrocosm-os/features/streamer

README.md

@@ -1,12 +1,12 @@
 
 <picture>
     <source srcset="./assets/macrocosmos-white.png"  media="(prefers-color-scheme: dark)">
-    <img src="macrocosmos-white.png">
+    <img src="assets/macrocosmos-white.png">
 </picture>
 
 <picture>
     <source srcset="./assets/macrocosmos-black.png"  media="(prefers-color-scheme: light)">
-    <img src="macrocosmos-black.png">
+    <img src="assets/macrocosmos-black.png">
 </picture>
 
 
@@ -75,28 +75,36 @@ EXPECTED_ACCESS_KEY=<ACCESS_KEY> pm2 start server.py --interpreter python3 --nam
 ## API Usage
 At present, the API provides two endpoints: `/chat` (live) and `/echo` (test). 
 
-`/chat` is used to chat with the network and receive a response. The endpoint requires a JSON payload with the following fields:
-- `k: int`: The number of responses to return
-- `timeout: float`: The time in seconds to wait for a response
-- `roles: List[str]`: The roles of the agents to query
-- `messages: List[str]`: The messages to send to the network
-- `prefer: str`: The preferred response to use as the default view. Should be one of `{'longest', 'shortest'}`
-
-Responses from the `/chat` endpoint are streamed back to the client as they are received from the network. Upon completion, the server will return a JSON response with the following fields:
-- `streamed_chunks: List[str]`: The streamed responses from the network
-- `streamed_chunks_timings: List[float]`: The time taken to receive each streamed response
-- `synapse: StreamPromptingSynapse`: The synapse used to query the network. This contains full context and metadata about the query.
-
-The `StreamPromptingSynapse` object contains the following fields:
-- `uid: int`: The unique identifier of the synapse
-- `completion: str`: The final response from the network
-- `timing: float`: The total time taken to receive the final response
+`/chat` is used to chat with the network and receive a response. It requires a JSON payload structured as per the QueryValidatorParams class.
+The request payload requires the following parameters encapsulated within the [`QueryValidatorParams`](./validators/base.py) data class:
+- `k_miners: int`: The number of miners from which to request responses.
+- `exclude: List[str]`: A list of roles or agents to exclude from querying.
+- `roles: List[str]`: The roles of the agents to query.
+- `messages: List[str]`: The messages to be sent to the network.
+- `timeout: int`: The time in seconds to wait for a response.
+- `prefer: str`: The preferred response format, can be either `'longest'` or `'shortest'`.
+- `request: Request`: The original request object encapsulating all request data.
+- `sampling_mode: str`: The mode of sampling to use, defaults to `"random"`. Can be either `"random"` or `"top_incentive"`
+
+Responses from the `/chat` endpoint are handled by two classes: `StreamChunk` and `StreamError`, with their attributes defined as follows:
+- `StreamChunk`:
+  - `delta: str`: The new chunk of response received.
+  - `finish_reason: Optional[str]`: The reason for the response completion, if applicable Can be `None`, `finished` or `error` (check StreamError below).
+  - `accumulated_chunks: List[str]`: All chunks of responses accumulated thus far.
+  - `accumulated_chunks_timings: List[float]`: Timing for each chunk received.
+  - `timestamp: str`: The timestamp at which the chunk was processed.
+  - `sequence_number: int`: A sequential identifier for the response part.
+  - `selected_uid: int`: The identifier for the selected response source.
+
+- `StreamError`:
+  - `error: str`: Description of the error occurred.
+  - `timestamp: str`: The timestamp of the error.
+  - `sequence_number: int`: A sequential identifier for the error.
+  - `finish_reason: str`: Always set to `'error'` to indicate an error completion.
 
 > Note: The API is subject to change as the project evolves.
 
-
 ## Testing
-
 To test the API locally, you can use the following curl command:
 
 ```bash

responses.py

@@ -1,27 +0,0 @@
-from pydantic import BaseModel, Field
-from typing import List, Dict, Any
-
-
-class TextStreamResponse(BaseModel):
-    streamed_chunks: List[str] = Field(
-        default_factory=list, description="List of streamed chunks."
-    )
-    streamed_chunks_timings: List[float] = Field(
-        default_factory=list, description="List of streamed chunks timings, in seconds."
-    )
-    uid: int = Field(0, description="UID of queried miner")
-    completion: str = Field(
-        "", description="The final completed string from the stream."
-    )
-    timing: float = Field(
-        0, description="Timing information of all request, in seconds."
-    )
-
-    def to_dict(self):
-        return {
-            "streamed_chunks": self.streamed_chunks,
-            "streamed_chunks_timings": self.streamed_chunks_timings,
-            "uid": self.uid,
-            "completion": self.completion,
-            "timing": self.timing,
-        }

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/base.py

@@ -13,6 +13,7 @@ class QueryValidatorParams:
     timeout: int
     prefer: str
     request: Request
+    sampling_mode: str
 
     @staticmethod
     def from_request(request: Request):
@@ -26,6 +27,7 @@ class QueryValidatorParams:
             timeout=data.get("timeout", 10),
             prefer=data.get("prefer", "longest"),
             request=request,
+            sampling_mode=data.get("sampling_mode", "random"),
         )
 
 

validators/sn1_validator_wrapper.py

@@ -1,182 +1,58 @@
-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
-
-
-@dataclass
-class ProcessedStreamResponse:
-    streamed_chunks: List[str]
-    streamed_chunks_timings: List[float]
-    synapse: StreamPromptingSynapse
+from .streamer import AsyncResponseDataStreamer
+from .validator_utils import get_top_incentive_uids
 
 
 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.
+    def sample_uids(self, params: QueryValidatorParams):    
+        if params.sampling_mode == "random":
             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 chunkis not a StreamPrompting synapse")
+            return uids
+        if params.sampling_mode == "top_incentive":
+            metagraph = self.validator.metagraph
+            vpermit_tao_limit = self.validator.config.neuron.vpermit_tao_limit
+            
+            top_uids = get_top_incentive_uids(metagraph, k=params.k_miners, vpermit_tao_limit=vpermit_tao_limit)
+            
+            return top_uids
 
     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 =  self.sample_uids(params)
+        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()

validators/validator_utils.py

@@ -0,0 +1,23 @@
+from typing import List
+from prompting.utils.uids import check_uid_availability
+
+
+def get_top_incentive_uids(metagraph, k: int, vpermit_tao_limit: int) -> List[int]:
+    miners_uids = list(map(int, filter(lambda uid: check_uid_availability(metagraph, uid, vpermit_tao_limit), metagraph.uids)))
+    
+    # Builds a dictionary of uids and their corresponding incentives
+    all_miners_incentives = {
+        "miners_uids": miners_uids,
+        "incentives": list(map(lambda uid: metagraph.I[uid], miners_uids))
+    }
+    
+    # Zip the uids and their corresponding incentives into a list of tuples
+    uid_incentive_pairs = list(zip(all_miners_incentives['miners_uids'], all_miners_incentives['incentives']))
+
+    # Sort the list of tuples by the incentive value in descending order
+    uid_incentive_pairs_sorted = sorted(uid_incentive_pairs, key=lambda x: x[1], reverse=True)
+
+    # Extract the top 10 uids
+    top_k_uids = [uid for uid, incentive in uid_incentive_pairs_sorted[:k]]
+    
+    return top_k_uids