Investigate with Otto

"""Triggered by Astro DAG trigger notifications to investigate a failed DAG run and post the diagnosis to Slack.

Required settings. Each can be set as an environment variable OR an Airflow
Variable of the same name (set BEFORE enabling the alert):
- ASTRO_ORGANIZATION_ID   (your Astro organization ID)
- ASTRO_DEPLOYMENT_ID     (deployment to run diagnoses against)
- ASTRO_API_TOKEN         (Astro API token with labs/v1 access)
- SLACK_WEBHOOK_URL       (Slack incoming webhook URL)
- SLACK_CHANNEL           (optional; overrides the webhook's default channel)

Set these in the Astro UI under Deployment → Environment (Environment Variables
or Airflow Variables). Missing any required value will raise
`ValueError: Missing required setting ...` at task runtime.

Expected dag_run.conf keys (sent by the Astro Alert trigger):
- alertId              (Astro alert ID; surfaced in the Slack message)
- alertType            (DAG_FAILURE — other alert types are ignored)
- dagName              (target DAG being triggered — this DAG; not used, the failed source DAG comes from `message`)
- message              (free-form alert text; the failed source DAG ID is parsed from `"... for DAG <name>"`)
- airflowDagRunId      (the failed source DAG's run ID, e.g. `scheduled__2026-06-03T23:40:00+00:00`)
- logFailureSummaries  (dict keyed by failed task ID; if exactly one key, it's used as the failed task ID)

Example payload (DAG_FAILURE alert from `always_failing_dag`):
    {
      "alertId": "cmpyms67v000m01pxkfaqoe5t",
      "alertType": "DAG_FAILURE",
      "dagName": "alert_investigation_agent",
      "message": "DAG run failed for DAG always_failing_dag",
      "airflowDagRunId": "scheduled__2026-06-03T23:40:00+00:00",
      "logFailureSummaries": {"fail_on_purpose": "..."}
    }
"""

from __future__ import annotations

import json
import logging
import os
import re
from typing import Any

import requests
from airflow.sdk import Variable, dag, task
from pendulum import datetime

ASTRO_API_BASE = "https://api.astronomer.io/labs/v1"
# This demo investigates DAG-failure alerts. Other alert types are logged
# and skipped.
ALLOWED_ALERT_TYPES = {"DAG_FAILURE"}

log = logging.getLogger(__name__)


def _get_setting(key: str) -> str:
    # Resolve a setting from an environment variable or an Airflow Variable of
    # the same name. Env var wins if both are set.
    env_value = os.environ.get(key, "")
    if env_value:
        return env_value

    try:
        return Variable.get(key)
    except Exception:
        return ""


def _get_required_setting(key: str) -> str:
    value = _get_setting(key)
    if value:
        return value
    raise ValueError(f"Missing required setting: '{key}' (set it as an env var or Airflow Variable)")


def _extract_source_dag_id(message: str) -> str:
    # Astro Alert payloads don't include the failed source DAG ID as a
    # top-level key (`dagName` in conf is the *target* DAG being triggered,
    # i.e. this one). Parse it out of the DAG_FAILURE alert `message`:
    #   "DAG run failed for DAG <dag_id>"
    match = re.search(r"for DAG\s+['\"]?([a-zA-Z0-9_.-]+)['\"]?", message, re.IGNORECASE)
    if match:
        return match.group(1)

    raise ValueError(
        "Could not determine the failed source DAG ID from the alert message. "
        f"Expected a 'for DAG <name>' phrase; got message={message!r}"
    )


def _extract_run_id(conf: dict[str, Any]) -> str:
    value = conf.get("airflowDagRunId")
    if value:
        return str(value)

    raise ValueError("Missing required dag_run.conf key: airflowDagRunId")


