268 lines
10 KiB
Python
268 lines
10 KiB
Python
"""Agent engine — orchestrates LLM calls, tool execution, and memory."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from typing import Any, Dict, Generator, List, Optional, Union
|
|
|
|
from app.agent.memory import BaseMemory, ConversationMemory, Message
|
|
from app.agent.tools import ToolRegistry, global_registry
|
|
from app.config import LLMProviderConfig
|
|
from app.llm.client import LLMClient
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _parse_tool_calls(response: Dict[str, Any]) -> List[Dict[str, Any]]:
|
|
"""Extract tool_calls from an LLM response dict."""
|
|
choices = response.get("choices", [])
|
|
if not choices:
|
|
return []
|
|
message = choices[0].get("message", {})
|
|
return message.get("tool_calls") or []
|
|
|
|
|
|
def _has_tool_calls(response: Dict[str, Any]) -> bool:
|
|
return bool(_parse_tool_calls(response))
|
|
|
|
|
|
def _extract_content(response: Dict[str, Any]) -> Optional[str]:
|
|
"""Pull the text content out of a response dict."""
|
|
choices = response.get("choices", [])
|
|
if not choices:
|
|
return None
|
|
return choices[0].get("message", {}).get("content")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Agent
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class Agent:
|
|
"""Full-featured agent loop with tool calling and memory.
|
|
|
|
Parameters
|
|
----------
|
|
llm_client : LLMClient
|
|
The LLM client used for completions.
|
|
tools : ToolRegistry, optional
|
|
Registry of available tools. Defaults to the global registry.
|
|
memory : BaseMemory, optional
|
|
Memory backend. Defaults to ``ConversationMemory(window_size=50)``.
|
|
system_prompt : str, optional
|
|
System-level instruction prepended to every request.
|
|
max_iterations : int
|
|
Safety cap on the tool-call loop to prevent infinite recursion.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
llm_client: LLMClient,
|
|
tools: Optional[ToolRegistry] = None,
|
|
memory: Optional[BaseMemory] = None,
|
|
system_prompt: Optional[str] = None,
|
|
max_iterations: int = 10,
|
|
) -> None:
|
|
self.llm = llm_client
|
|
self.tools = tools or global_registry()
|
|
self.memory = memory or ConversationMemory()
|
|
self.system_prompt = system_prompt or (
|
|
"You are a helpful assistant. You have access to tools that can "
|
|
"help you complete tasks. Use them when needed."
|
|
)
|
|
self.max_iterations = max_iterations
|
|
|
|
# -- public API -------------------------------------------------------------
|
|
|
|
def run(
|
|
self,
|
|
messages: Union[List[Dict[str, Any]], str],
|
|
*,
|
|
stream: bool = False,
|
|
) -> str:
|
|
"""Process messages and return the final assistant response.
|
|
|
|
Handles the full loop:
|
|
LLM call → detect tool calls → execute tools → feed results back
|
|
→ repeat until no more tool calls or max_iterations reached.
|
|
"""
|
|
if stream:
|
|
# Collect streamed chunks into one string.
|
|
return "".join(self.run_stream(messages))
|
|
|
|
# Normalise input --------------------------------------------------
|
|
msg_list = self._normalise_messages(messages)
|
|
|
|
# Build the initial conversation -----------------------------------
|
|
conversation: List[Dict[str, Any]] = []
|
|
if self.system_prompt:
|
|
conversation.append({"role": "system", "content": self.system_prompt})
|
|
conversation.extend(msg_list)
|
|
|
|
# Store user input in memory ---------------------------------------
|
|
for m in msg_list:
|
|
self.memory.add(Message.from_openai(m))
|
|
|
|
# Tool-call loop ---------------------------------------------------
|
|
for iteration in range(self.max_iterations):
|
|
logger.info("Agent iteration %d/%d", iteration + 1, self.max_iterations)
|
|
|
|
response = self.llm.chat(
|
|
conversation,
|
|
tools=self.tools.get_schema(),
|
|
)
|
|
|
|
# Check for tool calls
|
|
tool_calls = _parse_tool_calls(response)
|
|
if not tool_calls:
|
|
# No tool calls — this is the final answer
|
|
content = _extract_content(response) or ""
|
|
assistant_msg = Message(role="assistant", content=content)
|
|
self.memory.add(assistant_msg)
|
|
conversation.append({"role": "assistant", "content": content})
|
|
return content
|
|
|
|
# There are tool calls — execute them ---------------------------
|
|
assistant_msg = {
|
|
"role": "assistant",
|
|
"content": _extract_content(response),
|
|
"tool_calls": tool_calls,
|
|
}
|
|
self.memory.add(Message.from_openai(assistant_msg))
|
|
conversation.append(assistant_msg)
|
|
|
|
for tc in tool_calls:
|
|
tool_name = tc["function"]["name"]
|
|
raw_args = tc["function"]["arguments"]
|
|
try:
|
|
args = json.loads(raw_args) if isinstance(raw_args, str) else raw_args
|
|
except json.JSONDecodeError as exc:
|
|
args = {}
|
|
logger.error("Failed to parse tool args for %s: %s", tool_name, exc)
|
|
|
|
tool_call_id = tc.get("id", "call_unknown")
|
|
|
|
try:
|
|
result = self.tools.execute(tool_name, **args)
|
|
except Exception as exc:
|
|
result = f"Error executing {tool_name}: {exc}"
|
|
|
|
result_str = str(result) if not isinstance(result, str) else result
|
|
tool_response = {
|
|
"role": "tool",
|
|
"content": result_str,
|
|
"tool_call_id": tool_call_id,
|
|
}
|
|
self.memory.add(Message.from_openai(tool_response))
|
|
conversation.append(tool_response)
|
|
logger.debug(
|
|
"Tool %s → %s", tool_name, result_str[:200]
|
|
)
|
|
|
|
# If we exit the loop it means max_iterations was hit
|
|
logger.warning("Agent hit max_iterations (%d) — returning partial result", self.max_iterations)
|
|
last_content = _extract_content(response) or "(max iterations reached)"
|
|
assistant_msg = Message(role="assistant", content=last_content)
|
|
self.memory.add(assistant_msg)
|
|
return last_content
|
|
|
|
def run_stream(
|
|
self,
|
|
messages: Union[List[Dict[str, Any]], str],
|
|
) -> Generator[str, None, None]:
|
|
"""Stream the final assistant response.
|
|
|
|
Note: tool calls are handled silently in the background; only the
|
|
final assistant text chunk is yielded.
|
|
"""
|
|
msg_list = self._normalise_messages(messages)
|
|
|
|
conversation: List[Dict[str, Any]] = []
|
|
if self.system_prompt:
|
|
conversation.append({"role": "system", "content": self.system_prompt})
|
|
conversation.extend(msg_list)
|
|
|
|
for m in msg_list:
|
|
self.memory.add(Message.from_openai(m))
|
|
|
|
# Tool-call loop (non-streaming) until we reach the final turn
|
|
for iteration in range(self.max_iterations):
|
|
logger.info("Agent iteration %d/%d (stream)", iteration + 1, self.max_iterations)
|
|
|
|
response = self.llm.chat(
|
|
conversation,
|
|
tools=self.tools.get_schema(),
|
|
)
|
|
|
|
tool_calls = _parse_tool_calls(response)
|
|
if not tool_calls:
|
|
# Final answer — now stream it -------------------------------
|
|
streamed_chunks: List[str] = []
|
|
for chunk in self.llm.chat_stream(conversation):
|
|
streamed_chunks.append(chunk)
|
|
yield chunk
|
|
|
|
full_content = "".join(streamed_chunks)
|
|
self.memory.add(Message(role="assistant", content=full_content))
|
|
return
|
|
|
|
# Execute tool calls (same as non-streaming) --------------------
|
|
assistant_msg = {
|
|
"role": "assistant",
|
|
"content": _extract_content(response),
|
|
"tool_calls": tool_calls,
|
|
}
|
|
self.memory.add(Message.from_openai(assistant_msg))
|
|
conversation.append(assistant_msg)
|
|
|
|
for tc in tool_calls:
|
|
tool_name = tc["function"]["name"]
|
|
raw_args = tc["function"]["arguments"]
|
|
try:
|
|
args = json.loads(raw_args) if isinstance(raw_args, str) else raw_args
|
|
except json.JSONDecodeError:
|
|
args = {}
|
|
|
|
tool_call_id = tc.get("id", "call_unknown")
|
|
|
|
try:
|
|
result = self.tools.execute(tool_name, **args)
|
|
except Exception as exc:
|
|
result = f"Error executing {tool_name}: {exc}"
|
|
|
|
result_str = str(result) if not isinstance(result, str) else result
|
|
tool_response = {
|
|
"role": "tool",
|
|
"content": result_str,
|
|
"tool_call_id": tool_call_id,
|
|
}
|
|
self.memory.add(Message.from_openai(tool_response))
|
|
conversation.append(tool_response)
|
|
|
|
logger.warning("Agent stream hit max_iterations")
|
|
last_content = _extract_content(response) or ""
|
|
yield last_content
|
|
|
|
# -- private ----------------------------------------------------------------
|
|
|
|
@staticmethod
|
|
def _normalise_messages(
|
|
messages: Union[List[Dict[str, Any]], str],
|
|
) -> List[Dict[str, Any]]:
|
|
"""Accept a plain string or a list of message dicts."""
|
|
if isinstance(messages, str):
|
|
return [{"role": "user", "content": messages}]
|
|
return messages
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Alias so __init__.py can import Agent as AgentEngine if needed
|
|
# ---------------------------------------------------------------------------
|
|
AgentEngine = Agent
|