agent.py

from collections.abc import Callable
from datetime import date
from time import sleep
from typing import Iterator

from ollama import ChatResponse

from llm import BaseLLM
from memory.memory_manager import MemoryManager
from memory.summarizer import Summarizer
from models import History, Message
from tools.registry import registry
from tools.schema import tools_schema  # triggers imports and registration
from utils import comment_log, data_log, prompt_log


class Agent:
    """
    Conversational agent that builds prompts and streams responses from an LLM.
    """

    def __init__(self, llm: BaseLLM, memory: MemoryManager, summarizer: Summarizer):
        self.llm = llm
        self.memory = memory
        self.summarizer = summarizer
        ids = memory.sessions.init_session(model=llm.model)
        self.session_id = ids["session_id"]
        self.prev_session_id = ids["prev_session_id"]

        # INITIALIZATION RETRIEVAL:
        # The summaries of the exchanges of the previous session\n
        # (While these are semantic, vector db summaries, here they're used in an episodic way - retrieved by previous session id)
        self.prev_summaries = (
            memory.summaries.get_recent(
                n=15, where={"session_id": self.prev_session_id}
            )
            if self.prev_session_id
            else []
        )

    def available_tools(self) -> dict[str, Callable[..., str]]:
        return registry.available()

    def respond(self, short_term_mem: History, user_input: str) -> Iterator[str]:

        comment_log("RETRIEVAL:")
        comment_log("short-term memory from chat:")
        data_log(short_term_mem)
        comment_log("semantic memory from summaries collection:")
        # The summaries of the exchanges relevant to the user input\n
        # (Now thes summaries are used in an semantic way - retrieved by meaning)
        relevant_summaries = self.memory.summaries.read_list(user_input, 3)
        data_log(relevant_summaries)

        comment_log("PROMPT BUILDER:")
        messages: list[Message] = self.build_prompt(short_term_mem, self.prev_summaries)
        prompt_log(messages)

        comment_log("LLM:")

        comment_log("TOOL USE:")
        comment_log(
            "Run the inner tool use loop and use any tools needed to answer the question"
        )
        used_tools, response = self.tool_use_loop(messages)

        # Collect full response for summaries:
        response_parts: list[str] = []

        comment_log("Stream the final answer back to the user.")
        if used_tools:
            # response.message.content already has the entire answer, but instead of sending it all at once, it can still be streamed:
            for word in (response.message.content or "").split(" "):
                content = word + " "
                response_parts.append(content)
                yield content + " "
                sleep(0.02)

        else:
            # no tools used - stream from LLM
            # yield from self.llm.chat_stream(messages)
            for chunk in self.llm.chat_stream(messages):
                response_parts.append(chunk)
                yield chunk

        print("")

        comment_log("MEMORY WRITE:")
        self.memory.sessions.increment_message_count(self.session_id)

        comment_log("Summarizing conversation...")
        full_response = "".join(response_parts)
        self.summarizer.summarize(short_term_mem, full_response, self.session_id)

    def finish(self):
        self.memory.sessions.end_session(self.session_id)
        comment_log("Session ended:")
        data_log(self.memory.sessions.get_session(self.session_id))

    # Helper methods:

    def build_prompt(
        self, short_term_mem: History, prev_session_summaries: list[str]
    ) -> list[Message]:
        # For now, it's just a system prompt + the short-term memory from the chat history.
        # Eventually, this would be more complex - pulling from long-term memory, doing retrieval, etc.
        return [
            {
                "role": "system",
                "content": f"""
Today is {date.today()}.
Answer in 3 sentences or fewer.
Where tool results are unhelpful or insufficient, say so instead of guessing. Don't make up information!
For example, if news item results returned by the web_search tool are insufficient, say so instead of guessing. Don't make up news!
Cite sources inline, eg (Source Name, Date). Don't include model citation formats.
In the previous session, we spoke about the following, which may or may not be relevant to the current chat: {prev_session_summaries}
Always check these previous sessions for information for responses or answers to questions before resorting to tools.
                """,
            },
            *short_term_mem,
        ]

    def tool_use_loop(self, messages: list[Message]) -> tuple[bool, ChatResponse]:
        used_tools = False
        while (
            True
        ):  # <-- Tool use loop - keep going until the LLM doesn't call any more tools
            # response = self.llm.chat(messages, tools=list(available_tools().values()))  # type: ignore
            # Passing the actual Python functions as tools at runtime, (instead of the JSON schemas) is supported by Ollama.
            # (Pylance just can't see that from the type stubs.)
            # Ollama can introspect the function's docstring and type hints to build the schema automatically.
            # (So the docstring becomes both human documentation and what Ollama reads to build the schema behind the scenes.)

            response = self.llm.chat_sync(messages, tools=tools_schema)
            # But sometimes Ollama's auto-introspection isn't enough:
            # 1. Complex parameter types
            # 2. Nested/structured parameters
            # 3. The model keeps calling the tool wrong
            # 4. Switching away from Ollama - OpenAI, Anthropic, etc. don't do auto-introspection -they require explicit schemas.

            comment_log("decision to use tool")
            # LLMs are non-deterministic, so there'll be variation in the number of loops between identical queries.
            data_log(response)

            if response.message.tool_calls:
                used_tools = True
                comment_log("yes - use tool")
                # Fix: only append the assistant message if it has content:
                if response.message.content:
                    messages.append({
                        "role": "assistant",
                        "content": response.message.content,
                        # or "",  # content can be `None` when there are tool calls
                    })
                comment_log("choose tool")
                for tool_call in response.message.tool_calls:
                    comment_log(f"tool: {tool_call.function.name}")

                    fn = self.available_tools().get(tool_call.function.name)
                    if fn:
                        result = fn(**tool_call.function.arguments)
                        messages.append({
                            "role": "tool",
                            "content": str(
                                result  # or json.dumps(result) for structured data - content must be a string
                            ),
                            "name": tool_call.function.name,
                        })
                    else:
                        print(
                            f"[warn] model requested unknown tool: {tool_call.function.name}"
                        )
                # loop again with tool results
            else:
                comment_log("no - don't use tool")
                break  # no tool calls, final answer is in response.message.content

        return used_tools, response