Pipeline incident commander (Airflow)

"""
incident_commander.py
---------------------
Airflow on_failure_callback that triggers the `incident-commander`
DinoAI agent when a DAG task fails.

Attach to any DAG or individual task:

    # Whole DAG — fires on any task failure
    with DAG(
        "hourly_marts",
        on_failure_callback=trigger_incident_commander,
        ...
    ) as dag:
        ...

    # Single task — fires only when this task fails
    run_fct_orders = DbtRunOperator(
        task_id="fct_orders",
        on_failure_callback=trigger_incident_commander,
        ...
    )

Environment variables (required — set in Airflow Connections or Variables):
    PARADIME_API_ENDPOINT  – Paradime API endpoint
    PARADIME_API_KEY       – Paradime API key
    PARADIME_API_SECRET    – Paradime API secret

Optional:
    INCIDENT_SLACK_CHANNEL – Slack channel for incident posts (default: #incidents)
    SUSPECT_MODEL          – dbt™ model to focus triage on (default: extracted
                             from failed task_id if it matches a model name)
"""

from __future__ import annotations

import os
import threading
from typing import Any

from paradime import Paradime

AGENT_NAME     = "incident-commander"
SLACK_CHANNEL  = os.environ.get("INCIDENT_SLACK_CHANNEL", "#incidents")


def _build_client() -> Paradime:
    return Paradime(
        api_endpoint=os.environ["PARADIME_API_ENDPOINT"],
        api_key=os.environ["PARADIME_API_KEY"],
        api_secret=os.environ["PARADIME_API_SECRET"],
    )


def _build_incident_message(context: dict[str, Any]) -> str:
    """
    Build the incident trigger message from Airflow task context.
    Airflow passes a rich context dict to on_failure_callback — we extract
    the most useful fields for the triage agent.
    """
    dag_id        = context.get("dag").dag_id if context.get("dag") else "unknown"
    task_id       = context.get("task_instance").task_id if context.get("task_instance") else "unknown"
    run_id        = context.get("run_id", "unknown")
    execution_date = str(context.get("execution_date", "unknown"))
    exception     = str(context.get("exception", "unknown"))
    log_url       = context.get("task_instance").log_url if context.get("task_instance") else "unavailable"

    # Try to infer the suspect dbt™ model from the task ID or env override
    suspect_model = (
        os.environ.get("SUSPECT_MODEL")
        or task_id  # Airflow task IDs often match dbt™ model names directly
    )

    return (
        f"INCIDENT: Airflow task failure detected.\n\n"
        f"DAG:            {dag_id}\n"
        f"Failed task:    {task_id}\n"
        f"Run ID:         {run_id}\n"
        f"Execution date: {execution_date}\n"
        f"Exception:      {exception}\n"
        f"Log URL:        {log_url}\n"
        f"Suspect model:  {suspect_model}\n\n"
        f"Coordinate triage across log-analyzer, query-profiler, and "
        f"owner-notifier. Post a unified ROOT CAUSE / IMPACT / OWNER / "
        f"NEXT ACTION summary to {SLACK_CHANNEL}."
    )


def trigger_incident_commander(context: dict[str, Any]) -> None:
    """
    Airflow on_failure_callback. Fires asynchronously so it does not block
    Airflow's task runner thread while the agent session runs.

    Attach to a DAG:   on_failure_callback=trigger_incident_commander
    Attach to a task:  on_failure_callback=trigger_incident_commander
    """
    def _run() -> None:
        try:
            paradime = _build_client()
            message  = _build_incident_message(context)

            dag_id  = context.get("dag").dag_id if context.get("dag") else "unknown"
            task_id = context.get("task_instance").task_id if context.get("task_instance") else "unknown"

            print(f"[incident-commander] 🚨  Triggering triage for {dag_id}/{task_id} …")

            run = paradime.dinoai_agents.trigger_run_and_wait(
                agent=AGENT_NAME,
                message=message,
                slack_channel=SLACK_CHANNEL,
                timeout=1800,   # 30 minutes max
            )

            print(
                f"[incident-commander] ✅  Triage complete for {dag_id}/{task_id}.\n"
                f"{run.messages[-1].content}"
            )

        except Exception as exc:  # noqa: BLE001
            # Log but never re-raise — a callback failure must not mask the original error
            print(f"[incident-commander] ❌  Failed to trigger triage: {exc}")

    # Run in a daemon thread so Airflow doesn't wait for it
    thread = threading.Thread(target=_run, daemon=True)
    thread.start()

from airflow import DAG
from airflow.utils.dates import days_ago
from callbacks.incident_commander import trigger_incident_commander

with DAG(
    dag_id="hourly_marts",
    schedule_interval="@hourly",
    start_date=days_ago(1),
    on_failure_callback=trigger_incident_commander,   # ← fires on any task failure
    catchup=False,
) as dag:
    ...

