workflow.py

"""
LangGraph Workflow — Campaign Pacing Agent Pipeline.

This module defines the StateGraph that orchestrates the four-agent pipeline:

    fetch → monitor → [has alerts?] → diagnosis (LLM) → action (LLM + tools) →
                                        [has escalations?] → escalation (LLM) → END
                     ↘ END (no alerts)                    ↘ END (no escalations)

STATE FLOW:
    The PipelineState TypedDict flows through the graph. Each node reads from
    the state and returns updates. LangGraph merges the updates into the state
    automatically.

    1. fetch: Populates state["campaigns"] from the simulator/API.
    2. monitor: Populates state["alerts"] using rule-based threshold checks.
       Also filters out campaigns in cooldown from the actionable alerts.
    3. diagnosis: Populates state["diagnoses"] using Claude's reasoning.
    4. action: Populates state["actions_executed"], state["escalations"],
               state["action_reasoning"], and state["action_notifications"]
               using Claude with tool-calling.
    5. escalation: Populates state["notifications"] using Claude's writing.

COOLDOWN:
    After actions are taken on a campaign, that campaign enters a configurable
    cooldown period (default 2 minutes, set in config.yaml). During cooldown,
    the Monitor still reports alerts (for visibility), but those campaigns are
    excluded from Diagnosis and Action to let previous fixes take effect.

CONDITIONAL ROUTING:
    - After monitor: If alerts is empty → END (nothing to do).
    - After action: If escalations is empty → END (no bid changes needed).
    This prevents unnecessary LLM calls when everything is healthy.

COMPILATION:
    build_workflow() creates the graph definition.
    compile_workflow() compiles it into an executable app.
    The compiled app is invoked once per monitoring cycle with a fresh state.
"""

from typing import Any, TypedDict

from langgraph.graph import END, StateGraph

from agents.action import run_actions
from agents.diagnosis import run_diagnosis
from agents.escalation import run_escalation
from agents.monitor import run_monitor
from campaign_simulator import get_campaigns
from cooldown import tracker as cooldown_tracker


class PipelineState(TypedDict):
    """
    Shared state that flows through the LangGraph pipeline.

    Each key represents the output of one agent node. Nodes read from
    earlier keys and write to their own key. LangGraph handles the
    merging automatically.
    """
    # Raw campaign data from the simulator/API (set by fetch node).
    campaigns: list[dict]

    # Alert dicts from the Monitor agent (set by monitor node).
    alerts: list[dict]

    # Campaigns that are currently in cooldown (set by monitor node).
    cooldown_skipped: list[dict]

    # Diagnosis dicts from the Diagnosis agent — LLM-generated (set by diagnosis node).
    diagnoses: list[dict]

    # Actions the Action agent executed autonomously (set by action node).
    actions_executed: list[dict]

    # Bid-change recommendations requiring human approval (set by action node).
    escalations: list[dict]

    # LLM's reasoning trace from the Action agent (set by action node).
    action_reasoning: str

    # FYI notifications about autonomous actions (set by action node).
    action_notifications: list[dict]

    # Formatted notifications from the Escalation agent — LLM-written (set by escalation node).
    notifications: list[dict]


# ---------------------------------------------------------------------------
# Node Functions — Each corresponds to one step in the pipeline.
# They take the current state and return a dict of updates to merge.
# ---------------------------------------------------------------------------

def fetch_campaigns(state: PipelineState) -> dict[str, Any]:
    """
    Node 1: Fetch current campaign data from the simulator/API.

    This calls get_campaigns() which:
    - In the POC: Advances the simulator tick and returns drifted metrics.
    - In production: Would call real DSP API endpoints (DV360, TTD, Xandr).

    The returned data is the same dict schema either way — the rest of
    the pipeline doesn't know or care about the data source.
    """
    campaigns = get_campaigns()
    return {"campaigns": campaigns}


