Stateful Agents with LangGraph Checkpoints: Complete Guide

from dataclasses import dataclass, field
from typing import List, Dict, Any, Optional
from datetime import datetime, timezone
import json

@dataclass
class AgentState:
    """Agent state snapshot."""
    state_id: str
    session_id: str
    agent_id: str
    user_id: str
    created_at: str
    conversation_history: List[Dict[str, str]] = field(default_factory=list)
    context: Dict[str, Any] = field(default_factory=dict)
    tool_results: List[Dict[str, Any]] = field(default_factory=list)
    metadata: Dict[str, Any] = field(default_factory=dict)
    
    def add_message(self, role: str, content: str) -> None:
        """Add message to conversation history."""
        self.conversation_history.append({
            "role": role,
            "content": content,
            "timestamp": datetime.now(timezone.utc).isoformat(),
        })
    
    def set_context(self, key: str, value: Any) -> None:
        """Update context."""
        self.context[key] = value
    
    def get_context(self, key: str, default: Any = None) -> Any:
        """Get context value."""
        return self.context.get(key, default)
    
    def add_tool_result(self, tool: str, result: Any) -> None:
        """Record tool execution result."""
        self.tool_results.append({
            "tool": tool,
            "result": result,
            "timestamp": datetime.now(timezone.utc).isoformat(),
        })

class StateManager:
    """Manage agent state with checkpointing."""
    
    def __init__(self, storage_backend):
        self.storage = storage_backend
        self.active_states: Dict[str, AgentState] = {}
    
    async def create_state(
        self,
        session_id: str,
        agent_id: str,
        user_id: str,
    ) -> AgentState:
        """Create new agent state."""
        import uuid
        
        state = AgentState(
            state_id=str(uuid.uuid4()),
            session_id=session_id,
            agent_id=agent_id,
            user_id=user_id,
            created_at=datetime.now(timezone.utc).isoformat(),
        )
        self.active_states[session_id] = state
        
        # Persist to storage
        await self.checkpoint(state)
        
        return state
    
    async def get_state(self, session_id: str) -> Optional[AgentState]:
        """Get existing state."""
        # Check memory cache
        if session_id in self.active_states:
            return self.active_states[session_id]
        
        # Load from storage
        state = await self.storage.load_state(session_id)
        
        if state:
            self.active_states[session_id] = state
        
        return state
    
    async def checkpoint(self, state: AgentState) -> None:
        """Save state checkpoint."""
        await self.storage.save_state(state)
    
    async def delete_state(self, session_id: str) -> None:
        """Delete state."""
        # Remove from cache
        if session_id in self.active_states:
            del self.active_states[session_id]
        
        # Remove from storage
        await self.storage.delete_state(session_id)
    
    async def list_user_sessions(self, user_id: str) -> List[str]:
        """List all sessions for user."""
        return await self.storage.list_sessions(user_id)

# Usage
state_manager = StateManager(storage_backend=postgres_backend)

# Create new session
state = await state_manager.create_state(
    session_id="sess-12345",
    agent_id="support-agent",
    user_id="user-123",
)

# Update state
state.add_message("user", "I need help with my order")
state.set_context("current_intent", "order_help")

# Checkpoint (save)
await state_manager.checkpoint(state)

# Later: retrieve state
restored_state = await state_manager.get_state("sess-12345")
print(f"Conversation history: {len(restored_state.conversation_history)} messages")

from langgraph.graph import StateGraph, END
from langgraph.checkpoint.postgres import PostgresSaver
from typing import TypedDict, Annotated
import operator

# Define state schema
class AgentGraphState(TypedDict):
    """State for LangGraph agent."""
    messages: Annotated[List[Dict[str, str]], operator.add]
    context: Dict[str, Any]
    next_action: Optional[str]
    tool_results: List[Dict[str, Any]]

# Define agent nodes
async def process_user_input(state: AgentGraphState) -> AgentGraphState:
    """Process user input."""
    user_message = state["messages"][-1]["content"]
    
    # Extract intent
    intent = await classify_intent(user_message)
    
    state["context"]["intent"] = intent
    state["next_action"] = "select_tool"
    
    return state

