Agent Memory Systems: From Context to Persistent Storage › Code Implementation: Building a Memory-Enabled Agent

from datetime import datetime
from typing import Dict, Optional

import numpy as np


# We'll use a simple embedding function for demonstration.
# In production, use sentence-transformers or OpenAI embeddings.
def simple_embed(text: str, dim: int = 64) -> np.ndarray:
    """Simple bag-of-characters embedding for demonstration."""
    vec = np.zeros(dim)
    for i, char in enumerate(text.lower()):
        vec[ord(char) % dim] += 1
    # Normalize to unit length for cosine similarity
    norm = np.linalg.norm(vec)
    return vec / norm if norm > 0 else vec


class Memory:
    """Single memory entry with all metadata needed for retrieval scoring."""

    def __init__(
        self,
        content: str,
        memory_type: str = "episodic",
        importance: float = 1.0,
        metadata: Optional[Dict] = None,
    ):
        self.content = content
        self.memory_type = memory_type  # 'episodic', 'semantic', 'procedural'
        self.importance = importance
        self.created_at = datetime.now()
        self.embedding = simple_embed(content)
        self.metadata = metadata or {}
        self.access_count = 0
        self.last_accessed = datetime.now()

    def to_dict(self) -> Dict:
        return {
            "content": self.content[:100] + "..."
            if len(self.content) > 100
            else self.content,
            "type": self.memory_type,
            "importance": self.importance,
            "age_hours": (datetime.now() - self.created_at).total_seconds()
            / 3600,
            "access_count": self.access_count,
        }

from typing import Dict, List, Optional, Tuple


class AgentMemory:
    """Hierarchical memory system implementing working memory and long-term storage."""

    def __init__(self, capacity: int = 1000, context_size: int = 5):
        self.long_term_memory: List[Memory] = []
        self.working_memory: List[Memory] = []  # Recent, unconsolidated
        self.capacity = capacity
        self.context_size = context_size  # Items to keep in working memory

        # Retrieval scoring weights (must sum to 1.0)
        self.alpha = 0.6  # Similarity weight
        self.beta = 0.3  # Recency weight
        self.gamma = 0.1  # Importance weight

    def add_memory(
        self,
        content: str,
        memory_type: str = "episodic",
        importance: float = 1.0,
        metadata: Optional[Dict] = None,
    ):
        """Add a new memory to working memory, triggering consolidation if needed."""
        mem = Memory(content, memory_type, importance, metadata)
        self.working_memory.append(mem)

        # Consolidate working memory if it grows too large
        if len(self.working_memory) > self.context_size * 2:
            self._consolidate()

    def _consolidate(self):
        """Move older working memories to long-term storage."""
        # Preserve the most recent items in working memory
        to_archive = self.working_memory[: -self.context_size]
        self.working_memory = self.working_memory[-self.context_size :]

        for mem in to_archive:
            # Boost importance for frequently accessed memories before archiving
            if mem.access_count > 2:
                mem.importance *= 1.5
            self.long_term_memory.append(mem)

        # Evict from long-term memory if over capacity
        if len(self.long_term_memory) > self.capacity:
            self._evict_memories()

    def _evict_memories(self):
        """Remove lowest-scoring memories to respect capacity constraints."""
        now = datetime.now()
        scores = []

        for i, mem in enumerate(self.long_term_memory):
            age_days = (now - mem.created_at).days
            recency_score = np.exp(-0.1 * age_days)
            # Eviction score weights retention factors: importance most, then access, then recency
            retention_score = (
                0.3 * recency_score
                + 0.5 * min(mem.importance / 5.0, 1.0)
                + 0.2 * min(mem.access_count / 10.0, 1.0)
            )
            scores.append((retention_score, i))

        # Keep the highest-scoring memories (capacity - 100 to create buffer)
        scores.sort()
        keep_indices = set(idx for _, idx in scores[-(self.capacity - 100) :])
        self.long_term_memory = [
            m for i, m in enumerate(self.long_term_memory) if i in keep_indices
        ]

    def retrieve(self, query: str, k: int = 3) -> List[Tuple[Memory, float]]:
        """Retrieve top-k memories using multi-factor scoring."""
        query_vec = simple_embed(query)
        now = datetime.now()

        scored_memories = []
        all_memories = self.working_memory + self.long_term_memory

        for mem in all_memories:
            # Semantic similarity via cosine dot product on unit vectors
            similarity = np.dot(query_vec, mem.embedding)

            # Recency: exponential decay based on age in days
            age_days = (now - mem.created_at).total_seconds() / (24 * 3600)
            recency = np.exp(-0.1 * age_days)

            # Importance: normalized to [0, 1]
            importance = min(mem.importance / 5.0, 1.0)

            # Combined score from the multi-factor formula
            score = (
                self.alpha * similarity
                + self.beta * recency
                + self.gamma * importance
            )
            scored_memories.append((mem, score))

        scored_memories.sort(key=lambda x: x[1], reverse=True)

        # Update access statistics for retrieved memories
        for mem, _ in scored_memories[:k]:
            mem.access_count += 1
            mem.last_accessed = now

        return scored_memories[:k]

