hordelab/hub
2
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: fix agent prompt data flow, fix tools in openai endpoint, account for completion usage info

Browse Source
This commit is contained in:
Shivering Ocelot committed 2026-10-03 16:38:33 -05:00
1 parent 868e9b0545
commit 1bd1880d9b
4 files changed
+121 -64
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

backend/src/agent.py
+24 -23
View File
@@ -77,38 +77,39 @@ class Agent:
**kwargs,
},
)
out_stream = Stream(self.app)
def on_output_event(stream):
done_event = asyncio.Event()
out_stream = Stream(self.app)
tool_calls = []
tool_names = None
if 'tools' in kwargs:
tool_names = [t['function']['name'] for t in kwargs['tools'] if isinstance(t, dict)]
def on_chunk(chunk):
out_stream.feed(chunk)
stream.listen(on_chunk, lambda: None)
def on_done(data):
out_stream.close()
run.on('output', on_output_event)
run.on('done', on_done)
def on_inference(msg):
kwargs = {'content': msg.content}
# Change the output shape based on the presence of tools
if 'tools' in kwargs:
tool_calls = []
tool_names = [t['function']['name'] for t in kwargs['tools'] if isinstance(t, dict)]
def on_tool_call(call):
print('prompt call', call)
if call.tool_name in tool_names:
tool_calls.append(call)
run.on('tool-call', on_tool_call)
if msg.role == 'assistant':
msg.content.listen(on_chunk, lambda: None)
elif msg.role == 'tool' and tool_names:
if msg.content.tool_name in tool_names:
tool_calls.append(msg.content)
def on_done(ev):
done_event.set()
out_stream.close()
run.on('inference', on_inference)
run.on('done', on_done)
await self.run(run)
# Cannot stream tool calls out, wait for them
await out_stream.content_wait()
print('tool names' + repr(tool_names))
print('tool calls' + repr(tool_calls))
if tool_names:
await done_event.wait()
return {
'output': out_stream,
'tool_calls': tool_calls,
}
else:
await self.run(run)
return out_stream
@@ -157,9 +158,9 @@ class Agent:
kwargs['kind'] = "log"
kwargs['subkind'] = "tool-call"
kwargs['content'] = repr(msg.content)
else:
# Unknown message
kwargs['kind'] = "log"
else: return
new_message = await self.app.send_chat_message(
sender=self.id,
backend/src/agents/expert.py
+3 -2
View File
@@ -47,8 +47,9 @@ async def chat_completion_loop(app, run, messages, model_id=None, tools=None):
print(f"[agent] unknown tool name: {tc.tool_name}")
continue
tc.result = await tool.execute_async(**tc.arguments)
elif turn_response:
break
# By default, end if there's no calls
else: break
run.close()
backend/src/inference.py
+34 -7
View File
@@ -11,8 +11,10 @@ from tool import Tool, ToolSpec, ToolCall
#from llama_cpp import Llama
from openrouter import OpenRouter
# TODO: Collect usage data
TEST_MOCK = False
VALID_ROLES = ['user', 'assistant', 'system', 'tool', 'reasoning']
VALID_ROLES = ['user', 'assistant', 'system', 'tool', 'reasoning', 'usage']
CONTEXT_WINDOW = 10_000
# Background inference tasks are kept referenced here so they aren't GC'd mid-flight.
@@ -60,14 +62,24 @@ def format_openai(messages):
formatted.append(msg)
return msg
# For debugging
_stream_count = 0
_badrole_count = 0
_badroles = set()
for msg in messages:
content = msg.content
if isinstance(content, Stream):
_stream_count += 1
content = content.content
if msg.role == 'tool':
last = last_assistant_message()
tool_calls = last.get('tool_calls')
if tool_calls == None:
tool_calls = []
last['tool_calls'] = tool_calls
tc = msg.content
tc = content
if isinstance(tc, dict):
tc = ToolCall(
tool_name=tc['tool_name'],
@@ -94,10 +106,19 @@ def format_openai(messages):
elif msg.role == 'reasoning':
last = last_assistant_message()
reasoning = last.get('reasoning', '')
last['reasoning'] = reasoning + msg.content
else:
fmsg = {'role': msg.role, 'content': msg.content}
last['reasoning'] = reasoning + content
elif msg.role in ['user', 'assistant', 'system']:
fmsg = {'role': msg.role, 'content': content}
formatted.append(fmsg)
else:
_badrole_count += 1
_badroles.add(msg.role)
_warnings = []
if _stream_count: _warnings.append(f"{_stream_count} streams in message list")
if _badrole_count: _warnings.append(f"{_badrole_count} messages with bad roles {_badroles}")
if _warnings: print(f"[inference] WARNING: " + ", ".join(_warnings))
return formatted
@@ -350,8 +371,14 @@ async def openrouter_chat_completion(app, messages, model, tools=[], opts={}):
tool_call.tool_name += tc.function.name
if tc.function.arguments:
tool_call.arguments += tc.function.arguments
else:
print("[openrouter] no event choices: " + repr(event))
if event.usage:
usage = {}
for key in ['completion_tokens', 'prompt_tokens', 'total_tokens', 'cost']:
if hasattr(event.usage, key):
value = getattr(event.usage, key)
if value != None:
usage[key] = value
queue.put_nowait(Message('usage', usage))
if content_stream: content_stream.close()
if reasoning_stream: reasoning_stream.close()
backend/src/server/openai.py
+60 -32
View File
@@ -1,7 +1,7 @@
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 tool import ToolSpec, ToolParameter, ToolCall
from agent import AgentRun
from utils import encode_json
@@ -72,43 +72,43 @@ def attach_run_handlers(run, tool_names):
queue = asyncio.Queue()
loop = asyncio.get_event_loop()
def on_output(stream):
async def on_chunk(chunk):
async def on_content(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):
async def on_reasoning(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)
def on_inference(msg):
# assistant and reasoning are assumed to have a Stream
# tool is assumed to have a ToolCall
if msg.role == 'assistant':
msg.content.listen(on_content, on_eof)
elif msg.role == 'reasoning':
msg.content.listen(on_reasoning, on_eof)
elif msg.role == 'tool':
deferred = route_tool_call(msg.content, tool_names)
msg_type = 'deferred_call' if deferred else 'tool_call'
loop.create_task(queue.put((msg_type, tool_call)))
loop.create_task(queue.put((msg_type, msg.content)))
elif msg.role == 'usage':
loop.create_task(queue.put(('usage', msg.content)))
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('inference', on_inference)
run.on('done', on_done)
return queue
def format_chunk(agent_id, delta, finish_reason=None):
def format_chunk(agent_id, delta, finish_reason=None, usage=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,
"usage": usage,
"choices": [{"delta": delta, "index": 0, "finish_reason": finish_reason}]
}
@@ -124,8 +124,8 @@ def register(server, app):
# Extract required fields
agent_id = data.get("model")
openai_messages = data.get("messages", [])
openai_tools = data.get("tools", [])
openai_messages = data.get("messages") or []
openai_tools = data.get("tools") or []
stream_mode = data.get("stream", False)
if not agent_id or not openai_messages:
@@ -164,20 +164,36 @@ def register(server, app):
# message carries the output.
call.result = content
entry['resolved'] = True
messages.append({'role': 'tool', 'content': call.result, 'tool_call_id': call.id})
messages.append({'role': 'tool', 'content': call})
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.
# Rebuild the assistant's tool call turn from the request.
# Calls already in the cache (persistent conversation) are
# reused as-is. Client-declared function calls carry their
# full payload on the wire, so they are reconstructed from
# the message when the cache is gone. Server-originated
# (deferred) calls only ever sent an id + type, so they can
# only be recovered from the cache and are dropped otherwise.
calls = []
for tc in msg['tool_calls']:
entry = tool_call_cache.get(tc.get('id', ''))
if entry is None:
continue
call_id = tc.get('id', '')
entry = tool_call_cache.get(call_id)
if entry is not None:
calls.append(entry['call'])
continue
func = tc.get('function')
if tc.get('type') == 'function' and func:
call = ToolCall(
id=call_id,
tool_name=func.get('name', ''),
arguments=json.loads(func.get('arguments') or '{}'),
)
# Seeding the cache lets the following role='tool'
# message resolve the call and attach its result.
cache_tool_call(call, deferred=False)
calls.append(call)
messages.append({
'role': 'assistant',
'content': content if type(content) == str else None,
@@ -239,17 +255,19 @@ def register(server, app):
# 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,
"tool_policy": "ignore",
},
)
print("run input: " + repr(run.input))
queue = attach_run_handlers(run, tool_names)
await control.run(run)
usage = {'completion_tokens': 0, 'prompt_tokens': 0, 'total_tokens': 0, 'cost': 0}
# --- Streaming ---
if stream_mode:
async def generate():
nonlocal usage
try:
tool_call_count = 0
while True:
@@ -281,6 +299,9 @@ def register(server, app):
"type": "tool",
}]}
elif msg_type == 'usage':
usage = data
elif msg_type == 'eof':
# Determine finish_reason — was there a tool call?
finish_reason = "tool_calls" if any(
@@ -289,11 +310,15 @@ def register(server, app):
yield f"data: {json.dumps(format_chunk(agent_id, {}, finish_reason))}\n\n"
break
if delta:
yield f"data: {json.dumps(format_chunk(agent_id, delta))}\n\n"
else:
raise Exception(f"Unrecognized msg_type: {msg_type}. {repr(data)}")
if delta:
yield f"data: {json.dumps(format_chunk(agent_id, delta))}\n\n"
# Usage
yield f"data: {json.dumps(format_chunk(agent_id, {}, finish_reason, usage))}\n\n"
except Exception as e:
print("[openai] ERROR: " + str(e))
print(traceback.format_exc())
@@ -327,6 +352,8 @@ def register(server, app):
# 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"})
elif msg_type == 'usage':
usage = data
else:
raise Exception(f"Unrecognized msg_type: {msg_type}. {repr(data)}")
except Exception as e:
@@ -346,6 +373,7 @@ def register(server, app):
"object": "chat.completion",
"created": int(time.time()),
"model": agent_id,
"usage": usage,
"choices": [{
"message": message,
"finish_reason": finish_reason,