🎯 What You'll Learn Today

import uuid
from datetime import datetime
from typing import Annotated, Any, Literal, TypedDict
#!pip install langchain-core
#!pip install langgraph
import numexpr
from langchain_core.messages import (
    AIMessage,
    BaseMessage,
    HumanMessage,
    SystemMessage,
    ToolMessage,
)
from langchain_core.tools import tool
from langgraph.graph import END, START, StateGraph
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode

class State(TypedDict, total=False):
    """Advanced state container for multi-tool system.

    Notes:
        - messages tracks conversation
        - available_tools lists valid tools
        - tool_usage tracks execution counts
        - rate_limits sets usage boundaries
        - tool_name tracks current selection
        - tool_outputs stores results

    Attributes:
        messages: Conversation history
        available_tools: Valid tool list
        tool_usage: Tool usage tracking
        rate_limits: Tool usage limits
        tool_name: Active tool
        tool_outputs: Execution results
    """

    messages: Annotated[list[BaseMessage], add_messages]
    available_tools: list[str]
    tool_usage: dict[str, int]
    rate_limits: dict[str, int]
    tool_name: str | None
    tool_outputs: list[str]

@tool
def calculator(query: str) -> str:
    """Calculate mathematical expressions safely.

    Notes:
        - Uses numexpr for safe evaluation
        - Handles input cleaning
        - Returns formatted results

    Args:
        query: Mathematical expression

    Returns:
        Formatted calculation result
    """
    try:
        expression = query.replace("calculate ", "").replace("what is ", "").strip()
        result = numexpr.evaluate(expression).item()
        return f"{expression} = {result}"
    except Exception as e:
        return f"Error in calculation: {e!s}"

@tool
def check_weather(query: str) -> str:
    """Get weather information with proper formatting.

    Notes:
        - Handles query cleaning
        - Formats location properly
        - Includes timestamps

    Args:
        query: Location query

    Returns:
        Formatted weather information
    """
    try:
        location = (
            query.lower()
            .replace("what's the weather in ", "")
            .replace("weather in ", "")
        )
        current_time = datetime.now().strftime("%H:%M")
        return f"Weather in {location.title()} at {current_time}: 22°C, Sunny"
    except Exception as e:
        return f"Error checking weather: {e!s}"

def can_use_tool(state: State, tool_name: str) -> bool:
    """Check tool availability against rate limits.

    Notes:
        - Verifies tool existence
        - Checks usage counts
        - Enforces limits

    Args:
        state: Current state
        tool_name: Tool to check

    Returns:
        Availability status
    """
    if tool_name not in state["tool_usage"]:
        return True
    current_usage = state["tool_usage"].get(tool_name, 0)
    limit = state["rate_limits"].get(tool_name, float("inf"))
    return current_usage < limit

def tool_selector(state: State) -> State:
    """Select appropriate tool with rate limit checking.

    Notes:
        - Parses user intent
        - Checks rate limits
        - Provides feedback

    Args:
        state: Current state

    Returns:
        Updated state with selection
    """
    if not state.get("messages"):
        return {
            **state,
            "messages": [
                SystemMessage(
                    content="I am a helpful AI assistant that can use tools."
                ),
                AIMessage(content="How can I help you today?"),
            ],
        }

    # Get last human message
    last_message = None
    for msg in reversed(state["messages"]):
        if isinstance(msg, HumanMessage):
            last_message = msg.content.lower()
            break

    if not last_message:
        return state

    # Select tool with rate limit check
    if "weather" in last_message:
        tool_name = "check_weather"
        reason = "I'll check the weather information."
    elif any(
        word in last_message for word in ["calculate", "what is", "+", "-", "*", "/"]
    ):
        tool_name = "calculator"
        reason = "I'll help you with that calculation."
    else:
        return {
            **state,
            "messages": state["messages"]
            + [AIMessage(content="I'm not sure how to help with that.")],
        }

    # Check rate limit
    if not can_use_tool(state, tool_name):
        limit = state["rate_limits"].get(tool_name, 0)
        return {
            **state,
            "messages": state["messages"]
            + [
                AIMessage(
                    content=f"Sorry, the {tool_name} tool has reached its limit of {limit}."
                )
            ],
        }

    return {
        **state,
        "tool_name": tool_name,
        "messages": state["messages"] + [AIMessage(content=reason)],
    }