def monitor_node(state: PipelineState) -> dict[str, Any]:
    """
    Node 2: Run rule-based monitoring checks across all campaigns.

    This is intentionally NOT LLM-powered — threshold checks should be
    fast, deterministic, and free. See agents/monitor.py for the 11 checks.

    After detecting alerts, filters out campaigns that are in cooldown.
    Those campaigns still show alerts (for visibility in the Monitor panel),
    but are excluded from Diagnosis and Action to let previous fixes work.

    Returns alerts list. If empty (after cooldown filtering), graph short-circuits to END.
    """
    all_alerts = run_monitor(state["campaigns"])

    # Separate alerts into actionable vs cooldown-skipped.
    actionable_alerts = []
    cooldown_skipped = []

    for alert in all_alerts:
        cid = alert.get("campaign_id", "")
        if cooldown_tracker.is_in_cooldown(cid):
            remaining = cooldown_tracker.get_cooldown_remaining(cid)
            alert["cooldown_remaining"] = remaining
            cooldown_skipped.append(alert)
        else:
            actionable_alerts.append(alert)

    return {"alerts": actionable_alerts, "cooldown_skipped": cooldown_skipped}


def should_diagnose(state: PipelineState) -> str:
    """
    Conditional edge after Monitor: Should we invoke the LLM for diagnosis?

    If there are no alerts, there's nothing to diagnose — skip straight
    to END. This saves LLM API calls (and cost) on healthy cycles.
    """
    if state["alerts"]:
        return "diagnose"
    return "end"


def diagnosis_node(state: PipelineState) -> dict[str, Any]:
    """
    Node 3: LLM-powered root cause analysis.

    Claude analyzes the alerts, cross-references campaign data (optionally
    calling read tools for deeper inspection), and determines the primary
    root cause and contributing factors for each affected campaign.

    This is where the system becomes truly agentic — the LLM reasons about
    WHY campaigns are struggling, not just that they are.
    """
    diagnoses = run_diagnosis(state["alerts"], state["campaigns"])
    return {"diagnoses": diagnoses}


def action_node(state: PipelineState) -> dict[str, Any]:
    """
    Node 4: LLM with tool-calling executes corrective actions.

    Claude receives the diagnoses and decides what to do about each one.
    It enters an agentic tool-calling loop:
    - Inspects campaigns via read tools
    - Executes autonomous actions (audience expansion, creative rotation, etc.)
    - Recommends bid changes (via recommend_bid_change tool)
    - Checks results and decides if more actions are needed

    After execution, records cooldown for all acted-on campaigns and
    generates FYI notifications for the campaign manager about autonomous actions.

    This is the most complex node — it can make multiple tool calls per
    campaign, across multiple campaigns, in a single invocation.
    """
    result = run_actions(state["diagnoses"], state["campaigns"])

    # Record cooldown for every campaign that had actions taken.
    acted_campaign_ids = set()
    for action in result["actions_executed"]:
        cid = action.get("campaign_id", "")
        action_type = action.get("action", "unknown")
        if cid:
            cooldown_tracker.record_action(cid, action_type)
            acted_campaign_ids.add(cid)

    for esc in result["escalations"]:
        cid = esc.get("campaign_id", "")
        if cid:
            cooldown_tracker.record_action(cid, "recommend_bid_change")
            acted_campaign_ids.add(cid)

    # Build FYI notifications for campaign manager about autonomous actions.
    action_notifications = _build_action_notifications(
        result["actions_executed"], state["campaigns"]
    )

    return {
        "actions_executed": result["actions_executed"],
        "escalations": result["escalations"],
        "action_reasoning": result["llm_reasoning"],
        "action_notifications": action_notifications,
    }


def should_escalate(state: PipelineState) -> str:
    """
    Conditional edge after Action: Are there bid changes needing approval?

    If the Action agent didn't recommend any bid changes, we skip the
    Escalation agent entirely — no need to generate notifications.
    """
    if state["escalations"]:
        return "escalate"
    return "end"


