diff --git a/backend/src/agent.py b/backend/src/agent.py index 4f4910b..0780c34 100644 --- a/backend/src/agent.py +++ b/backend/src/agent.py @@ -77,9 +77,30 @@ class Agent: out_stream.close() run.on('output', on_output_event) run.on('done', on_done) - await self.run(run) - return out_stream + # 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) + 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)) + return { + 'output': out_stream, + 'tool_calls': tool_calls, + } + else: + await self.run(run) + return out_stream + async def received_message(self, message): chat = await self.app.get_chat(message.chat_id) diff --git a/backend/src/agents/expert.py b/backend/src/agents/expert.py index a8342eb..b98c1ab 100644 --- a/backend/src/agents/expert.py +++ b/backend/src/agents/expert.py @@ -57,11 +57,10 @@ async def chat_completion_loop(app, run, messages, model_id=None, tools=None): content = output_stream and await output_stream.content_wait() reasoning = reasoning_stream and await reasoning_stream.content_wait() - if content or reasoning: - msg = inference.LlamaMessage('assistant', content or "") - if reasoning: - msg.reasoning = reasoning - messages.append(msg) + if content: + messages.append(inference.Message('assistant', content)) + if reasoning: + messages.append(inference.Message('reasoning', reasoning)) if execute_tools and tool_calls: for tc in tool_calls: @@ -70,11 +69,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) - messages.append(inference.LlamaMessage('tool', tc)) + messages.append(inference.Message('tool', tc)) tool_calls.clear() - # Buffer time to let the provider rest before the next call - await asyncio.sleep(2) - elif content or not reasoning: + elif content or not reasoning or not execute_tools: break run.close() diff --git a/backend/src/inference.py b/backend/src/inference.py index 3c26800..8fe5ffc 100644 --- a/backend/src/inference.py +++ b/backend/src/inference.py @@ -5,22 +5,20 @@ import os import traceback import random from stream import Stream -from tool import format_tools_for_engine, build_tool_prompt_prefix, ToolSpec, ToolCall +from tool import Tool, ToolSpec, ToolCall #from llama_cpp import Llama from openrouter import OpenRouter -VALID_ROLES = ['user', 'assistant', 'system', 'tool'] +VALID_ROLES = ['user', 'assistant', 'system', 'tool', 'reasoning'] CONTEXT_WINDOW = 10_000 -# TODO: Rename to ChatMessage. Add methods for constructing from instruct/openai/etc formatted objects. -class LlamaMessage: - def __init__(self, role, content, reasoning=None): +class Message: + def __init__(self, role, content): if role not in VALID_ROLES: raise Exception("message role must be one of: " + ",".join(VALID_ROLES)) self.role = role self.content = content - self.reasoning = reasoning def to_dict(self): """Serialize to the dict format expected by llama-cpp / ollama.""" @@ -30,44 +28,68 @@ class LlamaMessage: return d def __repr__(self): - return f"LlamaMessage(role={self.role!r}, content={self.content!r}, reasoning={self.reasoning!r})" + return f"Message(role={self.role!r}, content={self.content!r})" -class Models: - llama_8b = "meta/meta-llama-3-8b-instruct" - llama_70b = "meta/meta-llama-3-70b-instruct" - qwen3_8b = "lucataco/qwen3-vl-8b-instruct:39e893666996acf464cff75688ad49ac95ef54e9f1c688fbc677330acc478e11" - qwen3_235b = "qwen/qwen3-235b-a22b-instruct-2507" +def format_tools(tools): + specs = [] + for t in tools: + if isinstance(t, dict): + specs.append(t) + elif isinstance(t, ToolSpec): + specs.append(t.format_openai()) + elif isinstance(t, Tool): + specs.append(t.get_spec().format_openai()) + return specs def format_openai(messages): + print('format openai ' + repr(messages)) formatted = [] - last_assistant_message = None + def last_assistant_message(): + if formatted and formatted[-1]['role'] == 'assistant': + return formatted[-1] + else: + msg = {'role': 'assistant', 'content': ''} + formatted.append(msg) + return msg + for msg in messages: if msg.role == 'tool': - if not last_assistant_message: - # Create one implicitly - last_assistant_message = {'role': 'assistant', 'content': ''} - tool_calls = last_assistant_message.get('tool_calls') + last = last_assistant_message() + tool_calls = last.get('tool_calls') if tool_calls == None: tool_calls = [] - last_assistant_message['tool_calls'] = tool_calls + last['tool_calls'] = tool_calls + tc = msg.content + if isinstance(tc, dict): + tc = ToolCall( + tool_name=tc['tool_name'], + arguments=tc['arguments'], + result=tc['result'], + id=tc['id'] + ) + # Skip non-resolved calls + if not tc.result: continue + # TODO: Fix call id if missing tool_calls.append({ - 'id': msg.content.id, + 'id': tc.id, 'type': 'function', 'function': { - 'name': msg.content.tool_name, - 'arguments': json.dumps(msg.content.arguments), + 'name': tc.tool_name, + 'arguments': json.dumps(tc.arguments), }, }) formatted.append({ 'role': 'tool', - 'tool_call_id': msg.content.id, - 'content': json.dumps(msg.content.result), + 'tool_call_id': tc.id, + 'content': json.dumps(tc.result), }) + 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} formatted.append(fmsg) - if msg.role == 'assistant': - last_assistant_message = fmsg return formatted @@ -116,10 +138,10 @@ def get_llama_cpp_model(app, model_file): # Generate a conversation based on agent inputs def build_messages(system_prompt="", instruction=None, input=None, messages=None): def normalize(message): - if isinstance(message, LlamaMessage): + if isinstance(message, Message): return message elif isinstance(message, dict): - return LlamaMessage(message['role'], message['content']) + return Message(message['role'], message['content']) else: raise Exception("Unsupported message type: " + repr(type(message))) @@ -139,9 +161,9 @@ def build_messages(system_prompt="", instruction=None, input=None, messages=None # Add default messages to list if system_prompt: - messages.insert(0, LlamaMessage("system", system_prompt)) + messages.insert(0, Message("system", system_prompt)) if prompt: - messages.append(LlamaMessage("user", prompt)) + messages.append(Message("user", prompt)) return messages @@ -205,7 +227,7 @@ async def local_chat_completion(app, messages, model, tools=[], opts={}, agent_r kwargs['repeat_penalty'] = repeat # Format tools for llama-cpp (OpenAI-compatible function format) - formatted_tools = format_tools_for_engine(tools, "llama-cpp") if tools else None + formatted_tools = format_tools(tools) if tools else None if formatted_tools: kwargs['tools'] = formatted_tools kwargs['tool_choice'] = 'auto' @@ -301,9 +323,10 @@ async def openrouter_chat_completion(app, messages, model, tools=[], opts={}, ag } # Pass tools in OpenAI-compatible format if provided - formatted_tools = format_tools_for_engine(tools, "openrouter") if tools else None + 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: diff --git a/backend/src/server/websocket.py b/backend/src/server/websocket.py index 6ace764..35e2182 100644 --- a/backend/src/server/websocket.py +++ b/backend/src/server/websocket.py @@ -294,18 +294,32 @@ def register(server, app, auth_manager): await send({'error': "instruction or input required"}) continue + # Shape of response changes based on options lol response = await app.send_agent_prompt( await agent_id_from_body(), body.get('instruction'), body.get('input'), body.get('options', {}) ) + print("prompt response", repr(response)) + tool_calls = None + + # Extract from dict + if isinstance(response, dict): + tool_calls = response['tool_calls'] + response = response['output'] + if isinstance(response, Stream): if body.get('stream'): response = {"streamId": response.id} else: response = await response.content_wait() - await send({'message': "prompt sent", 'response': response, 'nonce': nonce}) + + packet_out = {'message': "prompt sent", 'response': response, 'nonce': nonce} + if tool_calls: + packet_out['toolCalls'] = tool_calls + + await send(packet_out) elif action == "read-stream": stream_id = body['streamId'] diff --git a/backend/src/tool.py b/backend/src/tool.py index 463f031..c6e503d 100644 --- a/backend/src/tool.py +++ b/backend/src/tool.py @@ -218,24 +218,17 @@ class ToolCall: is_error: Whether the result represents an error. """ - def __init__(self, tool_name, arguments={}, id=None): + def __init__(self, tool_name, arguments={}, id=None, result=None): self.id = id or "" self.tool_name = tool_name self.arguments = arguments - self.result = None + self.result = result def __repr__(self): id_prefix = self.id + ', ' if self.id else '' return f"ToolCall({id_prefix}tool={self.tool_name!r}, args={self.arguments!r})" - def call_message(self): - raise Exception("TODO: Llama message for tool call") - - def result_message(self): - from inference import LlamaMessage - return LlamaMessage("tool", self.result) - # --------------------------------------------------------------------------- # Helpers to format tool specs per engine diff --git a/backend/src/utils.py b/backend/src/utils.py index c37d5f1..5013e5a 100644 --- a/backend/src/utils.py +++ b/backend/src/utils.py @@ -1,14 +1,36 @@ import json from stream import Stream -from tool import ToolSpec +from tool import ToolSpec, ToolCall +from inference import Message + +# TODO: Move to module 'format' +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)} + elif isinstance(value, ToolCall): + json = { + 'id': value.id, + 'tool_name': value.tool_name, + 'arguments': encode_value(value.arguments), + } + if value.result != None: + json['result'] = encode_value(value.result), + return json + elif isinstance(value, Stream): + return {'streamId': value.id} + else: + return value def encode_json(value): def encode_unknown(value): - if isinstance(value, ToolSpec): - return value.to_openai_dict() - elif isinstance(value, Stream): - return {'streamId': value.id} + # All complex values are transformed to a dict + value = encode_value(value) + # otherwise simply show the repr + if isinstance(value, dict): + return value else: return repr(value) return json.dumps(value, default=encode_unknown) \ No newline at end of file