runs black

common/utils.py

@@ -153,7 +153,6 @@ async def echo_stream(request: web.Request) -> web.StreamResponse:
     k = request_data.get("k", 1)
     message = "\n\n".join(request_data["messages"])
 
-
     echo_iterator = EchoAsyncIterator(message, k, delay=0.3)
     streamer = AsyncResponseDataStreamer(echo_iterator, selected_uid=0, delay=0.3)
 

server.py

@@ -2,7 +2,13 @@ import asyncio
 
 import bittensor as bt
 from aiohttp import web
-from aiohttp_apispec import docs, request_schema, response_schema, setup_aiohttp_apispec, validation_middleware
+from aiohttp_apispec import (
+    docs,
+    request_schema,
+    response_schema,
+    setup_aiohttp_apispec,
+    validation_middleware,
+)
 
 from common import utils
 from common.middlewares import api_key_middleware, json_parsing_middleware
@@ -10,11 +16,7 @@ from common.schemas import QueryChatSchema, StreamChunkSchema, StreamErrorSchema
 from validators import QueryValidatorParams, S1ValidatorAPI, ValidatorAPI
 
 
-@docs(
-    tags=["Prompting API"],
-    summary="Chat",
-    description="Chat endpoint."
-)
+@docs(tags=["Prompting API"], summary="Chat", description="Chat endpoint.")
 @request_schema(QueryChatSchema)
 @response_schema(StreamChunkSchema, 200)
 @response_schema(StreamErrorSchema, 400)
@@ -32,7 +34,7 @@ async def chat(request: web.Request) -> web.StreamResponse:
 @docs(
     tags=["Prompting API"],
     summary="Echo test",
-    description="Echo endpoint for testing purposes."
+    description="Echo endpoint for testing purposes.",
 )
 @request_schema(QueryChatSchema)
 @response_schema(StreamChunkSchema, 200)
@@ -45,7 +47,9 @@ class ValidatorApplication(web.Application):
     def __init__(self, validator_instance=None, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
-        self["validator"] = validator_instance if validator_instance else S1ValidatorAPI()
+        self["validator"] = (
+            validator_instance if validator_instance else S1ValidatorAPI()
+        )
 
         # Add middlewares to application
         self.add_routes(

validators/__init__.py

@@ -1,4 +1,4 @@
 from .base import QueryValidatorParams, ValidatorAPI, MockValidator
 from .sn1_validator_wrapper import S1ValidatorAPI
 from .streamer import AsyncResponseDataStreamer
-from .stream_manager import StreamManager
+from .stream_manager import StreamManager

validators/database.py

@@ -8,30 +8,32 @@ from .streamer import ProcessedStreamResponse
 class LogDatabase:
     def __init__(self, log_database_path: str):
         self.log_database_path = log_database_path
-        self.ensure_db_exists(log_database_path)        
-
+        self.ensure_db_exists(log_database_path)
 
     def ensure_db_exists(self, file_path):
         if not os.path.exists(file_path):
             # Create an empty JSONL file
-            with open(file_path, 'w') as file:
+            with open(file_path, "w") as file:
                 pass
             # TODO: change log to debug
             bt.logging.info(f"File '{file_path}' created.")
         else:
             bt.logging.info(f"File '{file_path}' already exists.")
 
-
     async def add_streams_to_db(self, stream_responses: ProcessedStreamResponse):
         bt.logging.info(f"Writing streams to the database...")
         try:
-            stream_responses_dict = [dict(stream_response) for stream_response in stream_responses]
-            await self.append_dicts_to_file(self.log_database_path, stream_responses_dict)
+            stream_responses_dict = [
+                dict(stream_response) for stream_response in stream_responses
+            ]
+            await self.append_dicts_to_file(
+                self.log_database_path, stream_responses_dict
+            )
         except Exception as e:
             bt.logging.error(f"Error while adding streams to the database: {e}")
-            raise e                
+            raise e
 
     async def append_dicts_to_file(self, file_path, dictionaries):
-        async with aiofiles.open(file_path, mode='a') as file:
+        async with aiofiles.open(file_path, mode="a") as file:
             for dictionary in dictionaries:
-                await file.write(json.dumps(dictionary) + '\n')
+                await file.write(json.dumps(dictionary) + "\n")

validators/sn1_validator_wrapper.py

@@ -10,11 +10,12 @@ from .streamer import AsyncResponseDataStreamer
 from .validator_utils import get_top_incentive_uids
 from .stream_manager import StreamManager
 
+
 class S1ValidatorAPI(ValidatorAPI):
     def __init__(self):
         self.validator = Validator()
 
-    def sample_uids(self, params: QueryValidatorParams):    
+    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 []
@@ -23,9 +24,11 @@ class S1ValidatorAPI(ValidatorAPI):
         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)
-            
+
+            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:
@@ -33,11 +36,13 @@ class S1ValidatorAPI(ValidatorAPI):
         # task_name = utils.guess_task_name(params.messages[-1])
 
         # Get the list of uids to query for this step.
-        uids =  self.sample_uids(params)
+        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"Sampling dendrite by {params.sampling_mode} with roles {params.roles} and messages {params.messages}")
+        bt.logging.info(
+            f"Sampling dendrite by {params.sampling_mode} with roles {params.roles} and messages {params.messages}"
+        )
 
         streams_responses = await self.validator.dendrite(
             axons=axons,
@@ -48,13 +53,14 @@ class S1ValidatorAPI(ValidatorAPI):
             deserialize=False,
             streaming=True,
         )