async def select_tool(state: AgentGraphState) -> AgentGraphState:
    """Select appropriate tool."""
    intent = state["context"].get("intent")
    
    if intent == "order_help":
        state["next_action"] = "get_order"
    elif intent == "refund":
        state["next_action"] = "process_refund"
    else:
        state["next_action"] = "general_response"
    
    return state

async def get_order(state: AgentGraphState) -> AgentGraphState:
    """Get order details."""
    # Extract order ID from conversation
    order_id = extract_order_id(state["messages"])
    
    # Call tool
    order = await fetch_order(order_id)
    
    state["tool_results"].append({
        "tool": "get_order",
        "result": order,
    })
    
    state["next_action"] = "generate_response"
    
    return state

async def generate_response(state: AgentGraphState) -> AgentGraphState:
    """Generate final response."""
    # Use LLM to generate response based on state
    response = await generate_llm_response(state)
    
    state["messages"].append({
        "role": "assistant",
        "content": response,
    })
    
    state["next_action"] = END
    
    return state

# Build graph
def create_agent_graph(checkpointer):
    """Create stateful agent graph."""
    
    graph = StateGraph(AgentGraphState)
    
    # Add nodes
    graph.add_node("process_input", process_user_input)
    graph.add_node("select_tool", select_tool)
    graph.add_node("get_order", get_order)
    graph.add_node("generate_response", generate_response)
    
    # Add edges
    graph.set_entry_point("process_input")
    graph.add_edge("process_input", "select_tool")
    
    # Conditional routing from select_tool
    graph.add_conditional_edges(
        "select_tool",
        lambda state: state["next_action"],
        {
            "get_order": "get_order",
            "process_refund": "process_refund",
            "general_response": "generate_response",
        },
    )
    
    graph.add_edge("get_order", "generate_response")
    graph.add_edge("generate_response", END)
    
    # Compile with checkpointer
    return graph.compile(checkpointer=checkpointer)

# Setup PostgreSQL checkpointer
from langgraph.checkpoint.postgres import PostgresSaver

checkpointer = PostgresSaver.from_conn_string(
    "postgresql://user:pass@localhost:5432/agents"
)

# Create agent
agent = create_agent_graph(checkpointer)

# Run with session
config = {"configurable": {"thread_id": "sess-12345"}}

# First turn
result = await agent.ainvoke(
    {
        "messages": [{"role": "user", "content": "I need help with order ORD-123"}],
        "context": {},
        "next_action": None,
        "tool_results": [],
    },
    config=config,
)

print(f"Response: {result['messages'][-1]['content']}")

# Second turn (state persisted from first turn)
result = await agent.ainvoke(
    {
        "messages": [{"role": "user", "content": "Can I get a refund?"}],
    },
    config=config,  # Same thread_id loads previous state
)

print(f"Response: {result['messages'][-1]['content']}")

from typing import Protocol
from collections import deque

class MemoryStrategy(Protocol):
    """Memory management strategy."""
    
    async def add_memory(self, key: str, value: Any) -> None:
        """Add to memory."""
        ...
    
    async def recall(self, query: str, limit: int = 5) -> List[Any]:
        """Recall relevant memories."""
        ...

