wip: refactor tool processing in openai endpoint

backend/src/agents/expert.py

@@ -108,12 +108,13 @@ class ExpertAgent(Agent):
             if using_agent_tools and tool_calls:
                 print("Will execute tool calls")
                 for tc in tool_calls:
-                    tool = next(t for t in tools if t.name == tc.tool_name)
+                    try:
-                    if not tool:
+                        tool = next(t for t in tools if t.name == tc.tool_name)
+                        if not tool: raise StopIteration()
+                        tc.result = await tool.execute_async(**tc.arguments)
+                        messages.append(inference.LlamaMessage('tool', tc))
+                    except StopIteration:
                         print(f"Bad agent tool name: {tc.name}")
-                        continue
-                    tc.result = await tool.execute_async(**tc.arguments)
-                    messages.append(inference.LlamaMessage('tool', tc))
                 tool_calls = []
                 # Buffer time to let openrouter rest?
                 await asyncio.sleep(2)

backend/src/inference.py

@@ -45,7 +45,8 @@ def format_openai(messages):
     for msg in messages:
         if msg.role == 'tool':
             if not last_assistant_message:
-                raise Exception("tool message before any assistant message")
+                # Create one implicitly
+                last_assistant_message = {'role': 'assistant', 'content': ''}
             tool_calls = last_assistant_message.get('tool_calls')
             if tool_calls == None:
                 tool_calls = []
@@ -314,8 +315,21 @@ async def openrouter_chat_completion(app, messages, model, tools=[], opts={}, ag
     # TODO: Separate as InferenceRun class
     content_stream = None
     reasoning_stream = None
+
+    last_channel = None
+    def open_channel(channel):
+        nonlocal last_channel, content_stream, reasoning_stream
+        if channel != last_channel:
+            if content_stream: content_stream.close()
+            if reasoning_stream: reasoning_stream.close()
+            content_stream = None
+            reasoning_stream = None
+        last_channel = channel
+
+
     async def message_loop(name, value=None):
         nonlocal content_stream, reasoning_stream, agent_run
+        open_channel(name) # Hack: use the message as the channel name
         if name == 'output':
             if content_stream == None:
                 content_stream = Stream(app)

backend/src/server/openai.py

@@ -1,10 +1,118 @@
-import json, time, traceback, asyncio
+import json, time, traceback, asyncio, uuid
 from quart import Response, request
 from quart_auth import login_required, current_user
 from tool import ToolSpec, ToolParameter
 from agent import AgentRun
 from utils import encode_json
 
+# ---------------------------------------------------------------------------
+# Tool call deferring
+#
+# Every tool call emitted by an agent run is cached here so follow-up requests
+# can reference it by call id. Calls that match a client-passed tool are
+# executed by the client; anything else targets a native tool and is deferred:
+# the client only sees an opaque deferred call (id + type "tool") and
+# execution happens on our side through the existing tool infrastructure.
+# ---------------------------------------------------------------------------
+
+# call id -> {"call": ToolCall, "time": insertion time,
+#             "deferred": native call awaiting our execution,
+#             "resolved": output available on call.result}
+tool_call_cache = {}
+
+
+def cache_tool_call(call, deferred):
+    """Store an emitted tool call in the tool call cache.
+
+    The call keeps its own output on ``call.result`` once resolved. Insertion
+    time is recorded so stale entries can be expired later.
+    """
+    if not call.id:
+        call.id = 'call_' + uuid.uuid4().hex
+    tool_call_cache[call.id] = {
+        'call': call,
+        'time': time.time(),
+        'deferred': deferred,
+        'resolved': False,
+    }
+    return call
+
+
+def route_tool_call(call, tool_names):
+    """Check an emitted tool call against the client-passed tool names.
+
+    Returns True when the call targets a native tool and must be deferred.
+    Either way the call is stored in the tool call cache.
+    """
+    deferred = call.tool_name not in tool_names
+    cache_tool_call(call, deferred)
+    return deferred
+
+
+async def execute_deferred(function_call):
+    """Execute a deferred (native) tool call.
+
+    The output must be stored on the call itself (``function_call.result``).
+
+    TODO: resolve the ToolSpec for ``function_call.tool_name`` from the
+    agent's native tools and dispatch through the existing tool
+    infrastructure (ToolSpec.execute_async), recording the result or error on
+    the call. Deferred execution needs more design work; stubbed for now.
+    """
+    pass
+
+
+def attach_run_handlers(run, tool_names):
+    """Wire the run's output events into a single queue.
+
+    The queue receives ('content'|'reasoning'|'tool_call'|'deferred_call', data)
+    tuples and terminates with ('eof', None). Tool calls are routed (client vs
+    native/deferred) here so both response modes share the same logic.
+    """
+    queue = asyncio.Queue()
+    loop = asyncio.get_event_loop()
+
+    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)
+
+    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):
+        deferred = route_tool_call(tool_call, tool_names)
+        msg_type = 'deferred_call' if deferred else 'tool_call'
+        loop.create_task(queue.put((msg_type, tool_call)))
+
+    def on_done(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)
+
+    return queue
+
+
+def format_chunk(agent_id, delta, finish_reason=None):
+    """Build one OpenAI-format streaming chunk body."""
+    return {
+        "id": "chatcmpl-" + str(time.time()),
+        "object": "chat.completion.chunk",
+        "created": int(time.time()),
+        "model": agent_id,
+        "choices": [{"delta": delta, "index": 0, "finish_reason": finish_reason}]
+    }
+
+
 def register(server, app):
     # https://github.com/hrgarber/openai-compatible-api/blob/main/app/main.py
     @server.post('/api/openai/v1/chat/completions')
