Kaushik's Notes

Memory

7 min read

What is Agent Memory?

AI agent memory is the persistent cognitive architecture that allows agents to accumulate knowledge, maintain contextual awareness, and adapt their behavior based on historical interactions and learned experiences.

  • Memory makes Agents: Reliable, Believable and Capable

What is Memory Management?

Memory management in agentic syustems refers to the systematic organization, persistence, and retrieval of different types of information that AI agents need to function effictively across interactions and sessions.

Core Components of Agent Memory Management

1. Generation

2. Storage

3. Retrieval

4. Integration

5. Updating

6. Deletion/Forgetting

Different Forms of Memory

Persona

  • Description: Stores agent identity information, including personality traits, roles, expertise domains, and communication styles.
  • Contents: Name, role, goals, background, and vector embeddings for semantic retrieval.
  • Usage: Provides consistent identity for agents across interactions and sessions.
  • Schema: Includes Persona_id, name, role, goals, background fields with embedding vectors

ToolBox

  • Description: Stores tool definitions, metadata, parameter schemas, and embeddings for function capabilities.
  • Contents: Tool names, descriptions, parameter specifications, and vector embeddings
  • Usage: Enables semantic discovery and execution of external functions by agents
  • Schema: Includes tool_id, name, function metadata, parameters, and embedding vectors.

Conversation Memory

  • Description: Stores historical exchanges between users and agents
  • Contents: Sequential turns with roles, content, timestamps, and conversation identifiers.
  • Usage: Provides context for ongoing conversations and enables coherent multi-turn interactions
  • Schema: Includes memory_id, conversation_id, role, content, timestamp fields

Workflow Memory

  • Description: Stores multi-step process information and state tracking
  • Contents: Workflow definitions, current state, transition history, and execution context
  • Usage: Supports complex, multi-stage operations that span multiple agent interactions
  • Schema: Includes workflow_id, stages, current_stage, history, and context information

Agent Memory

  • Description: Stores agent-specific state, decision patterns, and learned behaviors
  • Contents: Agent preferences, behavioral patterns, decision history, and adaptation parameters
  • Usage: Enables agents to maintain consistent behavior patterns and learn from past decisions
  • Schema: Includes agent_id, behavior_patterns, decision_history, preferences, adaptation_metrics

Episodic Memory

  • Description: Stores specific events and experiences as discrete episodes with temporal context
  • Contents: Event sequences, context information, outcomes, and temporal relationships
  • Usage: Enables agents to recall specific past events and learn from similar situations
  • Schema: Includes episode_id, timestamp, context, actions, outcomes, emotional_state

Long-term Memory

  • Description: Stores persistent knowledge, patterns, and insights that transcend individual sessions
  • Contents: Accumulated knowledge, learned patterns, strategic insights, and domain expertise
  • Usage: Provides foundational knowledge base that persists across agent lifecycles
  • Schema: Includes memory_id, knowledge_type, content, confidence_score, last_updated

Entity Memory

  • Description: Stores information about specific entities (people, objects, concepts) encountered by the agent
  • Contents: Entity attributes, relationships, interaction history, and contextual associations
  • Usage: Enables agents to maintain detailed knowledge about recurring entities and their properties
  • Schema: Includes entity_id, entity_type, attributes, relationships, interaction_history

Memory Implementation in LangGraph

1. Checkpointing for Conversation Memory

LangGraph provides built-in checkpointing mechanisms for maintaining conversation state:

from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.graph import StateGraph, END
from typing import TypedDict
 
class ConversationState(TypedDict):
    messages: list
    user_id: str
    conversation_id: str
    context: dict
 
# Initialize memory with SQLite checkpointer
memory = SqliteSaver.from_conn_string(":memory:")
 
def conversation_node(state: ConversationState):
    # Process conversation with memory context
    return {"messages": updated_messages}
 
# Build graph with memory
workflow = StateGraph(ConversationState)
workflow.add_node("conversation", conversation_node)
workflow.set_entry_point("conversation")
workflow.add_edge("conversation", END)
 
