Logo of AI Product Engineer
AI Product Engineer
Tutorial Image: LangGraph Tutorial: Enhanced State Management for Multi-Tool Agents - Unit 2.2 Exercise 1

LangGraph Tutorial: Enhanced State Management for Multi-Tool Agents - Unit 2.2 Exercise 1

Learn how to implement enhanced state management for multi-tool agents in LangGraph. This tutorial covers creating a state structure for tool usage tracking, rate limiting, and type-safe updates, ensuring precise control over tool execution and clear status monitoring.

🎯 What You'll Learn Today

LangGraph Tutorial: Enhanced State Management for Multi-Tool Agents - Unit 2.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 demonstrates how to implement an advanced state management system for agents that can use multiple tools, with features like rate limiting and usage tracking.

Key Concepts Covered

  1. Advanced State Management
  2. Tool Usage Tracking
  3. Rate Limiting Implementation
  4. Type-Safe State Handling
from typing import Annotated, Any, TypedDict
#!pip install langchain-core
#!pip install langgraph
from langchain_core.messages import BaseMessage
from langgraph.graph.message import add_messages

Step 1: State Definition

We define our enhanced state structure for multi-tool operations.

Why This Matters

Advanced state structure is crucial because

  1. Enables controlled tool access
  2. Maintains usage history
  3. Enforces rate limits
  4. Supports type safety

Debug Tips

  1. State Structure:

    • Verify field initialization
    • Check type annotations
    • Monitor optional fields
    • Validate tool configurations
class State(TypedDict, total=False):
    """Enhanced state container for multi-tool operations.

    This state implementation uses TypedDict with total=False for:
    1. Optional field support
    2. Backward compatibility
    3. Type safety maintenance
    4. Clean state updates

    Attributes:
        messages: Conversation history with add_messages annotation
        available_tools: List of accessible tools
        tool_usage: Usage count per tool
        rate_limits: Maximum uses per tool
        tool_name: Currently selected tool (optional)
        tool_outputs: Tool execution results (optional)
    """

    messages: Annotated[list[BaseMessage], add_messages]
    available_tools: list[Any]
    tool_usage: dict[str, int]
    rate_limits: dict[str, int]
    tool_name: str | None
    tool_outputs: list[str]
class ToolLimitExceeded(Exception):
    """Exception for tool usage limit violations."""

    pass

Step 2: State Initialization

We implement the state initialization logic.

Why This Matters

Proper initialization is crucial because

  1. Ensures consistent starting state
  2. Sets up tool availability
  3. Configures rate limits
  4. Enables clean tracking

Debug Tips

  1. Initialization:

    • Verify default values
    • Check rate limit setup
    • Monitor tool registration
    • Validate state structure
def initialize_state(
    available_tools: list[str] = ["calculator", "weather"],
    rate_limits: dict[str, int] | None = None,
) -> State:
    """Initialize state with tools and limits.

    Args:
        available_tools: Tools to make available
        rate_limits: Optional tool usage limits

    Returns:
        Initialized State object
    """
    if rate_limits is None:
        rate_limits = {
            "calculator": 3,
            "weather": 1,
            "search": 2,
        }

    for tool in available_tools:
        if tool not in rate_limits:
            rate_limits[tool] = 1

    return {
        "messages": [],
        "available_tools": available_tools,
        "tool_usage": {tool: 0 for tool in available_tools},
        "rate_limits": rate_limits,
    }

Step 3: Tool Usage Management

We implement tool usage checking and tracking.

Why This Matters

Usage management is crucial because

  1. Prevents tool overuse
  2. Maintains rate limits
  3. Tracks usage patterns
  4. Enables monitoring

Debug Tips

  1. Usage Tracking:

    • Monitor usage counts
    • Verify limit checks
    • Track state updates
    • Check error handling
def can_use_tool(state: State, tool_name: str) -> bool:
    """Check tool availability based on usage and limits.

    Args:
        state: Current state
        tool_name: Tool to check

    Returns:
        Whether tool can be used
    """
    if tool_name not in state["available_tools"]:
        raise KeyError(f"Tool '{tool_name}' is not available")

    current_usage = state["tool_usage"][tool_name]
    limit = state["rate_limits"][tool_name]
    return current_usage < limit
def use_tool(state: State, tool_name: str) -> State:
    """Record tool usage in state.

    Args:
        state: Current state
        tool_name: Tool being used

    Returns:
        Updated state
    """
    if not can_use_tool(state, tool_name):
        limit = state["rate_limits"][tool_name]
        raise ToolLimitExceeded(
            f"Tool '{tool_name}' has reached its limit of {limit} uses"
        )

    new_state = state.copy()
    new_state["tool_usage"] = state["tool_usage"].copy()
    new_state["tool_usage"][tool_name] += 1
    new_state["tool_name"] = tool_name

    return new_state

Step 4: Status Reporting

We implement tool status monitoring.

Why This Matters

Status reporting is crucial because

  1. Enables usage monitoring
  2. Facilitates debugging
  3. Supports decision making
  4. Maintains transparency

Debug Tips

  1. Status Tracking:

    • Verify calculations
    • Check format consistency
    • Monitor updates
    • Validate accessibility
def get_tool_status(state: State) -> dict[str, dict[str, Any]]:
    """Get complete tool status summary.

    Args:
        state: Current state

    Returns:
        Tool status information
    """
    status = {}
    for tool in state["available_tools"]:
        uses_left = state["rate_limits"][tool] - state["tool_usage"][tool]
        status[tool] = {
            "current_usage": state["tool_usage"][tool],
            "limit": state["rate_limits"][tool],
            "uses_remaining": uses_left,
            "available": uses_left > 0,
        }
    return status
def demonstrate_usage():
    """Demonstrate the state management system."""
    state = initialize_state(
        available_tools=["calculator", "weather", "search"],
        rate_limits={"calculator": 2, "weather": 1, "search": 3},
    )

    print("\nInitial tool status:")
    print_tool_status(state)

    try:
        state = use_tool(state, "calculator")
        print("\nAfter using calculator once:")
        print_tool_status(state)

        state = use_tool(state, "calculator")
        print("\nAfter using calculator twice:")
        print_tool_status(state)

        state = use_tool(state, "calculator")
    except ToolLimitExceeded as e:
        print(f"\nError: {e}")
def print_tool_status(state: State) -> None:
    """Print readable tool status."""
    status = get_tool_status(state)
    for tool, info in status.items():
        print(
            f"{tool}: {info['current_usage']}/{info['limit']} uses "
            f"({info['uses_remaining']} remaining)"
        )

if __name__ == "__main__":
    demonstrate_usage()

Common Pitfalls

  1. Not copying state during updates
  2. Missing tool validation
  3. Incorrect limit tracking
  4. Poor error handling

Key Takeaways

  1. State Design: TypedDict ensures type safety
  2. Usage Tracking: Clean tracking prevents overuse
  3. Rate Limiting: Proper limits maintain control
  4. Status Monitoring: Clear reporting enables oversight

Next Steps

  1. Add tool priority system
  2. Implement cooldown periods
  3. Add usage statistics
  4. Enhance error handling
  5. Add state persistence

Expected Output

Initial tool status

calculator: 0/2 uses (2 remaining)
weather: 0/1 uses (1 remaining)
search: 0/3 uses (3 remaining)

After using calculator once

calculator: 1/2 uses (1 remaining)
weather: 0/1 uses (1 remaining)
search: 0/3 uses (3 remaining)
Error: Tool 'calculator' has reached its limit of 2 uses

Rod Rivera

🇬🇧 Chapter