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

This commit is contained in:
Shivering Ocelot committed 2026-10-02 05:55:18 -05:00
1 parent 2ba8096a3c
commit 659b7497fb
6 files changed
+127 -57

No files matched your search

+23 -2
View File
@@ -77,9 +77,30 @@ class Agent:
out_stream.close() out_stream.close()
run.on('output', on_output_event) run.on('output', on_output_event)
run.on('done', on_done) 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): async def received_message(self, message):
chat = await self.app.get_chat(message.chat_id) chat = await self.app.get_chat(message.chat_id)
+6 -9
View File
@@ -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() content = output_stream and await output_stream.content_wait()
reasoning = reasoning_stream and await reasoning_stream.content_wait() reasoning = reasoning_stream and await reasoning_stream.content_wait()
if content or reasoning: if content:
msg = inference.LlamaMessage('assistant', content or "") messages.append(inference.Message('assistant', content))
if reasoning: if reasoning:
msg.reasoning = reasoning messages.append(inference.Message('reasoning', reasoning))
messages.append(msg)
if execute_tools and tool_calls: if execute_tools and tool_calls:
for tc in 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}") print(f"[agent] unknown tool name: {tc.tool_name}")
continue continue
tc.result = await tool.execute_async(**tc.arguments) tc.result = await tool.execute_async(**tc.arguments)
messages.append(inference.LlamaMessage('tool', tc)) messages.append(inference.Message('tool', tc))
tool_calls.clear() tool_calls.clear()
# Buffer time to let the provider rest before the next call elif content or not reasoning or not execute_tools:
await asyncio.sleep(2)
elif content or not reasoning:
break break
run.close() run.close()
+54 -31
View File
@@ -5,22 +5,20 @@ import os
import traceback import traceback
import random import random
from stream import Stream 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 llama_cpp import Llama
from openrouter import OpenRouter from openrouter import OpenRouter
VALID_ROLES = ['user', 'assistant', 'system', 'tool'] VALID_ROLES = ['user', 'assistant', 'system', 'tool', 'reasoning']
CONTEXT_WINDOW = 10_000 CONTEXT_WINDOW = 10_000
# TODO: Rename to ChatMessage. Add methods for constructing from instruct/openai/etc formatted objects. class Message:
class LlamaMessage: def __init__(self, role, content):
def __init__(self, role, content, reasoning=None):
if role not in VALID_ROLES: if role not in VALID_ROLES:
raise Exception("message role must be one of: " + ",".join(VALID_ROLES)) raise Exception("message role must be one of: " + ",".join(VALID_ROLES))
self.role = role self.role = role
self.content = content self.content = content
self.reasoning = reasoning
def to_dict(self): def to_dict(self):
"""Serialize to the dict format expected by llama-cpp / ollama.""" """Serialize to the dict format expected by llama-cpp / ollama."""
@@ -30,44 +28,68 @@ class LlamaMessage:
return d return d
def __repr__(self): 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: def format_tools(tools):
llama_8b = "meta/meta-llama-3-8b-instruct" specs = []
llama_70b = "meta/meta-llama-3-70b-instruct" for t in tools:
qwen3_8b = "lucataco/qwen3-vl-8b-instruct:39e893666996acf464cff75688ad49ac95ef54e9f1c688fbc677330acc478e11" if isinstance(t, dict):
qwen3_235b = "qwen/qwen3-235b-a22b-instruct-2507" 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): def format_openai(messages):
print('format openai ' + repr(messages))
formatted = [] 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: for msg in messages:
if msg.role == 'tool': if msg.role == 'tool':
if not last_assistant_message: last = last_assistant_message()
# Create one implicitly tool_calls = last.get('tool_calls')
last_assistant_message = {'role': 'assistant', 'content': ''}
tool_calls = last_assistant_message.get('tool_calls')
if tool_calls == None: if tool_calls == None:
tool_calls = [] 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({ tool_calls.append({
'id': msg.content.id, 'id': tc.id,
'type': 'function', 'type': 'function',
'function': { 'function': {
'name': msg.content.tool_name, 'name': tc.tool_name,
'arguments': json.dumps(msg.content.arguments), 'arguments': json.dumps(tc.arguments),
}, },
}) })
formatted.append({ formatted.append({
'role': 'tool', 'role': 'tool',
'tool_call_id': msg.content.id, 'tool_call_id': tc.id,
'content': json.dumps(msg.content.result), 'content': json.dumps(tc.result),
}) })
elif msg.role == 'reasoning':
last = last_assistant_message()
reasoning = last.get('reasoning', '')
last['reasoning'] = reasoning + msg.content
else: else:
fmsg = {'role': msg.role, 'content': msg.content} fmsg = {'role': msg.role, 'content': msg.content}
formatted.append(fmsg) formatted.append(fmsg)
if msg.role == 'assistant':
last_assistant_message = fmsg
return formatted return formatted
@@ -116,10 +138,10 @@ def get_llama_cpp_model(app, model_file):
# Generate a conversation based on agent inputs # Generate a conversation based on agent inputs
def build_messages(system_prompt="", instruction=None, input=None, messages=None): def build_messages(system_prompt="", instruction=None, input=None, messages=None):
def normalize(message): def normalize(message):
if isinstance(message, LlamaMessage): if isinstance(message, Message):
return message return message
elif isinstance(message, dict): elif isinstance(message, dict):
return LlamaMessage(message['role'], message['content']) return Message(message['role'], message['content'])
else: else:
raise Exception("Unsupported message type: " + repr(type(message))) 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 # Add default messages to list
if system_prompt: if system_prompt:
messages.insert(0, LlamaMessage("system", system_prompt)) messages.insert(0, Message("system", system_prompt))
if prompt: if prompt:
messages.append(LlamaMessage("user", prompt)) messages.append(Message("user", prompt))
return messages return messages
@@ -205,7 +227,7 @@ async def local_chat_completion(app, messages, model, tools=[], opts={}, agent_r
kwargs['repeat_penalty'] = repeat kwargs['repeat_penalty'] = repeat
# Format tools for llama-cpp (OpenAI-compatible function format) # 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: if formatted_tools:
kwargs['tools'] = formatted_tools kwargs['tools'] = formatted_tools
kwargs['tool_choice'] = 'auto' 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 # 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: if formatted_tools:
kwargs['tools'] = formatted_tools kwargs['tools'] = formatted_tools
print(f"Tools: {tools} and {formatted_tools}")
temp = agent_run.input.get("temperature") temp = agent_run.input.get("temperature")
if temp is not None: if temp is not None:
+15 -1
View File
@@ -294,18 +294,32 @@ def register(server, app, auth_manager):
await send({'error': "instruction or input required"}) await send({'error': "instruction or input required"})
continue continue
# Shape of response changes based on options lol
response = await app.send_agent_prompt( response = await app.send_agent_prompt(
await agent_id_from_body(), await agent_id_from_body(),
body.get('instruction'), body.get('instruction'),
body.get('input'), body.get('input'),
body.get('options', {}) 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 isinstance(response, Stream):
if body.get('stream'): if body.get('stream'):
response = {"streamId": response.id} response = {"streamId": response.id}
else: else:
response = await response.content_wait() 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": elif action == "read-stream":
stream_id = body['streamId'] stream_id = body['streamId']
+2 -9
View File
@@ -218,24 +218,17 @@ class ToolCall:
is_error: Whether the result represents an error. 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.id = id or ""
self.tool_name = tool_name self.tool_name = tool_name
self.arguments = arguments self.arguments = arguments
self.result = None self.result = result
def __repr__(self): def __repr__(self):
id_prefix = self.id + ', ' if self.id else '' id_prefix = self.id + ', ' if self.id else ''
return f"ToolCall({id_prefix}tool={self.tool_name!r}, args={self.arguments!r})" 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 # Helpers to format tool specs per engine
+27 -5
View File
@@ -1,14 +1,36 @@
import json import json
from stream import Stream 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_json(value):
def encode_unknown(value): def encode_unknown(value):
if isinstance(value, ToolSpec): # All complex values are transformed to a dict
return value.to_openai_dict() value = encode_value(value)
elif isinstance(value, Stream): # otherwise simply show the repr
return {'streamId': value.id} if isinstance(value, dict):
return value
else: else:
return repr(value) return repr(value)
return json.dumps(value, default=encode_unknown) return json.dumps(value, default=encode_unknown)