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
+164 -96

No files matched your search

+46 -20
View File
@@ -4,6 +4,7 @@ import inspect
import time import time
from temporary import generate_tmp_id from temporary import generate_tmp_id
from stream import Stream
# Tool related classes must be moved to a different module # Tool related classes must be moved to a different module
@@ -53,7 +54,7 @@ class Agent:
}, },
) )
await self.run(run) await self.run(run)
return run.stream return run.get_output_stream()
async def received_message(self, message): async def received_message(self, message):
chat_view = await self.app.get_chat_view(message.chat_id) chat_view = await self.app.get_chat_view(message.chat_id)
@@ -77,18 +78,22 @@ class Agent:
) )
await self.run(run) await self.run(run)
new_message = await self.app.send_chat_message(
sender=self.id, async def on_output_event(data):
chat_id=message.chat_id, new_message = await self.app.send_chat_message(
thread_id=message.thread_id, sender=self.id,
content=run.stream, chat_id=message.chat_id,
) thread_id=message.thread_id,
run.emit_event('chat-message', { content=data.stream,
'messageId': new_message['id'], )
'chatId': message.chat_id, run.emit_event('chat-message', {
'threadId': message.thread_id, 'messageId': new_message['id'],
}) 'chatId': message.chat_id,
return run.stream 'threadId': message.thread_id,
})
run.on('output', on_output_event)
return run
def _role_for_message(self, message): def _role_for_message(self, message):
sender = message.sender sender = message.sender
@@ -138,11 +143,12 @@ class AgentRun:
# I/O envelope (no generic data bag) # I/O envelope (no generic data bag)
self.input = input or {} self.input = input or {}
from stream import Stream
self.stream = Stream(app)
self.events = [] self.events = []
self.event_listeners = {} self.event_listeners = {}
# Default Streams
self.output = None self.output = None
self.reasoning = None
# Cancellation tracking # Cancellation tracking
self.task = None # For asyncio tasks self.task = None # For asyncio tasks
@@ -156,7 +162,6 @@ class AgentRun:
"userId": self.user_id, "userId": self.user_id,
"startTime": self.start_time, "startTime": self.start_time,
"input": self.input, "input": self.input,
"streamId": self.stream.id,
}) })
return self return self
@@ -185,7 +190,10 @@ class AgentRun:
self.events.append(event) self.events.append(event)
listeners = self.event_listeners.get(event_name, []) listeners = self.event_listeners.get(event_name, [])
for cb in listeners: 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 # Broadcast event to websocket listeners
self.app.emit_event("run-event", { self.app.emit_event("run-event", {
@@ -193,6 +201,24 @@ class AgentRun:
"name": event_name, "name": event_name,
"data": data "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): def add_tool_call(self, tool_name, arguments):
self.emit_event("tool_call", {"name": tool_name, "arguments": arguments}) self.emit_event("tool_call", {"name": tool_name, "arguments": arguments})
@@ -261,6 +287,6 @@ class AgentRun:
return return
self.done = True self.done = True
self.end_time = int(time.time()) self.end_time = int(time.time())
self.output = self.stream.content if self.output: self.output.close()
self.stream.close() if self.reasoning: self.reasoning.close()
self.emit_event("done", {"output": self.output, "agent_id": self.agent_id}) 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, agent_run=run,
) )
run.stream.listen(lambda chunk: None, lambda: run.close()) run.get_output_stream().listen(lambda chunk: None, lambda: run.close())
+55 -22
View File
@@ -3,8 +3,9 @@ import asyncio
import json import json
import os import os
import ollama import ollama
import traceback
from stream import Stream 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 llama_cpp import Llama
from openrouter import OpenRouter from openrouter import OpenRouter
@@ -115,11 +116,6 @@ async def chat_completion(app,
model_engine = db_model['engine'] model_engine = db_model['engine']
model_name = db_model['model'] 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": if model_engine == "llama-cpp":
result = await local_chat_completion(app, messages, model=model_name, tools=tools, opts=opts, agent_run=agent_run) result = await local_chat_completion(app, messages, model=model_name, tools=tools, opts=opts, agent_run=agent_run)
elif model_engine == "replicate": 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): async def openrouter_chat_completion(app, messages, model, tools=[], opts={}, agent_run=None):
if agent_run: if not agent_run:
content_stream = agent_run.stream raise Exception("agent_run required")
else:
content_stream = Stream(app)
async def worker(): async def worker():
print("START OPENROUTER INFERENCE") print("START OPENROUTER INFERENCE")
@@ -284,6 +278,12 @@ async def openrouter_chat_completion(app, messages, model, tools=[], opts={}, ag
try: try:
stream_response = await app.openrouter_client.chat.send_async(**kwargs) 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: async for event in stream_response:
if agent_run and agent_run.cancelled: if agent_run and agent_run.cancelled:
print(f"Cancelling openrouter inference for run {agent_run.id}") 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: if event.choices and len(event.choices) > 0:
delta = event.choices[0].delta 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") print("END OPENROUTER INFERENCE")
task = asyncio.create_task(worker()) task = asyncio.create_task(worker())
if agent_run: agent_run.task = task
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): async def replicate_completion(app, messages, model=None, tools=[], opts={}, agent_run=None):
+60 -51
View File
@@ -7,6 +7,7 @@ import time
import json import json
import asyncio import asyncio
import nest_asyncio import nest_asyncio
import traceback
from stream import Stream from stream import Stream
from tool import ToolSpec, ToolParameter from tool import ToolSpec, ToolParameter
from agent import AgentRun from agent import AgentRun
@@ -15,6 +16,14 @@ from agent import AgentRun
# relative to the module file???? wtf # relative to the module file???? wtf
frontend_path = '../../frontend/dist' 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): async def start_server(app):
nest_asyncio.apply() nest_asyncio.apply()
@@ -128,7 +137,7 @@ async def start_server(app):
nonlocal send_queue nonlocal send_queue
while True: while True:
msg = await send_queue.get() msg = await send_queue.get()
await websocket.send(json.dumps(msg)) await websocket.send(encode_json(msg))
async def send(msg): async def send(msg):
await send_queue.put(msg) await send_queue.put(msg)
@@ -403,69 +412,56 @@ async def start_server(app):
loop = asyncio.get_event_loop() loop = asyncio.get_event_loop()
# Listen to the stream for both content and tool_call chunks # 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(): def on_output(stream):
await queue.put(('eof', None)) 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 # Start the agent run
await control.run(run) await control.run(run)
async def generate(): async def generate():
try: try:
tool_call_count = 0
while True: while True:
msg_type, data = await queue.get() msg_type, data = await queue.get()
delta = None
if msg_type == 'content': if msg_type == 'content':
chunk_body = { delta = {'content': data}
"id": "chatcmpl-" + str(time.time()), elif msg_type == 'reasoning':
"object": "chat.completion.chunk", delta = {'reasoning': data}
"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"
elif msg_type == 'tool_call': elif msg_type == 'tool_call':
# Forward the tool call delta (already in OpenAI format from llama-cpp) tool_call_count = tool_call_count + 1
tc_entry = {"index": data.get('index', 0)} arguments = encode_json(data.arguments) if data.arguments else '{}'
if data.get('id'): tc_entry = {"index": tool_call_count, "function": {'name': data.tool_name, 'arguments': arguments}}
tc_entry['id'] = data['id'] if data.id:
if 'function' in data: tc_entry['id'] = data.id
func_part = {} delta = {"tool_calls": [tc_entry]}
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"
elif msg_type == 'eof': elif msg_type == 'eof':
# Determine finish_reason — was there a tool call? # 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" yield f"data: {json.dumps(chunk_body)}\n\n"
break 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: except Exception as e:
print("[openai] ERROR: " + str(e)) print("[openai] ERROR: " + str(e))
print(traceback.format_exc())
yield f"data: {json.dumps({'error': str(e)})}\n\n" yield f"data: {json.dumps({'error': str(e)})}\n\n"
return Response(generate(), content_type="text/event-stream") return Response(generate(), content_type="text/event-stream")
@@ -492,7 +501,7 @@ async def start_server(app):
# --- Non-streaming --- # --- Non-streaming ---
else: else:
await control.run(run) await control.run(run)
await run.stream.content_wait() await run.get_output_stream().content_wait()
# Collect any tool calls that were emitted # Collect any tool calls that were emitted
tool_events = [e['data'] for e in run.events if e['name'] == 'tool_call'] 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. 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.id = id or ""
self.tool_name = tool_name self.tool_name = tool_name
self.arguments = arguments or {} self.arguments = arguments
self.result = None self.result = None