add middlewares + api refactoring

api.py

@@ -1,107 +0,0 @@
-
-import json
-import asyncio
-
-import traceback
-import bittensor as bt
-
-import utils
-
-from typing import List
-from neurons.validator import Validator
-from prompting.forward import handle_response
-from prompting.dendrite import DendriteResponseEvent
-from prompting.protocol import PromptingSynapse, StreamPromptingSynapse
-from prompting.utils.uids import get_random_uids
-
-from aiohttp import web
-
-from aiohttp.web_response import Response
-
-
-async def single_response(validator: Validator, messages: List[str], roles: List[str], k: int = 5, timeout: float = 3.0, exclude: List[int] = None, prefer: str = 'longest') -> Response:
-
-    try:
-        # Guess the task name of current request
-        task_name = utils.guess_task_name(messages[-1])
-
-        # Get the list of uids to query for this step.
-        uids = get_random_uids(validator, k=k, exclude=exclude or []).tolist()
-        axons = [validator.metagraph.axons[uid] for uid in uids]
-
-        # Make calls to the network with the prompt.
-        bt.logging.info(f'Calling dendrite')
-        responses = await validator.dendrite(
-            axons=axons,
-            synapse=PromptingSynapse(roles=roles, messages=messages),
-            timeout=timeout,
-        )
-
-        bt.logging.info(f"Creating DendriteResponseEvent:\n {responses}")
-        # Encapsulate the responses in a response event (dataclass)
-        response_event = DendriteResponseEvent(responses, uids)
-
-        # 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=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 {single_response.__name__}:\n{traceback.format_exc()}')
-        return Response(status=500, reason="Internal error")
-
-
-async def stream_response(validator: Validator, messages: List[str], roles: List[str], k: int = 5, timeout: float = 3.0, exclude: List[int] = None, prefer: str = 'longest') -> web.StreamResponse:
-
-    try:
-        # Guess the task name of current request
-        task_name = utils.guess_task_name(messages[-1])
-
-        # Get the list of uids to query for this step.
-        uids = get_random_uids(validator, k=k, exclude=exclude or []).tolist()
-        axons = [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 validator.dendrite(
-            axons=axons,
-            synapse=StreamPromptingSynapse(roles=roles, messages=messages),
-            timeout=timeout,
-            deserialize=False,
-            streaming=True,
-        )
-
-        # Prepare the task for handling stream responses
-        handle_stream_responses_task = asyncio.create_task(
-            handle_response(responses=dict(zip(uids, streams_responses)))
-        )
-
-        stream_results = await handle_stream_responses_task
-
-        responses = [stream_result.synapse for stream_result in stream_results]
-        bt.logging.info(f"Creating DendriteResponseEvent:\n {responses}")
-        # Encapsulate the responses in a response event (dataclass)
-        response_event = DendriteResponseEvent(responses, uids)
-
-        # 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=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 {single_response.__name__}:\n{traceback.format_exc()}')
-        return Response(status=500, reason="Internal error")

forward.py

@@ -1,244 +1,244 @@
-import time
-import sys
-import asyncio
-import numpy as np
-import bittensor as bt
-import traceback
-from typing import List, Dict, Awaitable
-from prompting.agent import HumanAgent
-from prompting.dendrite import DendriteResponseEvent
-from prompting.conversation import create_task
-from prompting.protocol import StreamPromptingSynapse
-from prompting.rewards import RewardResult
-from prompting.utils.uids import get_random_uids
-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)
-    return result
-
-@async_log
-async def execute_dendrite_call(dendrite_call):
-    responses = await dendrite_call
-    return responses
-
-@dataclass
-class StreamResult:
-    synapse: StreamPromptingSynapse = None
-    exception: BaseException = None
-    uid: int = None
-
-
-async def process_response(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}")
-
-        if chunk is not None:
-            synapse = chunk  # last object yielded is the synapse itself with completion filled
-
-            # Assuming chunk holds the last value yielded which should be a synapse
-            if isinstance(synapse, StreamPromptingSynapse):
-                return synapse
-
-        bt.logging.debug(
-            f"Synapse is not StreamPromptingSynapse. Miner uid {uid} completion set to '' "
-        )
-    except Exception as e:
-        # bt.logging.error(f"Error in generating reference or handling responses: {e}", exc_info=True)
-        traceback_details = traceback.format_exc()
-        bt.logging.error(
-            f"Error in generating reference or handling responses for uid {uid}: {e}\n{traceback_details}"
-        )
-
-        failed_synapse = StreamPromptingSynapse(
-            roles=["user"], messages=["failure"], completion=""
-        )
-
-        return failed_synapse
-
-
-@async_log
-async def handle_response(responses: Dict[int, Awaitable]) -> List[StreamResult]:
-    """The handle_response function is responsible for creating asyncio tasks around acquiring streamed miner chunks
-    and processing them asynchronously. It then pairs the results with their original UIDs and returns a list of StreamResults.
-
-    Args:
-        responses (Dict[int, Awaitable]): Responses contains awaitables that are used to acquire streamed miner chunks.
-
-    Raises:
-        ValueError
-
-    Returns:
-        List[StreamResult]: DataClass containing the synapse, exception, and uid
-    """
-    tasks_with_uid = [
-        (uid, responses[uid]) for uid, _ in responses.items()
-    ]  # Pair UIDs with their tasks
-
-    # 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
-    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 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
-        elif isinstance(result, BaseException):
-            failed_synapse = StreamPromptingSynapse(
-                roles=["user"], messages=["failure"], completion=""
-            )
-            mapped_results.append(
-                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
-        else:
-            bt.logging.error(f"Unexpected result type for UID {uid}: {result}")
-            raise ValueError(f"Unexpected result type for UID {uid}: {result}")
-
-    return mapped_results
-
-
-@async_log
-async def generate_reference(agent: HumanAgent):
-    loop = asyncio.get_running_loop()
-    result = await loop.run_in_executor(
-        None, agent.task.generate_reference, agent.llm_pipeline
-    )
-    return result
-
-
-def log_stream_results(stream_results: List[StreamResult]):
-    failed_responses = [
-        response for response in stream_results if response.exception is not None
-    ]
-    empty_responses = [
-        response
-        for response in stream_results
-        if response.exception is None and response.synapse.completion == ""
-    ]
-    non_empty_responses = [
-        response
-        for response in stream_results
-        if response.exception is None and response.synapse.completion != ""
-    ]
-
-    bt.logging.info(f"Total of non_empty responses: ({len(non_empty_responses)})")
-    bt.logging.info(f"Total of empty responses: ({len(empty_responses)})")
-    bt.logging.info(
-        f"Total of failed responses: ({len(failed_responses)}):\n {failed_responses}"
-    )
-
-    for failed_response in failed_responses:
-        formatted_exception = serialize_exception_to_string(failed_response.exception)
-        bt.logging.error(
-            f"Failed response for uid {failed_response.uid}: {formatted_exception}"
-        )
-
-
-async def run_step(
-    self, agent: HumanAgent, k: int, timeout: float, exclude: list = None
-):
-    """Executes a single step of the agent, which consists of:
-    - Getting a list of uids to query
-    - Querying the network
-    - Rewarding the network
-    - Updating the scores
-    - Logging the event
-
-    Args:
-        agent (HumanAgent): The agent to run the step for.
-        k (int): The number of uids to query.
-        timeout (float): The timeout for the queries.
-        exclude (list, optional): The list of uids to exclude from the query. Defaults to [].
-    """
-
-    bt.logging.debug("run_step", agent.task.name)
-
-    # Record event start time.
-    start_time = time.time()
-    # Get the list of uids to query for this step.
-    uids = get_random_uids(self, k=k, exclude=exclude or []).to(self.device)
-    uids_cpu = uids.cpu().tolist()
-
-    axons = [self.metagraph.axons[uid] for uid in uids]
-
-    # Directly call dendrite and process responses in parallel
-    streams_responses = await self.dendrite(
-        axons=axons,
-        synapse=StreamPromptingSynapse(roles=["user"], messages=[agent.challenge]),
-        timeout=timeout,
-        deserialize=False,
-        streaming=True,
-    )
-
-    # Prepare the task for handling stream responses
-    handle_stream_responses_task = asyncio.create_task(
-        handle_response(responses=dict(zip(uids_cpu, streams_responses)))
-    )
-
-    if not agent.task.static_reference:
-        reference_generation_task = generate_reference(agent)
-        _, stream_results = await asyncio.gather(
-            reference_generation_task, handle_stream_responses_task
-        )
-    else:
-        stream_results = await handle_stream_responses_task
-
-    log_stream_results(stream_results)
-
-    all_synapses_results = [stream_result.synapse for stream_result in stream_results]
-
-    # Encapsulate the responses in a response event (dataclass)
-    response_event = DendriteResponseEvent(
-        responses=all_synapses_results, uids=uids, timeout=timeout
-    )
-
-    bt.logging.info(f"Created DendriteResponseEvent:\n {response_event}")
-    # Reward the responses and get the reward result (dataclass)
-    # This contains a list of RewardEvents but can be exported as a dict (column-wise) for logging etc
-    reward_result = RewardResult(
-        self.reward_pipeline,
-        agent=agent,
-        response_event=response_event,
-        device=self.device,
-    )
-    bt.logging.info(f"Created RewardResult:\n {reward_result}")
-
-    # The original idea was that the agent is 'satisfied' when it gets a good enough response (e.g. reward critera is met, such as ROUGE>threshold)
-    agent.update_progress(
-        top_reward=reward_result.rewards.max(),
-        top_response=response_event.completions[reward_result.rewards.argmax()],
-    )
-
-    self.update_scores(reward_result.rewards, uids)
-
-    stream_results_uids = [stream_result.uid for stream_result in stream_results]
-    stream_results_exceptions = [
-        serialize_exception_to_string(stream_result.exception)
-        for stream_result in stream_results
-    ]
-    # Log the step event.
-    event = {
-        "block": self.block,
-        "step_time": time.time() - start_time,
-        "stream_results_uids": stream_results_uids,
-        "stream_results_exceptions": stream_results_exceptions,
-        **agent.__state_dict__(full=self.config.neuron.log_full),
-        **reward_result.__state_dict__(full=self.config.neuron.log_full),
-        **response_event.__state_dict__(),
-    }
-
-    return event
+# import time
+# import sys
+# import asyncio
+# import numpy as np
+# import bittensor as bt
+# import traceback
+# from typing import List, Dict, Awaitable
+# from prompting.agent import HumanAgent
+# from prompting.dendrite import DendriteResponseEvent
+# from prompting.conversation import create_task
+# from prompting.protocol import StreamPromptingSynapse
+# from prompting.rewards import RewardResult
+# from prompting.utils.uids import get_random_uids
+# 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)
+#     return result
+
+# @async_log
+# async def execute_dendrite_call(dendrite_call):
+#     responses = await dendrite_call
+#     return responses
+
+# @dataclass
+# class StreamResult:
+#     synapse: StreamPromptingSynapse = None
+#     exception: BaseException = None
+#     uid: int = None
+
+
+# async def process_response(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}")
+
+#         if chunk is not None:
+#             synapse = chunk  # last object yielded is the synapse itself with completion filled
+
+#             # Assuming chunk holds the last value yielded which should be a synapse
+#             if isinstance(synapse, StreamPromptingSynapse):
+#                 return synapse
+
+#         bt.logging.debug(
+#             f"Synapse is not StreamPromptingSynapse. Miner uid {uid} completion set to '' "
+#         )
+#     except Exception as e:
+#         # bt.logging.error(f"Error in generating reference or handling responses: {e}", exc_info=True)
+#         traceback_details = traceback.format_exc()
+#         bt.logging.error(
+#             f"Error in generating reference or handling responses for uid {uid}: {e}\n{traceback_details}"
+#         )
+
+#         failed_synapse = StreamPromptingSynapse(
+#             roles=["user"], messages=["failure"], completion=""
+#         )
+
+#         return failed_synapse
+
+
+# @async_log
+# async def handle_response(responses: Dict[int, Awaitable]) -> List[StreamResult]:
+#     """The handle_response function is responsible for creating asyncio tasks around acquiring streamed miner chunks
+#     and processing them asynchronously. It then pairs the results with their original UIDs and returns a list of StreamResults.
+
+#     Args:
+#         responses (Dict[int, Awaitable]): Responses contains awaitables that are used to acquire streamed miner chunks.
+
+#     Raises:
+#         ValueError
+
+#     Returns:
+#         List[StreamResult]: DataClass containing the synapse, exception, and uid
+#     """
+#     tasks_with_uid = [
+#         (uid, responses[uid]) for uid, _ in responses.items()
+#     ]  # Pair UIDs with their tasks
+
+#     # 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
+#     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 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
+#         elif isinstance(result, BaseException):
+#             failed_synapse = StreamPromptingSynapse(
+#                 roles=["user"], messages=["failure"], completion=""
+#             )
+#             mapped_results.append(
+#                 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
+#         else:
+#             bt.logging.error(f"Unexpected result type for UID {uid}: {result}")
+#             raise ValueError(f"Unexpected result type for UID {uid}: {result}")
+
+#     return mapped_results
+
+
+# @async_log
+# async def generate_reference(agent: HumanAgent):
+#     loop = asyncio.get_running_loop()
+#     result = await loop.run_in_executor(
+#         None, agent.task.generate_reference, agent.llm_pipeline
+#     )
+#     return result
+
+
+# def log_stream_results(stream_results: List[StreamResult]):
+#     failed_responses = [
+#         response for response in stream_results if response.exception is not None
+#     ]
+#     empty_responses = [
+#         response
+#         for response in stream_results
+#         if response.exception is None and response.synapse.completion == ""
+#     ]
+#     non_empty_responses = [
+#         response
+#         for response in stream_results
+#         if response.exception is None and response.synapse.completion != ""
+#     ]
+
+#     bt.logging.info(f"Total of non_empty responses: ({len(non_empty_responses)})")
+#     bt.logging.info(f"Total of empty responses: ({len(empty_responses)})")
+#     bt.logging.info(
+#         f"Total of failed responses: ({len(failed_responses)}):\n {failed_responses}"
+#     )
+
+#     for failed_response in failed_responses:
+#         formatted_exception = serialize_exception_to_string(failed_response.exception)
+#         bt.logging.error(
+#             f"Failed response for uid {failed_response.uid}: {formatted_exception}"
+#         )
+
+
+# async def run_step(
+#     self, agent: HumanAgent, k: int, timeout: float, exclude: list = None
+# ):
+#     """Executes a single step of the agent, which consists of:
+#     - Getting a list of uids to query
+#     - Querying the network
+#     - Rewarding the network
+#     - Updating the scores
+#     - Logging the event
+
+#     Args:
+#         agent (HumanAgent): The agent to run the step for.
+#         k (int): The number of uids to query.
+#         timeout (float): The timeout for the queries.
+#         exclude (list, optional): The list of uids to exclude from the query. Defaults to [].
+#     """
+
+#     bt.logging.debug("run_step", agent.task.name)
+
+#     # Record event start time.
+#     start_time = time.time()
+#     # Get the list of uids to query for this step.
+#     uids = get_random_uids(self, k=k, exclude=exclude or []).to(self.device)
+#     uids_cpu = uids.cpu().tolist()
+
+#     axons = [self.metagraph.axons[uid] for uid in uids]
+
+#     # Directly call dendrite and process responses in parallel
+#     streams_responses = await self.dendrite(
+#         axons=axons,
+#         synapse=StreamPromptingSynapse(roles=["user"], messages=[agent.challenge]),
+#         timeout=timeout,
+#         deserialize=False,
+#         streaming=True,
+#     )
+
+#     # Prepare the task for handling stream responses
+#     handle_stream_responses_task = asyncio.create_task(
+#         handle_response(responses=dict(zip(uids_cpu, streams_responses)))
+#     )
+
+#     if not agent.task.static_reference:
+#         reference_generation_task = generate_reference(agent)
+#         _, stream_results = await asyncio.gather(
+#             reference_generation_task, handle_stream_responses_task
+#         )
+#     else:
+#         stream_results = await handle_stream_responses_task
+
+#     log_stream_results(stream_results)
+
+#     all_synapses_results = [stream_result.synapse for stream_result in stream_results]
+
+#     # Encapsulate the responses in a response event (dataclass)
+#     response_event = DendriteResponseEvent(
+#         responses=all_synapses_results, uids=uids, timeout=timeout
+#     )
+
+#     bt.logging.info(f"Created DendriteResponseEvent:\n {response_event}")
+#     # Reward the responses and get the reward result (dataclass)
+#     # This contains a list of RewardEvents but can be exported as a dict (column-wise) for logging etc
+#     reward_result = RewardResult(
+#         self.reward_pipeline,
+#         agent=agent,
+#         response_event=response_event,
+#         device=self.device,
+#     )
+#     bt.logging.info(f"Created RewardResult:\n {reward_result}")
+
+#     # The original idea was that the agent is 'satisfied' when it gets a good enough response (e.g. reward critera is met, such as ROUGE>threshold)
+#     agent.update_progress(
+#         top_reward=reward_result.rewards.max(),
+#         top_response=response_event.completions[reward_result.rewards.argmax()],
+#     )
+
+#     self.update_scores(reward_result.rewards, uids)
+
+#     stream_results_uids = [stream_result.uid for stream_result in stream_results]
+#     stream_results_exceptions = [
+#         serialize_exception_to_string(stream_result.exception)
+#         for stream_result in stream_results
+#     ]
+#     # Log the step event.
+#     event = {
+#         "block": self.block,
+#         "step_time": time.time() - start_time,
+#         "stream_results_uids": stream_results_uids,
+#         "stream_results_exceptions": stream_results_exceptions,
+#         **agent.__state_dict__(full=self.config.neuron.log_full),
+#         **reward_result.__state_dict__(full=self.config.neuron.log_full),
+#         **response_event.__state_dict__(),
+#     }
+
+#     return event

middlewares.py

@@ -0,0 +1,34 @@
+import os
+import json
+import bittensor as bt
+from aiohttp.web import Response
+
+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}")
+
+        # 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
+
+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")
+
+        # Continue to the next handler if JSON is successfully parsed
+        return await handler(request)
+    return middleware_handler

