feat: (WIP) improve openrouter tool parsing and openai tool encoding
This commit is contained in:
1 parent
7058cf0771
commit
17d3041a60
5 files changed
+164
-96
No files matched your search
+46
-20
@@ -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", {
|
||||||
@@ -194,6 +202,24 @@ class AgentRun:
|
|||||||
"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})
|
||||||
@@ -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
@@ -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
@@ -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?
|
||||||
@@ -483,8 +479,21 @@ 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
@@ -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
|
||||||
|
|
||||||
|
|||||||
Reference in new issue
Block a user