Episodic and semantic memory: remember turns, and remember facts

"""
Episodic and semantic memory: remember turns, and remember facts

Two stores, because they answer different questions.

  Episodic  what was said, when. Searchable transcript of past turns.
  Semantic  what is true. Extracted, deduplicated facts about the user.

    Turn -> store episode -> extract facts -> upsert semantic
    Query -> retrieve facts (+ episodes only if needed) -> answer

The failure this prevents is subtle. Loading whole past conversations into
context to "remember" the user is expensive and actively harmful: old
transcripts pull the model toward whatever was being discussed then, not what
is being asked now. Facts are small, current, and steer nothing.

So the default retrieval here is facts only. Episodes are fetched when the
question is genuinely about what happened rather than about what is true.

Run: python episodic_semantic.py
"""

from typing import Literal

from dotenv import load_dotenv
from openai import OpenAI
from pydantic import BaseModel, Field

load_dotenv()

client = OpenAI()

MODEL = "gpt-5.4-mini"

# --------------------------------------------------------------
# Two stores. In production these are a vector index and a table;
# the shape of the pattern is the same.
# --------------------------------------------------------------

EPISODES: list[dict[str, str]] = []
FACTS: dict[str, str] = {}


class ExtractedFact(BaseModel):
    key: str = Field(
        description="Stable snake_case identifier, e.g. 'preferred_language'."
    )
    value: str = Field(description="The fact itself, as short as possible.")


class Extraction(BaseModel):
    facts: list[ExtractedFact] = Field(
        description=(
            "Durable facts about the user or their situation. Empty when the "
            "turn contains nothing worth remembering. Most turns do not."
        )
    )


class Retrieval(BaseModel):
    needs_episodes: bool = Field(
        description=(
            "True only when the question is about what was previously said or "
            "done, rather than about what is currently true."
        )
    )
    why: str


def remember(user_text: str, assistant_text: str) -> list[ExtractedFact]:
    """Append the episode verbatim, then distil anything durable into facts."""
    EPISODES.append({"user": user_text, "assistant": assistant_text})

    completion = client.beta.chat.completions.parse(
        model=MODEL,
        messages=[
            {
                "role": "system",
                "content": (
                    "Extract durable facts about the user from this exchange. "
                    "A durable fact is still true next month. Preferences, "
                    "constraints, and stable context qualify. Questions asked, "
                    "one-off requests, and anything about the current task do "
                    "not. Return an empty list when nothing qualifies."
                ),
            },
            {"role": "user", "content": f"User: {user_text}\nAssistant: {assistant_text}"},
        ],
        response_format=Extraction,
    )
    extracted = completion.choices[0].message.parsed.facts
    for fact in extracted:
        # Upsert, so a changed preference replaces rather than accumulates.
        # Append-only fact stores end up holding both the old answer and the
        # new one, with nothing to say which is current.
        FACTS[fact.key] = fact.value
    return extracted


def decide_retrieval(question: str) -> Retrieval:
    completion = client.beta.chat.completions.parse(
        model=MODEL,
        messages=[
            {
                "role": "system",
                "content": (
                    "Decide whether answering needs the past transcript, or "
                    "whether the known facts are enough. Transcripts are "
                    "expensive and drag the answer toward old topics, so "
                    "prefer facts unless the question is genuinely about what "
                    "was said or done."
                ),
            },
            {"role": "user", "content": question},
        ],
        response_format=Retrieval,
    )
    return completion.choices[0].message.parsed


def answer(question: str) -> str:
    plan = decide_retrieval(question)
    print(f"  retrieval: facts{' + episodes' if plan.needs_episodes else ' only'}")
    print(f"      why  : {plan.why}")

    known = "\n".join(f"- {k}: {v}" for k, v in FACTS.items()) or "(nothing known)"
    context = f"Known facts:\n{known}"
    if plan.needs_episodes:
        transcript = "\n".join(
            f"User: {e['user']}\nAssistant: {e['assistant']}" for e in EPISODES
        )
        context += f"\n\nPast turns:\n{transcript}"

    completion = client.chat.completions.create(
        model=MODEL,
        messages=[
            {
                "role": "system",
                "content": (
                    "Answer using the context provided. If it does not contain "
                    "the answer, say you do not know rather than guessing."
                ),
            },
            {"role": "user", "content": f"{context}\n\nQuestion: {question}"},
        ],
    )
    return completion.choices[0].message.content or ""


if __name__ == "__main__":
    print("\n> seeding two turns")
    for user_text, assistant_text in [
        ("I run a small ceramics shop on Shopify.", "Noted, thanks."),
        ("Keep any tooling under 200 a month, that is my ceiling.", "Understood."),
    ]:
        new_facts = remember(user_text, assistant_text)
        print(f"  learned: {[(f.key, f.value) for f in new_facts] or 'nothing durable'}")

    for question in [
        "What platform am I on, and what is my budget?",
        "What exactly did I say when I mentioned the budget?",
    ]:
        print(f"\n> {question}")
        print(f"  answer: {answer(question)}")