workflow.py

"""
workflow.py
Multi-agent workflow built with LangGraph. Each disaster request
flows through four agent nodes:

    intake_agent -> classification_agent -> urgency_agent -> allocation_agent

Each node updates a shared state (CyrexState) that is finally
persisted to the database by the caller. Splitting classification and
urgency into separate nodes (even though a single LLM call currently
produces both) keeps the graph faithful to the concept note's
"multi-agent workflow" design and makes it easy to later swap in a
dedicated urgency-scoring model or additional business rules without
touching the rest of the pipeline.
"""

from typing import TypedDict, Optional
from langgraph.graph import StateGraph, END

from llm_client import classify_emergency
from resources import pick_team


class CyrexState(TypedDict, total=False):
    name: str
    contact: str
    location: str
    description: str
    num_people: int

    disaster_type: str
    urgency_level: str
    urgency_score: int
    reasoning: str

    assigned_team: str
    status: str


def intake_agent(state: CyrexState) -> CyrexState:
    """Validates / normalizes the raw citizen submission."""
    state["description"] = state.get("description", "").strip()
    state["num_people"] = int(state.get("num_people") or 1)
    state["status"] = "Pending"
    return state


def classification_agent(state: CyrexState) -> CyrexState:
    """Calls the LLM (or fallback) to classify disaster type + urgency in one pass."""
    result = classify_emergency(state["description"], state["num_people"])
    state["disaster_type"] = result["disaster_type"]
    state["urgency_level"] = result["urgency_level"]
    state["urgency_score"] = result["urgency_score"]
    state["reasoning"] = result["reasoning"]
    return state


def urgency_agent(state: CyrexState) -> CyrexState:
    """
    Applies additional business rules on top of the LLM's urgency
    assessment -- e.g. boosting priority for large groups -- and
    settles the final urgency level/score.
    """
    score = state.get("urgency_score", 40)

    if state.get("num_people", 1) >= 10 and score < 90:
        score = min(score + 10, 100)
        state["reasoning"] = (state.get("reasoning", "") +
                               " Priority boosted due to large number of people affected.")

    if score >= 80:
        level = "Critical"
    elif score >= 60:
        level = "High"
    elif score >= 35:
        level = "Medium"
    else:
        level = "Low"

    state["urgency_score"] = score
    state["urgency_level"] = level
    return state


def allocation_agent(state: CyrexState) -> CyrexState:
    """Assigns the most suitable available rescue team."""
    team = pick_team(state.get("disaster_type", "Other"), state.get("urgency_level", "Medium"))
    state["assigned_team"] = team
    state["status"] = "Team Assigned"
    return state


def build_graph():
    graph = StateGraph(CyrexState)
    graph.add_node("intake", intake_agent)
    graph.add_node("classify", classification_agent)
    graph.add_node("assess_urgency", urgency_agent)
    graph.add_node("allocate", allocation_agent)

    graph.set_entry_point("intake")
    graph.add_edge("intake", "classify")
    graph.add_edge("classify", "assess_urgency")
    graph.add_edge("assess_urgency", "allocate")
    graph.add_edge("allocate", END)

    return graph.compile()


_compiled_graph = None


def get_graph():
    global _compiled_graph
    if _compiled_graph is None:
        _compiled_graph = build_graph()
    return _compiled_graph


def process_request(name: str, contact: str, location: str,
                     description: str, num_people: int = 1) -> dict:
    """Runs a new citizen request through the full Cyrex agent workflow."""
    initial_state: CyrexState = {
        "name": name,
        "contact": contact,
        "location": location,
        "description": description,
        "num_people": num_people,
    }
    graph = get_graph()
    final_state = graph.invoke(initial_state)
    return dict(final_state)