diff --git a/backend/src/agent.py b/backend/src/agent.py index 6a4e907..03d2db2 100644 --- a/backend/src/agent.py +++ b/backend/src/agent.py @@ -77,38 +77,39 @@ class Agent: **kwargs, }, ) + + done_event = asyncio.Event() out_stream = Stream(self.app) - - def on_output_event(stream): - 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) - - # Change the output shape based on the presence of tools + tool_calls = [] + tool_names = None 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) - 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)) + def on_chunk(chunk): + out_stream.feed(chunk) + def on_inference(msg): + kwargs = {'content': msg.content} + + 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) + + 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" + + # Unknown message + else: return new_message = await self.app.send_chat_message( sender=self.id, diff --git a/backend/src/agents/expert.py b/backend/src/agents/expert.py index ff4f1a0..b3f97e5 100644 --- a/backend/src/agents/expert.py +++ b/backend/src/agents/expert.py @@ -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() diff --git a/backend/src/inference.py b/backend/src/inference.py index 1f1fb77..0804b29 100644 --- a/backend/src/inference.py +++ b/backend/src/inference.py @@ -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() diff --git a/backend/src/server/openai.py b/backend/src/server/openai.py index 79dd237..7172952 100644 --- a/backend/src/server/openai.py +++ b/backend/src/server/openai.py @@ -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): - 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) + async def on_content(chunk): + await queue.put(('content', chunk)) + async def on_reasoning(chunk): + await queue.put(('reasoning', chunk)) + async def on_eof(): pass - 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_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, 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: + call_id = tc.get('id', '') + entry = tool_call_cache.get(call_id) + if entry is not None: + calls.append(entry['call']) continue - calls.append(entry['call']) + + 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,