๐ฏ What You'll Learn Today
LangGraph Tutorial: Advanced Message Processing with State Management - Unit 1.2 Exercise 2
Try It Yourself
๐ข Joint Initiative
This tutorial is part of a collaboration between AIPE and 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
- State Management with Multiple Fields
- Message Processing Logic
- Type Safety with TypedDict
- Field Preservation in State Updates
- 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
- Better conversation context tracking
- More sophisticated AI agent behaviors
- Improved error recovery through state preservation
- 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
-
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)
-
State Management:
- Ensure all required fields are present in state updates
- Validate summary updates maintain proper format
- Monitor window_size compliance
-
Common Errors:
- KeyError: Usually indicates missing state fields
- TypeError: Often means incorrect message type usage
- IndexError: Check empty message list handling
Key Takeaways
-
State Management:
- Always preserve all state fields during updates
- Use proper type annotations for safety
- Maintain consistency in field updates
-
Message Processing:
- Handle empty states explicitly
- Process messages with full context
- Implement clear processing logic
Common Pitfalls
- Forgetting to preserve non-message state fields
- Incorrect message type usage
- Missing empty state handling
- Improper annotation usage
Next Steps
- Add error handling and recovery
- Implement summary generation logic
- Add window size enforcement
- 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
-
Enhanced State Management:
- Add message timestamps for temporal tracking
- Implement message priority queuing Example use case: Time-sensitive conversation handling
-
Advanced Processing:
- Add message validation logic
- Implement custom message types Scenario: Multi-modal conversation handling
Expected Output
Initial message: Hello!