runs black

forward.py

@@ -15,17 +15,22 @@ 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)
+    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
@@ -82,20 +87,19 @@ async def handle_response(responses: Dict[int, Awaitable]) -> List[StreamResult]
         (uid, responses[uid]) for uid, _ in responses.items()
     ]  # Pair UIDs with their tasks
 
-    #Start tasks, preserving order and their associated UIDs
+    # 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
+    # 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 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
+        # 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=""
@@ -104,7 +108,7 @@ async def handle_response(responses: Dict[int, Awaitable]) -> List[StreamResult]
                 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
+        # 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}")

middlewares.py

@@ -3,29 +3,31 @@ import json
 import bittensor as bt
 from aiohttp.web import Request, Response, middleware
 
-EXPECTED_ACCESS_KEY = os.environ.get('EXPECTED_ACCESS_KEY')
+EXPECTED_ACCESS_KEY = os.environ.get("EXPECTED_ACCESS_KEY")
+
 
 @middleware
-async def api_key_middleware(request: Request, handler):    
+async def api_key_middleware(request: Request, handler):
     # Logging the request
     bt.logging.info(f"Handling {request.method} request to {request.path}")
 
     # Check access key
     access_key = request.headers.get("api_key")
     if EXPECTED_ACCESS_KEY is not None and access_key != EXPECTED_ACCESS_KEY:
-        bt.logging.error(f'Invalid access key: {access_key}')
+        bt.logging.error(f"Invalid access key: {access_key}")
         return Response(status=401, reason="Invalid access key")
 
     # Continue to the next handler if the API key is valid
-    return await handler(request)    
+    return await handler(request)
+
 
 @middleware
-async def json_parsing_middleware(request: Request, handler):    
+async def json_parsing_middleware(request: Request, handler):
     try:
         # Parsing JSON data from the request
-        request['data'] = await request.json()
+        request["data"] = await request.json()
     except json.JSONDecodeError as e:
-        bt.logging.error(f'Invalid JSON data: {str(e)}')
+        bt.logging.error(f"Invalid JSON data: {str(e)}")
         return Response(status=400, text="Invalid JSON")
 
     # Continue to the next handler if JSON is successfully parsed

server.py

