minor integration adjustments

middlewares.py

@@ -1,34 +1,32 @@
 import os
 import json
 import bittensor as bt
-from aiohttp.web import Response
+from aiohttp.web import Request, Response, middleware
 
 EXPECTED_ACCESS_KEY = os.environ.get('EXPECTED_ACCESS_KEY')
 
-async def api_key_middleware(app, handler):
-    async def middleware_handler(request):
-        # Logging the request
-        bt.logging.info(f"Handling {request.method} request to {request.path}")
+@middleware
+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}')
-            return Response(status=401, reason="Invalid access key")
+    # 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}')
+        return Response(status=401, reason="Invalid access key")
 
-        # Continue to the next handler if the API key is valid
-        return await handler(request)
-    return middleware_handler
+    # Continue to the next handler if the API key is valid
+    return await handler(request)    
 
-async def json_parsing_middleware(app, handler):
-    async def middleware_handler(request):
-        try:
-            # Parsing JSON data from the request
-            request['data'] = await request.json()
-        except json.JSONDecodeError as e:
-            bt.logging.error(f'Invalid JSON data: {str(e)}')
-            return Response(status=400, text="Invalid JSON")
+@middleware
+async def json_parsing_middleware(request: Request, handler):    
+    try:
+        # Parsing JSON data from the request
+        request['data'] = await request.json()
+    except json.JSONDecodeError as 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
-        return await handler(request)
-    return middleware_handler
+    # Continue to the next handler if JSON is successfully parsed
+    return await handler(request)

server.py

@@ -34,8 +34,6 @@ EXPECTED_ACCESS_KEY="hey-michal" python app.py --neuron.model_id mock --wallet.n
 ```
 add --mock to test the echo stream
 """
-@api_key_middleware
-@json_parsing_middleware
 async def chat(request: web.Request) -> Response:
     """
     Chat endpoint for the validator.
@@ -43,7 +41,7 @@ async def chat(request: web.Request) -> Response:
     request_data = request['data']
     params = QueryValidatorParams.from_dict(request_data)    
     # TODO: SET STREAM AS DEFAULT
-    stream = request_data.get('stream', False)        
+    stream = request_data.get('stream', True)        
     
     # Access the validator from the application context
     validator: ValidatorAPI = request.app['validator']
@@ -52,29 +50,29 @@ async def chat(request: web.Request) -> Response:
     return response
 
 
-@api_key_middleware
-@json_parsing_middleware
 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.middlewares.append(api_key_middleware)
-        self.middlewares.append(json_parsing_middleware)
-        
+        # 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()

validators/__init__.py

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

validators/sn1_validator_wrapper.py

@@ -1,14 +1,13 @@
 import json
 import utils
+import torch
 import traceback
 import bittensor as bt
-import asyncio
-from prompting.forward import handle_response
 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 base import QueryValidatorParams, ValidatorAPI
+from .base import QueryValidatorParams, ValidatorAPI
 from aiohttp.web_response import Response, StreamResponse
 from deprecated import deprecated
 
@@ -16,7 +15,7 @@ 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:
@@ -37,7 +36,7 @@ class S1ValidatorAPI(ValidatorAPI):
 
             bt.logging.info(f"Creating DendriteResponseEvent:\n {responses}")
             # Encapsulate the responses in a response event (dataclass)
-            response_event = DendriteResponseEvent(responses, uids)
+            response_event = DendriteResponseEvent(responses, torch.LongTensor(uids), params.timeout)
 
             # convert dict to json
             response = response_event.__state_dict__()