agent/graph.py

"""
LangGraph Graph Definition — compiles the scheduling workflow into an executable graph.
"""
from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import MemorySaver
from agent.state import SchedulingState
from agent.nodes import (
    validate_inputs,
    check_availability,
    calculate_travel,
    rank_and_present,
    book_appointment,
    send_notifications,
    no_match,
)


def should_continue_after_availability(state: SchedulingState) -> str:
    """Routing function after availability check.

    Returns 'calculate_travel' if providers are available, 'no_match' otherwise.
    """
    if state.get("available_providers"):
        return "calculate_travel"
    return "no_match"


def build_scheduling_graph():
    """Build and compile the scheduling LangGraph.

    Graph flow:
        validate_inputs → check_availability → [conditional]
            → calculate_travel → rank_and_present → (INTERRUPT for user selection)
            → book_appointment → send_notifications → END
        or:
            → no_match → END

    The graph uses interrupt_before=["book_appointment"] to pause for human-in-the-loop
    confirmation. The Streamlit UI should:
        1. Run the graph until it hits the interrupt
        2. Display ranked_matches to the user
        3. Get user selection
        4. Update state with user_selection
        5. Resume the graph

    Returns:
        Compiled LangGraph
    """
    workflow = StateGraph(SchedulingState)

    # Add nodes
    workflow.add_node("validate_inputs", validate_inputs)
    workflow.add_node("check_availability", check_availability)
    workflow.add_node("calculate_travel", calculate_travel)
    workflow.add_node("rank_and_present", rank_and_present)
    workflow.add_node("book_appointment", book_appointment)
    workflow.add_node("send_notifications", send_notifications)
    workflow.add_node("no_match", no_match)

    # Set entry point
    workflow.set_entry_point("validate_inputs")

    # Add edges
    workflow.add_edge("validate_inputs", "check_availability")
    workflow.add_conditional_edges(
        "check_availability",
        should_continue_after_availability,
        {
            "calculate_travel": "calculate_travel",
            "no_match": "no_match",
        }
    )
    workflow.add_edge("calculate_travel", "rank_and_present")
    workflow.add_edge("rank_and_present", "book_appointment")
    workflow.add_edge("book_appointment", "send_notifications")
    workflow.add_edge("send_notifications", END)
    workflow.add_edge("no_match", END)

    # Compile with interrupt for human-in-the-loop
    # The graph pauses before book_appointment so the UI can show options and get user selection
    graph = workflow.compile(checkpointer=MemorySaver(), interrupt_before=["book_appointment"])

    return graph