def execute_with_tool_node(state: State, tools: list[Any]) -> State:
    """Execute tools using ToolNode with proper formatting.

    Notes:
        - Handles rate limits
        - Updates usage stats
        - Formats messages properly

    Args:
        state: Current state
        tools: Available tools

    Returns:
        Updated state with results
    """
    if not state.get("tool_name"):
        return {**state, "tool_outputs": []}

    # Check rate limits
    current_usage = state["tool_usage"].get(state["tool_name"], 0)
    limit = state["rate_limits"].get(state["tool_name"], float("inf"))
    if current_usage >= limit:
        return {
            **state,
            "tool_outputs": [
                f"Tool {state['tool_name']} has reached its limit of {limit}"
            ],
        }

    try:
        tool_node = ToolNode(tools=tools)

        # Get last human message
        message = None
        for msg in reversed(state["messages"]):
            if isinstance(msg, HumanMessage):
                message = msg.content.lower()
                if state["tool_name"] == "calculator":
                    message = message.replace("what's", "").replace("what is", "")
                    message = message.replace("calculate", "").replace("?", "").strip()
                elif state["tool_name"] == "check_weather":
                    message = message.replace("what's the weather in", "")
                    message = message.replace("weather in", "").replace("?", "").strip()
                break

        if not message:
            return state

        # Create proper tool message
        tool_message = AIMessage(
            content="",
            tool_calls=[
                {
                    "name": state["tool_name"],
                    "args": {"query": message},
                    "id": str(uuid.uuid4()),
                    "type": "tool_call",
                }
            ],
        )

        # Execute tool
        result = tool_node.invoke({"messages": [tool_message]})

        if result and "messages" in result and result["messages"]:
            tool_output = result["messages"][0].content
        else:
            tool_output = "No result from tool"

        # Update usage count
        current_usage = state["tool_usage"].copy()
        current_usage[state["tool_name"]] = current_usage.get(state["tool_name"], 0) + 1

        return {**state, "tool_outputs": [tool_output], "tool_usage": current_usage}

    except Exception as e:
        return {**state, "tool_outputs": [f"Error executing tool: {e!s}"]}

def result_processor(state: State) -> State:
    """Process tool results into conversation messages.

    Notes:
        - Creates tool messages
        - Formats AI responses
        - Handles conversation flow

    Args:
        state: Current state

    Returns:
        Updated state with processed results
    """
    if not state.get("tool_outputs"):
        return state

    tool_output = state["tool_outputs"][-1]

    # Create properly formatted messages
    tool_message = ToolMessage(
        content=tool_output,
        tool_call_id=str(uuid.uuid4()),
        name=state.get("tool_name", "unknown_tool"),
    )

    ai_message = AIMessage(
        content=f"Here's what I found: {tool_output}. Is there anything else?"
    )

    return {**state, "messages": state["messages"] + [tool_message, ai_message]}

def get_next_step(state: State) -> Literal["continue", "end"]:
    """Determine next conversation step.

    Notes:
        - Checks error conditions
        - Verifies rate limits
        - Handles completion

    Args:
        state: Current state

    Returns:
        Next step identifier
    """
    # Check errors
    if state.get("tool_outputs"):
        last_output = state["tool_outputs"][-1]
        if "Error" in last_output:
            return "end"

    # Check rate limits
    for tool, usage in state.get("tool_usage", {}).items():
        limit = state.get("rate_limits", {}).get(tool, float("inf"))
        if usage >= limit:
            return "end"

    return "end"  # End after one tool use for simplicity