-        
+
         # Creates a streamer from the selected stream
         stream_manager = StreamManager()
-        selected_stream = await stream_manager.process_streams(params.request, streams_responses, uids)
-        
-        return selected_stream
+        selected_stream = await stream_manager.process_streams(
+            params.request, streams_responses, uids
+        )
 
+        return selected_stream
 
     async def query_validator(self, params: QueryValidatorParams) -> Response:
         return await self.get_stream_response(params)

validators/stream_manager.py

@@ -9,20 +9,37 @@ from aiohttp.web import Request
 class StreamManager:
     def __init__(self, log_database_path: str = "requests_db.jsonl"):
         self.log_database = LogDatabase(log_database_path)
-    
-    async def process_streams(self, request:Request, streams_responses: List[AsyncIterator], stream_uids: List[int]):
+
+    async def process_streams(
+        self,
+        request: Request,
+        streams_responses: List[AsyncIterator],
+        stream_uids: List[int],
+    ):
         lock = asyncio.Lock()
-        
-        streamers = [AsyncResponseDataStreamer(async_iterator=stream, selected_uid=stream_uid, lock=lock) for stream, stream_uid in zip(streams_responses, stream_uids)]            
-        completed_streams = await asyncio.gather(*[streamer.stream(request) for streamer in streamers])
-        
+
+        streamers = [
+            AsyncResponseDataStreamer(
+                async_iterator=stream, selected_uid=stream_uid, lock=lock
+            )
+            for stream, stream_uid in zip(streams_responses, stream_uids)
+        ]
+        completed_streams = await asyncio.gather(
+            *[streamer.stream(request) for streamer in streamers]
+        )
+
         lock.release()
         bt.logging.info(f"Streams from uids: {stream_uids} processing completed.")
-        
 
         await self.log_database.add_streams_to_db(completed_streams)
         # Gets the first stream that acquired the lock, meaning the first stream that was able to return a non-empty chunk
-        _, selected_stream = next(((streamer, completed_stream) for streamer, completed_stream in zip(streamers, completed_streams) if streamer.lock_acquired), None)
-        
-        return selected_stream 
-        
+        _, selected_stream = next(
+            (
+                (streamer, completed_stream)
+                for streamer, completed_stream in zip(streamers, completed_streams)
+                if streamer.lock_acquired
+            ),
+            None,
+        )
+
+        return selected_stream

validators/streamer.py

@@ -18,7 +18,7 @@ class StreamChunk(BaseModel):
     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)
@@ -28,17 +28,25 @@ class StreamError(BaseModel):
     error: str
     timestamp: str
     sequence_number: int
-    finish_reason: str = 'error'
+    finish_reason: str = "error"
 
     def encode(self, encoding: str) -> bytes:
         data = json.dumps(self.dict(), indent=4)
         return data.encode(encoding)
 
+
 ProcessedStreamResponse = Union[StreamChunk, StreamError]
 
+
 class AsyncResponseDataStreamer:
-    def __init__(self, async_iterator: AsyncIterator, selected_uid:int, lock: asyncio.Lock, delay: float = 0.1):
-        self.async_iterator = async_iterator        
+    def __init__(
+        self,
+        async_iterator: AsyncIterator,
+        selected_uid: int,
+        lock: asyncio.Lock,
+        delay: float = 0.1,
+    ):
+        self.async_iterator = async_iterator
         self.delay = delay
         self.selected_uid = selected_uid
         self.accumulated_chunks: List[str] = []
@@ -47,36 +55,43 @@ class AsyncResponseDataStreamer:
         self.sequence_number: int = 0
         self.lock = lock
         self.lock_acquired = False
-        
-    
-    def ensure_response_is_created(self, initiated_response: web.StreamResponse) -> web.StreamResponse:
+
+    def ensure_response_is_created(
+        self, initiated_response: web.StreamResponse
+    ) -> web.StreamResponse:
         # Creates response if it was not created
         if initiated_response == None:
             initiated_response = web_response.StreamResponse(status=200, reason="OK")
             initiated_response.headers["Content-Type"] = "application/json"
             return initiated_response
-        
+
         return initiated_response
