Building Long-Term Memory for AI Agents That Actually Works

from datetime import datetime
from typing import List, Optional
import hashlib

class EpisodicMemory:
    def __init__(self, db, vector_store, summarizer):
        self.db = db
        self.vector_store = vector_store
        self.summarizer = summarizer
    
    async def store_conversation(self, user_id: str, conversation: List[dict],
                                  metadata: dict) -> str:
        """Store a conversation with compression for older content."""
        
        conversation_id = hashlib.sha256(
            f"{user_id}:{datetime.now().isoformat()}".encode()
        ).hexdigest()[:16]
        
        # Store raw conversation
        await self.db.insert("conversations", {
            "id": conversation_id,
            "user_id": user_id,
            "messages": conversation,
            "timestamp": datetime.now(),
            "metadata": metadata
        })
        
        # Generate searchable summary
        summary = await self.summarizer.summarize_conversation(conversation)
        
        # Extract key points for vector search
        key_points = await self.summarizer.extract_key_points(conversation)
        
        # Index for retrieval
        for i, point in enumerate(key_points):
            embedding = await self.vector_store.embed(point["text"])
            await self.vector_store.upsert(
                id=f"{conversation_id}:{i}",
                vector=embedding,
                metadata={
                    "user_id": user_id,
                    "conversation_id": conversation_id,
                    "type": "episodic",
                    "point_type": point["type"],  # decision, preference, fact, etc.
                    "timestamp": datetime.now().isoformat(),
                    "summary": point["text"]
                }
            )
        
        return conversation_id
    
    async def retrieve_relevant_episodes(self, user_id: str, query: str,
                                          limit: int = 5) -> List[dict]:
        """Retrieve conversations relevant to current query."""
        
        query_embedding = await self.vector_store.embed(query)
        
        results = await self.vector_store.search(
            vector=query_embedding,
            filter={"user_id": user_id, "type": "episodic"},
            top_k=limit * 2  # Over-fetch then dedupe by conversation
        )
        
        # Dedupe by conversation, keeping highest scoring point per conversation
        conversations = {}
        for result in results:
            conv_id = result.metadata["conversation_id"]
            if conv_id not in conversations or result.score > conversations[conv_id]["score"]:
                conversations[conv_id] = {
                    "conversation_id": conv_id,
                    "score": result.score,
                    "relevant_point": result.metadata["summary"],
                    "timestamp": result.metadata["timestamp"]
                }
        
        # Get full conversations for top results
        top_conversations = sorted(
            conversations.values(), 
            key=lambda x: x["score"], 
            reverse=True
        )[:limit]
        
        return [
            await self.db.get("conversations", c["conversation_id"])
            for c in top_conversations
        ]

class SemanticMemory:
    """Long-term factual memory about users and contexts."""
    
    def __init__(self, db, vector_store, extractor):
        self.db = db
        self.vector_store = vector_store
        self.extractor = extractor
    
    async def extract_and_store_facts(self, user_id: str, 
                                       conversation: List[dict]) -> List[str]:
        """Extract durable facts from conversation and store them."""
        
        # Use LLM to extract facts
        extraction = await self.extractor.extract(conversation, schema={
            "preferences": "List of user preferences mentioned",
            "facts": "Factual information about user/company",
            "constraints": "Limitations or requirements mentioned",
            "relationships": "People/companies mentioned and their relationships",
            "decisions": "Decisions made or commitments given"
        })
        
        stored_ids = []
        
        for fact_type, facts in extraction.items():
            for fact in facts:
                # Check if this fact updates existing knowledge
                existing = await self._find_related_facts(user_id, fact)
                
                if existing:
                    # Update existing fact
                    updated_fact = await self._merge_facts(existing, fact)
                    fact_id = await self._update_fact(existing["id"], updated_fact)
                else:
                    # Store new fact
                    fact_id = await self._store_new_fact(user_id, fact_type, fact)
                
                stored_ids.append(fact_id)
        
        return stored_ids
    
    async def _store_new_fact(self, user_id: str, fact_type: str, 
                               fact: dict) -> str:
        """Store a new semantic memory."""
        
        fact_id = generate_id()
        
        # Store in database
        await self.db.insert("semantic_memory", {
            "id": fact_id,
            "user_id": user_id,
            "type": fact_type,
            "content": fact["content"],
            "confidence": fact.get("confidence", 0.8),
            "source_conversations": [fact.get("source_conversation_id")],
            "created_at": datetime.now(),
            "updated_at": datetime.now(),
            "access_count": 0,
            "last_accessed": None
        })
        
        # Index for retrieval
        embedding = await self.vector_store.embed(fact["content"])
        await self.vector_store.upsert(
            id=fact_id,
            vector=embedding,
            metadata={
                "user_id": user_id,
                "type": "semantic",
                "fact_type": fact_type,
                "content": fact["content"]
            }
        )
        
        return fact_id
    
    async def retrieve_relevant_facts(self, user_id: str, query: str,
                                       fact_types: Optional[List[str]] = None,
                                       limit: int = 10) -> List[dict]:
        """Retrieve facts relevant to query."""
        
        query_embedding = await self.vector_store.embed(query)
        
        filter_conditions = {"user_id": user_id, "type": "semantic"}
        if fact_types:
            filter_conditions["fact_type"] = {"$in": fact_types}
        
        results = await self.vector_store.search(
            vector=query_embedding,
            filter=filter_conditions,
            top_k=limit
        )
        
        # Update access patterns (for memory consolidation)
        fact_ids = [r.id for r in results]
        await self._update_access_patterns(fact_ids)
        
        return [
            {
                "content": r.metadata["content"],
                "type": r.metadata["fact_type"],
                "relevance": r.score
            }
            for r in results
        ]

