🎯 What You'll Learn Today

LangGraph Tutorial: Advanced Message Processing with State Management - Unit 1.2 Exercise 2

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 demonstrates how to implement sophisticated message processing in LangGraph using TypedDict for state management and proper annotations for message handling. We'll create a message processor that maintains conversation context and demonstrates state field preservation.

Key Concepts Covered

1. State Management with Multiple Fields
2. Message Processing Logic
3. Type Safety with TypedDict
4. Field Preservation in State Updates
5. Proper Message Annotation Usage

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

Step 1: Enhanced State Definition

We define a state structure that combines message handling with additional context fields for richer conversation management.

class State(TypedDict):
    """Advanced state container for rich conversation context.

    This state implementation demonstrates three key concepts:
    1. Message management with proper LangGraph annotations
    2. Conversation summarization capability
    3. Configurable context window management

    Attributes:
        messages: List of conversation messages with LangGraph's add_messages
                 annotation for proper message propagation
        summary: Running summary of the conversation context
        window_size: Number of previous messages to maintain in context

    Note:
        The add_messages annotation is crucial - it enables LangGraph to properly
        handle message additions and updates throughout the conversation flow.
    """

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

Why This Matters

Advanced state management with multiple fields enables

1. Better conversation context tracking
2. More sophisticated AI agent behaviors
3. Improved error recovery through state preservation
4. Flexible configuration of conversation parameters

Step 2: Message Processing Implementation

We implement a stateful message processor that demonstrates proper state management and field preservation.

def process_message(state: State) -> State:
    """Process messages while maintaining conversation context and state.

    This processor demonstrates several advanced concepts:
    1. Full state preservation across updates
    2. Conditional response generation
    3. Context-aware message handling

    The processing follows this flow:
    1. Handle empty state initialization
    2. Process existing messages with context
    3. Preserve non-message state fields

    Args:
        state: Current conversation state containing messages,
              summary, and configuration

    Returns:
        State: Updated state with new messages and preserved fields

    Example:
        >>> initial_state = {"messages": [], "summary": "", "window_size": 3}
        >>> new_state = process_message(initial_state)
        >>> print(new_state["messages"][0].content)
        "Hello!"
    """
    # Initialize state if empty
    if not state["messages"]:
        return {
            "messages": [HumanMessage(content="Hello!")],
            "summary": "",
            "window_size": 3,
        }

    # Process last message with context
    last_message = state["messages"][-1]
    if last_message.content == "Hello!":
        return {
            "messages": [HumanMessage(content="How are you?")],
            "summary": state["summary"],
            "window_size": state["window_size"],
        }

    # Maintain state if no conditions met
    return state

Debug Tips

1. Message Processing Issues:

   • Print state before and after processing to verify field preservation
   • Check message list length matches window_size configuration
   • Verify message content types (HumanMessage vs AIMessage)

2. State Management:

   • Ensure all required fields are present in state updates
   • Validate summary updates maintain proper format
   • Monitor window_size compliance

3. Common Errors:

   • KeyError: Usually indicates missing state fields
   • TypeError: Often means incorrect message type usage
   • IndexError: Check empty message list handling

Key Takeaways

1. State Management:

   • Always preserve all state fields during updates
   • Use proper type annotations for safety
   • Maintain consistency in field updates

2. Message Processing:

   • Handle empty states explicitly
   • Process messages with full context
   • Implement clear processing logic

Common Pitfalls

1. Forgetting to preserve non-message state fields
2. Incorrect message type usage
3. Missing empty state handling
4. Improper annotation usage

Next Steps

1. Add error handling and recovery
2. Implement summary generation logic
3. Add window size enforcement
4. Integrate with external message processors

Example Usage

initial_state = {"messages": [], "summary": "", "window_size": 3}

processed_state = process_message(initial_state)

print(f"Initial message: {processed_state['messages'][0].content}")

Variations and Extensions

1. Enhanced State Management:

   • Add message timestamps for temporal tracking
   • Implement message priority queuing Example use case: Time-sensitive conversation handling

2. Advanced Processing:

   • Add message validation logic
   • Implement custom message types Scenario: Multi-modal conversation handling

Expected Output

Initial message: Hello!

Rod Rivera

🇬🇧 Chapter