class ConversationBufferMemory:
    """Store recent conversation turns."""
    
    def __init__(self, max_turns: int = 10):
        self.max_turns = max_turns
        self.buffer = deque(maxlen=max_turns)
    
    async def add_memory(self, role: str, content: str) -> None:
        """Add message to buffer."""
        self.buffer.append({
            "role": role,
            "content": content,
            "timestamp": datetime.now(timezone.utc).isoformat(),
        })
    
    async def recall(self, query: str = None, limit: int = 10) -> List[Dict[str, str]]:
        """Recall recent messages."""
        return list(self.buffer)[-limit:]
    
    def to_dict(self) -> Dict[str, Any]:
        """Serialize for checkpointing."""
        return {
            "max_turns": self.max_turns,
            "buffer": list(self.buffer),
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "ConversationBufferMemory":
        """Deserialize from checkpoint."""
        memory = cls(max_turns=data["max_turns"])
        memory.buffer = deque(data["buffer"], maxlen=data["max_turns"])
        return memory

class SummaryMemory:
    """Store conversation summaries."""
    
    def __init__(self):
        self.summaries = []
        self.current_conversation = []
        self.summary_threshold = 10  # Summarize every 10 messages
    
    async def add_memory(self, role: str, content: str) -> None:
        """Add message and summarize if needed."""
        self.current_conversation.append({
            "role": role,
            "content": content,
        })
        
        if len(self.current_conversation) >= self.summary_threshold:
            await self._summarize()
    
    async def _summarize(self) -> None:
        """Summarize current conversation."""
        from openai import AsyncOpenAI
        
        client = AsyncOpenAI()
        
        conversation_text = "\n".join(
            f"{msg['role']}: {msg['content']}"
            for msg in self.current_conversation
        )
        
        response = await client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{
                "role": "user",
                "content": f"Summarize this conversation concisely:\n\n{conversation_text}",
            }],
        )
        
        summary = response.choices[0].message.content
        
        self.summaries.append({
            "summary": summary,
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "num_messages": len(self.current_conversation),
        })
        
        # Clear current conversation
        self.current_conversation = []
    
    async def recall(self, query: str = None, limit: int = 5) -> List[str]:
        """Recall recent summaries."""
        return [s["summary"] for s in self.summaries[-limit:]]

class VectorMemory:
    """Store memories with semantic search."""
    
    def __init__(self, vector_store):
        self.vector_store = vector_store
    
    async def add_memory(self, content: str, metadata: Dict[str, Any]) -> None:
        """Add memory with embedding."""
        from openai import AsyncOpenAI
        
        client = AsyncOpenAI()
        
        # Embed content
        response = await client.embeddings.create(
            model="text-embedding-3-small",
            input=content,
        )
        
        embedding = response.data[0].embedding
        
        # Store in vector DB
        await self.vector_store.upsert({
            "content": content,
            "embedding": embedding,
            "metadata": metadata,
            "timestamp": datetime.now(timezone.utc).isoformat(),
        })
    
    async def recall(self, query: str, limit: int = 5) -> List[Dict[str, Any]]:
        """Recall semantically similar memories."""
        from openai import AsyncOpenAI
        
        client = AsyncOpenAI()
        
        # Embed query
        response = await client.embeddings.create(
            model="text-embedding-3-small",
            input=query,
        )
        
        query_embedding = response.data[0].embedding
        
        # Search vector store
        results = await self.vector_store.search(
            query_embedding,
            limit=limit,
        )
        
        return results

# Usage: Combine memory patterns
class HybridMemory:
    """Combine multiple memory strategies."""
    
    def __init__(self):
        self.buffer = ConversationBufferMemory(max_turns=5)
        self.summary = SummaryMemory()
        self.vector = VectorMemory(vector_store=pinecone_client)
    
    async def add_memory(self, role: str, content: str, metadata: Dict[str, Any] = None) -> None:
        """Add to all memory systems."""
        await self.buffer.add_memory(role, content)
        await self.summary.add_memory(role, content)
        
        if metadata:
            await self.vector.add_memory(content, metadata)
    
    async def recall(self, query: str, strategy: str = "vector") -> List[Any]:
        """Recall using specified strategy."""
        if strategy == "buffer":
            return await self.buffer.recall()
        elif strategy == "summary":
            return await self.summary.recall()
        else:  # vector
            return await self.vector.recall(query)

# Checkpoint hybrid memory
memory = HybridMemory()

# Add memories
await memory.add_memory("user", "I ordered item XYZ last week", {"topic": "order"})
await memory.add_memory("assistant", "Let me check on that order for you.")

# Recall relevant context
relevant = await memory.recall("What did I order?", strategy="vector")

