diff --git a/backend/src/agent.py b/backend/src/agent.py index 93fe6d2..6a4e907 100644 --- a/backend/src/agent.py +++ b/backend/src/agent.py @@ -144,13 +144,28 @@ class Agent: }, ) - async def on_output_event(stream): - print("[agent] run output " + repr(stream)) + async def on_inference(msg): + kwargs = {'content': msg.content} + + if msg.role == 'assistant': + # kind and subkind are default + pass + elif msg.role == 'reasoning': + kwargs['kind'] = "log" + kwargs['subkind'] = "reasoning" + elif msg.role == 'tool': + kwargs['kind'] = "log" + kwargs['subkind'] = "tool-call" + kwargs['content'] = repr(msg.content) + else: + # Unknown message + kwargs['kind'] = "log" + new_message = await self.app.send_chat_message( sender=self.id, chat_id=message.chat_id, thread_id=message.thread_id, - content=stream, + **kwargs ) run.emit_event('chat-message', { 'messageId': new_message['id'], @@ -158,41 +173,7 @@ class Agent: 'threadId': message.thread_id, }) - async def on_reasoning_event(stream): - print("[agent] run output " + repr(stream)) - new_message = await self.app.send_chat_message( - sender=self.id, - chat_id=message.chat_id, - thread_id=message.thread_id, - content=stream, - kind="log", - subkind="reasoning", - ) - run.emit_event('chat-message', { - 'messageId': new_message['id'], - 'chatId': message.chat_id, - 'threadId': message.thread_id, - }) - - async def on_toolcall_event(stream): - print("[agent] run output " + repr(stream)) - new_message = await self.app.send_chat_message( - sender=self.id, - chat_id=message.chat_id, - thread_id=message.thread_id, - content=stream, - kind="log", - subkind="tool-call", - ) - run.emit_event('chat-message', { - 'messageId': new_message['id'], - 'chatId': message.chat_id, - 'threadId': message.thread_id, - }) - - await run.on_async('reasoning', on_reasoning_event) - await run.on_async('output', on_output_event) - await run.on_async('tool-call', on_toolcall_event) + await run.on_async('inference', on_inference) await self.run(run) return run diff --git a/backend/src/agents/expert.py b/backend/src/agents/expert.py index b98c1ab..ff4f1a0 100644 --- a/backend/src/agents/expert.py +++ b/backend/src/agents/expert.py @@ -20,47 +20,25 @@ async def chat_completion_loop(app, run, messages, model_id=None, tools=None): tools = list(tools or []) execute_tools = run.input.get('tool_policy', 'auto') == 'auto' - tool_calls = [] - output_stream = None - reasoning_stream = None - inference_end = asyncio.Event() - - def on_output(stream): - nonlocal output_stream - output_stream = stream - - def on_reasoning(stream): - nonlocal reasoning_stream - reasoning_stream = stream - - run.on('output', on_output) - run.on('reasoning', on_reasoning) - run.on('tool-call', tool_calls.append) - run.on('inference-end', lambda _: inference_end.set()) - while True: - # Reset per-iteration state - output_stream = None - reasoning_stream = None - inference_end.clear() - - await inference.chat_completion( + tool_calls = [] + turn_response = False + generator = inference.chat_completion( app, model_id=model_id, messages=messages, tools=tools, - agent_run=run, + opts=run.input, ) - await inference_end.wait() + async for msg in generator: + messages.append(msg) + if msg.role == 'tool': + tool_calls.append(msg.content) + if msg.role == 'assistant': + turn_response = True - content = output_stream and await output_stream.content_wait() - reasoning = reasoning_stream and await reasoning_stream.content_wait() - - if content: - messages.append(inference.Message('assistant', content)) - if reasoning: - messages.append(inference.Message('reasoning', reasoning)) + run.emit_event('inference', msg) if execute_tools and tool_calls: for tc in tool_calls: @@ -69,9 +47,7 @@ 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) - messages.append(inference.Message('tool', tc)) - tool_calls.clear() - elif content or not reasoning or not execute_tools: + elif turn_response: break run.close() diff --git a/backend/src/inference.py b/backend/src/inference.py index 76e9b5a..1f1fb77 100644 --- a/backend/src/inference.py +++ b/backend/src/inference.py @@ -4,15 +4,28 @@ import json import os import traceback import random +import inspect from stream import Stream from tool import Tool, ToolSpec, ToolCall #from llama_cpp import Llama from openrouter import OpenRouter +TEST_MOCK = False VALID_ROLES = ['user', 'assistant', 'system', 'tool', 'reasoning'] CONTEXT_WINDOW = 10_000 +# Background inference tasks are kept referenced here so they aren't GC'd mid-flight. +_background_tasks = set() + +def spawn(coro): + task = asyncio.create_task(coro) + _background_tasks.add(task) + task.add_done_callback(_background_tasks.discard) + return task + +# These reflect "emitted messages" from any actor, individual turn parts, rather than entire turns. +# role acts as a type and a sender simultaneously, and content is dynamic class Message: def __init__(self, role, content): if role not in VALID_ROLES: @@ -21,11 +34,7 @@ class Message: self.content = content def to_dict(self): - """Serialize to the dict format expected by llama-cpp / ollama.""" - d = {"role": self.role, "content": self.content} - if self.reasoning: - d["reasoning"] = self.reasoning - return d + return {"role": self.role, "content": self.content} def __repr__(self): return f"Message(role={self.role!r}, content={self.content!r})" @@ -42,7 +51,6 @@ def format_tools(tools): return specs def format_openai(messages): - print('format openai ' + repr(messages)) formatted = [] def last_assistant_message(): if formatted and formatted[-1]['role'] == 'assistant': @@ -103,9 +111,12 @@ def safe_queue(callback): queue = asyncio.Queue() loop = asyncio.get_running_loop() async def read_loop(): + nonlocal callback while True: args = await queue.get() - await callback(*args) + task = callback(*args) + if inspect.isawaitable(task): + await task loop.create_task(read_loop()) def send(*args): @@ -173,232 +184,200 @@ async def chat_completion(app, messages=None, tools=[], opts={}, - agent_run=None ): - # TESTING OVERRIDE - # return await mock_completion(app, agent_run) db_model = (await app.get_model(model_id)).value model_engine = db_model['engine'] model_name = db_model['model'] - if model_engine == "llama-cpp": - result = await local_chat_completion(app, messages, model=model_name, tools=tools, opts=opts, agent_run=agent_run) - elif model_engine == "openrouter": - result = await openrouter_chat_completion(app, messages, model=model_name, tools=tools, opts=opts, agent_run=agent_run) - else: - raise Exception("Unsupported model: " + repr(model)) + print(f"[inference] chat completion: {model_engine} '{model_name}', {len(messages)} messages, {len(tools)} tools") - if agent_run: - return None - return result + if TEST_MOCK: + generator = mock_completion(app) + elif model_engine == "llama-cpp": + generator = local_chat_completion(app, messages, model=model_name, tools=tools, opts=opts) + elif model_engine == "openrouter": + generator = openrouter_chat_completion(app, messages, model=model_name, tools=tools, opts=opts) + else: + raise Exception("Unsupported model engine: " + repr(model_engine)) + + async for message in generator: + yield message + + print(f"[inference] chat completion end: {model_engine} '{model_name}'") # llama.cpp -async def local_chat_completion(app, messages, model, tools=[], opts={}, agent_run=None): - if not agent_run: - raise Exception("agent_run required") - - content_stream = None - async def message_loop(name, value=None): - nonlocal content_stream, agent_run - if name == 'output': - if content_stream == None: - content_stream = Stream(app) - agent_run.emit_event('output', content_stream) - content_stream.feed(value) - elif name == 'close': - if content_stream: - content_stream.close() - agent_run.emit_event('inference-end', {}) - send_message = safe_queue(message_loop) - - def worker(): - llm = get_llama_cpp_model(app, model) - - kwargs = {} - - temp = opts.get("temperature") - if temp is not None: - kwargs['temperature'] = temp - - repeat = opts.get("repetition_penalty") - if repeat is not None: - kwargs['repeat_penalty'] = repeat - - # Format tools for llama-cpp (OpenAI-compatible function format) - formatted_tools = format_tools(tools) if tools else None - if formatted_tools: - kwargs['tools'] = formatted_tools - kwargs['tool_choice'] = 'auto' - - # With stream=True, returns an iterator that infers tokens on demand (openai compatible) - llm_output = llm.create_chat_completion( - messages=[msg.to_dict() for msg in messages], - stream=True, - **kwargs - ) - - for chunk in llm_output: - # Check for cancellation - if agent_run and agent_run.cancelled: - print(f"Cancelling local inference for run {agent_run.id}") - break - - delta = chunk['choices'][0]['delta'] - if 'content' in delta: - send_message('output', delta['content']) - # Handle tool_call deltas if the model emits them - if 'tool_calls' in delta: - # Forward tool call chunks as JSON strings for the consumer - send_message("tool-call", { - "id": tc.get("id"), - "function": tc.get("function"), - }) - - send_message('close') - - thread = threading.Thread(target=worker) - if agent_run: - agent_run.thread = thread - thread.start() - - if agent_run: - return None - return content_stream - - -async def openrouter_chat_completion(app, messages, model, tools=[], opts={}, agent_run=None): - if not agent_run: - raise Exception("agent_run required") - - # 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) - agent_run.emit_event('output', content_stream) - content_stream.feed(value) - if name == 'reasoning': - if reasoning_stream == None: - reasoning_stream = Stream(app) - agent_run.emit_event('reasoning', reasoning_stream) - reasoning_stream.feed(value) - if name == 'tool-call': - agent_run.emit_event('tool-call', value) - elif name == 'close': - if content_stream: content_stream.close() - if reasoning_stream: reasoning_stream.close() - agent_run.emit_event('inference-end', {}) - emit = safe_queue(message_loop) +async def local_chat_completion(app, messages, model, tools=[], opts={}): + queue = asyncio.Queue() + loop = asyncio.get_running_loop() + emit = safe_queue(queue.put_nowait) async def worker(): - print("START OPENROUTER INFERENCE") - if not app.openrouter_client: - raise Exception("OPENROUTER_API_KEY environment variable is not set") - - kwargs = { - 'model': model, - 'messages': format_openai(messages), - 'stream': True, - # Options - 'timeout_ms': 45_000, - 'retries': 3, - } - - # Pass tools in OpenAI-compatible format if provided - formatted_tools = format_tools(tools) if tools else None - if formatted_tools: - kwargs['tools'] = formatted_tools - print(f"Tools: {tools} and {formatted_tools}") - - temp = agent_run.input.get("temperature") - if temp is not None: - kwargs['temperature'] = temp - - reasoning = agent_run.input.get("reasoning_effort") - if reasoning is not None: - kwargs['reasoning'] = {"effort": reasoning} - - stream_response = await app.openrouter_client.chat.send_async(**kwargs) - - # OpenAI format streams tool calls in parts - # We could, but I won't assume any tool call is complete until the message is over - tool_calls = [] - - async for event in stream_response: - if agent_run and agent_run.cancelled: - stream_response.cancel() - break - - if event.choices and len(event.choices) > 0: - delta = event.choices[0].delta - - if delta.content: - emit('output', delta.content) - if delta.reasoning: - emit('reasoning', delta.reasoning) - if delta.tool_calls: - print("[openrouter] tool call:" + repr(delta.tool_calls)) - for tc in delta.tool_calls: - if tc.index >= len(tool_calls): - tool_calls.append(ToolCall(tool_name="", arguments="", id="")) - tool_call = tool_calls[tc.index] - - if tc.id: - tool_call.id += tc.id - - if tc.function: - if tc.function.name: - 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)) - - for tc in tool_calls: - try: - arguments = json.loads(tc.arguments) - tc.arguments = arguments - emit('tool-call', tc) - except Exception: - print("[openrouter] Bad tool call arguments for " + repr(tc)) - print(traceback.format_exc()) - - emit('close') - print("END OPENROUTER INFERENCE") - - def handle_done(task): + stream = None try: - task.result() - except asyncio.CancelledError: - print("[openrouter] inference cancelled") - except Exception as e: - print(f"[openrouter] Worker error: {e}") + llm = get_llama_cpp_model(app, model) + + kwargs = {} + + temp = opts.get("temperature") + if temp is not None: + kwargs['temperature'] = temp + + repeat = opts.get("repetition_penalty") + if repeat is not None: + kwargs['repeat_penalty'] = repeat + + # Format tools for llama-cpp (OpenAI-compatible function format) + formatted_tools = format_tools(tools) if tools else None + if formatted_tools: + kwargs['tools'] = formatted_tools + kwargs['tool_choice'] = 'auto' + + # With stream=True, returns an iterator that infers tokens on demand (openai compatible) + iterator = await loop.run_in_executor(None, lambda: llm.create_chat_completion( + messages=[msg.to_dict() for msg in messages], + stream=True, + **kwargs + )) + + sentinel = object() + while True: + chunk = await loop.run_in_executor(None, next, iterator, sentinel) + if chunk is sentinel: + break + + delta = chunk['choices'][0]['delta'] + if 'content' in delta: + if stream is None: + stream = Stream(app) + emit(Message('assistant', stream)) + stream.feed(delta['content']) + if 'tool_calls' in delta: + for tc in delta['tool_calls']: + emit(Message('tool', tc)) + except Exception: print(traceback.format_exc()) + finally: + if stream: stream.close() + emit(None) - task = asyncio.create_task(worker()) - task.add_done_callback(handle_done) - agent_run.task = task + spawn(worker()) + + while True: + message = await queue.get() + if message is None: + break + yield message -async def mock_completion(app, agent_run=None): +async def openrouter_chat_completion(app, messages, model, tools=[], opts={}): + if not app.openrouter_client: + raise Exception("OPENROUTER_API_KEY environment variable is not set") + + queue = asyncio.Queue() + + async def worker(): + 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 + + try: + kwargs = { + 'model': model, + 'messages': format_openai(messages), + 'stream': True, + # Options + 'timeout_ms': 45_000, + 'retries': 3, + } + + # Pass tools in OpenAI-compatible format if provided + formatted_tools = format_tools(tools) if tools else None + if formatted_tools: + kwargs['tools'] = formatted_tools + + temp = opts.get("temperature") + if temp is not None: + kwargs['temperature'] = temp + + reasoning = opts.get("reasoning_effort") + if reasoning is not None: + kwargs['reasoning'] = {"effort": reasoning} + + stream_response = await app.openrouter_client.chat.send_async(**kwargs) + + # OpenAI format streams tool calls in parts + # We could, but I won't assume any tool call is complete until the message is over + tool_calls = [] + + async for event in stream_response: + if event.choices and len(event.choices) > 0: + delta = event.choices[0].delta + + if delta.reasoning: + open_channel('reasoning') + if reasoning_stream is None: + reasoning_stream = Stream(app) + queue.put_nowait(Message('reasoning', reasoning_stream)) + reasoning_stream.feed(delta.reasoning) + if delta.content: + open_channel('output') + if content_stream is None: + content_stream = Stream(app) + queue.put_nowait(Message('assistant', content_stream)) + content_stream.feed(delta.content) + if delta.tool_calls: + for tc in delta.tool_calls: + if tc.index >= len(tool_calls): + tool_calls.append(ToolCall(tool_name="", arguments="", id="")) + tool_call = tool_calls[tc.index] + + if tc.id: + tool_call.id += tc.id + + if tc.function: + if tc.function.name: + 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 content_stream: content_stream.close() + if reasoning_stream: reasoning_stream.close() + + for tc in tool_calls: + try: + tc.arguments = json.loads(tc.arguments) + queue.put_nowait(Message('tool', tc)) + except Exception: + print("[openrouter] Bad tool call arguments for " + repr(tc)) + print(traceback.format_exc()) + except Exception: + print(traceback.format_exc()) + finally: + queue.put_nowait(None) + + spawn(worker()) + + while True: + message = await queue.get() + if message is None: + break + yield message + + +async def mock_completion(app): VERBS = ['fix','scan','sync','deploy','parse','render','cache','trace', 'merge','purge','index','probe','flush','bind','warm','drain'] OBJECTS = ['nodes','packets','tokens','frames','shards','clusters', @@ -406,20 +385,28 @@ async def mock_completion(app, agent_run=None): 'handles','slots','cursors','segments'] DELAY = 0.3 - content_stream = Stream(app) - if agent_run: - agent_run.emit_event('output', content_stream) + queue = asyncio.Queue() + emit = safe_queue(queue.put_nowait) - async def content_task (): - await asyncio.sleep(DELAY) - content_stream.feed(random.choice(VERBS)) - await asyncio.sleep(DELAY) - content_stream.feed(f" {random.randint(1, 99)} ") - await asyncio.sleep(DELAY) - content_stream.feed(random.choice(VERBS)) - await asyncio.sleep(DELAY) - content_stream.close() - agent_run.emit_event('inference-end', {}) + async def worker(): + stream = Stream(app) + emit(Message('assistant', stream)) + try: + await asyncio.sleep(DELAY) + stream.feed(random.choice(VERBS)) + await asyncio.sleep(DELAY) + stream.feed(f" {random.randint(1, 99)} ") + await asyncio.sleep(DELAY) + stream.feed(random.choice(VERBS)) + await asyncio.sleep(DELAY) + finally: + stream.close() + emit(None) - asyncio.create_task(content_task()) - if not agent_run: return content_stream + spawn(worker()) + + while True: + message = await queue.get() + if message is None: + break + yield message diff --git a/backend/src/server/websocket.py b/backend/src/server/websocket.py index 49851c1..f1a2885 100644 --- a/backend/src/server/websocket.py +++ b/backend/src/server/websocket.py @@ -1,4 +1,4 @@ -import json, asyncio +import json, asyncio, traceback from quart import websocket from utils import encode_json from stream import Stream # isinstance check in send-agent-prompt @@ -40,10 +40,14 @@ def register(server, app, auth_manager): send_queue = asyncio.Queue() async def sender_impl(): + # This thread MUST NOT crash nonlocal send_queue while True: - msg = await send_queue.get() - await websocket.send(encode_json(msg)) + try: + msg = await send_queue.get() + await websocket.send(encode_json(msg)) + except Exception as err: + print(traceback.format_exc()) async def send(msg): await send_queue.put(msg) @@ -72,6 +76,7 @@ def register(server, app, auth_manager): action = req['action'] body = req.get('body', {}) nonce = req.get('nonce', None) + print(f"[ws] action @{user_id} {action}") async def agent_id_from_body (allow_none=False): if 'agentId' in body: diff --git a/backend/src/utils.py b/backend/src/utils.py index 5013e5a..894d90d 100644 --- a/backend/src/utils.py +++ b/backend/src/utils.py @@ -9,7 +9,7 @@ def encode_value(value): if isinstance(value, ToolSpec): return value.to_openai_dict() elif isinstance(value, Message): - return {'type': value.type, 'content': encode_value(value.content)} + return {'role': value.role, 'content': encode_value(value.content)} elif isinstance(value, ToolCall): json = { 'id': value.id,