🎯 What You'll Learn Today

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

class State(TypedDict):
    """Manages the state of our routing system.

    This state implementation introduces three key concepts:
    1. Annotated message lists for proper message handling
    2. Classification tracking for routing decisions
    3. Confidence scoring for uncertainty management

    Attributes:
        messages: List of conversation messages with LangGraph's add_messages
                 annotation for proper message management
        classification: Category of the current message (e.g., greeting, help)
        confidence: Confidence score for the classification (0.0 to 1.0)
    """

    messages: Annotated[list[BaseMessage], add_messages]
    classification: str
    confidence: float

def classifier_node(state: State) -> State:
    """Classifies incoming messages into categories.

    This node demonstrates several advanced concepts:
    1. Defensive programming with empty message handling
    2. Confidence-based classification
    3. Fallback mechanisms for unknown inputs
    4. Structured state updates

    Args:
        state: Current system state containing messages

    Returns:
        Updated state with classification and confidence
    """
    if not state["messages"]:
        return {
            "messages": state["messages"],
            "classification": "unknown",
            "confidence": 0.0,
        }

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

    # Classification logic
    if "hello" in message or "hi" in message:
        return {
            "messages": state["messages"],
            "classification": "greeting",
            "confidence": 0.9,
        }
    elif "help" in message or "support" in message:
        return {
            "messages": state["messages"],
            "classification": "help",
            "confidence": 0.8,
        }
    elif "bye" in message or "goodbye" in message:
        return {
            "messages": state["messages"],
            "classification": "farewell",
            "confidence": 0.9,
        }

    return {
        "messages": state["messages"],
        "classification": "unknown",
        "confidence": 0.1,
    }

def response_node_1(state: State) -> State:
    """Handles greeting messages with appropriate responses.

    Args:
        state: Current state with classification

    Returns:
        State with greeting response added
    """
    return {
        "messages": [AIMessage(content="Hello! How can I assist you today?")],
        "classification": state["classification"],
        "confidence": state["confidence"],
    }

def response_node_2(state: State) -> State:
    """Handles help requests with supportive responses.

    Args:
        state: Current state with classification

    Returns:
        State with help response added
    """
    return {
        "messages": [
            AIMessage(content="I'm here to help! What do you need assistance with?")
        ],
        "classification": state["classification"],
        "confidence": state["confidence"],
    }

def response_node_3(state: State) -> State:
    """Handles unknown or low confidence cases with clarification requests.

    Args:
        state: Current state with classification

    Returns:
        State with clarification response added
    """
    return {
        "messages": [
            AIMessage(
                content="I'm not quite sure what you mean. Could you please rephrase that?"
            )
        ],
        "classification": state["classification"],
        "confidence": state["confidence"],
    }

def get_next_node(state: State) -> str:
    """Routes messages to appropriate response nodes.

    This function demonstrates LangGraph's conditional routing by:
    1. Evaluating confidence thresholds
    2. Mapping classifications to handlers
    3. Implementing fallback logic

    Args:
        state: Current state with classification and confidence

    Returns:
        str: Name of the next node to route to
    """
    if state["confidence"] > 0.7:
        classification_routes = {"greeting": "response_1", "help": "response_2"}
        return classification_routes.get(state["classification"], "response_3")

    return "response_3"

Create and configure the graph

graph = StateGraph(State)

graph.add_node("classifier", classifier_node)
graph.add_node("response_1", response_node_1)
graph.add_node("response_2", response_node_2)
graph.add_node("response_3", response_node_3)
Configure routing
graph.add_edge(START, "classifier")
graph.add_conditional_edges(
    "classifier",
    get_next_node,
    {
        "response_1": "response_1",
        "response_2": "response_2",
        "response_3": "response_3",
    },
)
Add terminal edges
graph.add_edge("response_1", END)
graph.add_edge("response_2", END)
graph.add_edge("response_3", END)
Compile the graph

chain = graph.compile()

def process_message(message: str) -> None:
    """Process a message through the routing system.

    Args:
        message: Input message to process
    """
    initial_state = {
        "messages": [AIMessage(content=message)],
        "classification": "",
        "confidence": 0.0,
    }

    result = chain.invoke(initial_state)

    print(f"\nInput: {message}")
    print("Response:", result["messages"][-1].content)
    print(
        f"Classification: {result['classification']} (confidence: {result['confidence']})"
    )

if __name__ == "__main__":
    print("Testing the Message Routing System:")
    print("==================================")
    test_messages = [
        "Hello there!",  # Should trigger greeting response

        "I need help with something",  # Should trigger help response

        "What's the weather like?",  # Should trigger unknown response

        "goodbye",  # Should trigger farewell response

    ]
    ## for message in test_messages
    process_message(message)