Memory &
Knowledge

cd agent-lab && source .venv/bin/activate
touch memory_agent.py

import os
import json
import anthropic
from datetime import datetime, timezone
from dotenv import load_dotenv

load_dotenv()
client = anthropic.Anthropic()

MEMORY_FILE = "episodes.json"
MAX_TURNS   = 6   # sliding window: keep last N user+assistant pairs
TOKEN_WARN  = 8000 # log a warning when input tokens exceed this


def load_episodes() -> list[dict]:
    """Load past episodes from the JSON file, or return empty list."""
    if not os.path.exists(MEMORY_FILE):
        return []
    with open(MEMORY_FILE) as f:
        return json.load(f)


def save_episode(goal: str, turns: int, outcome: str) -> None:
    """Append the current session to the episodic memory file."""
    episodes = load_episodes()
    episodes.append({
        "id": f"ep_{datetime.now(timezone.utc).strftime('%Y%m%d_%H%M%S')}",
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "user_goal": goal,
        "turns": turns,
        "outcome": outcome,
    })
    with open(MEMORY_FILE, "w") as f:
        json.dump(episodes, f, indent=2)
    print(f"[memory] episode saved → {MEMORY_FILE}")

def trim_messages(messages: list[dict]) -> list[dict]:
    """Keep only the last MAX_TURNS user+assistant pairs.
    Inserts a notice when history is trimmed so the model is aware."""
    # A "pair" is one user + one assistant message = 2 items
    max_messages = MAX_TURNS * 2
    if len(messages) <= max_messages:
        return messages

    trimmed = messages[-max_messages:]
    notice = {
        "role": "user",
        "content": (
            "[System note: earlier conversation turns were trimmed to fit "
            "the context window. Only the last few turns are shown.]"
        )
    }
    print(f"[memory] context trimmed: keeping last {MAX_TURNS} turns")
    return [notice] + trimmed

def build_system_prompt(episodes: list[dict]) -> str:
    """Build the system prompt, injecting past episode summaries."""
    base = (
        "You are a helpful assistant with memory of past conversations. "
        "You answer questions clearly and concisely."
    )
    if not episodes:
        return base

    # Inject the last 3 episodes as context
    recent = episodes[-3:]
    history_lines = [
        f"- [{ep['timestamp'][:10]}] Goal: {ep['user_goal']} "
        f"({ep['turns']} turns, {ep['outcome']})"
        for ep in recent
    ]
    history_block = "\n".join(history_lines)
    return (
        f"{base}\n\n"
        "Past sessions (most recent first):\n"
        f"{history_block}\n\n"
        "Use this context to give continuity to the user across sessions."
    )

def run_chat() -> None:
    episodes = load_episodes()
    system   = build_system_prompt(episodes)
    messages: list[dict] = []
    turn_count = 0
    first_message = ""

    print("\n=== Memory Agent ===")
    print(f"Loaded {len(episodes)} past episode(s).")
    print('Type your message. Type "quit" to exit and save.\n')

    while True:
        user_input = input("You: ").strip()
        if not user_input:
            continue
        if user_input.lower() == "quit":
            break

        if not first_message:
            first_message = user_input

        messages.append({"role": "user", "content": user_input})
        messages = trim_messages(messages)

        response = client.messages.create(
            model="claude-opus-4-6",  # check docs.anthropic.com/en/docs/about-claude/models
            max_tokens=512,
            system=system,
            messages=messages
        )

        input_tokens = response.usage.input_tokens
        if input_tokens > TOKEN_WARN:
            print(f"[memory] ⚠ high token usage: {input_tokens} input tokens")

        reply = response.content[0].text
        messages.append({"role": "assistant", "content": reply})
        turn_count += 1

        print(f"Agent [{input_tokens} tok]: {reply}\n")

    # Save episode on exit
    if turn_count > 0:
        save_episode(
            goal=first_message,
            turns=turn_count,
            outcome="completed"
        )


if __name__ == "__main__":
    run_chat()

python memory_agent.py

=== Memory Agent ===
Loaded 0 past episode(s).
Type your message. Type "quit" to exit and save.

You: My name is Alex and I'm building a RAG pipeline in Python.
Agent [312 tok]: Great to meet you, Alex! Happy to help with your RAG pipeline...

You: What are the main chunking strategies I should consider?
Agent [498 tok]: The three main strategies are fixed-size, sentence-boundary, and semantic...

You: quit
[memory] episode saved → episodes.json

=== Memory Agent ===
Loaded 1 past episode(s).
Type your message. Type "quit" to exit and save.

You: What was I working on last time?
Agent [441 tok]: Based on our previous session, you're Alex and you were building
a RAG pipeline in Python. You asked about chunking strategies — we covered
fixed-size, sentence-boundary, and semantic chunking. Would you like to continue?

def trim_messages_with_summary(messages: list[dict]) -> list[dict]:
    """Summarise old turns before trimming them."""
    max_messages = MAX_TURNS * 2
    if len(messages) <= max_messages:
        return messages

    old_turns  = messages[:-max_messages]
    keep_turns = messages[-max_messages:]

    # Summarise the turns being dropped
    summary_prompt = (
        "Summarise these conversation turns in 2-3 sentences, "
        "preserving key facts, decisions, and user preferences:\n\n"
        + "\n".join(
            f"{m['role'].upper()}: {m['content'] if isinstance(m['content'], str) else str(m['content'])}"
            for m in old_turns
        )
    )
    summary_resp = client.messages.create(
        model="claude-opus-4-6",
        max_tokens=200,
        messages=[{"role": "user", "content": summary_prompt}]
    )
    summary = summary_resp.content[0].text

    notice = {
        "role": "user",
        "content": f"[Earlier turns summarised]: {summary}"
    }
    print(f"[memory] summarised {len(old_turns)} old turns")
    return [notice] + keep_turns