Add WIP refactor

api.py

@@ -0,0 +1,107 @@
+
+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

@@ -16,10 +16,10 @@ from prompting.utils.misc import async_log, serialize_exception_to_string
 from dataclasses import dataclass
 
 @async_log
-async def generate_reference(agent):    
+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    
+    return result
 
 @async_log
 async def execute_dendrite_call(dendrite_call):
@@ -199,8 +199,8 @@ async def run_step(
 
     log_stream_results(stream_results)
 
-    all_synapses_results = [stream_result.synapse for stream_result in 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

server.py

@@ -2,20 +2,12 @@
 
 
 import os
-import re
 import time
 import asyncio
 import json
-import traceback
 import bittensor as bt
 
-from collections import Counter
-
 from neurons.validator import Validator
-from prompting.dendrite import DendriteResponseEvent
-from prompting.protocol import PromptingSynapse
-from prompting.utils.uids import get_random_uids
-from prompting.rewards import DateRewardModel, FloatDiffModel
 from aiohttp import web
 from aiohttp.web_response import Response
 
@@ -43,7 +35,7 @@ EXPECTED_ACCESS_KEY="hey-michal" pm2 start app.py --interpreter python3 --name a
 
 basic testing
 ```
-EXPECTED_ACCESS_KEY="hey-michal" python app.py --neuron.model_id mock --wallet.name sn1 --wallet.hotkey v1 --netuid 1 --neuron.tasks math --neuron.task_p 1 --neuron.device cpu 
+EXPECTED_ACCESS_KEY="hey-michal" python app.py --neuron.model_id mock --wallet.name sn1 --wallet.hotkey v1 --netuid 1 --neuron.tasks math --neuron.task_p 1 --neuron.device cpu
 ```
 add --mock to test the echo stream
 """
@@ -51,102 +43,7 @@ add --mock to test the echo stream
 EXPECTED_ACCESS_KEY = os.environ.get('EXPECTED_ACCESS_KEY')
 
 validator = None
-reward_models = {
-    'date_qa': DateRewardModel(),
-    'math': FloatDiffModel(),
-}
 
-def completion_is_valid(completion: str):
-    """
-    Get the completion statuses from the completions.
-    """
-    patt = re.compile(r'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')
-    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'):
-    """
-    Ensemble completions from multiple models.
-    # TODO: Measure agreement
-    # TODO: Figure out how to mitigate the cabal effect (large groups will appear to be more credible)
-    # TODO: Reward pipeline
-    """
-    if not completions:
-        return None
-    
-    
-    answer = None
-    if task_name in ('qa', 'summarization'):
-        # No special handling for QA or summarization
-        supporting_completions = completions
-    
-    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}")
-        valid_date_indices = [i for i, d in enumerate(dates) if d]
-        valid_completions = [completions[i] for i in valid_date_indices]
-        valid_dates = [dates[i] for i in valid_date_indices]
-        dates = [f"{d[0].strftime('%-d %B')} {d[1]}" for d in valid_dates]
-        if not dates:
-            return None
-        
-        counter = Counter(dates)
-        most_common, count = counter.most_common()[0]
-        answer = most_common
-        if count == 1:
-            supporting_completions = valid_completions
-        else:
-            supporting_completions = [c for i, c in enumerate(valid_completions) if dates[i]==most_common]
-        
-    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))
-        vals = [val for val in vals if val]
-        if not vals:
-            return None
-        
-        most_common, count = Counter(dates).most_common()[0]
-        bt.logging.info(f"Most common value: {most_common}, count: {count}")
-        answer = most_common
-        if count == 1:
-            supporting_completions = completions
-        else:
-            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':
-        preferred_completion = sorted(supporting_completions, key=len)[-1]
-    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)
-    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'
-    }
-    
-def guess_task_name(challenge: str):
-    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'),
-    }
-    for task_name, patt in categories.items():
-        if patt.search(challenge):
-            return task_name
-
-    return 'qa'
 
 async def chat(request: web.Request) -> Response:
     """
@@ -178,50 +75,18 @@ async def chat(request: web.Request) -> Response:
         return Response(status=400)
 
     bt.logging.info(f'Request data: {request_data}')
-    k = request_data.get('k', 10)
-    exclude = request_data.get('exclude', [])
-    timeout = request_data.get('timeout', 10)
-    prefer = request_data.get('prefer', 'longest')
-    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.
-        uids = get_random_uids(validator, k=k, exclude=exclude or []).to(validator.device)
-        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=request_data['roles'], messages=request_data['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(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'] = 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")
+
+    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")
@@ -238,7 +103,7 @@ async def echo_stream(request):
     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)    
+    timeout = request_data.get('timeout', 0.2)
     message = '\n\n'.join(request_data['messages'])
 
     # Create a StreamResponse
@@ -251,7 +116,7 @@ async def echo_stream(request):
         for word in message.split():
             chunk = f'{word} '
             await response.write(chunk.encode('utf-8'))
-            completion += chunk 
+            completion += chunk
             time.sleep(timeout)
             bt.logging.info(f"Echoed: {chunk}")
 
@@ -269,7 +134,7 @@ async def echo_stream(request):
         "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'))
 
@@ -277,6 +142,7 @@ async def echo_stream(request):
     await response.write_eof()
     return response
 
+
 class ValidatorApplication(web.Application):
     def __init__(self, *a, **kw):
         super().__init__(*a, **kw)
@@ -296,14 +162,12 @@ bt.logging.info(validator_app)
 def main(run_aio_app=True, test=False) -> None:
 
     loop = asyncio.get_event_loop()
-
-    # port = validator.metagraph.axons[validator.uid].port
     port = 10000
     if run_aio_app:
         try:
             web.run_app(validator_app, port=port, loop=loop)
         except KeyboardInterrupt:
-            bt.logging.info("Keyboard interrupt detected. Exiting validator.")
+            bt.logging.warning("Keyboard interrupt detected. Exiting validator.")
         finally:
             pass
 

test.py

@@ -0,0 +1,16 @@
+import pytest
+
+
+
+def test_query_network():
+    pass
+
+def test_filter_completions():
+    pass
+
+
+def test_guess_task_name():
+    pass
+
+def test_ensemble_completions():
+    pass

utils.py

@@ -0,0 +1,109 @@
+import re
+import bittensor as bt
+
+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"]
+
+reward_models = {
+    'date_qa': DateRewardModel(),
+    'math': FloatDiffModel(),
+}
+
+def completion_is_valid(completion: str):
+    """
+    Get the completion statuses from the completions.
+    """
+    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):
+        return False
+    return True
+
+
+def ensemble_result(completions: list, task_name: str, prefer: str = 'longest'):
+    """
+    Ensemble completions from multiple models.
+    # TODO: Measure agreement
+    # TODO: Figure out how to mitigate the cabal effect (large groups will appear to be more credible)
+    # TODO: Reward pipeline
+    """
+    if not completions:
+        return None
+
+    answer = None
+    if task_name in ('qa', 'summarization'):
+        # No special handling for QA or summarization
+        supporting_completions = completions
+
+    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}")
+        valid_date_indices = [i for i, d in enumerate(dates) if d]
+        valid_completions = [completions[i] for i in valid_date_indices]
+        valid_dates = [dates[i] for i in valid_date_indices]
+        dates = [f"{d[0].strftime('%-d %B')} {d[1]}" for d in valid_dates]
+        if not dates:
+            return None
+
+        counter = Counter(dates)
+        most_common, count = counter.most_common()[0]
+        answer = most_common
+        if count == 1:
+            supporting_completions = valid_completions
+        else:
+            supporting_completions = [c for i, c in enumerate(valid_completions) if dates[i]==most_common]
+
+    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))
+        vals = [val for val in vals if val]
+        if not vals:
+            return None
+
+        most_common, count = Counter(dates).most_common()[0]
+        bt.logging.info(f"Most common value: {most_common}, count: {count}")
+        answer = most_common
+        if count == 1:
+            supporting_completions = completions
+        else:
+            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':
+        preferred_completion = sorted(supporting_completions, key=len)[-1]
+    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)
+    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'
+    }
+
+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'),
+    }
+    for task_name, patt in categories.items():
+        if patt.search(challenge):
+            return task_name
+
+    return 'qa'