LangChain to LangGraph Migration: Enterprise Stateful AI Orchestration Guide

# Hybrid Migration Wrapper Pattern
from langchain.agents import initialize_agent
from langgraph.graph import StateGraph
from typing import TypedDict, Dict, Any

class HybridAgentState(TypedDict):
    input: str
    output: str
    metadata: Dict[str, Any]

class HybridMigrationWrapper:
    def __init__(self, langchain_agent, enable_langgraph=False):
        self.langchain_agent = langchain_agent
        self.enable_langgraph = enable_langgraph
        self.graph = self._build_langgraph() if enable_langgraph else None
    
    def invoke(self, state: HybridAgentState) -> HybridAgentState:
        # Feature flag controlled migration
        if self.enable_langgraph and self.graph:
            try:
                # Attempt LangGraph execution
                result = self.graph.invoke(state)
                return result
            except Exception as e:
                # Fallback to LangChain on failure
                print(f"LangGraph failed, falling back to LangChain: {e}")
                return self._langchain_fallback(state)
        else:
            # Original LangChain execution
            return self._langchain_fallback(state)
    
    def _langchain_fallback(self, state: HybridAgentState) -> HybridAgentState:
        result = self.langchain_agent.run(state["input"])
        return {"input": state["input"], "output": result, "metadata": {"engine": "langchain"}}
    
    def _build_langgraph(self):
        # Build equivalent LangGraph workflow
        workflow = StateGraph(HybridAgentState)
        # Add nodes and edges for LangGraph implementation
        return workflow.compile()

# BEFORE: Deprecated LangChain initialize_agent
from langchain.agents import initialize_agent, AgentType
from langchain.chat_models import ChatOpenAI
from langchain.tools import Tool

def create_langchain_agent():
    llm = ChatOpenAI(temperature=0, model="gpt-4")
    tools = [
        Tool(name="calculator", func=lambda x: eval(x), description="Perform calculations"),
        Tool(name="search", func=lambda x: f"Search results for: {x}", description="Search knowledge base")
    ]
    
    # Deprecated initialize_agent API
    agent = initialize_agent(
        tools=tools,
        llm=llm,
        agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
        verbose=True,
        return_intermediate_steps=True
    )
    return agent

# AFTER: Modern LangGraph StateGraph Implementation
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import create_react_agent
from typing import TypedDict, List, Dict, Any
from langchain_openai import ChatOpenAI
from langchain.tools import tool

class AgentState(TypedDict):
    input: str
    output: str
    intermediate_steps: List[Dict[str, Any]]
    metadata: Dict[str, Any]

@tool
def calculator(expression: str) -> str:
    """Perform mathematical calculations."""
    try:
        result = eval(expression)
        return f"Calculation result: {result}"
    except Exception as e:
        return f"Calculation error: {str(e)}"

@tool
def search(query: str) -> str:
    """Search knowledge base for information."""
    return f"Search results for query '{query}': Found relevant information in company knowledge base."

def create_langgraph_agent():
    llm = ChatOpenAI(temperature=0, model="gpt-4")
    tools = [calculator, search]
    
    # Create stateful React agent with built-in checkpointing
    agent = create_react_agent(llm, tools)
    
    # Wrap in StateGraph for additional control
    workflow = StateGraph(AgentState)
    workflow.add_node("agent", agent)
    workflow.add_edge("agent", END)
    workflow.set_entry_point("agent")
    
    # Compile with memory for state persistence
    from langgraph.checkpoint.memory import MemorySaver
    memory = MemorySaver()
    app = workflow.compile(checkpointer=memory)
    
    return app

# Usage comparison
print("LangChain (Deprecated):")
langchain_agent = create_langchain_agent()
result = langchain_agent.run("Calculate 15% of 240 and search for financial formulas")
print(result)

print("\nLangGraph (Modern):")
langgraph_agent = create_langgraph_agent()
config = {"configurable": {"thread_id": "enterprise-thread-001"}}
result = langgraph_agent.invoke({
    "input": "Calculate 15% of 240 and search for financial formulas"
}, config)
print(result["output"])