def _extract_task_id(conf: dict[str, Any]) -> str | None:
    # DAG_FAILURE payloads include a `logFailureSummaries` dict keyed by the
    # failed task ID(s). If exactly one task failed, use it.
    summaries = conf.get("logFailureSummaries")
    if isinstance(summaries, dict) and len(summaries) == 1:
        return next(iter(summaries))

    return None


def _call_investigation_agent(
    organization_id: str,
    deployment_id: str,
    api_token: str,
    dag_id: str,
    run_id: str,
    task_id: str | None,
) -> dict[str, Any]:
    auth_headers = {"Authorization": f"Bearer {api_token}"}
    deployment_base = (
        f"{ASTRO_API_BASE}/organizations/{organization_id}"
        f"/observability/deployments/{deployment_id}/dag-failure-diagnosis/runs"
    )

    start_response = requests.post(
        deployment_base,
        headers={**auth_headers, "Content-Type": "application/json"},
        json={
            "dagId": dag_id,
            "runId": run_id,
            **({"taskId": task_id} if task_id else {}),
        },
        timeout=30,
    )
    start_response.raise_for_status()
    diagnosis_run_id = start_response.json()["runId"]

    response = requests.get(
        f"{deployment_base}/{diagnosis_run_id}/events",
        headers={**auth_headers, "Accept": "text/event-stream"},
        stream=True,
        timeout=(30, 300),
    )
    response.raise_for_status()

    event_type: str | None = None
    data_lines: list[str] = []
    text_chunks: list[str] = []
    diagnosis: dict[str, Any] | None = None

    def flush_event() -> None:
        nonlocal event_type, data_lines, diagnosis
        if not event_type:
            data_lines = []
            return

        payload = "\n".join(data_lines)
        if event_type == "rca_diagnosis" and payload:
            diagnosis = json.loads(payload)
        elif event_type == "text_delta" and payload:
            text_chunks.append(json.loads(payload).get("text", ""))
        elif event_type == "error" and payload:
            raise RuntimeError(json.loads(payload).get("message", "Investigation Agent returned an error"))

        event_type = None
        data_lines = []

    for raw_line in response.iter_lines(decode_unicode=True):
        if raw_line is None:
            continue

        line = raw_line.rstrip("\r")
        if not line:
            flush_event()
            if diagnosis:
                break
            continue

        if line.startswith(":"):
            continue
        if line.startswith("event:"):
            event_type = line.split(":", 1)[1].strip()
            continue
        if line.startswith("data:"):
            data_lines.append(line.split(":", 1)[1].lstrip())

    flush_event()

    if diagnosis:
        return diagnosis
    if text_chunks:
        return {"title": f"Investigation for {dag_id}", "summary": "".join(text_chunks).strip()}

    raise RuntimeError("Investigation Agent returned no diagnosis payload")


def _format_slack_message(
    diagnosis: dict[str, Any],
    alert_id: str,
    alert_type: str,
    dag_id: str,
    run_id: str,
    task_id: str | None,
) -> dict[str, Any]:
    title = diagnosis.get("title") or f"Investigation for {dag_id}"
    summary = diagnosis.get("summary") or "No summary returned."
    suggested_fix = diagnosis.get("suggested_fix") or "No suggested fix returned."
    severity = diagnosis.get("severity") or "UNKNOWN"
    priority = diagnosis.get("priority") or "UNKNOWN"
    root_cause_type = diagnosis.get("root_cause_type") or "UNKNOWN"
    confidence = diagnosis.get("confidence")

    fields = [
        {"type": "mrkdwn", "text": f"*Alert ID*\n`{alert_id}`"},
        {"type": "mrkdwn", "text": f"*Alert Type*\n`{alert_type}`"},
        {"type": "mrkdwn", "text": f"*DAG*\n`{dag_id}`"},
        {"type": "mrkdwn", "text": f"*Run ID*\n`{run_id}`"},
        {"type": "mrkdwn", "text": f"*Severity*\n`{severity}`"},
        {"type": "mrkdwn", "text": f"*Root Cause Type*\n`{root_cause_type}`"},
    ]
    if task_id:
        fields.append({"type": "mrkdwn", "text": f"*Task ID*\n`{task_id}`"})
    if confidence is not None:
        fields.append({"type": "mrkdwn", "text": f"*Confidence*\n`{confidence}`"})
    if priority != "UNKNOWN":
        fields.append({"type": "mrkdwn", "text": f"*Priority*\n`{priority}`"})

    blocks: list[dict[str, Any]] = [
        {"type": "header", "text": {"type": "plain_text", "text": title[:150]}},
        {"type": "section", "fields": fields[:10]},
        {"type": "section", "text": {"type": "mrkdwn", "text": f"*Summary*\n{summary[:2900]}"}},
        {"type": "section", "text": {"type": "mrkdwn", "text": f"*Suggested Fix*\n{suggested_fix[:2900]}"}},
    ]

    text = (
        f"[{severity}] {title}\n"
        f"Root cause type: {root_cause_type}\n"
        f"Summary: {summary}\n\n"
        f"Suggested fix: {suggested_fix}"
    )

    return {"text": text[:4000], "blocks": blocks}


