feat: (WIP) improve openrouter tool parsing and openai tool encoding

backend/src/agent.py

@@ -4,6 +4,7 @@ import inspect
 import time
 
 from temporary import generate_tmp_id
+from stream import Stream
 
 # Tool related classes must be moved to a different module
 
@@ -53,7 +54,7 @@ class Agent:
             },
         )
         await self.run(run)
-        return run.stream
+        return run.get_output_stream()
 
     async def received_message(self, message):
         chat_view = await self.app.get_chat_view(message.chat_id)
@@ -77,18 +78,22 @@ class Agent:
         )
 
         await self.run(run)
-        new_message = await self.app.send_chat_message(
-            sender=self.id,
-            chat_id=message.chat_id,
-            thread_id=message.thread_id,
-            content=run.stream,
-        )
-        run.emit_event('chat-message', {
-            'messageId': new_message['id'],
-            'chatId': message.chat_id,
-            'threadId': message.thread_id,
-        })
-        return run.stream
+
+        async def on_output_event(data):
+            new_message = await self.app.send_chat_message(
+                sender=self.id,
+                chat_id=message.chat_id,
+                thread_id=message.thread_id,
+                content=data.stream,
+            )
+            run.emit_event('chat-message', {
+                'messageId': new_message['id'],
+                'chatId': message.chat_id,
+                'threadId': message.thread_id,
+            })
+
+        run.on('output', on_output_event)
+        return run
 
     def _role_for_message(self, message):
         sender = message.sender
@@ -138,11 +143,12 @@ class AgentRun:
 
         # I/O envelope (no generic data bag)
         self.input = input or {}
-        from stream import Stream
-        self.stream = Stream(app)
         self.events = []
         self.event_listeners = {}
+
+        # Default Streams
         self.output = None
+        self.reasoning = None
         
         # Cancellation tracking
         self.task = None  # For asyncio tasks
@@ -156,7 +162,6 @@ class AgentRun:
             "userId": self.user_id,
             "startTime": self.start_time,
             "input": self.input,
-            "streamId": self.stream.id,
         })
         
         return self
@@ -185,7 +190,10 @@ class AgentRun:
         self.events.append(event)
         listeners = self.event_listeners.get(event_name, [])
         for cb in listeners:
-            cb(data)
+            try:
+                cb(data)
+            except Exception as e:
+                print(f"ERROR: handling event '{event_name}' in run {self.id}\n" + str(e))
         
         # Broadcast event to websocket listeners
         self.app.emit_event("run-event", {
@@ -194,6 +202,24 @@ class AgentRun:
             "data": data
         })
     
+    def get_output_stream(self):
+        if not self.output:
+            self.output = Stream(self.app)
+            self.emit_event('output', self.output)
+        return self.output
+    
+    def feed_output(self, chunk):
+        stream = self.get_output_stream()
+        if chunk:
+            self.output.feed(chunk)
+    
+    def feed_reasoning(self, chunk):
+        if not self.reasoning:
+            self.reasoning = Stream(self.app)
+            self.emit_event('reasoning', self.reasoning)
+        if chunk:
+            self.reasoning.feed(chunk)
+
     def add_tool_call(self, tool_name, arguments):
         self.emit_event("tool_call", {"name": tool_name, "arguments": arguments})
 
@@ -261,6 +287,6 @@ class AgentRun:
             return
         self.done = True
         self.end_time = int(time.time())
-        self.output = self.stream.content
-        self.stream.close()
-        self.emit_event("done", {"output": self.output, "agent_id": self.agent_id})
+        if self.output: self.output.close()
+        if self.reasoning: self.reasoning.close()
+        self.emit_event("done", {"output": self.output.content, "agent_id": self.agent_id})

backend/src/agents/expert.py

@@ -45,4 +45,4 @@ class ExpertAgent(Agent):
             agent_run=run,
         )
 
-        run.stream.listen(lambda chunk: None, lambda: run.close())
+        run.get_output_stream().listen(lambda chunk: None, lambda: run.close())

backend/src/inference.py

@@ -3,8 +3,9 @@ import asyncio
 import json
 import os
 import ollama
+import traceback
 from stream import Stream
-from tool import format_tools_for_engine, build_tool_prompt_prefix, ToolSpec
+from tool import format_tools_for_engine, build_tool_prompt_prefix, ToolSpec, ToolCall
 
 from llama_cpp import Llama
 from openrouter import OpenRouter
