src/graph.py

from __future__ import annotations

from typing import TypedDict, List, Dict

from langgraph.graph import StateGraph, END

from .agents import (
    QiddiyaState,
    orchestrator_node,
    wait_time_predictor_node,
    route_optimizer_node,
    experience_writer_node,
    critic_node,
    reflection_node,
)


class GraphState(TypedDict, total=False):
    user_request: Dict
    logs: List[str]
    wait_time_forecast: Dict[str, int] | None
    raw_plan: Dict | None
    final_plan: Dict | None
    critique: str | None
    reflection_round: int
    refined_attraction_ids: List[str] | None


QiddiyaState = GraphState


def _should_reflect(state: QiddiyaState) -> str:
    critique = state.get("critique") or ""
    reflection_round = int(state.get("reflection_round", 0))
    if critique and reflection_round < 2:
        return "reflect"
    return "finish"


def build_qiddiya_graph() -> StateGraph:
    graph = StateGraph(GraphState)

    graph.add_node("orchestrator", orchestrator_node)
    graph.add_node("wait_time", wait_time_predictor_node)
    graph.add_node("route", route_optimizer_node)
    graph.add_node("guide", experience_writer_node)
    graph.add_node("critic", critic_node)
    graph.add_node("reflection", reflection_node)

    graph.set_entry_point("orchestrator")
    graph.add_edge("orchestrator", "wait_time")
    graph.add_edge("wait_time", "route")
    graph.add_edge("route", "guide")
    graph.add_edge("guide", "critic")

    graph.add_conditional_edges(
        "critic",
        _should_reflect,
        {
            "reflect": "reflection",
            "finish": END,
        },
    )
    graph.add_edge("reflection", "route")

    return graph