Logo of AI Product Engineer
AI Product Engineer
Tutorial Image: LangGraph Tutorial: Advanced State Management with Extended Fields - Unit 1.2 Exercise 1

LangGraph Tutorial: Advanced State Management with Extended Fields - Unit 1.2 Exercise 1

Explore advanced state management in LangGraph with extended fields for context tracking and memory control. Learn how to enhance your state structure for robust conversational agents.

🎯 What You'll Learn Today

LangGraph Tutorial: Advanced State Management with Extended Fields - Unit 1.2 Exercise 1

This tutorial is also available in Google Colab here or for download here

Joint Initiative: This tutorial is part of a collaboration between AI Product Engineer and the Nebius Academy.

This tutorial explores advanced state management in LangGraph by extending the state structure beyond simple message tracking. We'll learn how to create a sophisticated state that maintains conversation context and controls memory usage.

Key Concepts Covered

  1. Extended State Fields
  2. Message Annotation
  3. Memory Window Control
  4. Context Management

Let's break it down step by step:

from typing import Annotated, TypedDict
!pip install langchain-core
!pip install langgraph
from langchain_core.messages import BaseMessage
from langgraph.graph.message import add_messages

Step 1: Enhanced State Definition

We define a more sophisticated State class that goes beyond basic message tracking to include conversation context and memory management.

class State(TypedDict):
    """Enhanced state container with context management capabilities.

    This implementation demonstrates advanced state management by tracking:
    1. Message history with proper LangGraph annotations
    2. Conversation summaries for context retention
    3. Memory window control for efficient processing

    Attributes:
        messages: List of conversation messages with LangGraph's add_messages
                 annotation for proper message handling
        summary: Running summary of the conversation context
        window_size: Control parameter for message history retention

    Note:
        The add_messages annotation is crucial for proper message
        handling in LangGraph, ensuring correct state updates while
        the additional fields enable sophisticated conversation management.
    """

    messages: Annotated[list[BaseMessage], add_messages]
    summary: str
    window_size: int

Step 2: State Initialization

Demonstrate proper state initialization with all required fields.

Initialize state with default values

state: State = {
    "messages": [],  # Empty message history

    "summary": "",  # No initial summary

    "window_size": 3,  # Keep last 3 messages in memory

}
Display initial state
print(state)  # Output: {'messages': [], 'summary': '', 'window_size': 3}

Key Takeaways

  1. Extended State: State can track more than just messages
  2. Context Management: Additional fields enable sophisticated tracking
  3. Memory Control: Window size parameter enables memory management
  4. Type Safety: TypedDict ensures proper field typing

Common Pitfalls to Avoid

  1. Forgetting to initialize all required fields
  2. Using incorrect types for fields
  3. Neglecting to update all state fields consistently
  4. Mishandling the add_messages annotation

Next Steps

  • Implement summary updates
  • Add message window management
  • Create context-aware processing
  • Add state validation logic

Example Extended Usage

from langchain_core.messages import HumanMessage
Add a message
state['messages'].append(HumanMessage(content="Hello!"))
Update summary
state['summary'] = "Greeting initiated"
Manage message window
if len(state['messages']) > state['window_size']:
   state['messages'] = state['messages'][-state['window_size']:]

Rod Rivera

🇬🇧 Chapter