feat(backend): tool support in send agent prompt, regularize inference messages (wip)
This commit is contained in:
1 parent
2ba8096a3c
commit
659b7497fb
6 files changed
+127
-57
No files matched your search
+23
-2
@@ -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)
|
||||||
|
|||||||
@@ -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
@@ -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:
|
||||||
|
|||||||
@@ -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
@@ -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
@@ -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)
|
||||||
Reference in new issue
Block a user