feat: initial release - Hunzi agent framework v0.1.0
This commit is contained in:
@@ -0,0 +1,267 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user