@@ -34,53 +34,53 @@ EXPECTED_ACCESS_KEY="hey-michal" python app.py --neuron.model_id mock --wallet.n
 ```
 add --mock to test the echo stream
 """
+
+
 async def chat(request: web.Request) -> Response:
     """
     Chat endpoint for the validator.
-    """        
-    params = QueryValidatorParams.from_request(request)    
+    """
+    params = QueryValidatorParams.from_request(request)
 
-    
     # Access the validator from the application context
-    validator: ValidatorAPI = request.app['validator']
-    
+    validator: ValidatorAPI = request.app["validator"]
+
     response = await validator.query_validator(params)
     return response
 
 
-async def echo_stream(request, request_data):    
-    request_data = request['data']
+async def echo_stream(request, request_data):
+    request_data = request["data"]
     return await utils.echo_stream(request_data)
 
 
-
 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()
-        
-        # Add middlewares to application                
-        self.add_routes([
-            web.post('/chat/', chat),
-            web.post('/echo/', echo_stream)
-        ])
+
+        self["validator"] = (
+            validator_instance if validator_instance else S1ValidatorAPI()
+        )
+
+        # Add middlewares to application
+        self.add_routes([web.post("/chat/", chat), web.post("/echo/", echo_stream)])
         self.setup_middlewares()
         # TODO: Enable rewarding and other features
-        
+
     def setup_middlewares(self):
         self.middlewares.append(json_parsing_middleware)
         self.middlewares.append(api_key_middleware)
-        
+
+
 def main(run_aio_app=True, test=False) -> None:
     loop = asyncio.get_event_loop()
     port = 10000
     if run_aio_app:
         # Instantiate the application with the actual validator
         bt.logging.info("Starting validator application.")
-        validator_app = ValidatorApplication()        
-        bt.logging.success(f'Validator app initialized successfully', validator_app)
-        
+        validator_app = ValidatorApplication()
+        bt.logging.success(f"Validator app initialized successfully", validator_app)
+
         try:
             web.run_app(validator_app, port=port, loop=loop)
         except KeyboardInterrupt:
@@ -88,5 +88,6 @@ def main(run_aio_app=True, test=False) -> None:
         finally:
             pass
 
+
 if __name__ == "__main__":
     main()

test.py

@@ -1,10 +1,10 @@
 import pytest
 
 
-
 def test_query_network():
     pass
 
+
 def test_filter_completions():
     pass
 
@@ -12,5 +12,6 @@ def test_filter_completions():
 def test_guess_task_name():
     pass
 
+
 def test_ensemble_completions():
-    pass
+    pass

utils.py

@@ -6,13 +6,26 @@ from aiohttp import web
 from collections import Counter
 from prompting.rewards import DateRewardModel, FloatDiffModel
 
-UNSUCCESSFUL_RESPONSE_PATTERNS = ["I'm sorry", "unable to", "I cannot", "I can't", "I am unable", "I am sorry", "I can not", "don't know", "not sure", "don't understand", "not capable"]
+UNSUCCESSFUL_RESPONSE_PATTERNS = [
+    "I'm sorry",
+    "unable to",
+    "I cannot",
+    "I can't",
+    "I am unable",
+    "I am sorry",
+    "I can not",
+    "don't know",
+    "not sure",
+    "don't understand",
+    "not capable",
+]
 
 reward_models = {
-    'date_qa': DateRewardModel(),
-    'math': FloatDiffModel(),
+    "date_qa": DateRewardModel(),
+    "math": FloatDiffModel(),
 }
 
+
 def completion_is_valid(completion: str):
     """
     Get the completion statuses from the completions.
@@ -20,13 +33,15 @@ def completion_is_valid(completion: str):
     if not completion.strip():
         return False
 
-    patt = re.compile(r'\b(?:' + '|'.join(UNSUCCESSFUL_RESPONSE_PATTERNS) + r')\b', re.IGNORECASE)
-    if not len(re.findall(r'\w+',completion)) or patt.search(completion):
+    patt = re.compile(
+        r"\b(?:" + "|".join(UNSUCCESSFUL_RESPONSE_PATTERNS) + r")\b", re.IGNORECASE
+    )
+    if not len(re.findall(r"\w+", completion)) or patt.search(completion):
         return False
     return True
 
 
-def ensemble_result(completions: list, task_name: str, prefer: str = 'longest'):
+def ensemble_result(completions: list, task_name: str, prefer: str = "longest"):
     """
     Ensemble completions from multiple models.
     # TODO: Measure agreement
@@ -37,11 +52,11 @@ def ensemble_result(completions: list, task_name: str, prefer: str = 'longest'):
         return None
 
     answer = None
-    if task_name in ('qa', 'summarization'):
+    if task_name in ("qa", "summarization"):
         # No special handling for QA or summarization
         supporting_completions = completions
 
-    elif task_name == 'date_qa':
+    elif task_name == "date_qa":
         # filter the completions to be the ones that contain valid dates and if there are multiple dates, select the most common one (with support > 1)
         dates = list(map(reward_models[task_name].parse_dates_from_text, completions))
         bt.logging.info(f"Unprocessed dates: {dates}")
@@ -58,9 +73,11 @@ def ensemble_result(completions: list, task_name: str, prefer: str = 'longest'):
         if count == 1:
             supporting_completions = valid_completions
         else:
-            supporting_completions = [c for i, c in enumerate(valid_completions) if dates[i]==most_common]
+            supporting_completions = [
+                c for i, c in enumerate(valid_completions) if dates[i] == most_common
+            ]
 
-    elif task_name == 'math':
+    elif task_name == "math":
         # filter the completions to be the ones that contain valid numbers and if there are multiple values, select the most common one (with support > 1)
         # TODO: use the median instead of the most common value
         vals = list(map(reward_models[task_name].extract_number, completions))
@@ -74,57 +91,67 @@ def ensemble_result(completions: list, task_name: str, prefer: str = 'longest'):
         if count == 1:
             supporting_completions = completions
         else:
-            supporting_completions = [c for i, c in enumerate(completions) if vals[i]==most_common]
-
+            supporting_completions = [
+                c for i, c in enumerate(completions) if vals[i] == most_common
+            ]
 
     bt.logging.info(f"Supporting completions: {supporting_completions}")
-    if prefer == 'longest':
+    if prefer == "longest":
         preferred_completion = sorted(supporting_completions, key=len)[-1]
-    elif prefer == 'shortest':
+    elif prefer == "shortest":
         preferred_completion = sorted(supporting_completions, key=len)[0]
-    elif prefer == 'most_common':
-        preferred_completion = max(set(supporting_completions), key=supporting_completions.count)
+    elif prefer == "most_common":
+        preferred_completion = max(
+            set(supporting_completions), key=supporting_completions.count
+        )
     else:
         raise ValueError(f"Unknown ensemble preference: {prefer}")
 
     return {
-        'completion': preferred_completion,
-        'accepted_answer': answer,
-        'support': len(supporting_completions),
-        'support_indices': [completions.index(c) for c in supporting_completions],
-        'method': f'Selected the {prefer.replace("_", " ")} completion'
+        "completion": preferred_completion,
+        "accepted_answer": answer,
+        "support": len(supporting_completions),
+        "support_indices": [completions.index(c) for c in supporting_completions],
+        "method": f'Selected the {prefer.replace("_", " ")} completion',
     }
 
+
 def guess_task_name(challenge: str):
     # TODO: use a pre-trained classifier to guess the task name
     categories = {
-        'summarization': re.compile('summar|quick rundown|overview'),
-        'date_qa': re.compile('exact date|tell me when|on what date|on what day|was born?|died?'),
-        'math': re.compile('math|solve|solution| sum |problem|geometric|vector|calculate|degrees|decimal|factorial'),
+        "summarization": re.compile("summar|quick rundown|overview"),
+        "date_qa": re.compile(
+            "exact date|tell me when|on what date|on what day|was born?|died?"
+        ),
+        "math": re.compile(
+            "math|solve|solution| sum |problem|geometric|vector|calculate|degrees|decimal|factorial"
+        ),
     }
     for task_name, patt in categories.items():
         if patt.search(challenge):
             return task_name
 
-    return 'qa'
+    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)
-    message = '\n\n'.join(request_data['messages'])
+    k = request_data.get("k", 1)
+    exclude = request_data.get("exclude", [])
+    timeout = request_data.get("timeout", 0.2)
+    message = "\n\n".join(request_data["messages"])
 
     # Create a StreamResponse
-    response = web.StreamResponse(status=200, reason='OK', headers={'Content-Type': 'text/plain'})
+    response = web.StreamResponse(
+        status=200, reason="OK", headers={"Content-Type": "text/plain"}
+    )
     await response.prepare()
 
-    completion = ''
+    completion = ""
     # 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'))
+            chunk = f"{word} "
+            await response.write(chunk.encode("utf-8"))
             completion += chunk
             time.sleep(timeout)
             bt.logging.info(f"Echoed: {chunk}")
@@ -132,21 +159,23 @@ async def echo_stream(request_data: dict):
     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": {}
-    })
+    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": {},
+        }
+    )
 
     # 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(f"\n\nJSON_RESPONSE_BEGIN:\n{json_chunk}".encode("utf-8"))
 
     # Finalize the response
     await response.write_eof()
-    return response
+    return response

validators/__init__.py

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

validators/base.py

@@ -3,6 +3,7 @@ from typing import List
 from dataclasses import dataclass
 from aiohttp.web import Response, Request
 
+
 @dataclass
 class QueryValidatorParams:
     k_miners: int
@@ -12,28 +13,28 @@ class QueryValidatorParams:
     timeout: int
     prefer: str
     request: Request
-    
+
     @staticmethod
     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'),
-            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"),
+            request=request,
         )
 
+
 class ValidatorAPI(ABC):
     @abstractmethod
-    async def query_validator(self, params:QueryValidatorParams) -> Response:
+    async def query_validator(self, params: QueryValidatorParams) -> Response:
         pass
-    
-    
-class MockValidator(ValidatorAPI):    
-    async def query_validator(self, params:QueryValidatorParams) -> Response:
+
+
+class MockValidator(ValidatorAPI):
+    async def query_validator(self, params: QueryValidatorParams) -> Response:
         ...
-        

validators/sn1_validator_wrapper.py

@@ -13,23 +13,27 @@ from .base import QueryValidatorParams, ValidatorAPI
 from aiohttp.web_response import Response, StreamResponse
 from deprecated import deprecated
 
+
 class S1ValidatorAPI(ValidatorAPI):
     def __init__(self):
-        self.validator = Validator()    
-    
+        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:
+    @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()
+            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')
+            bt.logging.info(f"Calling dendrite")
             responses = await self.validator.dendrite(
                 axons=axons,
                 synapse=PromptingSynapse(roles=params.roles, messages=params.messages),
@@ -38,89 +42,113 @@ class S1ValidatorAPI(ValidatorAPI):
 
             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)
+            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["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)
+            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))
+            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()}')
+            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, uid: int, async_generator: Awaitable):
+
+    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'))
-                    
+                    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()}')
+            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'))
+            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:
+
+    async def get_stream_response(self, params: QueryValidatorParams) -> StreamResponse:
         response = StreamResponse(status=200, reason="OK")
-        response.headers['Content-Type'] = 'application/json'
+        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()
+            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')
+            bt.logging.info(f"Calling dendrite")
             streams_responses = await self.validator.dendrite(
                 axons=axons,
-                synapse=StreamPromptingSynapse(roles=params.roles, messages=params.messages),
+                synapse=StreamPromptingSynapse(
+                    roles=params.roles, messages=params.messages
+                ),
                 timeout=params.timeout,
                 deserialize=False,
                 streaming=True,
             )
-            
-            tasks = [self.process_response(uid, res) for uid, res in dict(zip(uids, streams_responses))]                        
+
+            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      
+
+            # 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()}')
+            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'))
+            await response.write(json.dumps({"error": str(e)}).encode("utf-8"))
         finally:
             await response.write_eof()  # Ensure to close the response properly
 
         return response
 
-
-    async def query_validator(self, params:QueryValidatorParams) -> Response:        
+    async def query_validator(self, params: QueryValidatorParams) -> Response:
         # TODO: SET STREAM AS DEFAULT
-        stream = params.request.get('stream', False)        
-        
+        stream = params.request.get("stream", False)
+
         if stream:
             return await self.get_stream_response(params)
         else:
             # DEPRECATED
             return await self.get_response(params)
-