Logo of AI Product Engineer
AI Product Engineer
Tutorial Image: LangGraph Tutorial: Tool State Setup with Mock Integration - Unit 2.1 Exercise 1

LangGraph Tutorial: Tool State Setup with Mock Integration - Unit 2.1 Exercise 1

Learn how to set up a tool-enabled state in LangGraph with mock API integration. This tutorial covers configuration management using Pydantic, state initialization with TypedDict, and tool setup for seamless execution, providing a secure and modular foundation for LangGraph applications.

🎯 What You'll Learn Today

LangGraph Tutorial: Tool State Setup with Mock Integration - Unit 2.1 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 set up a tool-enabled state in LangGraph. We'll use a mock API key for demonstration purposes, making this code self-contained and easy to run while teaching proper practices.

Learning Objectives

  1. Understanding state management in LangGraph
  2. Working with configuration using Pydantic
  3. Setting up external tools (using Tavily as an example)
  4. Proper type hinting and annotations

Part 1: Mock Configuration

In practice, you would use a real API key stored securely. For this educational example, we use a mock key.

MOCK_TAVILY_KEY = "mock_tavily_api_key_12345"

Part 2: Required Imports

We import necessary components from

  • typing: For type hints
  • langchain: For message and tool functionality
  • pydantic: For configuration management
import os
from typing import Annotated, Any, TypedDict
!pip install langchain-core
!pip install langgraph
!pip install pydantic-settings
from langchain_community.tools import TavilySearchResults
from langchain_core.messages import BaseMessage
from langgraph.graph.message import add_messages
from pydantic_settings import BaseSettings

Part 3: Configuration Management

We use Pydantic's BaseSettings for configuration. This class would typically load from environment variables, but for our example, it falls back to our mock key.

class Settings(BaseSettings):
    """Configuration for our tool-enabled agent.

    This class demonstrates:
    1. Using Pydantic for configuration management
    2. Providing default values (our mock key)
    3. Environment variable integration

    In production:
    - The API key would come from environment variables
    - You'd use proper secret management
    - You'd have additional configuration options
    """

    tavily_api_key: str = MOCK_TAVILY_KEY  # Default to mock key if no env var

Part 4: State Definition

The State class defines our application's state structure. It uses TypedDict to ensure type safety and proper structure.

class State(TypedDict):
    """State that tracks tool usage and results.

    This state tracks three key elements:
    1. messages: Conversation history (with special handling)
    2. tool_calls: Record of when tools are called
    3. tool_outputs: Results from tool executions

    The Annotated type on messages tells LangGraph to use
    add_messages for handling updates to this field.
    """

    messages: Annotated[list[BaseMessage], add_messages]
    tool_calls: list[dict]
    tool_outputs: list[Any]

Part 5: Tool and State Setup

Now we put everything together

  1. Initialize settings
  2. Set up the tool
  3. Create initial state
def setup_tool_environment():
    """Set up the tool environment with mock integration.

    This function:
    1. Initializes settings
    2. Sets up environment variables
    3. Creates the tool instance
    4. Initializes the state

    Returns:
        tuple: Containing the settings, tool, and initial state
    """
    # Initialize settings
    settings = Settings()

    # In a real application, you'd verify the API key here
    os.environ["TAVILY_API_KEY"] = settings.tavily_api_key

    # Create tool instance
    tavily_tool = TavilySearchResults()

    # Initialize state
    initial_state: State = {"messages": [], "tool_calls": [], "tool_outputs": []}

    return settings, tavily_tool, initial_state

Part 6: Usage Example

Demonstrate how to use the setup and verify everything is working.

def demonstrate_setup():
    """Run a demonstration of the tool setup."""
    # Set up the environment
    settings, tool, state = setup_tool_environment()

    # Print setup results
    print("Tool State Setup Results:")
    print("------------------------")
    print(f"Settings configured: {bool(settings)}")
    print(f"Tool initialized: {bool(tool)}")
    print(f"State created: {bool(state)}")
    print("\nInitial State Structure:")
    print(f"Messages: {len(state['messages'])}")
    print(f"Tool Calls: {len(state['tool_calls'])}")
    print(f"Tool Outputs: {len(state['tool_outputs'])}")

if __name__ == "__main__":
    demonstrate_setup()

Key Learning Points

  1. Configuration Management:

    • Using Pydantic for structured configuration
    • Handling API keys (mock version for learning)
    • Environment variable integration

  2. State Structure:

    • TypedDict for type-safe state definition
    • Special handling for messages with Annotated
    • Tracking tool usage and results

  3. Tool Setup:

    • Proper initialization sequence
    • Environment variable management
    • State initialization

  4. Best Practices:

    • Type hints throughout
    • Clear documentation
    • Modular structure
    • Error handling considerations

Next Steps for Students

  1. Replace mock key with a real Tavily API key
  2. Add error handling for API key verification
  3. Extend the state with additional fields
  4. Add more tool configurations
  5. Implement tool usage tracking

Note: In a production environment, you would:

  1. Use proper secret management
  2. Add error handling
  3. Validate API keys
  4. Add logging
  5. Implement security measures

Rod Rivera

🇬🇧 Chapter