wip: refactor tool processing in openai endpoint
This commit is contained in:
1 parent
fd0d0b994b
commit
13e89b63fc
3 files changed
+236
-107
No files matched your search
@@ -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)
|
||||
|
||||
@@ -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
@@ -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,
|
||||
}]
|
||||
}
|
||||
|
||||
|
||||
Reference in new issue
Block a user