🎯 What You'll Learn Today

from typing import Annotated, Literal, TypedDict
!pip install langchain-core
!pip install langgraph
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
from langgraph.graph import END, START, StateGraph
from langgraph.graph.message import add_messages

class State(TypedDict):
    """State container for agent graph system.

    Notes:
        - messages tracks conversation flow
        - current_step tracks execution state
        - tool_name tracks selected tools
        - tool_output stores execution results

    Attributes:
        messages: Conversation history with proper typing
        current_step: Current execution phase
        tool_name: Currently selected tool
        tool_output: Latest tool result
    """

    messages: Annotated[list[BaseMessage], add_messages]
    current_step: str
    tool_name: str | None
    tool_output: str | None

def tool_selector(state: State) -> State:
    """Select appropriate tools based on conversation.

    Notes:
        - Handles initial message
        - Implements tool selection logic
        - Updates state immutably

    Args:
        state: Current conversation state

    Returns:
        Updated state with tool selection
    """
    if not state.get("messages"):
        return {
            **state,
            "messages": [HumanMessage(content="What would you like to know?")],
            "current_step": "waiting_for_input",
        }

    last_message = state["messages"][-1].content.lower()

    # Tool selection logic
    if "weather" in last_message:
        tool = "weather_tool"
        response = "I'll check the weather for you."
    elif "calculate" in last_message:
        tool = "calculator_tool"
        response = "I'll help you with that calculation."
    else:
        tool = "search_tool"
        response = "I'll search for that information."

    return {
        **state,
        "tool_name": tool,
        "messages": state["messages"] + [AIMessage(content=response)],
        "current_step": "tool_selected",
    }

def tool_executor(state: State) -> State:
    """Execute selected tool and manage results.

    Notes:
        - Verifies tool selection
        - Simulates tool execution
        - Updates state with results

    Args:
        state: Current conversation state

    Returns:
        Updated state with execution results
    """
    tool_name = state.get("tool_name")
    if not tool_name:
        return state

    # Simulate tool execution
    if tool_name == "weather_tool":
        output = "The weather is sunny and 22°C"
    elif tool_name == "calculator_tool":
        output = "The result is 42"
    else:
        output = "Here's what I found from the search"

    return {**state, "tool_output": output, "current_step": "tool_executed"}

def result_processor(state: State) -> State:
    """Process and format tool execution results.

    Notes:
        - Verifies tool output
        - Formats response
        - Updates conversation

    Args:
        state: Current conversation state

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

    response = f"Here's what I found: {tool_output}"

    return {
        **state,
        "messages": state["messages"] + [AIMessage(content=response)],
        "current_step": "result_processed",
    }

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

    Notes:
        - Checks current step
        - Determines continuation
        - Handles termination

    Args:
        state: Current conversation state

    Returns:
        Next step identifier
    """
    return "end" if state.get("current_step") == "result_processed" else "continue"

def create_agent_graph() -> StateGraph:
    """Create complete agent graph configuration.

    Notes:
        - Initializes graph
        - Adds all nodes
        - Configures edges
        - Sets up routing

    Returns:
        Configured StateGraph
    """
    graph = StateGraph(State)

    # Add nodes
    graph.add_node("tool_selector", tool_selector)
    graph.add_node("tool_executor", tool_executor)
    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

def demonstrate_graph():
    """Demonstrate agent graph functionality."""
    graph = create_agent_graph()
    chain = graph.compile()

    # Initialize state
    state = {
        "messages": [HumanMessage(content="What's the weather like?")],
        "current_step": "start",
        "tool_name": None,
        "tool_output": None,
    }

    print("Agent Graph Demonstration")
    print("========================")

    result = chain.invoke(state)

    print("\nFinal Conversation:")
    for msg in result["messages"]:
        speaker = "Human" if isinstance(msg, HumanMessage) else "Assistant"
        print(f"{speaker}: {msg.content}")

Human: What's the weather like?
Assistant: I'll check the weather for you.
Assistant: Here's what I found: The weather is sunny and 22°C

if __name__ == "__main__":
    demonstrate_graph()