import re


class SummarizingAgent:
    """Agent with memory capabilities including importance scoring and summarization."""

    def __init__(self):
        self.memory = AgentMemory(capacity=100, context_size=3)
        self.conversation_turns = 0

    def interact(self, user_input: str, system_response: str = None):
        """Process a conversation turn, storing both sides with estimated importance."""
        self.memory.add_memory(
            content=f"User: {user_input}",
            memory_type="episodic",
            importance=self._estimate_importance(user_input),
        )

        if system_response:
            self.memory.add_memory(
                content=f"Assistant: {system_response}",
                memory_type="episodic",
                importance=1.0,
            )

        self.conversation_turns += 1

        # Every 5 turns, create a compressed summary of recent context
        if self.conversation_turns % 5 == 0:
            self._create_summary()

    def _estimate_importance(self, text: str) -> float:
        """Assign importance scores using heuristic signals."""
        importance = 1.0

        # Questions indicate information gaps that should be preserved
        if "?" in text:
            importance += 1.0

        # Task-framing language suggests high-stakes information
        if any(
            word in text.lower()
            for word in ["need", "must", "should", "task", "goal"]
        ):
            importance += 1.5

        # Personal information is useful for personalization
        if any(
            word in text.lower()
            for word in ["i like", "i prefer", "my name", "i am"]
        ):
            importance += 2.0

        return min(importance, 5.0)

    def _create_summary(self):
        """Compress recent context into a single high-importance summary memory."""
        recent = self.memory.working_memory[-5:]
        if len(recent) < 3:
            return

        # In production, this would call an LLM with a summarization prompt.
        # Here we extract salient terms as a proxy for summarization.
        topics = set()
        for mem in recent:
            words = re.findall(r'"([^"]*)"|\b([A-Z][a-z]+)\b', mem.content)
            topics.update(w for pair in words for w in pair if w)

        summary = f"Summary: Discussed {', '.join(sorted(topics)[:3])}. Key points from last 5 turns."

        self.memory.add_memory(
            content=summary,
            memory_type="semantic",
            importance=2.5,
            metadata={"is_summary": True, "turns_covered": 5},
        )

    def recall(self, query: str) -> str:
        """Retrieve and format relevant memories for a given query."""
        memories = self.memory.retrieve(query, k=3)

        if not memories:
            return "No relevant memories found."

        results = []
        for mem, score in memories:
            results.append(
                f"[{mem.memory_type}, score={score:.2f}] {mem.content[:80]}..."
            )
        return "\n".join(results)

# Initialize agent
agent = SummarizingAgent()

# Session 1: Initial project setup conversation
session1_inputs = [
    (
        "We need to build a Python API for our e-commerce platform",
        "I'll help you design a Python API for e-commerce. What specific features do you need?",
    ),
    (
        "It needs to handle user authentication with JWT tokens",
        "JWT authentication is a good choice. I'll note we need secure token handling.",
    ),
    (
        "Also implement a shopping cart with Redis caching",
        "Redis caching for the shopping cart will ensure fast performance.",
    ),
    (
        "My name is Sarah and I prefer FastAPI over Flask",
        "Noted, Sarah. FastAPI is excellent for this use case with its async support.",
    ),
]

for user_msg, agent_msg in session1_inputs:
    agent.interact(user_msg, agent_msg)

# Simulate time passing with less important conversation
for i in range(20):
    agent.interact(f"Filler conversation point {i}", "Acknowledged.")

# Session 2: User returns with a specific question
query = "How should we handle the JWT refresh strategy?"
recall_result = agent.recall(query)

# Session 3: Checking personal preferences
query2 = "What framework should I use for the web API?"
recall_result2 = agent.recall(query2)

def analyze_memory(agent):
    """Analyze the distribution of memories across tiers and types."""
    stats = {}
    stats["working_size"] = len(agent.memory.working_memory)
    stats["long_term_size"] = len(agent.memory.long_term_memory)

    types = {}
    for mem in agent.memory.long_term_memory + agent.memory.working_memory:
        types[mem.memory_type] = types.get(mem.memory_type, 0) + 1
    stats["type_distribution"] = types

    all_mems = agent.memory.long_term_memory + agent.memory.working_memory
    top_accessed = sorted(all_mems, key=lambda x: x.access_count, reverse=True)[
        :3
    ]
    stats["top_accessed"] = [
        (mem.content[:60], mem.access_count)
        for mem in top_accessed
        if mem.access_count > 0
    ]

    return stats


memory_stats = analyze_memory(agent)