From 17d3041a609625d15cee1fa1b25b0930ffdee673 Mon Sep 17 00:00:00 2001 From: Shivering Ocelot Date: Thu, 4 Jun 2026 11:27:42 -0500 Subject: [PATCH] feat: (WIP) improve openrouter tool parsing and openai tool encoding --- backend/src/agent.py | 66 ++++++++++++++------- backend/src/agents/expert.py | 2 +- backend/src/inference.py | 77 +++++++++++++++++------- backend/src/server.py | 111 +++++++++++++++++++---------------- backend/src/tool.py | 4 +- 5 files changed, 164 insertions(+), 96 deletions(-) diff --git a/backend/src/agent.py b/backend/src/agent.py index db08672..87b84ed 100644 --- a/backend/src/agent.py +++ b/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", { @@ -193,6 +201,24 @@ class AgentRun: "name": event_name, "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}) diff --git a/backend/src/agents/expert.py b/backend/src/agents/expert.py index 660c11f..ec07dd6 100644 --- a/backend/src/agents/expert.py +++ b/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()) diff --git a/backend/src/inference.py b/backend/src/inference.py index 68f8d7b..ace4239 100644 --- a/backend/src/inference.py +++ b/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): diff --git a/backend/src/server.py b/backend/src/server.py index 04669af..c9fe86d 100644 --- a/backend/src/server.py +++ b/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? @@ -482,9 +478,22 @@ 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'] diff --git a/backend/src/tool.py b/backend/src/tool.py index e3d4837..b0b98e0 100644 --- a/backend/src/tool.py +++ b/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