-    
-        
-    async def write_to_stream(self, request: web.Request, initiated_response: web.StreamResponse, stream_chunk: StreamChunk, lock: asyncio.Lock) -> web.StreamResponse:
-        # Try to acquire the lock and sets the lock_acquired flag. Only the stream that acquires the lock should write to the response            
-        if lock.locked() == False:            
+
+    async def write_to_stream(
+        self,
+        request: web.Request,
+        initiated_response: web.StreamResponse,
+        stream_chunk: StreamChunk,
+        lock: asyncio.Lock,
+    ) -> web.StreamResponse:
+        # Try to acquire the lock and sets the lock_acquired flag. Only the stream that acquires the lock should write to the response
+        if lock.locked() == False:
             self.lock_acquired = await lock.acquire()
-        
+
         if initiated_response == None and self.lock_acquired:
             initiated_response = self.ensure_response_is_created(initiated_response)
             # Prepare and send the headers
             await initiated_response.prepare(request)
-            
+
         if self.lock_acquired:
-            await initiated_response.write(stream_chunk.encode('utf-8'))
+            await initiated_response.write(stream_chunk.encode("utf-8"))
         else:
-            bt.logging.info(f"Stream of uid {stream_chunk.selected_uid} was not the first to return, skipping...")
-        
+            bt.logging.info(
+                f"Stream of uid {stream_chunk.selected_uid} was not the first to return, skipping..."
+            )
+
         return initiated_response
-        
-    
+
     async def stream(self, request: web.Request) -> ProcessedStreamResponse:
         # response = web_response.StreamResponse(status=200, reason="OK")
         # response.headers["Content-Type"] = "application/json"
@@ -86,55 +101,61 @@ class AsyncResponseDataStreamer:
             start_time = time.time()
             client_response: web.Response = None
             final_response: ProcessedStreamResponse
-            
+
             async for chunk in self.async_iterator:
-                 if isinstance(chunk, str):
+                if isinstance(chunk, str):
                     # 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.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)                    
+                    new_response_state = self._create_chunk_response(
+                        concatenated_chunks
+                    )
                     # Writes the new response state to the response
-                    client_response = await self.write_to_stream(request, client_response, new_response_state, self.lock)
-                    #await response.write(new_response_state.encode('utf-8'))
-                                
+                    client_response = await self.write_to_stream(
+                        request, client_response, new_response_state, self.lock
+                    )
+                    # await response.write(new_response_state.encode('utf-8'))
+
             if chunk is not None and isinstance(chunk, StreamPromptingSynapse):
-                if len(self.accumulated_chunks) == 0:                    
-                    self.accumulated_chunks.append(chunk.completion)                    
+                if len(self.accumulated_chunks) == 0:
+                    self.accumulated_chunks.append(chunk.completion)
                     self.accumulated_chunks_timings.append(time.time() - start_time)
-                
+
                 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 
+                synapse = chunk
                 final_response = self._create_chunk_response(synapse.completion)
-                
-                if synapse.completion:                
-                    client_response = await self.write_to_stream(request, client_response, final_response, self.lock)
+
+                if synapse.completion:
+                    client_response = await self.write_to_stream(
+                        request, client_response, final_response, self.lock
+                    )
             else:
                 raise ValueError("Stream did not return a valid synapse.")
-                                        
+
         except Exception as e:
             bt.logging.error(
                 f"Encountered an error while processing stream for uid {self.selected_uid} get_stream_response:\n{traceback.format_exc()}"
             )
             error_response = self._create_error_response(str(e))
             final_response = error_response
-            
+
             # Only the stream that acquires the lock should write the error response
             if self.lock_acquired:
                 self.ensure_response_is_created(client_response)
                 client_response.set_status(500, reason="Internal error")
-                client_response.write(error_response.encode('utf-8'))            
+                client_response.write(error_response.encode("utf-8"))
         finally:
             # Only the stream that acquires the lock should close the response
-            if self.lock_acquired:                
+            if self.lock_acquired:
                 self.ensure_response_is_created(client_response)
                 # Ensure to close the response properly
                 await client_response.write_eof()
-                
+
             return final_response
 
     def _create_chunk_response(self, chunk: str) -> StreamChunk:
@@ -151,7 +172,7 @@ class AsyncResponseDataStreamer:
             accumulated_chunks_timings=self.accumulated_chunks_timings,
             timestamp=self._current_timestamp(),
             sequence_number=self.sequence_number,
-            selected_uid=self.selected_uid
+            selected_uid=self.selected_uid,
         )
 
     def _create_error_response(self, error_message: str) -> StreamError:
@@ -164,7 +185,7 @@ class AsyncResponseDataStreamer:
         return StreamError(
             error=error_message,
             timestamp=self._current_timestamp(),
-            sequence_number=self.sequence_number
+            sequence_number=self.sequence_number,
         )
 
     def _current_timestamp(self) -> str:
@@ -173,4 +194,4 @@ class AsyncResponseDataStreamer:
 
         :return: Current timestamp as a string.
         """
-        return datetime.utcnow().isoformat()
+        return datetime.utcnow().isoformat()

validators/validator_utils.py

@@ -3,21 +3,33 @@ 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)))
-    
+    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))
+        "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']))
+    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)
+    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     
+
+    return top_k_uids