def create_multi_tool_agent():
    """Create integrated multi-tool agent system.

    Notes:
        - Configures all nodes
        - Sets up edges
        - Enables proper routing

    Returns:
        Compiled agent system
    """
    tools = [calculator, check_weather]
    graph = StateGraph(State)

    # Add nodes
    graph.add_node("tool_selector", tool_selector)
    graph.add_node("tool_executor", lambda state: execute_with_tool_node(state, tools))
    graph.add_node("result_processor", result_processor)

    # Configure edges
    graph.add_edge(START, "tool_selector")
    graph.add_edge("tool_selector", "tool_executor")
    graph.add_edge("tool_executor", "result_processor")

    # Add conditional routing
    graph.add_conditional_edges(
        "result_processor", get_next_step, {"continue": "tool_selector", "end": END}
    )

    return graph.compile()

def demonstrate_multi_tool_agent() -> State:
    """Demonstrate the complete system.

    Notes:
        - Tests multiple queries
        - Verifies rate limits
        - Shows conversation flow

    Returns:
        State: Final system state after all queries
    """
    agent = create_multi_tool_agent()

    queries = [
        "What's 2 + 2?",
        "What's the weather in Paris?",
        "Calculate 5 * 3",
        "What's the weather in London?",
    ]

    initial_state = {
        "messages": [],
        "available_tools": ["calculator", "check_weather"],
        "tool_usage": {"calculator": 0, "check_weather": 0},
        "rate_limits": {"calculator": 2, "check_weather": 1},
        "tool_outputs": [],
    }

    print("Multi-Tool Agent Demonstration")
    print("=============================")

    current_state = initial_state.copy()

    for query in queries:
        print(f"\nUser Query: {query}")

        query_state = {
            **current_state,
            "messages": current_state["messages"] + [HumanMessage(content=query)],
        }

        result = agent.invoke(query_state)
        current_state = {**result, "tool_usage": result["tool_usage"]}

        print("\nConversation:")
        for msg in result["messages"][-3:]:
            prefix = {
                HumanMessage: "Human",
                AIMessage: "Assistant",
                ToolMessage: "Tool",
                SystemMessage: "System",
            }.get(type(msg), "Unknown")
            print(f"{prefix}: {msg.content}")

        print("\nTool Usage:")
        for tool, usage in result["tool_usage"].items():
            limit = result["rate_limits"].get(tool, "unlimited")
            print(f"{tool}: {usage}/{limit}")

        print("-" * 50)

    return current_state

if __name__ == "__main__":
    demonstrate_multi_tool_agent()

User Query: What's 2 + 2?

# Conversation
Human: What's 2 + 2?
Assistant: I'll help you with that calculation.
Tool: 2 + 2 = 4
Assistant: Here's what I found: 2 + 2 = 4. Is there anything else?

## Tool Usage
calculator: 1/2
check_weather: 0/1

User Query: What's the weather in Paris?

## Conversation
Human: What's the weather in Paris?
Assistant: I'll check the weather information.
Tool: Weather in Paris at 14:30: 22°C, Sunny
Assistant: Here's what I found: Weather in Paris at 14:30: 22°C, Sunny. Is there anything else?

## Tool Usage
calculator: 1/2
check_weather: 1/1

User Query: Calculate 5 * 3

## Conversation
Human: Calculate 5 * 3
Assistant: I'll help you with that calculation.
Tool: 5 * 3 = 15
Assistant: Here's what I found: 5 * 3 = 15. Is there anything else?

## Tool Usage
calculator: 2/2
check_weather: 1/1

User Query: What's the weather in London?

## Conversation
Human: What's the weather in London?
Assistant: Sorry, the check_weather tool has reached its limit of 1.

## Tool Usage
calculator: 2/2
## check_weather: 1/1