Agent Memory & Context Budgeting

# script_id: day_046_agent_memory/layered_memory
from collections import deque
from dataclasses import dataclass, field

@dataclass
class AgentMemory:
    """Three-tier memory for an agent."""
    # Short-term: recent turns, bounded so it can't grow forever
    short_term: deque = field(default_factory=lambda: deque(maxlen=10))
    # Working: facts that should survive short-term trimming, this session
    working: dict = field(default_factory=dict)
    # Long-term: a handle to durable storage (DB/vector store); stubbed here
    long_term_store: dict = field(default_factory=dict)

    def add_turn(self, role: str, content: str):
        self.short_term.append({"role": role, "content": content})

    def remember_fact(self, key: str, value: str):
        """Promote something important to working memory so trimming won't lose it."""
        self.working[key] = value

    def persist(self, user_id: str):
        """Flush working memory to long-term storage (a real DB in production)."""
        self.long_term_store.setdefault(user_id, {}).update(self.working)

    def load(self, user_id: str):
        """Rehydrate working memory from long-term storage at session start."""
        self.working.update(self.long_term_store.get(user_id, {}))


mem = AgentMemory()
mem.load("user_42")                       # bring back what we knew about this user
mem.add_turn("user", "I'm on the Pro plan and hitting a rate limit.")
mem.remember_fact("plan", "Pro")          # survives even after 10+ more turns
mem.persist("user_42")                    # available next session
print(mem.working)                        # {'plan': 'Pro'}

# script_id: day_046_agent_memory/context_budget
# pip install tiktoken   (OpenAI's tokenizer; for Claude use client.messages.count_tokens)
import tiktoken

enc = tiktoken.encoding_for_model("gpt-4o")

def count(text: str) -> int:
    return len(enc.encode(text))

class ContextBudget:
    """Allocate a token window across zones and detect overflow."""
    def __init__(self, total: int, reserve_for_output: int):
        self.total = total
        self.reserve = reserve_for_output          # leave room to actually answer
        self.available = total - reserve_for_output

    def fits(self, system: str, history: list[dict], retrieved: str, query: str) -> bool:
        used = (count(system)
                + sum(count(m["content"]) for m in history)
                + count(retrieved)
                + count(query))
        return used <= self.available

    def trim_history(self, history: list[dict], system: str, retrieved: str, query: str) -> list[dict]:
        """Drop oldest turns until everything fits. Returns the kept turns."""
        fixed = count(system) + count(retrieved) + count(query)
        budget = self.available - fixed
        kept: list[dict] = []
        running = 0
        for msg in reversed(history):              # keep the most recent first
            cost = count(msg["content"])
            if running + cost > budget:
                break
            kept.append(msg)
            running += cost
        return list(reversed(kept))


budget = ContextBudget(total=200_000, reserve_for_output=8_000)
history = [{"role": "user", "content": "..."} for _ in range(50)]
kept = budget.trim_history(history, system="You are...", retrieved="docs...", query="next?")
print(f"kept {len(kept)} of {len(history)} turns")

# script_id: day_046_agent_memory/summarize_history
from openai import OpenAI
client = OpenAI()

def summarize_old_turns(turns: list[dict]) -> str:
    """Compress a batch of old turns into a short factual summary."""
    transcript = "\n".join(f"{m['role']}: {m['content']}" for m in turns)
    resp = client.chat.completions.create(
        model="gpt-4o-mini",                       # cheap model is fine for summarizing
        messages=[
            {"role": "system", "content": "Summarize this conversation in 2-3 sentences, "
                                           "preserving any facts, decisions, or user preferences."},
            {"role": "user", "content": transcript},
        ],
    )
    return resp.choices[0].message.content

# Strategy: every N turns, summarize the oldest half and keep the summary as one
# 'system' note, so the running context stays small without losing the thread.

# script_id: day_046_agent_memory/redis_working_memory
# Working memory in Redis — fast, shared across processes, expires on its own.
import json
import redis

r = redis.Redis(decode_responses=True)

def save_working(session_id: str, facts: dict, ttl_seconds: int = 3600):
    r.set(f"mem:{session_id}", json.dumps(facts), ex=ttl_seconds)

def load_working(session_id: str) -> dict:
    raw = r.get(f"mem:{session_id}")
    return json.loads(raw) if raw else {}