def _post_to_slack(webhook_url: str, channel: str | None, payload: dict[str, Any]) -> None:
    body = dict(payload)
    if channel:
        body["channel"] = channel

    response = requests.post(webhook_url, json=body, timeout=30)
    response.raise_for_status()


@dag(
    dag_id="alert_investigation_agent",
    start_date=datetime(2025, 1, 1),
    schedule=None,
    catchup=False,
    tags=["alerts", "investigation-agent", "slack"],
    default_args={"owner": "Astro", "retries": 0},
    doc_md=__doc__,
)
def alert_investigation_agent():
    @task
    def handle_alert(**context) -> None:
        dag_run = context.get("dag_run")
        conf = dict(dag_run.conf or {}) if dag_run else {}

        alert_id = str(conf.get("alertId", "unknown"))
        alert_type = str(conf.get("alertType", "unknown"))
        message = str(conf.get("message", ""))

        if alert_type not in ALLOWED_ALERT_TYPES:
            log.info(
                "Ignoring alert: alertType %r is not a DAG failure. "
                "This demo only investigates DAG-failure alerts.",
                alert_type,
            )
            return

        dag_id = _extract_source_dag_id(message)
        run_id = _extract_run_id(conf)
        task_id = _extract_task_id(conf)

        organization_id = _get_required_setting("ASTRO_ORGANIZATION_ID")
        deployment_id = _get_required_setting("ASTRO_DEPLOYMENT_ID")
        api_token = _get_required_setting("ASTRO_API_TOKEN")
        slack_webhook_url = _get_required_setting("SLACK_WEBHOOK_URL")
        slack_channel = _get_setting("SLACK_CHANNEL") or None

        diagnosis = _call_investigation_agent(
            organization_id=organization_id,
            deployment_id=deployment_id,
            api_token=api_token,
            dag_id=dag_id,
            run_id=run_id,
            task_id=task_id,
        )

        log.info(
            "Diagnosis completed for %s run %s: title=%s severity=%s root_cause_type=%s summary=%s suggested_fix=%s",
            dag_id,
            run_id,
            diagnosis.get("title") or f"Investigation for {dag_id}",
            diagnosis.get("severity") or "UNKNOWN",
            diagnosis.get("root_cause_type") or "UNKNOWN",
            diagnosis.get("summary") or "No summary returned.",
            diagnosis.get("suggested_fix") or "No suggested fix returned.",
        )

        slack_payload = _format_slack_message(
            diagnosis=diagnosis,
            alert_id=alert_id,
            alert_type=alert_type,
            dag_id=dag_id,
            run_id=run_id,
            task_id=task_id,
        )
        _post_to_slack(slack_webhook_url, slack_channel, slack_payload)

    handle_alert()


alert_investigation_agent()