feat(backend): tool support in send agent prompt, regularize inference messages (wip)

backend/src/agent.py

@@ -77,10 +77,31 @@ class Agent:
             out_stream.close()
         run.on('output', on_output_event)
         run.on('done', on_done)
+
+        # 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)
         thread = chat.get_thread(message.thread_id)

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 content:
+            messages.append(inference.Message('assistant', content))
         if reasoning:
-                msg.reasoning = reasoning
-            messages.append(msg)
+            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()

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:

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']

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

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
 
-def encode_json(value):
-    def encode_unknown(value):
+# 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):
+        # 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)