🎯 What You'll Learn Today
LangGraph Tutorial: Rate Limiting Implementation - Unit 2.2 Exercise 3
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 rate limiting in LangGraph applications to control tool usage and prevent abuse.
Key Concepts Covered
- State Management with TypedDict
- Rate Limit Tracking
- Usage Monitoring
- State Immutability Patterns
from typing import Annotated, TypedDict
!pip install langchain-core
!pip install langgraph
from langchain_core.messages import AIMessage, BaseMessage
from langgraph.graph.message import add_messagesStep 1: State Structure Definition
We define our state structure to track messages, tool usage, and rate limits.
Why This Matters
State structure is crucial because
- Enables tracking of tool usage across multiple interactions
- Provides a foundation for enforcing rate limits
- Maintains system stability through usage controls
Debug Tips
-
State Initialization Issues:
- Verify all dictionaries are properly initialized
- Check for correct type annotations
- Ensure rate_limits contains expected tools
class State(TypedDict):
"""State container with rate limiting capabilities.
Attributes:
messages: List of interaction messages
tool_usage: Dictionary tracking number of times each tool is used
rate_limits: Dictionary defining maximum uses allowed for each tool
"""
messages: Annotated[list[BaseMessage], add_messages]
tool_usage: dict[str, int]
rate_limits: dict[str, int]Step 2: Rate Limit Verification
Implement the logic to check if a tool has exceeded its rate limit.
Why This Matters
Rate limit checking is essential because
- Prevents abuse of expensive or limited resources
- Ensures fair system usage
- Protects external API quotas
Debug Tips
-
Rate Checking Issues:
- Add logging for limit checks
- Verify default values for new tools
- Check comparison logic
def check_rate_limit(state: State, tool_name: str) -> bool:
"""Check if tool has exceeded its rate limit.
Args:
state: Current application state
tool_name: Name of the tool to check
Returns:
bool: True if tool hasn't exceeded its limit, False otherwise
"""
usage = state["tool_usage"].get(tool_name, 0)
limit = state["rate_limits"].get(tool_name, float("inf"))
return usage < limitStep 3: Usage Tracking Implementation
Implement the logic to update tool usage counts and handle rate limit violations.
Why This Matters
Usage tracking is critical because
- Maintains accurate usage statistics
- Enables proactive rate limit enforcement
- Provides audit trail for tool usage
Debug Tips
-
State Update Issues:
- Verify state immutability
- Check dictionary updates
- Monitor message addition
def update_usage(state: State, tool_name: str) -> State:
"""Update tool usage counts and handle rate limit violations.
Args:
state: Current application state
tool_name: Name of the tool being used
Returns:
State: Updated state with new usage counts or rate limit message
"""
if check_rate_limit(state, tool_name):
return {
**state,
"tool_usage": {
**state["tool_usage"],
tool_name: state["tool_usage"].get(tool_name, 0) + 1,
},
}
return {
**state,
"messages": [AIMessage(content=f"Rate limit exceeded for {tool_name}")],
}Step 4: Implementation Example
Demonstrate practical usage of the rate limiting system.
Why This Matters
Example implementation helps
- Verify system functionality
- Demonstrate expected behavior
- Provide usage patterns
Debug Tips
-
Testing Issues:
- Test edge cases
- Verify limit boundaries
- Check message generation
Initialize state with example values
state = {
"messages": [],
"tool_usage": {"calculator": 2},
"rate_limits": {"calculator": 3},
}Test rate limit checking
calculator_allowed = check_rate_limit(state, "calculator")
print(f"Calculator usage allowed: {calculator_allowed}")Test usage update
updated_state = update_usage(state, "calculator")
print(f"Updated tool usage: {updated_state['tool_usage']}")Test rate limit exceeded
final_state = update_usage(updated_state, "calculator")
print(f"Final state messages: {final_state['messages']}")Common Pitfalls
- Forgetting to initialize tool_usage or rate_limits dictionaries
- Not handling new tools without defined limits
- Mutating state directly instead of creating new state
- Missing error handling for edge cases
Key Takeaways
- Rate limiting is essential for system stability
- Immutable state updates prevent side effects
- Clear error messages improve user experience
- Default values handle undefined tools gracefully
Next Steps
- Add persistent storage for usage tracking
- Implement time-based rate limiting
- Add usage analytics and reporting
- Implement rate limit recovery mechanisms
Expected Output
Calculator usage allowed: True
Updated tool usage: {'calculator': 3}
Final state messages: [AIMessage(content='Rate limit exceeded for calculator')]
Execute example if run directly
if __name__ == "__main__":
# Initialize state with example values
state = {
"messages": [],
"tool_usage": {"calculator": 2},
"rate_limits": {"calculator": 3},
}
# Test rate limit checking
calculator_allowed = check_rate_limit(state, "calculator")
print(f"Calculator usage allowed: {calculator_allowed}")
# Test usage update
updated_state = update_usage(state, "calculator")
print(f"Updated tool usage: {updated_state['tool_usage']}")
# Test rate limit exceeded
final_state = update_usage(updated_state, "calculator")
print(f"Final state messages: {final_state['messages']}")
🇬🇧 Chapter