def escalation_node(state: PipelineState) -> dict[str, Any]:
    """
    Node 5: LLM writes campaign manager notifications.

    For each bid-change recommendation, Claude writes a professional,
    data-driven notification that the campaign manager can act on.
    The notification includes the problem, recommendation, cost/benefit
    analysis, and consequences of inaction.
    """
    notifications = run_escalation(state["escalations"], state["campaigns"])
    return {"notifications": notifications}


def _build_action_notifications(actions: list[dict], campaigns: list[dict]) -> list[dict]:
    """
    Build FYI notifications to send to campaign managers about autonomous
    actions that were taken. These are informational (no approval needed)
    but keep the CM in the loop.

    Groups actions by campaign for a clean notification per campaign.
    """
    if not actions:
        return []

    # Build a campaign name lookup.
    name_map = {c["campaign_id"]: c["name"] for c in campaigns}

    # Group actions by campaign.
    by_campaign: dict[str, list[dict]] = {}
    for action in actions:
        cid = action.get("campaign_id", "unknown")
        by_campaign.setdefault(cid, []).append(action)

    notifications = []
    for cid, campaign_actions in by_campaign.items():
        action_lines = []
        for a in campaign_actions:
            action_type = a.get("action", "?")
            args = a.get("args", {})

            if action_type == "expand_audience":
                action_lines.append(f"Expanded audience: {args.get('segment_description', '?')}")
            elif action_type == "rotate_creative":
                action_lines.append(f"Rotated creative: {args.get('creative_description', '?')}")
            elif action_type == "adjust_frequency_cap":
                action_lines.append(f"Adjusted frequency cap to {args.get('new_cap', '?')}")
            elif action_type == "adjust_daily_budget":
                budget = args.get('new_daily_budget', 0)
                action_lines.append(f"Adjusted daily budget to ${budget:,.2f}")
            else:
                action_lines.append(f"{action_type}: {args}")

        notifications.append({
            "campaign_id": cid,
            "campaign_name": name_map.get(cid, cid),
            "type": "autonomous_action_fyi",
            "priority": "info",
            "actions_taken": action_lines,
            "message": (
                f"Automated actions taken on {name_map.get(cid, cid)}: "
                + "; ".join(action_lines)
                + ". These actions were executed autonomously and do not require approval. "
                + "Next evaluation after cooldown period."
            ),
        })

    return notifications


# ---------------------------------------------------------------------------
# Graph Construction
# ---------------------------------------------------------------------------

def build_workflow() -> StateGraph:
    """
    Build the LangGraph StateGraph with conditional routing.

    Graph structure:
        fetch → monitor → [condition] → diagnosis → action → [condition] → escalation → END
                              ↘ END                              ↘ END

    Returns:
        StateGraph: The uncompiled graph definition.
    """
    workflow = StateGraph(PipelineState)

    # Add all nodes to the graph.
    workflow.add_node("fetch", fetch_campaigns)
    workflow.add_node("monitor", monitor_node)
    workflow.add_node("diagnose", diagnosis_node)
    workflow.add_node("act", action_node)
    workflow.add_node("escalate", escalation_node)

    # Set the entry point — every cycle starts by fetching campaign data.
    workflow.set_entry_point("fetch")

    # Linear edge: fetch always leads to monitor.
    workflow.add_edge("fetch", "monitor")

    # Conditional edge: monitor → diagnosis (if alerts) or END (if clean).
    workflow.add_conditional_edges(
        "monitor",
        should_diagnose,
        {"diagnose": "diagnose", "end": END},
    )

    # Linear edge: diagnosis always leads to action.
    workflow.add_edge("diagnose", "act")

    # Conditional edge: action → escalation (if bid changes) or END.
    workflow.add_conditional_edges(
        "act",
        should_escalate,
        {"escalate": "escalate", "end": END},
    )

    # Linear edge: escalation leads to END.
    workflow.add_edge("escalate", END)

    return workflow


def compile_workflow():
    """
    Compile the workflow into an executable LangGraph app.

    The compiled app can be invoked with: app.invoke(initial_state)
    Returns the final state with all node outputs merged in.
    """
    return build_workflow().compile()