Memory Retrieval: Surfacing the Right Fact at the Right Time

import math
import re
from dataclasses import dataclass
 
 
@dataclass
class MemoryHit:
    id: str
    kind: str  # profile | episode
    text: str
    score: float
 
 
PROFILE_KEYS = ("preferred_language", "timezone", "tone")
 
 
def load_profile(rows: list[dict], user_id: str) -> list[MemoryHit]:
    hits = []
    for row in rows:
        if row["user_id"] != user_id or row.get("namespace", "profile") != "profile":
            continue
        if row["key"] in PROFILE_KEYS:
            hits.append(
                MemoryHit(
                    id=f"profile:{row['key']}",
                    kind="profile",
                    text=f"{row['key']}={row['value']}",
                    score=1.0,
                )
            )
    return hits
 
 
def tokenize(text: str) -> set[str]:
    return {t for t in re.findall(r"[a-z0-9_]+", text.lower()) if len(t) > 2}
 
 
def keyword_score(query: str, doc: str) -> float:
    q, d = tokenize(query), tokenize(doc)
    if not q or not d:
        return 0.0
    overlap = len(q & d)
    return overlap / math.sqrt(len(q) * len(d))
 
 
def retrieve_episodes(
    rows: list[dict],
    user_id: str,
    query: str,
    *,
    k: int = 3,
    min_score: float = 0.08,
) -> list[MemoryHit]:
    scored: list[MemoryHit] = []
    for row in rows:
        if row["user_id"] != user_id or row.get("namespace") != "episode":
            continue
        if row.get("expired"):
            continue
        text = row["content"]
        score = keyword_score(query, text)
        if score >= min_score:
            scored.append(MemoryHit(id=row["id"], kind="episode", text=text, score=score))
    scored.sort(key=lambda h: h.score, reverse=True)
    return scored[:k]
 
 
def format_memory_block(profile: list[MemoryHit], episodes: list[MemoryHit], max_chars: int = 800) -> str:
    lines = ["## Long-term memory (may be incomplete)"]
    if profile:
        lines.append("### Profile")
        for h in profile:
            lines.append(f"- ({h.id}) {h.text}")
    if episodes:
        lines.append("### Relevant past notes")
        for h in episodes:
            lines.append(f"- ({h.id}, score={h.score:.2f}) {h.text}")
    if len(lines) == 1:
        return ""
    block = "\n".join(lines)
    return block if len(block) <= max_chars else block[: max_chars - 20] + "\n...[truncated]"
 
 
def build_system_with_memory(
    base_system: str,
    store_rows: list[dict],
    user_id: str,
    user_message: str,
) -> str:
    profile = load_profile(store_rows, user_id)
    episodes = retrieve_episodes(store_rows, user_id, user_message, k=3)
    block = format_memory_block(profile, episodes)
    if not block:
        return base_system
    return base_system + "\n\n" + block + "\n\nUse memories only when relevant. Prefer fresher facts if they conflict."
 
 
# Demo data
rows = [
    {"user_id": "u1", "namespace": "profile", "key": "preferred_language", "value": "Portuguese"},
    {"user_id": "u1", "namespace": "profile", "key": "tone", "value": "concise"},
    {
        "id": "e1",
        "user_id": "u1",
        "namespace": "episode",
        "content": "User deploys the billing API to ECS on Fridays.",
        "expired": False,
    },
    {
        "id": "e2",
        "user_id": "u1",
        "namespace": "episode",
        "content": "User's dog is named Pixel.",
        "expired": False,
    },
]
 
system = build_system_with_memory(
    "You are a support agent.",
    rows,
    "u1",
    "Help me plan this Friday's billing API deploy.",
)
print(system)