requirements.txt

@@ -1,2 +1,3 @@
 git+https://github.com/opentensor/prompting.git@features/move-validator-into-prompting
-aiohttp
+aiohttp
+deprecated

server.py

@@ -1,16 +1,10 @@
-
-
-
-import os
-import time
 import asyncio
-import json
+import utils
 import bittensor as bt
-from collections import Counter
-from validator_wrapper import QueryValidatorParams, S1ValidatorWrapper
-from prompting.rewards import DateRewardModel, FloatDiffModel
 from aiohttp import web
 from aiohttp.web_response import Response
+from validators import S1ValidatorWrapper, QueryValidatorParams
+from middlewares import api_key_middleware, json_parsing_middleware
 
 """
 # test
@@ -41,135 +35,34 @@ EXPECTED_ACCESS_KEY="hey-michal" python app.py --neuron.model_id mock --wallet.n
 add --mock to test the echo stream
 """
 
-EXPECTED_ACCESS_KEY = os.environ.get('EXPECTED_ACCESS_KEY')
-
 validator = None
 
-
+@api_key_middleware
+@json_parsing_middleware
 async def chat(request: web.Request) -> Response:
     """
     Chat endpoint for the validator.
-
-    Required headers:
-    - api_key: The access key for the validator.
-
-    Required body:
-    - roles: The list of roles to query.
-    - messages: The list of messages to query.
-    Optional body:
-    - k: The number of nodes to query.
-    - exclude: The list of nodes to exclude from the query.
-    - timeout: The timeout for the query.
-    """
-
-    bt.logging.info(f'chat()')
-    # 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")
-
-    try:
-        request_data = await request.json()
-    except ValueError:
-        bt.logging.error(f'Invalid request data: {request_data}')
-        return Response(status=400)
-        
-    # try:
-    #     # Guess the task name of current request
-    #     task_name = guess_task_name(request_data['messages'][-1])
-                        
-    #     # Get the list of uids to query for this step.
-    #     params = QueryValidatorParams.from_dict(request_data)
-    #     response_event = await validator.query_validator(params)        
-                
-    #     # convert dict to json
-    #     response = response_event.__state_dict__()
-        
-    #     response['completion_is_valid'] = valid = list(map(completion_is_valid, response['completions']))
-    #     valid_completions = [response['completions'][i] for i, v in enumerate(valid) if v]
-
-    #     response['task_name'] = task_name
-    #     prefer = request_data.get('prefer', 'longest')
-    #     response['ensemble_result'] = ensemble_result(valid_completions, task_name=task_name, prefer=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 {chat.__name__}:\n{traceback.format_exc()}')
-    #     return Response(status=500, reason="Internal error")
-    bt.logging.info(f'Request data: {request_data}')
-
-    stream = request_data.get('stream', False)
-    if stream:
-        return stream_response(**request_data)
-    else:
-        return single_response(**request_data)
-
-
-
-
-async def echo_stream(request):
-
-    bt.logging.info(f'echo_stream()')
-    # 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")
-
-    try:
-        request_data = await request.json()
-    except ValueError:
-        bt.logging.error(f'Invalid request data: {request_data}')
-        return Response(status=400)
-
-    bt.logging.info(f'Request data: {request_data}')
-    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'})
-    await response.prepare(request)
-
-    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'))
-            completion += chunk
-            time.sleep(timeout)
-            bt.logging.info(f"Echoed: {chunk}")
-
-    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": {}
-    })
-
-    # Send the final JSON as part of the stream
-    await response.write(f"\n\nJSON_RESPONSE_BEGIN:\n{json_chunk}".encode('utf-8'))
-
-    # Finalize the response
-    await response.write_eof()
+    """    
+    request_data = request['data']
+    params = QueryValidatorParams.from_dict(request_data)    
+    # TODO: SET STREAM AS DEFAULT
+    stream = request_data.get('stream', False)        
+    response = await validator.query_validator(params, stream=stream)
     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, *a, **kw):
         super().__init__(*a, **kw)
+        self.middlewares.append(api_key_middleware)
+        self.middlewares.append(json_parsing_middleware)
+
         # TODO: Enable rewarding and other features
 
 

utils.py

@@ -1,8 +1,9 @@
 import re
 import bittensor as bt
-
+import time
+import json
+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"]
@@ -95,7 +96,6 @@ def ensemble_result(completions: list, task_name: str, prefer: str = 'longest'):
     }
 
 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'),
@@ -106,4 +106,47 @@ def guess_task_name(challenge: str):
         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'])
+
+    # Create a StreamResponse
+    response = web.StreamResponse(status=200, reason='OK', headers={'Content-Type': 'text/plain'})
+    await response.prepare()
+
+    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'))
+            completion += chunk
+            time.sleep(timeout)
+            bt.logging.info(f"Echoed: {chunk}")
+
+    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": {}
+    })
+
+    # Send the final JSON as part of the stream
+    await response.write(f"\n\nJSON_RESPONSE_BEGIN:\n{json_chunk}".encode('utf-8'))
+
+    # Finalize the response
+    await response.write_eof()
+    return response

validator_wrapper.py

@@ -1,70 +0,0 @@
-import bittensor as bt
-from prompting.validator import Validator
-from prompting.utils.uids import get_random_uids
-from prompting.protocol import PromptingSynapse
-from prompting.dendrite import DendriteResponseEvent
-from abc import ABC, abstractmethod
-from typing import List
-from dataclasses import dataclass
-
-@dataclass
-class QueryValidatorParams:
-    k_miners: int
-    exclude: List[str]
-    roles: List[str]
-    messages: List[str]
-    timeout: int
-    
-    @staticmethod
-    def from_dict(data: dict):
-        return QueryValidatorParams(            
-            k_miners=data.get('k', 10),
-            exclude=data.get('exclude', []),
-            roles=data['roles'],
-            messages=data['messages'],
-            timeout=data.get('timeout', 10)
-        )
-
-class ValidatorWrapper(ABC):
-    @abstractmethod
-    async def query_validator(self, params:QueryValidatorParams):
-        pass
-    
-    
-class S1ValidatorWrapper(ValidatorWrapper):
-    def __init__(self):
-        self.validator = Validator()    
-    
-    async def query_validator(self, params:QueryValidatorParams) -> DendriteResponseEvent:
-        # Get the list of uids to query for this step.
-        uids = get_random_uids(
-            self.validator,
-            k=params.k_miners,
-            exclude=params.exclude).to(self.validator.device)
-        
-        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,
-        )
-        
-        # Encapsulate the responses in a response event (dataclass)
-        bt.logging.info(f"Creating DendriteResponseEvent:\n {responses}")
-        response_event = DendriteResponseEvent(responses, uids, params.timeout)
-        return response_event
-        
-    
-class MockValidator(ValidatorWrapper):    
-    async def query_validator(self, params:QueryValidatorParams) -> DendriteResponseEvent:
-        ...
-        
-        
-        
-
-
-
-