feat: (WIP) improve openrouter tool parsing and openai tool encoding
This commit is contained in:
Shivering Ocelot committed
1 parent
7058cf0771
commit
17d3041a60
5 files changed
+164
-96
No files matched your search
+60
-51
@@ -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']
|
||||
|
||||
Reference in new issue
Block a user