372 lines
12 KiB
Python
372 lines
12 KiB
Python
"""Main agent for media library management."""
|
|
|
|
import json
|
|
import logging
|
|
from collections.abc import AsyncGenerator
|
|
from typing import Any
|
|
|
|
from infrastructure.persistence import get_memory
|
|
|
|
from .config import settings
|
|
from .prompts import PromptBuilder
|
|
from .registry import Tool, make_tools
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class Agent:
|
|
"""
|
|
AI agent for media library management.
|
|
|
|
Uses OpenAI-compatible tool calling API.
|
|
"""
|
|
|
|
def __init__(self, llm, max_tool_iterations: int = 5):
|
|
"""
|
|
Initialize the agent.
|
|
|
|
Args:
|
|
llm: LLM client with complete() method
|
|
max_tool_iterations: Maximum number of tool execution iterations
|
|
"""
|
|
self.llm = llm
|
|
self.tools: dict[str, Tool] = make_tools()
|
|
self.prompt_builder = PromptBuilder(self.tools)
|
|
self.max_tool_iterations = max_tool_iterations
|
|
|
|
def step(self, user_input: str) -> str:
|
|
"""
|
|
Execute one agent step with the user input.
|
|
|
|
This method:
|
|
1. Adds user message to memory
|
|
2. Builds prompt with history and context
|
|
3. Calls LLM, executing tools as needed
|
|
4. Returns final response
|
|
|
|
Args:
|
|
user_input: User's message
|
|
|
|
Returns:
|
|
Agent's final response
|
|
"""
|
|
memory = get_memory()
|
|
|
|
# Add user message to history
|
|
memory.stm.add_message("user", user_input)
|
|
memory.save()
|
|
|
|
# Build initial messages
|
|
system_prompt = self.prompt_builder.build_system_prompt()
|
|
messages: list[dict[str, Any]] = [{"role": "system", "content": system_prompt}]
|
|
|
|
# Add conversation history
|
|
history = memory.stm.get_recent_history(settings.max_history_messages)
|
|
messages.extend(history)
|
|
|
|
# Add unread events if any
|
|
unread_events = memory.episodic.get_unread_events()
|
|
if unread_events:
|
|
events_text = "\n".join(
|
|
[f"- {e['type']}: {e['data']}" for e in unread_events]
|
|
)
|
|
messages.append(
|
|
{"role": "system", "content": f"Background events:\n{events_text}"}
|
|
)
|
|
|
|
# Get tools specification for OpenAI format
|
|
tools_spec = self.prompt_builder.build_tools_spec()
|
|
|
|
# Tool execution loop
|
|
for _iteration in range(self.max_tool_iterations):
|
|
# Call LLM with tools
|
|
llm_result = self.llm.complete(messages, tools=tools_spec)
|
|
|
|
# Handle both tuple (response, usage) and dict response
|
|
if isinstance(llm_result, tuple):
|
|
response_message, usage = llm_result
|
|
else:
|
|
response_message = llm_result
|
|
|
|
# Check if there are tool calls
|
|
tool_calls = response_message.get("tool_calls")
|
|
|
|
if not tool_calls:
|
|
# No tool calls, this is the final response
|
|
final_content = response_message.get("content", "")
|
|
memory.stm.add_message("assistant", final_content)
|
|
memory.save()
|
|
return final_content
|
|
|
|
# Add assistant message with tool calls to conversation
|
|
messages.append(response_message)
|
|
|
|
# Execute each tool call
|
|
for tool_call in tool_calls:
|
|
tool_result = self._execute_tool_call(tool_call)
|
|
|
|
# Add tool result to messages
|
|
messages.append(
|
|
{
|
|
"tool_call_id": tool_call.get("id"),
|
|
"role": "tool",
|
|
"name": tool_call.get("function", {}).get("name"),
|
|
"content": json.dumps(tool_result, ensure_ascii=False),
|
|
}
|
|
)
|
|
|
|
# Max iterations reached, force final response
|
|
messages.append(
|
|
{
|
|
"role": "system",
|
|
"content": "Please provide a final response to the user without using any more tools.",
|
|
}
|
|
)
|
|
|
|
llm_result = self.llm.complete(messages)
|
|
if isinstance(llm_result, tuple):
|
|
final_message, usage = llm_result
|
|
else:
|
|
final_message = llm_result
|
|
|
|
final_response = final_message.get(
|
|
"content", "I've completed the requested actions."
|
|
)
|
|
memory.stm.add_message("assistant", final_response)
|
|
memory.save()
|
|
return final_response
|
|
|
|
def _execute_tool_call(self, tool_call: dict[str, Any]) -> dict[str, Any]:
|
|
"""
|
|
Execute a single tool call.
|
|
|
|
Args:
|
|
tool_call: OpenAI-format tool call dict
|
|
|
|
Returns:
|
|
Result dictionary
|
|
"""
|
|
function = tool_call.get("function", {})
|
|
tool_name = function.get("name", "")
|
|
|
|
try:
|
|
args_str = function.get("arguments", "{}")
|
|
args = json.loads(args_str)
|
|
except json.JSONDecodeError as e:
|
|
logger.error(f"Failed to parse tool arguments: {e}")
|
|
return {"error": "bad_args", "message": f"Invalid JSON arguments: {e}"}
|
|
|
|
# Validate tool exists
|
|
if tool_name not in self.tools:
|
|
available = list(self.tools.keys())
|
|
return {
|
|
"error": "unknown_tool",
|
|
"message": f"Tool '{tool_name}' not found",
|
|
"available_tools": available,
|
|
}
|
|
|
|
tool = self.tools[tool_name]
|
|
|
|
# Execute tool
|
|
try:
|
|
result = tool.func(**args)
|
|
return result
|
|
except KeyboardInterrupt:
|
|
# Don't catch KeyboardInterrupt - let it propagate
|
|
raise
|
|
except TypeError as e:
|
|
# Bad arguments
|
|
memory = get_memory()
|
|
memory.episodic.add_error(tool_name, f"bad_args: {e}")
|
|
return {"error": "bad_args", "message": str(e), "tool": tool_name}
|
|
except Exception as e:
|
|
# Other errors
|
|
memory = get_memory()
|
|
memory.episodic.add_error(tool_name, str(e))
|
|
return {"error": "execution_failed", "message": str(e), "tool": tool_name}
|
|
|
|
async def step_streaming(
|
|
self, user_input: str, completion_id: str, created_ts: int, model: str
|
|
) -> AsyncGenerator[dict[str, Any], None]:
|
|
"""
|
|
Execute agent step with streaming support for LibreChat.
|
|
|
|
Yields SSE chunks for tool calls and final response.
|
|
|
|
Args:
|
|
user_input: User's message
|
|
completion_id: Completion ID for the response
|
|
created_ts: Timestamp for the response
|
|
model: Model name
|
|
|
|
Yields:
|
|
SSE chunks in OpenAI format
|
|
"""
|
|
memory = get_memory()
|
|
|
|
# Add user message to history
|
|
memory.stm.add_message("user", user_input)
|
|
memory.save()
|
|
|
|
# Build initial messages
|
|
system_prompt = self.prompt_builder.build_system_prompt()
|
|
messages: list[dict[str, Any]] = [{"role": "system", "content": system_prompt}]
|
|
|
|
# Add conversation history
|
|
history = memory.stm.get_recent_history(settings.max_history_messages)
|
|
messages.extend(history)
|
|
|
|
# Add unread events if any
|
|
unread_events = memory.episodic.get_unread_events()
|
|
if unread_events:
|
|
events_text = "\n".join(
|
|
[f"- {e['type']}: {e['data']}" for e in unread_events]
|
|
)
|
|
messages.append(
|
|
{"role": "system", "content": f"Background events:\n{events_text}"}
|
|
)
|
|
|
|
# Get tools specification for OpenAI format
|
|
tools_spec = self.prompt_builder.build_tools_spec()
|
|
|
|
# Tool execution loop
|
|
for _iteration in range(self.max_tool_iterations):
|
|
# Call LLM with tools
|
|
llm_result = self.llm.complete(messages, tools=tools_spec)
|
|
|
|
# Handle both tuple (response, usage) and dict response
|
|
if isinstance(llm_result, tuple):
|
|
response_message, usage = llm_result
|
|
else:
|
|
response_message = llm_result
|
|
|
|
# Check if there are tool calls
|
|
tool_calls = response_message.get("tool_calls")
|
|
|
|
if not tool_calls:
|
|
# No tool calls, this is the final response
|
|
final_content = response_message.get("content", "")
|
|
memory.stm.add_message("assistant", final_content)
|
|
memory.save()
|
|
|
|
# Stream the final response
|
|
yield {
|
|
"id": completion_id,
|
|
"object": "chat.completion.chunk",
|
|
"created": created_ts,
|
|
"model": model,
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"delta": {"role": "assistant", "content": final_content},
|
|
"finish_reason": "stop",
|
|
}
|
|
],
|
|
}
|
|
return
|
|
|
|
# Stream tool calls
|
|
for tool_call in tool_calls:
|
|
function = tool_call.get("function", {})
|
|
tool_name = function.get("name", "")
|
|
tool_args = function.get("arguments", "{}")
|
|
|
|
# Yield chunk indicating tool call
|
|
yield {
|
|
"id": completion_id,
|
|
"object": "chat.completion.chunk",
|
|
"created": created_ts,
|
|
"model": model,
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"delta": {
|
|
"tool_calls": [
|
|
{
|
|
"index": 0,
|
|
"id": tool_call.get("id"),
|
|
"type": "function",
|
|
"function": {
|
|
"name": tool_name,
|
|
"arguments": tool_args,
|
|
},
|
|
}
|
|
]
|
|
},
|
|
"finish_reason": None,
|
|
}
|
|
],
|
|
}
|
|
|
|
# Add assistant message with tool calls to conversation
|
|
messages.append(response_message)
|
|
|
|
# Execute each tool call and stream results
|
|
for tool_call in tool_calls:
|
|
tool_result = self._execute_tool_call(tool_call)
|
|
function = tool_call.get("function", {})
|
|
tool_name = function.get("name", "")
|
|
|
|
# Add tool result to messages
|
|
messages.append(
|
|
{
|
|
"tool_call_id": tool_call.get("id"),
|
|
"role": "tool",
|
|
"name": tool_name,
|
|
"content": json.dumps(tool_result, ensure_ascii=False),
|
|
}
|
|
)
|
|
|
|
# Stream tool result as content
|
|
result_text = (
|
|
f"\n🔧 {tool_name}: {json.dumps(tool_result, ensure_ascii=False)}\n"
|
|
)
|
|
yield {
|
|
"id": completion_id,
|
|
"object": "chat.completion.chunk",
|
|
"created": created_ts,
|
|
"model": model,
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"delta": {"content": result_text},
|
|
"finish_reason": None,
|
|
}
|
|
],
|
|
}
|
|
|
|
# Max iterations reached, force final response
|
|
messages.append(
|
|
{
|
|
"role": "system",
|
|
"content": "Please provide a final response to the user without using any more tools.",
|
|
}
|
|
)
|
|
|
|
llm_result = self.llm.complete(messages)
|
|
if isinstance(llm_result, tuple):
|
|
final_message, usage = llm_result
|
|
else:
|
|
final_message = llm_result
|
|
|
|
final_response = final_message.get(
|
|
"content", "I've completed the requested actions."
|
|
)
|
|
memory.stm.add_message("assistant", final_response)
|
|
memory.save()
|
|
|
|
# Stream final response
|
|
yield {
|
|
"id": completion_id,
|
|
"object": "chat.completion.chunk",
|
|
"created": created_ts,
|
|
"model": model,
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"delta": {"content": final_response},
|
|
"finish_reason": "stop",
|
|
}
|
|
],
|
|
}
|