agent.py

"""agent.py -- CLI agent over the MemorySidecar, plus the surprise-vs-llm_extract A/B
token proof that is the heart of the pitch.

No provider import here: the LLM arrives as an injected client exposing
`chat(messages) -> result` with `.text` and `.total_tokens` (see llm_client.LLMClient).
That keeps the provider-abstraction contract and lets tests run with a fake, no creds.

Token accounting distinguishes two paths:
  - WRITE path: tokens spent DECIDING what to remember. surprise mode = 0 (neural gate);
    llm_extract mode = an LLM call per turn (the per-turn tax we eliminate).
  - READ path: tokens spent generating the answer. Identical in both modes.
The A/B claim is purely about the WRITE path.
"""
from __future__ import annotations

import json
from dataclasses import dataclass, field

from memory.sidecar import MemorySidecar, ObservationRecord

ANSWER_SYSTEM = (
    "You are a helpful assistant with long-term memory of the user. Use the MEMORY block "
    "(durable facts stored about the user) to personalize your answer. A fact tagged "
    "'(possibly outdated)' or '(uncertain)' may be stale -- prefer newer facts and hedge if it "
    "matters. If MEMORY is empty or does not contain what the user is asking about, say you don't "
    "have that information yet rather than guessing. Do not mention the memory mechanism itself."
)
EXTRACT_SYSTEM = (
    "You are a memory-extraction module. From the user's latest message, output the durable "
    "facts about the user worth remembering, one per line, or exactly 'NONE'. Terse; no commentary."
)
# Agentic read: the model's native function-calling drives WHEN and WHAT to recall. The WRITE stays the
# token-free surprise gate (so the zero-write-token thesis + the A/B result are untouched); this tool
# only reads. The model orchestrates; Flashbulb decides what to keep and forget.
TOOLS_SYSTEM = (
    "You are a helpful assistant with long-term memory of the user, accessed through the recall_memory "
    "tool. Before answering anything about the user, call recall_memory to check what you actually know. "
    "A fact the tool returns tagged '(possibly outdated)' or '(uncertain)' may be stale -- prefer newer "
    "facts and hedge if it matters. If recall_memory returns nothing relevant, say you don't have that "
    "information yet rather than guessing. Do not mention the memory mechanism itself."
)
RECALL_TOOL = {
    "type": "function",
    "function": {
        "name": "recall_memory",
        "description": ("Search your long-term memory of the user for facts relevant to a query. Returns "
                        "stored facts, some tagged '(possibly outdated)' or '(uncertain)'. Call this before "
                        "answering anything about the user."),
        "parameters": {
            "type": "object",
            "properties": {"query": {"type": "string",
                                     "description": "what to look up, e.g. 'where the user lives'"}},
            "required": ["query"],
        },
    },
}


@dataclass
class TurnResult:
    answer: str
    write_tokens: int  # provider tokens spent deciding what to store
    read_tokens: int   # provider tokens spent generating the answer
    stored: list[ObservationRecord] = field(default_factory=list)
    skipped: int = 0
    tool_calls: list = field(default_factory=list)  # agentic recalls the model made this turn (for the UI)


def _is_memory_candidate(text: str) -> bool:
    """Deterministic (0-token) 'extract' step for surprise mode: a durable fact is a
    non-empty declarative statement. Questions/requests are not facts about the user.
    This is a coarse heuristic, not an LLM call -- that is the whole point."""
    t = text.strip()
    return bool(t) and not t.endswith("?")