class MultiTurnAgent:
    """Agent with multi-turn conversation support."""
    
    def __init__(self, state_manager: StateManager):
        self.state_manager = state_manager
    
    async def handle_turn(
        self,
        session_id: str,
        user_message: str,
    ) -> Dict[str, Any]:
        """Handle single conversation turn."""
        # Get or create state
        state = await self.state_manager.get_state(session_id)
        
        if not state:
            state = await self.state_manager.create_state(
                session_id=session_id,
                agent_id="multi-turn-agent",
                user_id="user-123",  # Extract from request
            )
        
        # Add user message
        state.add_message("user", user_message)
        
        # Process with context
        response = await self._process_with_context(state, user_message)
        
        # Add assistant response
        state.add_message("assistant", response)
        
        # Checkpoint
        await self.state_manager.checkpoint(state)
        
        return {
            "response": response,
            "session_id": session_id,
            "turn_count": len(state.conversation_history) // 2,
        }
    
    async def _process_with_context(
        self,
        state: AgentState,
        user_message: str,
    ) -> str:
        """Process message with full conversation context."""
        from openai import AsyncOpenAI
        
        client = AsyncOpenAI()
        
        # Build context from conversation history
        messages = [
            {"role": msg["role"], "content": msg["content"]}
            for msg in state.conversation_history
        ]
        
        # Add current message
        messages.append({"role": "user", "content": user_message})
        
        # Call LLM with full context
        response = await client.chat.completions.create(
            model="gpt-4o",
            messages=messages,
        )
        
        return response.choices[0].message.content

# Usage
agent = MultiTurnAgent(state_manager)

# Turn 1
result1 = await agent.handle_turn(
    session_id="sess-123",
    user_message="I need help with my order",
)
print(f"Turn 1: {result1['response']}")

# Turn 2 (with context from turn 1)
result2 = await agent.handle_turn(
    session_id="sess-123",
    user_message="It's order number ORD-12345",
)
print(f"Turn 2: {result2['response']}")

# Turn 3
result3 = await agent.handle_turn(
    session_id="sess-123",
    user_message="Can I get a refund?",
)
print(f"Turn 3: {result3['response']}")

class StateRecovery:
    """Handle state recovery and rollback."""
    
    def __init__(self, state_manager: StateManager):
        self.state_manager = state_manager
    
    async def save_checkpoint(
        self,
        state: AgentState,
        checkpoint_name: str = "auto",
    ) -> str:
        """Save named checkpoint."""
        import uuid
        
        checkpoint_id = str(uuid.uuid4())
        
        await self.state_manager.storage.save_checkpoint(
            checkpoint_id=checkpoint_id,
            session_id=state.session_id,
            state=state,
            name=checkpoint_name,
        )
        
        return checkpoint_id
    
    async def restore_checkpoint(
        self,
        checkpoint_id: str,
    ) -> AgentState:
        """Restore from checkpoint."""
        state = await self.state_manager.storage.load_checkpoint(checkpoint_id)
        
        if not state:
            raise ValueError(f"Checkpoint {checkpoint_id} not found")
        
        # Restore to state manager
        self.state_manager.active_states[state.session_id] = state
        
        return state
    
    async def rollback_to_turn(
        self,
        session_id: str,
        turn_number: int,
    ) -> AgentState:
        """Rollback conversation to specific turn."""
        state = await self.state_manager.get_state(session_id)
        
        if not state:
            raise ValueError(f"Session {session_id} not found")
        
        # Truncate conversation history
        messages_to_keep = turn_number * 2  # user + assistant per turn
        state.conversation_history = state.conversation_history[:messages_to_keep]
        
        # Checkpoint rolled-back state
        await self.state_manager.checkpoint(state)
        
        return state
    
    async def list_checkpoints(
        self,
        session_id: str,
    ) -> List[Dict[str, Any]]:
        """List available checkpoints."""
        return await self.state_manager.storage.list_checkpoints(session_id)

# Usage
recovery = StateRecovery(state_manager)

# Save checkpoint before risky operation
state = await state_manager.get_state("sess-123")
checkpoint_id = await recovery.save_checkpoint(state, "before_refund")

try:
    # Attempt risky operation
    await process_refund(state)
except Exception as e:
    print(f"Operation failed: {e}")
    
    # Rollback to checkpoint
    state = await recovery.restore_checkpoint(checkpoint_id)
    print("State restored to before refund attempt")

# Rollback conversation
await recovery.rollback_to_turn("sess-123", turn_number=2)

# PostgreSQL backend
from typing import Protocol
import asyncpg
import json

class StateStorage(Protocol):
    """State storage interface."""
    
    async def save_state(self, state: AgentState) -> None:
        ...
    
    async def load_state(self, session_id: str) -> Optional[AgentState]:
        ...

