135 lines
4.9 KiB
Python
135 lines
4.9 KiB
Python
"""Core agent loop - the brain of CheddahBot."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import uuid
|
|
from typing import Generator
|
|
|
|
from .config import Config
|
|
from .db import Database
|
|
from .llm import LLMAdapter
|
|
from .router import build_system_prompt, format_messages_for_llm
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
MAX_TOOL_ITERATIONS = 10
|
|
|
|
|
|
class Agent:
|
|
def __init__(self, config: Config, db: Database, llm: LLMAdapter):
|
|
self.config = config
|
|
self.db = db
|
|
self.llm = llm
|
|
self.conv_id: str | None = None
|
|
self._memory = None # set by app after memory system init
|
|
self._tools = None # set by app after tool system init
|
|
|
|
def set_memory(self, memory):
|
|
self._memory = memory
|
|
|
|
def set_tools(self, tools):
|
|
self._tools = tools
|
|
|
|
def ensure_conversation(self) -> str:
|
|
if not self.conv_id:
|
|
self.conv_id = uuid.uuid4().hex[:12]
|
|
self.db.create_conversation(self.conv_id)
|
|
return self.conv_id
|
|
|
|
def new_conversation(self) -> str:
|
|
self.conv_id = uuid.uuid4().hex[:12]
|
|
self.db.create_conversation(self.conv_id)
|
|
return self.conv_id
|
|
|
|
def respond(self, user_input: str, files: list | None = None) -> Generator[str, None, None]:
|
|
"""Process user input and yield streaming response text."""
|
|
conv_id = self.ensure_conversation()
|
|
|
|
# Store user message
|
|
self.db.add_message(conv_id, "user", user_input)
|
|
|
|
# Build system prompt
|
|
memory_context = ""
|
|
if self._memory:
|
|
memory_context = self._memory.get_context(user_input)
|
|
|
|
tools_schema = []
|
|
tools_description = ""
|
|
if self._tools:
|
|
tools_schema = self._tools.get_tools_schema()
|
|
tools_description = self._tools.get_tools_description()
|
|
|
|
system_prompt = build_system_prompt(
|
|
identity_dir=self.config.identity_dir,
|
|
memory_context=memory_context,
|
|
tools_description=tools_description,
|
|
)
|
|
|
|
# Load conversation history
|
|
history = self.db.get_messages(conv_id, limit=self.config.memory.max_context_messages)
|
|
messages = format_messages_for_llm(system_prompt, history, self.config.memory.max_context_messages)
|
|
|
|
# Agent loop: LLM call → tool execution → repeat
|
|
for iteration in range(MAX_TOOL_ITERATIONS):
|
|
full_response = ""
|
|
tool_calls = []
|
|
|
|
for chunk in self.llm.chat(messages, tools=tools_schema or None, stream=True):
|
|
if chunk["type"] == "text":
|
|
full_response += chunk["content"]
|
|
yield chunk["content"]
|
|
elif chunk["type"] == "tool_use":
|
|
tool_calls.append(chunk)
|
|
|
|
# If no tool calls, we're done
|
|
if not tool_calls:
|
|
if full_response:
|
|
self.db.add_message(conv_id, "assistant", full_response, model=self.llm.current_model)
|
|
break
|
|
|
|
# Store assistant message with tool calls
|
|
self.db.add_message(
|
|
conv_id, "assistant", full_response,
|
|
tool_calls=[{"name": tc["name"], "input": tc["input"]} for tc in tool_calls],
|
|
model=self.llm.current_model,
|
|
)
|
|
|
|
# Execute tools
|
|
if self._tools:
|
|
messages.append({"role": "assistant", "content": full_response or "I'll use some tools to help with that."})
|
|
|
|
for tc in tool_calls:
|
|
yield f"\n\n🔧 **Using tool: {tc['name']}**\n"
|
|
try:
|
|
result = self._tools.execute(tc["name"], tc.get("input", {}))
|
|
except Exception as e:
|
|
result = f"Tool error: {e}"
|
|
yield f"```\n{result[:2000]}\n```\n\n"
|
|
|
|
self.db.add_message(conv_id, "tool", result, tool_result=tc["name"])
|
|
messages.append({"role": "user", "content": f'[Tool "{tc["name"]}" result]\n{result}'})
|
|
else:
|
|
# No tool system configured - just mention tool was requested
|
|
if full_response:
|
|
self.db.add_message(conv_id, "assistant", full_response, model=self.llm.current_model)
|
|
for tc in tool_calls:
|
|
yield f"\n(Tool requested: {tc['name']} - tool system not yet initialized)\n"
|
|
break
|
|
else:
|
|
yield "\n(Reached maximum tool iterations)"
|
|
|
|
# Check if memory flush is needed
|
|
if self._memory:
|
|
msg_count = self.db.count_messages(conv_id)
|
|
if msg_count > self.config.memory.flush_threshold:
|
|
self._memory.auto_flush(conv_id)
|
|
|
|
def respond_to_prompt(self, prompt: str) -> str:
|
|
"""Non-streaming response for scheduled tasks / internal use."""
|
|
result_parts = []
|
|
for chunk in self.respond(prompt):
|
|
result_parts.append(chunk)
|
|
return "".join(result_parts)
|