class WorkingMemory:
    """Session-level memory with automatic decay."""
    
    def __init__(self, max_tokens: int = 4000, decay_rate: float = 0.95):
        self.max_tokens = max_tokens
        self.decay_rate = decay_rate
        self.sessions: Dict[str, SessionMemory] = {}
    
    def get_session(self, session_id: str) -> "SessionMemory":
        if session_id not in self.sessions:
            self.sessions[session_id] = SessionMemory(
                self.max_tokens, self.decay_rate
            )
        return self.sessions[session_id]

class SessionMemory:
    def __init__(self, max_tokens: int, decay_rate: float):
        self.max_tokens = max_tokens
        self.decay_rate = decay_rate
        self.items: List[MemoryItem] = []
        self.scratchpad: dict = {}  # For agent reasoning state
    
    def add(self, content: str, importance: float = 0.5, 
            memory_type: str = "observation"):
        """Add item to working memory with importance score."""
        
        item = MemoryItem(
            content=content,
            importance=importance,
            type=memory_type,
            timestamp=datetime.now(),
            access_count=0
        )
        
        self.items.append(item)
        self._enforce_capacity()
    
    def get_context(self, query: Optional[str] = None) -> str:
        """Get working memory as context string."""
        
        # Decay old items
        self._apply_decay()
        
        # Sort by relevance if query provided, else by recency
        if query:
            items = self._rank_by_relevance(query)
        else:
            items = sorted(self.items, key=lambda x: x.timestamp, reverse=True)
        
        # Build context within token limit
        context_parts = []
        current_tokens = 0
        
        for item in items:
            item_tokens = estimate_tokens(item.content)
            if current_tokens + item_tokens > self.max_tokens:
                break
            
            context_parts.append(f"[{item.type}] {item.content}")
            current_tokens += item_tokens
        
        return "\n".join(context_parts)
    
    def _apply_decay(self):
        """Decay importance of old items."""
        now = datetime.now()
        for item in self.items:
            age_hours = (now - item.timestamp).total_seconds() / 3600
            decay_factor = self.decay_rate ** age_hours
            item.importance *= decay_factor
    
    def _enforce_capacity(self):
        """Remove lowest importance items when over capacity."""
        while self._total_tokens() > self.max_tokens * 1.5:
            # Remove lowest importance item
            min_item = min(self.items, key=lambda x: x.importance)
            self.items.remove(min_item)

