🎯 What You'll Learn Today

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

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

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,
    }

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

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()

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

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