# Production MemorySaver Configuration
from langgraph.checkpoint.memory import MemorySaver
from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.checkpoint.postgres import PostgresSaver
import os
from datetime import timedelta

class EnterpriseMemoryConfig:
    """Production-ready MemorySaver configurations for enterprise deployments."""
    
    @staticmethod
    def development_config():
        """Local development with SQLite persistence."""
        return SqliteSaver.from_conn_string("sqlite:///./agent_memory.db")
    
    @staticmethod
    def production_config():
        """High-availability production with PostgreSQL."""
        connection_string = os.getenv("DATABASE_URL", "postgresql://user:pass@localhost:5432/ai_agents")
        return PostgresSaver.from_conn_string(connection_string)
    
    @staticmethod
    def compliance_config():
        """Compliance-mode with encryption and audit logging."""
        from cryptography.fernet import Fernet
        import logging
        
        # Initialize encryption
        encryption_key = os.getenv("MEMORY_ENCRYPTION_KEY", Fernet.generate_key())
        cipher = Fernet(encryption_key)
        
        # Configure audit logging
        audit_logger = logging.getLogger("agent_memory_audit")
        audit_logger.setLevel(logging.INFO)
        
        # Custom MemorySaver with encryption
        class EncryptedPostgresSaver(PostgresSaver):
            def put(self, config, checkpoint, metadata):
                # Encrypt checkpoint data before storage
                encrypted_checkpoint = cipher.encrypt(str(checkpoint).encode())
                super().put(config, encrypted_checkpoint, metadata)
                audit_logger.info(f"Checkpoint stored for thread: {config.get('thread_id', 'unknown')}")
            
            def get(self, config):
                # Decrypt checkpoint data after retrieval
                encrypted_data = super().get(config)
                if encrypted_data:
                    decrypted_data = cipher.decrypt(encrypted_data).decode()
                    audit_logger.info(f"Checkpoint retrieved for thread: {config.get('thread_id', 'unknown')}")
                    return eval(decrypted_data)  # Convert back to dict
                return None
        
        connection_string = os.getenv("DATABASE_URL")
        return EncryptedPostgresSaver.from_conn_string(connection_string)

# Implementation with enterprise configuration
def create_enterprise_agent():
    from langgraph.prebuilt import create_react_agent
    from langchain_openai import ChatOpenAI
    
    llm = ChatOpenAI(temperature=0, model="gpt-4")
    tools = [calculator, search]  # Tools from previous example
    
    # Select memory configuration based on environment
    environment = os.getenv("ENVIRONMENT", "development")
    if environment == "production":
        memory = EnterpriseMemoryConfig.production_config()
    elif environment == "compliance":
        memory = EnterpriseMemoryConfig.compliance_config()
    else:
        memory = EnterpriseMemoryConfig.development_config()
    
    # Create agent with enterprise memory
    agent = create_react_agent(llm, tools, checkpointer=memory)
    return agent

# Conversation threading example
def demonstrate_conversation_threading():
    agent = create_enterprise_agent()
    
    # Customer service conversation with persistence
    thread_config = {"configurable": {"thread_id": "customer_12345"}}
    
    # First interaction
    result1 = agent.invoke({
        "input": "My name is Sarah Johnson and I'm having issues with order #ABC123"
    }, thread_config)
    print(f"Agent remembers: {result1['output']}")
    
    # Second interaction (hours or days later)
    result2 = agent.invoke({
        "input": "What's the status of my order again?"
    }, thread_config)
    print(f"Context preserved: {result2['output']}")
    
    # Audit trail for compliance
    return f"Conversation thread customer_12345 has full audit trail and state persistence"

# Human-in-the-Loop Approval Workflow
from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import MemorySaver
from typing import TypedDict, Optional
import asyncio
from datetime import datetime, timedelta

class ApprovalState(TypedDict):
    request_type: str
    request_data: dict
    risk_score: float
    approver_id: Optional[str]
    approval_status: str  # "pending", "approved", "rejected", "timeout"
    approval_timestamp: Optional[str]
    requester_justification: str
    approver_comments: Optional[str]

