code.py

"""Memory Management — short-term buffer, long-term facts, and a summarizer.

Three tiers, because they have different lifetimes:

  working  - the last few turns, verbatim. Cheap, precise, bounded.
  summary  - older turns, compressed. Keeps continuity without the token cost.
  facts    - durable things about the user that should survive any summary.

The interesting part is eviction. Context windows are finite, so the real
question isn't "how do I remember" but "what do I drop first". Here verbatim
turns age into a summary, while extracted facts are never dropped.
"""

from langchain_core.messages import AIMessage, HumanMessage, SystemMessage

from resources.agent import llm
from resources.helper import show_response

WORKING_LIMIT = 4  # messages kept verbatim before the oldest get summarized

summarize_prompt = (
    "Compress this conversation excerpt into 2 sentences of context. "
    "Keep names, preferences, and decisions. Drop pleasantries.\n\n{excerpt}"
)

extract_prompt = (
    "Extract durable facts about the user from this message - preferences, "
    "constraints, identity. One per line, no bullets. "
    "If there are none, reply NONE.\n\n{message}"
)


class Memory:
    """Holds the three tiers and decides what gets evicted when."""

    def __init__(self):
        self.working = []  # recent messages, verbatim
        self.summary = ""  # everything older, compressed
        self.facts = []  # durable, never evicted

    def remember(self, message):
        """Add a message, extracting facts and compacting if we're over limit."""
        self.working.append(message)

        # Facts are pulled out of user messages before those messages can age
        # into the summary, where details get lost.
        if isinstance(message, HumanMessage):
            found = llm.invoke(extract_prompt.format(message=message.content))
            if "NONE" not in found.content:
                new = [f.strip() for f in found.content.split("\n") if f.strip()]
                self.facts.extend(new)
                print(f"  [memory] +{len(new)} fact(s)")

        if len(self.working) > WORKING_LIMIT:
            self._compact()

    def _compact(self):
        """Age the oldest verbatim turns into the running summary."""
        evicted, self.working = self.working[:2], self.working[2:]
        excerpt = "\n".join(f"{m.type}: {m.content}" for m in evicted)
        if self.summary:
            excerpt = f"Earlier summary: {self.summary}\n\n{excerpt}"

        self.summary = llm.invoke(summarize_prompt.format(excerpt=excerpt)).content
        print(f"  [memory] compacted 2 messages -> summary ({len(self.summary)} chars)")

    def as_messages(self, question):
        """Assemble the three tiers into the prompt for one call."""
        context = []
        if self.facts:
            context.append("Known about the user:\n" + "\n".join(self.facts))
        if self.summary:
            context.append("Earlier conversation:\n" + self.summary)

        system = SystemMessage(
            "You are a helpful assistant.\n\n" + "\n\n".join(context)
            if context
            else "You are a helpful assistant."
        )
        return [system] + self.working + [HumanMessage(question)]


if __name__ == "__main__":
    memory = Memory()

    turns = [
        "Hi! I'm Priya, and I'm allergic to shellfish.",
        "I'm planning a dinner party for six people next Friday.",
        "Two of my guests are vegetarian.",
        "What's the weather usually like in Goa in December?",
        "Can you suggest a main course for the party?",
    ]

    for question in turns:
        print(f"\n{'=' * 60}\nUSER: {question}\n{'=' * 60}")
        reply = llm.invoke(memory.as_messages(question))
        memory.remember(HumanMessage(question))
        memory.remember(AIMessage(reply.content))
        show_response(reply)

    # The final turn is the test: by now the shellfish allergy has aged out of
    # the verbatim buffer, so a correct answer proves the fact store worked.
    print(f"\n{'=' * 60}\nMEMORY STATE\n{'=' * 60}")
    print(f"working ({len(memory.working)} msgs): "
          f"{[m.content[:40] for m in memory.working]}")
    print(f"\nsummary: {memory.summary}")
    print(f"\nfacts:\n  " + "\n  ".join(memory.facts))