class PostgreSQLStorage:
    """PostgreSQL storage backend."""
    
    def __init__(self, connection_string: str):
        self.connection_string = connection_string
        self.pool = None
    
    async def initialize(self) -> None:
        """Initialize connection pool."""
        self.pool = await asyncpg.create_pool(self.connection_string)
        
        # Create tables
        await self._create_tables()
    
    async def _create_tables(self) -> None:
        """Create database tables."""
        async with self.pool.acquire() as conn:
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS agent_states (
                    session_id VARCHAR(255) PRIMARY KEY,
                    state_id VARCHAR(255),
                    agent_id VARCHAR(255),
                    user_id VARCHAR(255),
                    conversation_history JSONB,
                    context JSONB,
                    tool_results JSONB,
                    metadata JSONB,
                    created_at TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)
            
            await conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_user_id 
                ON agent_states(user_id)
            """)
    
    async def save_state(self, state: AgentState) -> None:
        """Save state to PostgreSQL."""
        async with self.pool.acquire() as conn:
            await conn.execute("""
                INSERT INTO agent_states (
                    session_id, state_id, agent_id, user_id,
                    conversation_history, context, tool_results, metadata,
                    created_at
                ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
                ON CONFLICT (session_id) 
                DO UPDATE SET
                    conversation_history = $5,
                    context = $6,
                    tool_results = $7,
                    metadata = $8,
                    updated_at = CURRENT_TIMESTAMP
            """,
                state.session_id,
                state.state_id,
                state.agent_id,
                state.user_id,
                json.dumps(state.conversation_history),
                json.dumps(state.context),
                json.dumps(state.tool_results),
                json.dumps(state.metadata),
                state.created_at,
            )
    
    async def load_state(self, session_id: str) -> Optional[AgentState]:
        """Load state from PostgreSQL."""
        async with self.pool.acquire() as conn:
            row = await conn.fetchrow(
                "SELECT * FROM agent_states WHERE session_id = $1",
                session_id,
            )
            
            if not row:
                return None
            
            return AgentState(
                state_id=row["state_id"],
                session_id=row["session_id"],
                agent_id=row["agent_id"],
                user_id=row["user_id"],
                created_at=row["created_at"].isoformat(),
                conversation_history=json.loads(row["conversation_history"]),
                context=json.loads(row["context"]),
                tool_results=json.loads(row["tool_results"]),
                metadata=json.loads(row["metadata"]),
            )
    
    async def delete_state(self, session_id: str) -> None:
        """Delete state."""
        async with self.pool.acquire() as conn:
            await conn.execute(
                "DELETE FROM agent_states WHERE session_id = $1",
                session_id,
            )
    
    async def list_sessions(self, user_id: str) -> List[str]:
        """List sessions for user."""
        async with self.pool.acquire() as conn:
            rows = await conn.fetch(
                "SELECT session_id FROM agent_states WHERE user_id = $1 ORDER BY updated_at DESC",
                user_id,
            )
            
            return [row["session_id"] for row in rows]

# Redis backend for high-throughput
import redis.asyncio as redis

class RedisStorage:
    """Redis storage backend (fast, ephemeral)."""
    
    def __init__(self, redis_url: str):
        self.redis_url = redis_url
        self.client = None
    
    async def initialize(self) -> None:
        """Initialize Redis client."""
        self.client = await redis.from_url(self.redis_url)
    
    async def save_state(self, state: AgentState) -> None:
        """Save state to Redis."""
        key = f"agent_state:{state.session_id}"
        
        # Serialize state
        data = json.dumps({
            "state_id": state.state_id,
            "session_id": state.session_id,
            "agent_id": state.agent_id,
            "user_id": state.user_id,
            "created_at": state.created_at,
            "conversation_history": state.conversation_history,
            "context": state.context,
            "tool_results": state.tool_results,
            "metadata": state.metadata,
        })
        
        # Save with TTL (24 hours)
        await self.client.setex(key, 86400, data)
    
    async def load_state(self, session_id: str) -> Optional[AgentState]:
        """Load state from Redis."""
        key = f"agent_state:{session_id}"
        
        data = await self.client.get(key)
        
        if not data:
            return None
        
        state_dict = json.loads(data)
        
        return AgentState(**state_dict)

# Usage: Choose backend based on requirements
# PostgreSQL: Persistent, auditable, supports complex queries
# Redis: Fast, ephemeral (with TTL), high throughput

postgres_storage = PostgreSQLStorage("postgresql://user:pass@localhost/agents")
await postgres_storage.initialize()

# or

redis_storage = RedisStorage("redis://localhost:6379")
await redis_storage.initialize()

state_manager = StateManager(storage_backend=postgres_storage)

class OptimizedStateManager(StateManager):
    """Optimized state manager with caching and compression."""
    
    def __init__(self, storage_backend, cache_backend):
        super().__init__(storage_backend)
        self.cache = cache_backend
        self.compression_enabled = True
    
    async def get_state(self, session_id: str) -> Optional[AgentState]:
        """Get state with caching."""
        # Check memory cache
        if session_id in self.active_states:
            return self.active_states[session_id]
        
        # Check distributed cache (Redis)
        cached = await self.cache.get(f"state:{session_id}")
        if cached:
            state = self._deserialize(cached)
            self.active_states[session_id] = state
            return state
        
        # Load from storage
        state = await self.storage.load_state(session_id)
        
        if state:
            # Cache for fast retrieval
            await self.cache.setex(
                f"state:{session_id}",
                3600,  # 1 hour TTL
                self._serialize(state),
            )
            
            self.active_states[session_id] = state
        
        return state
    
    async def checkpoint(self, state: AgentState) -> None:
        """Checkpoint with compression."""
        # Compress large conversation histories
        if self.compression_enabled and len(state.conversation_history) > 20:
            state = self._compress_history(state)
        
        # Save to storage
        await super().checkpoint(state)
        
        # Update cache
        await self.cache.setex(
            f"state:{state.session_id}",
            3600,
            self._serialize(state),
        )
    
    def _compress_history(self, state: AgentState) -> AgentState:
        """Compress old conversation history."""
        # Keep recent 10 messages, summarize older ones
        if len(state.conversation_history) > 10:
            recent = state.conversation_history[-10:]
            old = state.conversation_history[:-10]
            
            # Store summary in context
            state.context["history_summary"] = f"Previous conversation ({len(old)} messages)"
            
            state.conversation_history = recent
        
        return state
    
    def _serialize(self, state: AgentState) -> str:
        """Serialize state."""
        import json
        return json.dumps(state.__dict__)
    
    def _deserialize(self, data: str) -> AgentState:
        """Deserialize state."""
        import json
        return AgentState(**json.loads(data))

# Usage
optimized_manager = OptimizedStateManager(
    storage_backend=postgres_storage,
    cache_backend=redis_client,
)

from dataclasses import dataclass

@dataclass
class AgentStateV2(AgentState):
    """Version 2 of agent state with new fields."""
    version: int = 2
    tags: List[str] = field(default_factory=list)
    sentiment_history: List[str] = field(default_factory=list)

class StateMigrator:
    """Migrate state between versions."""
    
    async def migrate_v1_to_v2(self, state_v1: AgentState) -> AgentStateV2:
        """Migrate from v1 to v2."""
        state_v2 = AgentStateV2(
            state_id=state_v1.state_id,
            session_id=state_v1.session_id,
            agent_id=state_v1.agent_id,
            user_id=state_v1.user_id,
            created_at=state_v1.created_at,
            conversation_history=state_v1.conversation_history,
            context=state_v1.context,
            tool_results=state_v1.tool_results,
            metadata=state_v1.metadata,
            version=2,
            tags=[],  # New field
            sentiment_history=[],  # New field
        )
        
        return state_v2
    
    async def detect_version(self, state_dict: Dict[str, Any]) -> int:
        """Detect state version."""
        return state_dict.get("version", 1)

# Usage
migrator = StateMigrator()

# Load state
state_dict = await storage.load_state_dict("sess-123")
version = await migrator.detect_version(state_dict)

if version == 1:
    state_v1 = AgentState(**state_dict)
    state_v2 = await migrator.migrate_v1_to_v2(state_v1)
    
    # Save migrated state
    await storage.save_state(state_v2)