def risk_assessment_node(state: ApprovalState) -> ApprovalState:
    """AI-powered risk assessment for financial transactions."""
    risk_factors = [
        state["request_data"].get("amount", 0) > 50000,  # High amount
        state["request_data"].get("currency") != "USD",   # Foreign currency
        state["request_data"].get("urgent", False)       # Urgent request
    ]
    
    risk_score = sum(risk_factors) * 0.33  # Simple scoring
    return {**state, "risk_score": risk_score}

def approval_decision_node(state: ApprovalState) -> ApprovalState:
    """Determine if human approval is required based on risk score."""
    if state["risk_score"] > 0.5:
        return {**state, "approval_status": "pending"}
    else:
        return {**state, "approval_status": "auto_approved"}

def human_approval_node(state: ApprovalState) -> ApprovalState:
    """Interrupt workflow for human approval with timeout handling."""
    print(f"🔄 Approval required for {state['request_type']}")
    print(f"Risk Score: {state['risk_score']:.2f}")
    print(f"Request: {state['request_data']}")
    print(f"Justification: {state['requester_justification']}")
    
    # Simulate human approval interface
    # In production, this would integrate with Slack, Teams, or custom UI
    try:
        # Wait for human input with timeout
        print("⏳ Waiting for approval (30 seconds timeout)...")
        approval_result = asyncio.run(wait_for_human_approval(timeout=30))
        return {**state, **approval_result}
    except asyncio.TimeoutError:
        print("⚠️ Approval timeout - escalating to manager")
        return {**state, "approval_status": "timeout", "approver_comments": "Auto-escalated due to timeout"}

async def wait_for_human_approval(timeout: int) -> dict:
    """Simulate human approval interface with timeout."""
    try:
        # In production, integrate with enterprise notification systems
        # For demo, simulate approval after delay
        await asyncio.sleep(5)  # Simulate approver response time
        
        # Simulate approval decision (in reality, comes from UI/API)
        return {
            "approval_status": "approved",
            "approver_id": "manager_001",
            "approver_comments": "Approved after review",
            "approval_timestamp": datetime.now().isoformat()
        }
    except asyncio.TimeoutError:
        raise

def final_processing_node(state: ApprovalState) -> ApprovalState:
    """Complete processing based on approval decision."""
    if state["approval_status"] == "approved":
        print(f"✅ Request approved by {state['approver_id']}")
        # Process the approved request
        return {**state, "processed": True}
    elif state["approval_status"] == "rejected":
        print(f"❌ Request rejected: {state['approver_comments']}")
        return {**state, "processed": False}
    else:
        print(f"⚠️ Request status: {state['approval_status']}")
        return {**state, "processed": False}

# Build approval workflow
workflow = StateGraph(ApprovalState)
workflow.add_node("risk_assessment", risk_assessment_node)
workflow.add_node("approval_decision", approval_decision_node)
workflow.add_node("human_approval", human_approval_node)
workflow.add_node("final_processing", final_processing_node)

workflow.add_edge("risk_assessment", "approval_decision")
workflow.add_edge("approval_decision", "human_approval")
workflow.add_edge("human_approval", "final_processing")
workflow.add_edge("final_processing", END)
workflow.set_entry_point("risk_assessment")

# Compile with memory for audit trail
memory = MemorySaver()
approval_app = workflow.compile(checkpointer=memory)

# Example usage
def demonstrate_approval_workflow():
    """Demonstrate human-in-the-loop approval workflow."""
    config = {"configurable": {"thread_id": "approval_001"}}
    
    # High-risk financial transaction request
    result = approval_app.invoke({
        "request_type": "wire_transfer",
        "request_data": {
            "amount": 75000,
            "currency": "EUR",
            "recipient": "International Supplier",
            "urgent": True
        },
        "requester_justification": "Critical supplier payment for production materials"
    }, config)
    
    print(f"Final approval status: {result['approval_status']}")
    print(f"Complete audit trail available for thread: approval_001")
    return result

# Enterprise Security Configuration
import os
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
import base64
from datetime import datetime, timedelta

