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

This commit is contained in:
Shivering Ocelot committed 2026-06-04 11:27:42 -05:00
1 parent 7058cf0771
commit 17d3041a60
5 files changed
+149 -81

No files matched your search

+35 -9
View File
@@ -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)
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=run.stream,
content=data.stream,
)
run.emit_event('chat-message', {
'messageId': new_message['id'],
'chatId': message.chat_id,
'threadId': message.thread_id,
})
return run.stream
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:
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})
+1 -1
View File
@@ -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())
+54 -21
View File
@@ -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,23 +292,56 @@ 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
async def replicate_completion(app, messages, model=None, tools=[], opts={}, agent_run=None):
if model is None:
+57 -48
View File
@@ -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
def on_output(stream):
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():
# 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']
+2 -2
View File
@@ -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