wip: refactor tool processing in openai endpoint

This commit is contained in:
Shivering Ocelot committed 2026-09-28 14:28:59 -05:00
1 parent fd0d0b994b
commit 13e89b63fc
3 files changed
+236 -107

No files matched your search

+4 -3
View File
@@ -108,12 +108,13 @@ class ExpertAgent(Agent):
if using_agent_tools and tool_calls:
print("Will execute tool calls")
for tc in tool_calls:
try:
tool = next(t for t in tools if t.name == tc.tool_name)
if not tool:
print(f"Bad agent tool name: {tc.name}")
continue
if not tool: raise StopIteration()
tc.result = await tool.execute_async(**tc.arguments)
messages.append(inference.LlamaMessage('tool', tc))
except StopIteration:
print(f"Bad agent tool name: {tc.name}")
tool_calls = []
# Buffer time to let openrouter rest?
await asyncio.sleep(2)
+15 -1
View File
@@ -45,7 +45,8 @@ def format_openai(messages):
for msg in messages:
if msg.role == 'tool':
if not last_assistant_message:
raise Exception("tool message before any assistant message")
# Create one implicitly
last_assistant_message = {'role': 'assistant', 'content': ''}
tool_calls = last_assistant_message.get('tool_calls')
if tool_calls == None:
tool_calls = []
@@ -314,8 +315,21 @@ async def openrouter_chat_completion(app, messages, model, tools=[], opts={}, ag
# TODO: Separate as InferenceRun class
content_stream = None
reasoning_stream = None
last_channel = None
def open_channel(channel):
nonlocal last_channel, content_stream, reasoning_stream
if channel != last_channel:
if content_stream: content_stream.close()
if reasoning_stream: reasoning_stream.close()
content_stream = None
reasoning_stream = None
last_channel = channel
async def message_loop(name, value=None):
nonlocal content_stream, reasoning_stream, agent_run
open_channel(name) # Hack: use the message as the channel name
if name == 'output':
if content_stream == None:
content_stream = Stream(app)
+217 -103
View File
@@ -1,10 +1,118 @@
import json, time, traceback, asyncio
import json, time, traceback, asyncio, uuid
from quart import Response, request
from quart_auth import login_required, current_user
from tool import ToolSpec, ToolParameter
from agent import AgentRun
from utils import encode_json
# ---------------------------------------------------------------------------
# Tool call deferring
#
# Every tool call emitted by an agent run is cached here so follow-up requests
# can reference it by call id. Calls that match a client-passed tool are
# executed by the client; anything else targets a native tool and is deferred:
# the client only sees an opaque deferred call (id + type "tool") and
# execution happens on our side through the existing tool infrastructure.
# ---------------------------------------------------------------------------
# call id -> {"call": ToolCall, "time": insertion time,
# "deferred": native call awaiting our execution,
# "resolved": output available on call.result}
tool_call_cache = {}
def cache_tool_call(call, deferred):
"""Store an emitted tool call in the tool call cache.
The call keeps its own output on ``call.result`` once resolved. Insertion
time is recorded so stale entries can be expired later.
"""
if not call.id:
call.id = 'call_' + uuid.uuid4().hex
tool_call_cache[call.id] = {
'call': call,
'time': time.time(),
'deferred': deferred,
'resolved': False,
}
return call
def route_tool_call(call, tool_names):
"""Check an emitted tool call against the client-passed tool names.
Returns True when the call targets a native tool and must be deferred.
Either way the call is stored in the tool call cache.
"""
deferred = call.tool_name not in tool_names
cache_tool_call(call, deferred)
return deferred
async def execute_deferred(function_call):
"""Execute a deferred (native) tool call.
The output must be stored on the call itself (``function_call.result``).
TODO: resolve the ToolSpec for ``function_call.tool_name`` from the
agent's native tools and dispatch through the existing tool
infrastructure (ToolSpec.execute_async), recording the result or error on
the call. Deferred execution needs more design work; stubbed for now.
"""
pass
def attach_run_handlers(run, tool_names):
"""Wire the run's output events into a single queue.
The queue receives ('content'|'reasoning'|'tool_call'|'deferred_call', data)
tuples and terminates with ('eof', None). Tool calls are routed (client vs
native/deferred) here so both response modes share the same logic.
"""
queue = asyncio.Queue()
loop = asyncio.get_event_loop()
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)
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):
deferred = route_tool_call(tool_call, tool_names)
msg_type = 'deferred_call' if deferred else 'tool_call'
loop.create_task(queue.put((msg_type, tool_call)))
def on_done(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)
return queue
def format_chunk(agent_id, delta, finish_reason=None):
"""Build one OpenAI-format streaming chunk body."""
return {
"id": "chatcmpl-" + str(time.time()),
"object": "chat.completion.chunk",
"created": int(time.time()),
"model": agent_id,
"choices": [{"delta": delta, "index": 0, "finish_reason": finish_reason}]
}
def register(server, app):
# https://github.com/hrgarber/openai-compatible-api/blob/main/app/main.py
@server.post('/api/openai/v1/chat/completions')
@@ -31,9 +139,6 @@ def register(server, app):
messages = []
for msg in openai_messages:
role = msg['role']
if role not in ['user', 'assistant', 'system']:
message_errors.append("Bad role: " + role)
continue
content = msg.get('content', None)
if type(content) == list:
@@ -43,6 +148,47 @@ def register(server, app):
if part['type'] == 'text':
content += part['text']
if role == 'tool':
# Tool result: scan the tool call cache by call id.
entry = tool_call_cache.get(msg.get('tool_call_id', ''))
if entry is None:
# Not in the cache: non existent or expired, drop silently.
continue
call = entry['call']
if not entry['resolved']:
if entry['deferred']:
# Native call: execute with the existing tool infrastructure.
await execute_deferred(call)
else:
# Client tool: the client already executed it, the
# message carries the output.
call.result = content
entry['resolved'] = True
messages.append({'role': 'tool', 'content': call.result, 'tool_call_id': call.id})
continue
if role == 'assistant' and msg.get('tool_calls'):
# Rebuild the assistant's tool call turn from the cache so
# unknown or expired calls are silently dropped and the agent
# receives the original call objects, which carry their own
# output once resolved.
calls = []
for tc in msg['tool_calls']:
entry = tool_call_cache.get(tc.get('id', ''))
if entry is None:
continue
calls.append(entry['call'])
messages.append({
'role': 'assistant',
'content': content if type(content) == str else None,
'tool_calls': calls,
})
continue
if role not in ['user', 'assistant', 'system']:
message_errors.append("Bad role: " + role)
continue
if type(content) != str:
message_errors.append("Bad content")
continue
@@ -75,6 +221,11 @@ def register(server, app):
parameters=parameters,
))
# Client-passed tool names: emitted tool calls that match one of these
# are sent back for the client to execute, anything else is a native
# tool and gets deferred.
tool_names = {spec.name for spec in tools}
# --- Get agent control and create AgentRun directly ---
control = await app.get_agent_control(agent_id)
@@ -85,47 +236,19 @@ def register(server, app):
input={
"messages": messages,
"tools": tools,
# Never execute tools inside the run: emit them instead so
# the handler can send client tools back and defer native
# ones.
"defer_tool_execution": True,
},
)
print("run input: " + repr(run.input))
# --- Streaming ---
if stream_mode:
queue = asyncio.Queue()
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):
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)
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
queue = attach_run_handlers(run, tool_names)
await control.run(run)
# --- Streaming ---
if stream_mode:
async def generate():
try:
tool_call_count = 0
@@ -141,36 +264,33 @@ def register(server, app):
elif msg_type == 'tool_call':
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]}
delta = {"tool_calls": [{
"index": tool_call_count,
"id": data.id,
"type": "function",
"function": {'name': data.tool_name, 'arguments': arguments},
}]}
elif msg_type == 'deferred_call':
tool_call_count = tool_call_count + 1
# Deferred calls stay opaque to the client: only
# the call id and type go on the wire.
delta = {"tool_calls": [{
"index": tool_call_count,
"id": data.id,
"type": "tool",
}]}
elif msg_type == 'eof':
# Determine finish_reason — was there a tool call?
finish_reason = "tool_calls" if any(
e['name'] == 'tool_call' for e in run.events
) else "stop"
chunk_body = {
"id": "chatcmpl-" + str(time.time()),
"object": "chat.completion.chunk",
"created": int(time.time()),
"model": agent_id,
"choices": [{"delta": {}, "index": 0, "finish_reason": finish_reason}]
}
yield f"data: {json.dumps(chunk_body)}\n\n"
yield f"data: {json.dumps(format_chunk(agent_id, {}, finish_reason))}\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"
yield f"data: {json.dumps(format_chunk(agent_id, delta))}\n\n"
else:
raise Exception(f"Unrecognized msg_type: {msg_type}. {repr(data)}")
@@ -181,31 +301,45 @@ def register(server, app):
return Response(generate(), content_type="text/event-stream")
# --- Non-streaming ---
# --- Non-streaming: drain the same event pipeline ---
else:
await control.run(run)
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']
if tool_events:
# Aggregate tool calls: combine partial argument chunks by id
tc_map = {}
for evt in tool_events:
tid = evt.get('id', '')
if tid not in tc_map:
tc_map[tid] = {
"id": tid,
try:
content = ""
tool_calls = []
while True:
msg_type, data = await queue.get()
if msg_type == 'eof':
break
elif msg_type == 'content':
content += data
elif msg_type == 'reasoning':
pass # not surfaced in non-streaming responses
elif msg_type == 'tool_call':
tool_calls.append({
"id": data.id,
"type": "function",
"function": {
"name": evt.get('function', {}).get('name', ''),
"arguments": evt.get('function', {}).get('arguments', ''),
}
}
"name": data.tool_name,
"arguments": encode_json(data.arguments) if data.arguments else '{}',
},
})
elif msg_type == 'deferred_call':
# Deferred calls stay opaque to the client: only the
# call id and type go on the wire.
tool_calls.append({"id": data.id, "type": "tool"})
else:
tc_map[tid]['function']['arguments'] += \
evt.get('function', {}).get('arguments', '')
raise Exception(f"Unrecognized msg_type: {msg_type}. {repr(data)}")
except Exception as e:
print("[openai] ERROR: " + str(e))
print(traceback.format_exc())
return {"error": str(e)}, 500
if tool_calls:
message = {"role": "assistant", "content": None, "tool_calls": tool_calls}
finish_reason = "tool_calls"
else:
message = {"role": "assistant", "content": content}
finish_reason = "stop"
return {
"id": "chatcmpl-" + str(time.time()),
@@ -213,28 +347,8 @@ def register(server, app):
"created": int(time.time()),
"model": agent_id,
"choices": [{
"message": {
"role": "assistant",
"content": None,
"tool_calls": list(tc_map.values()),
},
"finish_reason": "tool_calls",
}]
}
else:
# No tool calls: return text content
content = run.stream.content
return {
"id": "chatcmpl-" + str(time.time()),
"object": "chat.completion",
"created": int(time.time()),
"model": agent_id,
"choices": [{
"message": {
"role": "assistant",
"content": content,
},
"finish_reason": "stop",
"message": message,
"finish_reason": finish_reason,
}]
}