# Compile with memory checkpointing
app = workflow.compile(checkpointer=memory)
 
# Use with thread_id for conversation continuity
config = {"configurable": {"thread_id": "user_123_conversation"}}
result = app.invoke(input_data, config=config)

2. Custom Memory Storage Implementation

import sqlite3
import json
from datetime import datetime
from typing import Dict, List, Optional
 
class AgentMemoryManager:
    def __init__(self, db_path: str):
        self.db_path = db_path
        self.init_database()
    
    def init_database(self):
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # Conversation memory table
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS conversation_memory (
                id INTEGER PRIMARY KEY,
                conversation_id TEXT,
                role TEXT,
                content TEXT,
                timestamp DATETIME,
                metadata JSON
            )
        ''')
        
        # Entity memory table
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS entity_memory (
                id INTEGER PRIMARY KEY,
                entity_id TEXT UNIQUE,
                entity_type TEXT,
                attributes JSON,
                last_updated DATETIME
            )
        ''')
        
        # Episodic memory table
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS episodic_memory (
                id INTEGER PRIMARY KEY,
                episode_id TEXT UNIQUE,
                context JSON,
                actions JSON,
                outcomes JSON,
                timestamp DATETIME
            )
        ''')
        
        conn.commit()
        conn.close()
    
    def store_conversation(self, conversation_id: str, role: str, 
                          content: str, metadata: Dict = None):
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            INSERT INTO conversation_memory 
            (conversation_id, role, content, timestamp, metadata)
            VALUES (?, ?, ?, ?, ?)
        ''', (conversation_id, role, content, datetime.now(), 
              json.dumps(metadata or {})))
        
        conn.commit()
        conn.close()
    
    def retrieve_conversation_history(self, conversation_id: str, 
                                    limit: int = 50) -> List[Dict]:
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            SELECT role, content, timestamp, metadata 
            FROM conversation_memory 
            WHERE conversation_id = ? 
            ORDER BY timestamp DESC 
            LIMIT ?
        ''', (conversation_id, limit))
        
        results = cursor.fetchall()
        conn.close()
        
        return [{
            'role': row[0],
            'content': row[1],
            'timestamp': row[2],
            'metadata': json.loads(row[3])
        } for row in results]

3. Vector-Based Semantic Memory

from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_core.documents import Document
 
class SemanticMemoryManager:
    def __init__(self, persist_directory: str):
        self.embeddings = OpenAIEmbeddings()
        self.vectorstore = Chroma(
            persist_directory=persist_directory,
            embedding_function=self.embeddings
        )
    
    def store_memory(self, content: str, memory_type: str, 
                    metadata: Dict = None):
        doc = Document(
            page_content=content,
            metadata={
                "memory_type": memory_type,
                "timestamp": datetime.now().isoformat(),
                **(metadata or {})
            }
        )
        self.vectorstore.add_documents([doc])
    
    def retrieve_relevant_memories(self, query: str, 
                                 memory_type: str = None,
                                 k: int = 5) -> List[Document]:
        filter_dict = {"memory_type": memory_type} if memory_type else None
        return self.vectorstore.similarity_search(
            query, k=k, filter=filter_dict
        )

Memory Integration Patterns

1. Memory-Augmented Agent Node

def memory_augmented_agent_node(state: AgentState):
    memory_manager = state["memory_manager"]
    query = state["current_query"]
    
    # Retrieve relevant memories
    conversation_history = memory_manager.retrieve_conversation_history(
        state["conversation_id"]
    )
    
    relevant_memories = memory_manager.retrieve_relevant_memories(
        query, memory_type="episodic"
    )
    
    # Augment context with memories
    context = {
        "conversation_history": conversation_history,
        "relevant_experiences": relevant_memories,
        "current_query": query
    }
    
    # Process with memory context
    response = llm.invoke(context)
    
    # Store new memory
    memory_manager.store_conversation(
        state["conversation_id"], 
        "assistant", 
        response.content
    )
    
    return {"response": response, "updated_context": context}

2. Multi-Memory Retrieval

def multi_memory_retrieval_node(state: AgentState):
    query = state["query"]
    memory_managers = state["memory_managers"]
    
    # Retrieve from multiple memory types
    memories = {}
    
    # Conversation memory
    memories["conversation"] = memory_managers["conversation"].retrieve(
        state["conversation_id"]
    )
    
    # Entity memory
    memories["entities"] = memory_managers["entity"].retrieve_entities(
        query
    )
    
    # Episodic memory
    memories["episodes"] = memory_managers["episodic"].retrieve_similar_episodes(
        query
    )
    
    # Long-term knowledge
    memories["knowledge"] = memory_managers["knowledge"].retrieve_knowledge(
        query
    )
    
    return {"retrieved_memories": memories}

Memory Lifecycle Management

1. Memory Consolidation

def consolidate_memories(memory_manager: AgentMemoryManager):
    """Consolidate short-term memories into long-term storage"""
    
    # Retrieve recent episodic memories
    recent_episodes = memory_manager.get_recent_episodes(days=7)
    
    # Identify patterns and important events
    patterns = identify_patterns(recent_episodes)
    important_events = classify_importance(recent_episodes)
    
    # Store consolidated knowledge
    for pattern in patterns:
        memory_manager.store_long_term_knowledge(pattern)
    
    # Promote important episodes to long-term storage
    for event in important_events:
        memory_manager.promote_to_long_term(event)

2. Memory Cleanup and Forgetting

def cleanup_memories(memory_manager: AgentMemoryManager):
    """Remove outdated or irrelevant memories"""
    
    # Remove old conversation history (keep only recent)
    memory_manager.cleanup_conversations(keep_days=30)
    
    # Remove low-confidence episodic memories
    memory_manager.remove_low_confidence_episodes(threshold=0.3)
    
    # Compress similar entities
    memory_manager.merge_similar_entities(similarity_threshold=0.9)

Best Practices for Memory Management

1. Memory Hierarchy Design

  • Immediate Memory: Current conversation context
  • Short-term Memory: Recent interactions and temporary context
  • Long-term Memory: Persistent knowledge and patterns
  • Meta-Memory: Memory about memory (what the agent remembers)

2. Retrieval Strategies

  • Recency: Prioritize recent memories
  • Relevance: Use semantic similarity for retrieval
  • Importance: Weight memories by significance
  • Frequency: Consider how often memories are accessed

3. Privacy and Security

  • Data Encryption: Encrypt sensitive memory content
  • Access Control: Implement role-based memory access
  • Retention Policies: Automatic cleanup of personal data
  • Anonymization: Remove identifying information when possible

4. Performance Optimization

  • Indexing: Create efficient indexes for memory retrieval
  • Caching: Cache frequently accessed memories
  • Compression: Compress old memory data
  • Partitioning: Partition memories by user, time, or type

Advanced Memory Patterns

1. Hierarchical Memory

class HierarchicalMemory:
    def __init__(self):
        self.working_memory = {}
        self.short_term_memory = {}
        self.long_term_memory = {}
    
    def store(self, content, importance_level):
        if importance_level == "immediate":
            self.working_memory[content.id] = content
        elif importance_level == "session":
            self.short_term_memory[content.id] = content
        else:
            self.long_term_memory[content.id] = content

2. Associative Memory Networks

class AssociativeMemory:
    def __init__(self):
        self.memory_graph = nx.Graph()
    
    def add_memory(self, memory_id, content, associations):
        self.memory_graph.add_node(memory_id, content=content)
        for assoc_id, weight in associations:
            self.memory_graph.add_edge(memory_id, assoc_id, weight=weight)
    
    def retrieve_associated(self, memory_id, depth=2):
        return nx.ego_graph(self.memory_graph, memory_id, radius=depth)

This comprehensive memory management system enables agents to maintain context, learn from experiences, and provide more intelligent, personalized interactions across sessions.

Graph View

Backlinks