@@ -115,11 +116,6 @@ async def chat_completion(app,
     model_engine = db_model['engine']
     model_name = db_model['model']
 
-    # Warn when tools are passed to engines that do not support native tool calling
-    if tools:
-        if model_engine not in ("llama-cpp",):
-            print(f"WARNING: Tool calls not supported for engine '{model_engine}'. Tools will be ignored.")
-
     if model_engine == "llama-cpp":
         result = await local_chat_completion(app, messages, model=model_name, tools=tools, opts=opts, agent_run=agent_run)
     elif model_engine == "replicate":
@@ -260,10 +256,8 @@ async def ollama_chat_completion(app, messages, model, tools=[], opts={}, agent_
 
 
 async def openrouter_chat_completion(app, messages, model, tools=[], opts={}, agent_run=None):
-    if agent_run:
-        content_stream = agent_run.stream
-    else:
-        content_stream = Stream(app)
+    if not agent_run:
+        raise Exception("agent_run required")
 
     async def worker():
         print("START OPENROUTER INFERENCE")
@@ -284,6 +278,12 @@ async def openrouter_chat_completion(app, messages, model, tools=[], opts={}, ag
         try:
             stream_response = await app.openrouter_client.chat.send_async(**kwargs)
 
+            # OpenAI format streams tool calls in parts
+            # We could, but I won't assume any tool call is complete until the message is over
+            # TODO: tool_call_by_index = {}
+            tool_calls = []
+            last_tool_call = None
+
             async for event in stream_response:
                 if agent_run and agent_run.cancelled:
                     print(f"Cancelling openrouter inference for run {agent_run.id}")
@@ -292,22 +292,55 @@ async def openrouter_chat_completion(app, messages, model, tools=[], opts={}, ag
                     
                 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()
+                    if delta.content:
+                        agent_run.feed_output(delta.content)
+                    if delta.reasoning:
+                        agent_run.feed_reasoning(delta.reasoning)
+                    if delta.tool_calls:
+                        print("[openrouter] tool call:" + repr(delta.tool_calls)) 
+                        for tc in delta.tool_calls:
+                            if not last_tool_call:
+                                last_tool_call = ToolCall(None, arguments="")
+                                tool_calls.append(last_tool_call)
+
+                            if tc.id:
+                                if last_tool_call.id is None:
+                                    last_tool_call.id = tc.id
+                                elif last_tool_call.id != tc.id:
+                                    # New id = new tool
+                                    # TODO: WRONG! new id must append value to accumulating id. new index signals new tool
+                                    last_tool_call = ToolCall(None, arguments="", id=tc.id)
+                                    tool_calls.append(last_tool_call)
+
+                            if tc.function:
+                                if tc.function.name:
+                                    last_tool_call.tool_name = tc.function.name
+                                if tc.function.arguments:
+                                    last_tool_call.arguments += tc.function.arguments
+                else:
+                    print("[openrouter] no event choices: " + repr(event))
+
+            for tc in tool_calls:
+                try:
+                    arguments = json.loads(tc.arguments)
+                    tc.arguments = arguments
+                    agent_run.emit_event('tool-call', tc)
+                except Exception:
+                    print("[openrouter] Bad tool call arguments")
+                    print(traceback.format_exc())
+        except asyncio.CancelledError:
+            print(f"[openrouter] task cancelled for run {agent_run.id}")
+            raise
+        except Exception:
+            print(f"[openrouter] exception during inference")
+            print(traceback.format_exc())
+
+        agent_run.close()
         print("END OPENROUTER INFERENCE")
 
     task = asyncio.create_task(worker())
-    if agent_run:
-        agent_run.task = task
-
-    if agent_run:
-        return None
-    return content_stream
+    agent_run.task = task
 
 
 async def replicate_completion(app, messages, model=None, tools=[], opts={}, agent_run=None):

backend/src/server.py

@@ -7,6 +7,7 @@ import time
 import json
 import asyncio
 import nest_asyncio
+import traceback
 from stream import Stream
 from tool import ToolSpec, ToolParameter
 from agent import AgentRun
@@ -15,6 +16,14 @@ from agent import AgentRun
 # relative to the module file???? wtf
 frontend_path = '../../frontend/dist'
 
+def encode_json(value):
+    def encode_unknown(value):
+        if value is ToolSpec:
+            return value.to_openai_dict()
+        else:
+            return repr(value)
+    return json.dumps(value, default=encode_unknown)
+
 async def start_server(app):
     nest_asyncio.apply()
 
@@ -128,7 +137,7 @@ async def start_server(app):
             nonlocal send_queue
             while True:
                 msg = await send_queue.get()
-                await websocket.send(json.dumps(msg))
+                await websocket.send(encode_json(msg))
         async def send(msg):
             await send_queue.put(msg)
 
@@ -403,69 +412,56 @@ async def start_server(app):
             loop = asyncio.get_event_loop()
 
             # Listen to the stream for both content and tool_call chunks
-            async def on_chunk(chunk):
-                # Tool calls are fed as JSON strings: '{"tool_call": {...}}'
-                try:
-                    parsed = json.loads(chunk)
-                    if isinstance(parsed, dict) and 'tool_call' in parsed:
-                        tc = parsed['tool_call']
-                        await queue.put(('tool_call', tc))
-                        return
-                except (json.JSONDecodeError, TypeError):
-                    pass
-                # Regular content
-                await queue.put(('content', chunk))
 
-            async def on_eof():
-                await queue.put(('eof', None))
+            def on_output(stream):
+                async def on_chunk(chunk):
+                    await queue.put(('content', chunk))
+                async def on_eof():
+                    # on_done doesn't work?
+                    await queue.put(('eof', None))
+                stream.listen(on_chunk, on_eof)
 
-            run.stream.listen(on_chunk, on_eof)
+            def on_reasoning(stream):
+                async def on_chunk(chunk):
+                    await queue.put(('reasoning', chunk))
+                async def on_eof(): pass
+                stream.listen(on_chunk, on_eof)
+
+            def on_tool_call(tool_call):
+                print('agent tool call ' + repr(tool_call))
+                loop.create_task(queue.put(('tool_call', tool_call)))
+
+            def on_done(output):
+                print('agent done ' + repr(output))
+                loop.create_task(queue.put(('eof', None)))
+
+            run.on('output', on_output)
+            run.on('reasoning', on_reasoning)
+            run.on('tool-call', on_tool_call)
+            run.on('done', on_done)
 
             # Start the agent run
             await control.run(run)
 
             async def generate():
                 try:
+                    tool_call_count = 0
                     while True:
                         msg_type, data = await queue.get()
 
+                        delta = None
                         if msg_type == 'content':
-                            chunk_body = {
-                                "id": "chatcmpl-" + str(time.time()),
-                                "object": "chat.completion.chunk",
-                                "created": int(time.time()),
-                                "model": agent_id,
-                                "choices": [{"delta": {"content": data}, "index": 0, "finish_reason": None}]
-                            }
-                            yield f"data: {json.dumps(chunk_body)}\n\n"
+                            delta = {'content': data}
+                        elif msg_type == 'reasoning':
+                            delta = {'reasoning': data}
 
                         elif msg_type == 'tool_call':
-                            # Forward the tool call delta (already in OpenAI format from llama-cpp)
-                            tc_entry = {"index": data.get('index', 0)}
-                            if data.get('id'):
-                                tc_entry['id'] = data['id']
-                            if 'function' in data:
-                                func_part = {}
-                                fn = data['function']
-                                if fn.get('name'):
-                                    func_part['name'] = fn['name']
-                                if fn.get('arguments'):
-                                    func_part['arguments'] = fn['arguments']
-                                if func_part:
-                                    tc_entry['function'] = func_part
-
-                            chunk_body = {
-                                "id": "chatcmpl-" + str(time.time()),
-                                "object": "chat.completion.chunk",
-                                "created": int(time.time()),
-                                "model": agent_id,
-                                "choices": [{
-                                    "delta": {"tool_calls": [tc_entry]},
-                                    "index": 0,
-                                    "finish_reason": None
-                                }]
-                            }
-                            yield f"data: {json.dumps(chunk_body)}\n\n"
+                            tool_call_count = tool_call_count + 1
+                            arguments = encode_json(data.arguments) if data.arguments else '{}'
+                            tc_entry = {"index": tool_call_count, "function": {'name': data.tool_name, 'arguments': arguments}}
+                            if data.id:
+                                tc_entry['id'] = data.id
+                            delta = {"tool_calls": [tc_entry]}
 
                         elif msg_type == 'eof':
                             # Determine finish_reason — was there a tool call?
@@ -483,8 +479,21 @@ async def start_server(app):
                             yield f"data: {json.dumps(chunk_body)}\n\n"
                             break
                         
+                        if delta:
+                            chunk_body = {
+                                "id": "chatcmpl-" + str(time.time()),
+                                "object": "chat.completion.chunk",
+                                "created": int(time.time()),
+                                "model": agent_id,
+                                "choices": [{"delta": delta, "index": 0, "finish_reason": None}]
+                            }
+                            yield f"data: {json.dumps(chunk_body)}\n\n"
+                        else:
+                            raise Exception(f"Unrecognized msg_type: {msg_type}. {repr(data)}")
+
                 except Exception as e:
                     print("[openai] ERROR: " + str(e))
+                    print(traceback.format_exc())
                     yield f"data: {json.dumps({'error': str(e)})}\n\n"
 
             return Response(generate(), content_type="text/event-stream")
@@ -492,7 +501,7 @@ async def start_server(app):
         # --- Non-streaming ---
         else:
             await control.run(run)
-            await run.stream.content_wait()
+            await run.get_output_stream().content_wait()
 
             # Collect any tool calls that were emitted
             tool_events = [e['data'] for e in run.events if e['name'] == 'tool_call']

backend/src/tool.py

@@ -127,10 +127,10 @@ class ToolCall:
         is_error: Whether the result represents an error.
     """
 
-    def __init__(self, tool_name, arguments=None, id=None):
+    def __init__(self, tool_name, arguments={}, id=None):
         self.id = id or ""
         self.tool_name = tool_name
-        self.arguments = arguments or {}
+        self.arguments = arguments
 
         self.result = None