@@ -31,9 +139,6 @@ def register(server, app):
         messages = []
         for msg in openai_messages:
             role = msg['role']
-            if role not in ['user', 'assistant', 'system']:
-                message_errors.append("Bad role: " + role)
-                continue
             content = msg.get('content', None)
 
             if type(content) == list:
@@ -43,12 +148,53 @@ def register(server, app):
                     if part['type'] == 'text':
                         content += part['text']
 
+            if role == 'tool':
+                # Tool result: scan the tool call cache by call id.
+                entry = tool_call_cache.get(msg.get('tool_call_id', ''))
+                if entry is None:
+                    # Not in the cache: non existent or expired, drop silently.
+                    continue
+                call = entry['call']
+                if not entry['resolved']:
+                    if entry['deferred']:
+                        # Native call: execute with the existing tool infrastructure.
+                        await execute_deferred(call)
+                    else:
+                        # Client tool: the client already executed it, the
+                        # message carries the output.
+                        call.result = content
+                    entry['resolved'] = True
+                messages.append({'role': 'tool', 'content': call.result, 'tool_call_id': call.id})
+                continue
+
+            if role == 'assistant' and msg.get('tool_calls'):
+                # Rebuild the assistant's tool call turn from the cache so
+                # unknown or expired calls are silently dropped and the agent
+                # receives the original call objects, which carry their own
+                # output once resolved.
+                calls = []
+                for tc in msg['tool_calls']:
+                    entry = tool_call_cache.get(tc.get('id', ''))
+                    if entry is None:
+                        continue
+                    calls.append(entry['call'])
+                messages.append({
+                    'role': 'assistant',
+                    'content': content if type(content) == str else None,
+                    'tool_calls': calls,
+                })
+                continue
+
+            if role not in ['user', 'assistant', 'system']:
+                message_errors.append("Bad role: " + role)
+                continue
+
             if type(content) != str:
                 message_errors.append("Bad content")
                 continue
 
             messages.append({'role': role, 'content': content})
-        
+
         if message_errors:
             print("[openai] ERROR: bad messages input\n" + "\n".join([f"\t- {e}" for e in message_errors]))
 
@@ -75,6 +221,11 @@ def register(server, app):
                 parameters=parameters,
             ))
 
+        # Client-passed tool names: emitted tool calls that match one of these
+        # are sent back for the client to execute, anything else is a native
+        # tool and gets deferred.
+        tool_names = {spec.name for spec in tools}
+
         # --- Get agent control and create AgentRun directly ---
         control = await app.get_agent_control(agent_id)
 
@@ -85,47 +236,19 @@ def register(server, app):
             input={
                 "messages": messages,
                 "tools": tools,
+                # Never execute tools inside the run: emit them instead so
+                # the handler can send client tools back and defer native
+                # ones.
+                "defer_tool_execution": True,
             },
         )
         print("run input: " + repr(run.input))
 
+        queue = attach_run_handlers(run, tool_names)
+        await control.run(run)
+
         # --- Streaming ---
         if stream_mode:
-            queue = asyncio.Queue()
-            loop = asyncio.get_event_loop()
-
-            # Listen to the stream for both content and tool_call chunks
-
-            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)
-
-            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
@@ -141,36 +264,33 @@ def register(server, app):
                         elif msg_type == 'tool_call':
                             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}}