class EnterpriseSecurityConfig:
    """Production security configuration for LangGraph deployments."""
    
    def __init__(self):
        self.encryption_key = self._derive_key_from_password()
        self.cipher = Fernet(self.encryption_key)
    
    def _derive_key_from_password(self):
        """Derive encryption key from enterprise password manager."""
        password = os.getenv("ENTERPRISE_MASTER_KEY").encode()
        salt = os.getenv("KEY_SALT").encode()
        
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100000,
        )
        key = base64.urlsafe_b64encode(kdf.derive(password))
        return key
    
    def encrypt_state_data(self, data: dict) -> bytes:
        """Encrypt workflow state data before storage."""
        serialized_data = str(data).encode()
        encrypted_data = self.cipher.encrypt(serialized_data)
        return encrypted_data
    
    def decrypt_state_data(self, encrypted_data: bytes) -> dict:
        """Decrypt workflow state data after retrieval."""
        decrypted_data = self.cipher.decrypt(encrypted_data)
        return eval(decrypted_data.decode())
    
    def create_audit_log_entry(self, action: str, user_id: str, details: dict):
        """Create tamper-resistant audit log entry."""
        timestamp = datetime.utcnow().isoformat()
        log_entry = {
            "timestamp": timestamp,
            "action": action,
            "user_id": user_id,
            "details": details,
            "integrity_hash": self._calculate_integrity_hash(timestamp, action, user_id)
        }
        # In production, write to immutable audit system
        return log_entry
    
    def _calculate_integrity_hash(self, timestamp: str, action: str, user_id: str) -> str:
        """Calculate integrity hash for audit trail verification."""
        import hashlib
        data = f"{timestamp}:{action}:{user_id}:{self.encryption_key}".encode()
        return hashlib.sha256(data).hexdigest()

# Secure MemorySaver implementation
class SecureMemorySaver:
    """MemorySaver with enterprise-grade security features."""
    
    def __init__(self, security_config: EnterpriseSecurityConfig):
        self.security_config = security_config
        from langgraph.checkpoint.memory import MemorySaver
        self.base_saver = MemorySaver()
    
    def put(self, config, checkpoint, metadata):
        """Securely store checkpoint with encryption and audit trail."""
        # Encrypt checkpoint data
        encrypted_checkpoint = self.security_config.encrypt_state_data(checkpoint)
        
        # Store encrypted data
        self.base_saver.put(config, encrypted_checkpoint, metadata)
        
        # Create audit log entry
        user_id = config.get("user_id", "system")
        self.security_config.create_audit_log_entry(
            "CHECKPOINT_STORE",
            user_id,
            {"thread_id": config.get("thread_id"), "checkpoint_size": len(str(checkpoint))}
        )
    
    def get(self, config):
        """Retrieve and decrypt checkpoint data."""
        encrypted_data = self.base_saver.get(config)
        if encrypted_data:
            # Decrypt checkpoint data
            decrypted_checkpoint = self.security_config.decrypt_state_data(encrypted_data)
            
            # Create audit log for access
            user_id = config.get("user_id", "system")
            self.security_config.create_audit_log_entry(
                "CHECKPOINT_RETRIEVE",
                user_id,
                {"thread_id": config.get("thread_id")}
            )
            return decrypted_checkpoint
        return None

# GDPR-compliant data retention
class GDPRCompliantRetention:
    """Automated data lifecycle management for GDPR compliance."""
    
    def __init__(self, retention_days: int = 90):
        self.retention_days = retention_days
    
    def should_delete_data(self, created_timestamp: str) -> bool:
        """Determine if data should be deleted based on retention policy."""
        created_date = datetime.fromisoformat(created_timestamp)
        retention_date = created_date + timedelta(days=self.retention_days)
        return datetime.utcnow() > retention_date
    
    def create_deletion_log(self, user_id: str, data_type: str, deletion_reason: str):
        """Create audit log for data deletion events."""
        return {
            "timestamp": datetime.utcnow().isoformat(),
            "action": "DATA_DELETION",
            "user_id": user_id,
            "data_type": data_type,
            "reason": deletion_reason,
            "compliance_basis": "GDPR Article 17 - Right to erasure"
        }