"""
trigger_incident.py
-------------------
Manual trigger for the `incident-commander` DinoAI agent.

Usage:
    python scripts/trigger_incident.py \
        --schedule "hourly_marts" \
        --model fct_orders \
        --context "Failed 3 times in the last 30 minutes"

Environment variables (required):
    PARADIME_API_ENDPOINT  – Paradime API endpoint
    PARADIME_API_KEY       – Paradime API key
    PARADIME_API_SECRET    – Paradime API secret
"""

import argparse
import os
import sys

from paradime import Paradime

AGENT_NAME    = "incident-commander"
SLACK_CHANNEL = os.environ.get("INCIDENT_SLACK_CHANNEL", "#incidents")


def _parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Trigger the incident-commander agent.")
    parser.add_argument("--schedule", required=True, help="Bolt schedule or DAG name that is failing")
    parser.add_argument("--model",    required=True, help="Suspect dbt™ model name (e.g. fct_orders)")
    parser.add_argument("--context",  default="",    help="Any additional context about the failure")
    return parser.parse_args()


def main() -> None:
    args = _parse_args()

    required = {
        "PARADIME_API_ENDPOINT": os.environ.get("PARADIME_API_ENDPOINT"),
        "PARADIME_API_KEY":      os.environ.get("PARADIME_API_KEY"),
        "PARADIME_API_SECRET":   os.environ.get("PARADIME_API_SECRET"),
    }
    missing = [k for k, v in required.items() if not v]
    if missing:
        print(f"ERROR: Missing env vars: {', '.join(missing)}", file=sys.stderr)
        raise SystemExit(1)

    paradime = Paradime(
        api_endpoint=required["PARADIME_API_ENDPOINT"],
        api_key=required["PARADIME_API_KEY"],
        api_secret=required["PARADIME_API_SECRET"],
    )

    message = (
        f"INCIDENT: schedule '{args.schedule}' is failing. "
        f"Suspect model is {args.model}. "
        f"{args.context}\n\n"
        f"Coordinate triage across log-analyzer, query-profiler, and "
        f"owner-notifier. Post a unified ROOT CAUSE / IMPACT / OWNER / "
        f"NEXT ACTION summary to {SLACK_CHANNEL}."
    )

    print(f"[incident-commander] 🚨  Triggering triage for {args.schedule}/{args.model} …")

    run = paradime.dinoai_agents.trigger_run_and_wait(
        agent=AGENT_NAME,
        message=message,
        slack_channel=SLACK_CHANNEL,
        timeout=1800,
    )

    print("\n" + "=" * 72)
    print("INCIDENT TRIAGE REPORT")
    print("=" * 72)
    print(run.messages[-1].content)
    print("=" * 72)


if __name__ == "__main__":
    main()

"""
watch_incident.py
-----------------
Triggers the incident commander and streams sub-agent progress in real time.
Each child session is polled independently so you see findings as they arrive
rather than waiting for all three to complete.
"""

import os
import time

from paradime import Paradime
from paradime.apis.dinoai_agents.types import DinoaiAgentRunStatus

paradime = Paradime(
    api_endpoint=os.environ["PARADIME_API_ENDPOINT"],
    api_key=os.environ["PARADIME_API_KEY"],
    api_secret=os.environ["PARADIME_API_SECRET"],
)

# Trigger without waiting — returns immediately with the parent session ID
result    = paradime.dinoai_agents.trigger_run(
    agent="incident-commander",
    message=(
        "INCIDENT: schedule 'hourly_marts' has failed 3 times in the last "
        "30 minutes. Suspect model is fct_orders. Coordinate the triage "
        "and post a unified summary in #incidents."
    ),
)
parent_id = result.agent_session_id
print(f"[commander] 🚨  Session started: {parent_id}\n")

# Poll the parent — as sub-agents are spawned, child_session_ids populates
seen_children: set[str] = set()

while True:
    parent = paradime.dinoai_agents.get_run(agent_session_id=parent_id)

    # Stream any newly spawned child sessions
    for child_id in (parent.child_session_ids or []):
        if child_id in seen_children:
            continue
        seen_children.add(child_id)

        child = paradime.dinoai_agents.get_run(agent_session_id=child_id)
        print(f"--- sub-agent spawned: {child_id} ({child.status.value}) ---")
        for msg in child.messages[-3:]:
            print(f"  [{msg.role}] {msg.content[:200]}")

    if parent.status in (DinoaiAgentRunStatus.COMPLETED, DinoaiAgentRunStatus.FAILED):
        break

    time.sleep(15)

print("\n" + "=" * 72)
print("FINAL TRIAGE REPORT")
print("=" * 72)
print(parent.messages[-1].content)