class MemoryAssembler:
    """Combines all memory types into optimal context for current query."""
    
    def __init__(self, episodic: EpisodicMemory, semantic: SemanticMemory,
                 max_context_tokens: int = 6000):
        self.episodic = episodic
        self.semantic = semantic
        self.max_tokens = max_context_tokens
    
    async def assemble_memory(self, user_id: str, session_id: str,
                               query: str, working_memory: WorkingMemory) -> str:
        """Assemble relevant memories for the current query."""
        
        # Budget allocation (tokens)
        budgets = {
            "working": int(self.max_tokens * 0.3),   # 30% for session context
            "semantic": int(self.max_tokens * 0.4),  # 40% for facts
            "episodic": int(self.max_tokens * 0.3),  # 30% for past conversations
        }
        
        # Retrieve from each memory type
        working_context = working_memory.get_session(session_id).get_context(query)
        
        semantic_facts = await self.semantic.retrieve_relevant_facts(
            user_id, query, limit=20
        )
        
        episodic_memories = await self.episodic.retrieve_relevant_episodes(
            user_id, query, limit=3
        )
        
        # Format and fit within budgets
        memory_sections = []
        
        # Working memory (most relevant for current task)
        if working_context:
            working_section = self._format_section(
                "CURRENT SESSION CONTEXT",
                working_context,
                budgets["working"]
            )
            memory_sections.append(working_section)
        
        # Semantic memory (facts about user)
        if semantic_facts:
            facts_text = "\n".join([
                f"- {f['content']}" for f in semantic_facts
            ])
            semantic_section = self._format_section(
                "KNOWN FACTS ABOUT USER",
                facts_text,
                budgets["semantic"]
            )
            memory_sections.append(semantic_section)
        
        # Episodic memory (relevant past conversations)
        if episodic_memories:
            episodes_text = self._format_episodes(episodic_memories)
            episodic_section = self._format_section(
                "RELEVANT PAST INTERACTIONS",
                episodes_text,
                budgets["episodic"]
            )
            memory_sections.append(episodic_section)
        
        return "\n\n".join(memory_sections)
    
    def _format_section(self, header: str, content: str, 
                        max_tokens: int) -> str:
        """Format a memory section within token budget."""
        
        truncated = truncate_to_tokens(content, max_tokens - 10)
        return f"=== {header} ===\n{truncated}"
    
    def _format_episodes(self, episodes: List[dict]) -> str:
        """Format episodic memories as summaries."""
        
        formatted = []
        for episode in episodes:
            date = episode["timestamp"].strftime("%Y-%m-%d")
            summary = episode.get("summary", "No summary available")
            formatted.append(f"[{date}] {summary}")
        
        return "\n".join(formatted)

class MemoryConsolidator:
    """Nightly job to consolidate and prune memories."""
    
    async def consolidate_user_memories(self, user_id: str):
        # 1. Merge similar semantic memories
        await self._merge_similar_facts(user_id)
        
        # 2. Promote frequently accessed episodic memories to semantic
        await self._promote_important_episodes(user_id)
        
        # 3. Compress old episodic memories
        await self._compress_old_episodes(user_id)
        
        # 4. Prune low-value memories
        await self._prune_low_value_memories(user_id)
    
    async def _promote_important_episodes(self, user_id: str):
        """Extract durable facts from frequently accessed conversations."""
        
        # Find conversations accessed multiple times
        important_conversations = await self.db.query(
            "SELECT * FROM conversations WHERE user_id = ? AND access_count > 3",
            [user_id]
        )
        
        for conv in important_conversations:
            if not conv.get("facts_extracted"):
                # Extract facts to semantic memory
                await self.semantic.extract_and_store_facts(
                    user_id, conv["messages"]
                )
                # Mark as processed
                await self.db.update(
                    "conversations", 
                    conv["id"], 
                    {"facts_extracted": True}
                )
    
    async def _compress_old_episodes(self, user_id: str):
        """Replace old conversations with summaries."""
        
        old_conversations = await self.db.query(
            '''SELECT * FROM conversations 
               WHERE user_id = ? 
               AND timestamp < ? 
               AND NOT compressed''',
            [user_id, datetime.now() - timedelta(days=30)]
        )
        
        for conv in old_conversations:
            # Generate summary
            summary = await self.summarizer.summarize_conversation(
                conv["messages"]
            )
            
            # Replace full conversation with summary
            await self.db.update("conversations", conv["id"], {
                "messages": None,  # Remove raw messages
                "summary": summary,
                "compressed": True
            })