class Agent:
    def __init__(self, sidecar: MemorySidecar, llm, write_policy: str = "surprise", use_tools: bool = False):
        if write_policy not in ("surprise", "llm_extract"):
            raise ValueError("write_policy must be 'surprise' or 'llm_extract'")
        self.sidecar = sidecar
        self.llm = llm
        self.write_policy = write_policy
        self.use_tools = use_tools  # agentic read via model function-calling (write stays token-free)
        self.total_write_tokens = 0
        self.total_read_tokens = 0

    # ---- write policies ----
    def _write_surprise(self, user_msg: str) -> tuple[list, int]:
        """The neural surprise gate decides. ZERO provider tokens."""
        if not _is_memory_candidate(user_msg):
            return [], 0
        rec = self.sidecar.write(user_msg)
        return ([rec] if rec else []), 0

    def _write_llm_extract(self, user_msg: str) -> tuple[list, int]:
        """Baseline (Mem0-style): call the LLM every turn to extract/decide -- the per-turn
        token tax. Both policies share the same store; only the DECISION cost differs."""
        res = self.llm.chat([
            {"role": "system", "content": EXTRACT_SYSTEM},
            {"role": "user", "content": user_msg},
        ])
        stored = []
        for line in res.text.splitlines():
            fact = line.strip("-*0123456789. ").strip()
            if fact and fact.upper() != "NONE":
                rec = self.sidecar.write(fact)
                if rec:
                    stored.append(rec)
        return stored, res.total_tokens

    # ---- turn ----
    def turn(self, user_msg: str) -> TurnResult:
        skips_before = len(self.sidecar.skips)
        if self.write_policy == "surprise":
            stored, write_tokens = self._write_surprise(user_msg)     # ZERO provider tokens
        else:
            stored, write_tokens = self._write_llm_extract(user_msg)
        skipped = len(self.sidecar.skips) - skips_before

        if self.use_tools:
            answer, read_tokens, tool_calls = self._answer_with_tools(user_msg)
        else:
            answer, read_tokens = self._answer_direct(user_msg)
            tool_calls = []

        self.total_write_tokens += write_tokens
        self.total_read_tokens += read_tokens
        return TurnResult(answer, write_tokens, read_tokens, stored, skipped, tool_calls=tool_calls)

    def _answer_direct(self, user_msg: str) -> tuple[str, int]:
        """Retrieve top-k memory, inject it, answer in one shot. The non-agentic read path (and the
        fallback for the agentic one)."""
        _recs, memblock = self.sidecar.read(user_msg, k=5)
        system = ANSWER_SYSTEM + (f"\n\nMEMORY:\n{memblock}" if memblock else "")
        try:
            res = self.llm.chat([
                {"role": "system", "content": system},
                {"role": "user", "content": user_msg},
            ])
            return res.text, res.total_tokens
        except Exception:
            # Answer-LLM unreachable: the memory still works, so answer straight from what was retrieved.
            fallback = f"From memory:\n{memblock}" if memblock else "Nothing in memory about that yet — tell me a fact."
            return fallback, 0

    def _answer_with_tools(self, user_msg: str) -> tuple[str, int, list]:
        """Agentic read: the model decides whether and what to recall via the recall_memory tool, then answers
        over the returned (confidence-tagged) memories. The write already happened token-free, so this is
        the read path only. Any tool-calling failure falls back to the direct answer -- the demo never
        breaks."""
        messages = [{"role": "system", "content": TOOLS_SYSTEM}, {"role": "user", "content": user_msg}]
        tool_calls_made: list = []
        try:
            res = self.llm.chat(messages, tools=[RECALL_TOOL])
            read_tokens = res.total_tokens
            tcs = getattr(res, "tool_calls", None)
            if not tcs:
                return res.text, read_tokens, tool_calls_made  # the model answered without needing memory
            messages.append({
                "role": "assistant", "content": res.text or None,
                "tool_calls": [{"id": tc.id, "type": "function",
                                "function": {"name": tc.function.name, "arguments": tc.function.arguments}}
                               for tc in tcs],
            })
            for tc in tcs:
                if tc.function.name == "recall_memory":
                    try:
                        query = json.loads(tc.function.arguments).get("query", user_msg)
                    except Exception:
                        query = user_msg
                    _recs, memblock = self.sidecar.read(query, k=5)
                    tool_calls_made.append({"name": "recall_memory", "query": query,
                                            "result": memblock or "(no relevant memories)"})
                    messages.append({"role": "tool", "tool_call_id": tc.id,
                                     "content": memblock or "(no relevant memories)"})
                else:
                    messages.append({"role": "tool", "tool_call_id": tc.id, "content": "(unknown tool)"})
            final = self.llm.chat(messages)
            return final.text, read_tokens + final.total_tokens, tool_calls_made
        except Exception:
            answer, read_tokens = self._answer_direct(user_msg)  # tool-calling unsupported/failed
            return answer, read_tokens, tool_calls_made