+                            delta = {"tool_calls": [{
-                            if data.id:
+                                "index": tool_call_count,
-                                tc_entry['id'] = data.id
+                                "id": data.id,
-                            delta = {"tool_calls": [tc_entry]}
+                                "type": "function",
+                                "function": {'name': data.tool_name, 'arguments': arguments},
+                            }]}
+
+                        elif msg_type == 'deferred_call':
+                            tool_call_count = tool_call_count + 1
+                            # Deferred calls stay opaque to the client: only
+                            # the call id and type go on the wire.
+                            delta = {"tool_calls": [{
+                                "index": tool_call_count,
+                                "id": data.id,
+                                "type": "tool",
+                            }]}
 
                         elif msg_type == 'eof':
                             # Determine finish_reason — was there a tool call?
                             finish_reason = "tool_calls" if any(
                                 e['name'] == 'tool_call' for e in run.events
                             ) else "stop"
-
+                            yield f"data: {json.dumps(format_chunk(agent_id, {}, finish_reason))}\n\n"
-                            chunk_body = {
-                                "id": "chatcmpl-" + str(time.time()),
-                                "object": "chat.completion.chunk",
-                                "created": int(time.time()),
-                                "model": agent_id,
-                                "choices": [{"delta": {}, "index": 0, "finish_reason": finish_reason}]
-                            }
-                            yield f"data: {json.dumps(chunk_body)}\n\n"
                             break
-                        
+
                         if delta:
-                            chunk_body = {
+                            yield f"data: {json.dumps(format_chunk(agent_id, delta))}\n\n"
-                                "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)}")
 
@@ -181,62 +301,56 @@ def register(server, app):
 
             return Response(generate(), content_type="text/event-stream")
 
-        # --- Non-streaming ---
+        # --- Non-streaming: drain the same event pipeline ---
         else:
-            await control.run(run)
+            try:
-            await run.get_output_stream().content_wait()
+                content = ""
-
+                tool_calls = []
-            # Collect any tool calls that were emitted
+                while True:
-            tool_events = [e['data'] for e in run.events if e['name'] == 'tool_call']
+                    msg_type, data = await queue.get()
-
+                    if msg_type == 'eof':
-            if tool_events:
+                        break
-                # Aggregate tool calls: combine partial argument chunks by id
+                    elif msg_type == 'content':
-                tc_map = {}
+                        content += data
-                for evt in tool_events:
+                    elif msg_type == 'reasoning':
-                    tid = evt.get('id', '')
+                        pass  # not surfaced in non-streaming responses
-                    if tid not in tc_map:
+                    elif msg_type == 'tool_call':
-                        tc_map[tid] = {
+                        tool_calls.append({
-                            "id": tid,
+                            "id": data.id,
                             "type": "function",
                             "function": {
-                                "name": evt.get('function', {}).get('name', ''),
+                                "name": data.tool_name,
-                                "arguments": evt.get('function', {}).get('arguments', ''),
+                                "arguments": encode_json(data.arguments) if data.arguments else '{}',
-                            }
+                            },
-                        }
+                        })
+                    elif msg_type == 'deferred_call':
+                        # Deferred calls stay opaque to the client: only the
+                        # call id and type go on the wire.
+                        tool_calls.append({"id": data.id, "type": "tool"})
                     else:
-                        tc_map[tid]['function']['arguments'] += \
+                        raise Exception(f"Unrecognized msg_type: {msg_type}. {repr(data)}")
-                            evt.get('function', {}).get('arguments', '')
+            except Exception as e:
+                print("[openai] ERROR: " + str(e))
+                print(traceback.format_exc())
+                return {"error": str(e)}, 500
 
-                return {
+            if tool_calls:
-                    "id": "chatcmpl-" + str(time.time()),
+                message = {"role": "assistant", "content": None, "tool_calls": tool_calls}
-                    "object": "chat.completion",
+                finish_reason = "tool_calls"
-                    "created": int(time.time()),
-                    "model": agent_id,
-                    "choices": [{
-                        "message": {
-                            "role": "assistant",
-                            "content": None,
-                            "tool_calls": list(tc_map.values()),
-                        },
-                        "finish_reason": "tool_calls",
-                    }]
-                }
             else:
-                # No tool calls: return text content
+                message = {"role": "assistant", "content": content}
-                content = run.stream.content
+                finish_reason = "stop"
-                return {
+
-                    "id": "chatcmpl-" + str(time.time()),
+            return {
-                    "object": "chat.completion",
+                "id": "chatcmpl-" + str(time.time()),
-                    "created": int(time.time()),
+                "object": "chat.completion",
-                    "model": agent_id,
+                "created": int(time.time()),
-                    "choices": [{
+                "model": agent_id,
-                        "message": {
+                "choices": [{
-                            "role": "assistant",
+                    "message": message,
-                            "content": content,
+                    "finish_reason": finish_reason,
-                        },
+                }]
-                        "finish_reason": "stop",
+            }
-                    }]
-                }
 
     @server.get("/api/openai/v1/models")
     @login_required