feat: task cancellation (wip)

backend/src/inference.py

@@ -85,7 +85,7 @@ def build_messages(system_prompt="", instruction=None, input=None, messages=None
 
     # Add default messages to list
     if system_prompt:
-        messages.insert(0, LlamaMessage("system", system))
+        messages.insert(0, LlamaMessage("system", system_prompt))
     if prompt:
         messages.append(LlamaMessage("user", prompt))
 
@@ -176,6 +176,11 @@ async def local_chat_completion(app, messages, model, tools=[], opts={}, agent_r
         )
 
         for chunk in llm_output:
+            # Check for cancellation
+            if agent_run and agent_run.cancelled:
+                print(f"Cancelling local inference for run {agent_run.id}")
+                break
+                
             delta = chunk['choices'][0]['delta']
             if 'content' in delta:
                 chunk_text = delta['content']
@@ -195,6 +200,8 @@ async def local_chat_completion(app, messages, model, tools=[], opts={}, agent_r
         content_stream.close()
 
     thread = threading.Thread(target=worker)
+    if agent_run:
+        agent_run.thread = thread
     thread.start()
 
     if agent_run:
@@ -224,16 +231,28 @@ async def ollama_chat_completion(app, messages, model, tools=[], opts={}, agent_
         if formatted_tools:
             kwargs['tools'] = formatted_tools
 
-        stream_response = await ollama.AsyncClient().chat(**kwargs)
+        try:
+            stream_response = await ollama.AsyncClient().chat(**kwargs)
 
-        async for part in stream_response:
-            if part.message.content:
-                content_stream.feed(part.message.content)
+            async for part in stream_response:
+                if agent_run and agent_run.cancelled:
+                    print(f"Cancelling ollama inference for run {agent_run.id}")
+                    break
+                    
+                if part.message.content:
+                    content_stream.feed(part.message.content)
+        except asyncio.CancelledError:
+            print(f"Ollama task cancelled for run {agent_run.id}")
+            raise
 
         content_stream.close()
         print("END OLLAMA INFERENCE")
 
-    asyncio.create_task(worker())
+    task = asyncio.create_task(worker())
+    if agent_run:
+        agent_run.task = task
+    else:
+        task.start()
 
     if agent_run:
         return None
@@ -262,18 +281,29 @@ async def openrouter_chat_completion(app, messages, model, tools=[], opts={}, ag
         if formatted_tools:
             kwargs['tools'] = formatted_tools
 
-        stream_response = await app.openrouter_client.chat.send_async(**kwargs)
+        try:
+            stream_response = await app.openrouter_client.chat.send_async(**kwargs)
 
-        async for event in stream_response:
-            if event.choices and len(event.choices) > 0:
-                delta = event.choices[0].delta
-                if delta.content:
-                    content_stream.feed(delta.content)
+            async for event in stream_response:
+                if agent_run and agent_run.cancelled:
+                    print(f"Cancelling openrouter inference for run {agent_run.id}")
+                    # TODO: Need to cancell?
+                    break
+                    
+                if event.choices and len(event.choices) > 0:
+                    delta = event.choices[0].delta
+                    if delta.content:
+                        content_stream.feed(delta.content)
+        except asyncio.CancelledError:
+            print(f"OpenRouter task cancelled for run {agent_run.id if agent_run else 'unknown'}")
+            raise
 
         content_stream.close()
         print("END OPENROUTER INFERENCE")
 
-    asyncio.create_task(worker())
+    task = asyncio.create_task(worker())
+    if agent_run:
+        agent_run.task = task
 
     if agent_run:
         return None
@@ -318,21 +348,28 @@ async def replicate_base_completion(app, model, input, agent_run=None):
             'replicate_client': app.replicate_client,
             'model': model,
             'content_stream': content_stream,
-            'input': input
+            'input': input,
+            'agent_run': agent_run
         }
     )
     thread.start()
     #response = await content_stream.content_wait()
     #return response.lstrip()
     if agent_run:
+        agent_run.thread = thread
         return None
     return content_stream
 
 
-def run_replicate_inference_loop(replicate_client, model, input, content_stream):
+def run_replicate_inference_loop(replicate_client, model, input, content_stream, agent_run=None):
     print("START INFERENCE")
     completion = replicate_client.run(model, input=input, stream=True)
     for event in completion:
+        if agent_run and agent_run.cancelled:
+            print(f"Cancelling replicate inference for run {agent_run.id}")
+            break
+            # TODO: does completion need explicit cancellation?
+            
         if event:
             content_stream.feed(str(event))
     print("END INFERENCE")