You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by vi...@apache.org on 2026/05/27 14:26:31 UTC
(airflow) branch aip99-aipdev-example created (now abbc745d606)
This is an automated email from the ASF dual-hosted git repository.
vikramkoka pushed a change to branch aip99-aipdev-example
in repository https://gitbox.apache.org/repos/asf/airflow.git
at abbc745d606 Add AIP progress tracker example DAG for common.ai provider
This branch includes the following new commits:
new abbc745d606 Add AIP progress tracker example DAG for common.ai provider
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
(airflow) 01/01: Add AIP progress tracker example DAG for common.ai provider
Posted by vi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
vikramkoka pushed a commit to branch aip99-aipdev-example
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit abbc745d606be86254e7fc39e6d39144e1971b1c
Author: Vikram Koka <vi...@astronomer.io>
AuthorDate: Wed May 27 07:26:06 2026 -0700
Add AIP progress tracker example DAG for common.ai provider
Demonstrates LLMOperator with Dynamic Task Mapping, structured output, UsageLimits, and HITL approval — no framework dependency.
The DAG gathers AIP specs from the Confluence wiki and searches GitHub for related PRs, then uses an LLM to assess each AIP's
progress and synthesize a cross-AIP report for maintainer review.
---
.../example_dags/example_aip_progress_tracker.py | 475 +++++++++++++++++++++
1 file changed, 475 insertions(+)
diff --git a/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_aip_progress_tracker.py b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_aip_progress_tracker.py
new file mode 100644
index 00000000000..114b56302f6
--- /dev/null
+++ b/providers/common/ai/src/airflow/providers/common/ai/example_dags/example_aip_progress_tracker.py
@@ -0,0 +1,475 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+AIP progress tracker -- multi-source data fusion with common.ai operators.
+
+Demonstrates Dynamic Task Mapping, structured LLM output, cost-controlled
+synthesis, and HITL approval using only ``LLMOperator`` -- no LlamaIndex or
+LangChain dependency required.
+
+For each active Airflow Improvement Proposal the Dag gathers evidence from
+two sources (Confluence spec text, GitHub PRs and commits), asks an LLM to
+assess spec-vs-implementation progress, then synthesizes a cross-AIP report
+for maintainer review.
+
+``example_aip_progress_tracker`` (manual trigger):
+
+.. code-block:: text
+
+ fetch_aip_list (@task)
+ → gather_aip_evidence (@task, mapped ×N AIPs)
+ → format_analysis_prompt (@task, mapped ×N)
+ → analyze_aip (LLMOperator, mapped ×N)
+ → collect_analyses (@task)
+ → synthesize_report (LLMOperator, with UsageLimits)
+ → review_report (ApprovalOperator)
+
+**What this makes visible that a notebook hides:**
+
+* Each AIP investigation is a named, logged task instance with its own
+ retry behaviour -- not a loop iteration buried inside one cell.
+* If the GitHub API is rate-limited for one AIP, only that mapped
+ instance retries; the others preserve their XCom results.
+* The synthesis step's inputs and token budget are fully auditable.
+* A maintainer reviews the report before it goes to the dev list.
+
+Before running:
+
+1. Create an LLM connection named ``pydanticai_default`` (or the value of
+ ``LLM_CONN_ID``) for your chosen model provider.
+2. Set ``USE_SAMPLE_DATA = False`` in the DAG file to fetch live data
+ from the Apache Confluence wiki and GitHub API.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+import urllib.parse
+import urllib.request
+from datetime import timedelta
+
+from pydantic import BaseModel
+from pydantic_ai.usage import UsageLimits
+
+from airflow.providers.common.ai.operators.llm import LLMOperator
+from airflow.providers.common.compat.sdk import dag, task
+from airflow.providers.standard.operators.hitl import ApprovalOperator
+
+# ---------------------------------------------------------------------------
+# Configuration
+# ---------------------------------------------------------------------------
+
+LLM_CONN_ID = "pydanticai_default"
+
+# Confluence wiki -- the AIP listing page is public, no auth required.
+# https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals
+CONFLUENCE_BASE_URL = "https://cwiki.apache.org/confluence"
+AIP_LISTING_PAGE_ID = "89066602"
+GITHUB_REPO = "apache/airflow"
+
+# When True the Dag runs on built-in sample data (self-contained, no network).
+# Set to False to fetch live data from Confluence and GitHub.
+USE_SAMPLE_DATA = True
+
+# ---------------------------------------------------------------------------
+# Structured output model -- enforces a schema on the per-AIP LLM response
+# ---------------------------------------------------------------------------
+
+# [START aip_tracker_structured_output]
+
+
+class AIPStatus(BaseModel):
+ """Per-AIP analysis produced by the LLM."""
+
+ aip_number: int
+ title: str
+ spec_summary: str
+ implementation_status: str
+ key_prs: list[str]
+ blockers: list[str]
+ next_steps: list[str]
+ completion_pct: int
+
+
+# [END aip_tracker_structured_output]
+
+# ---------------------------------------------------------------------------
+# Sample data -- replace with Confluence / GitHub API calls for production
+# ---------------------------------------------------------------------------
+
+SAMPLE_AIPS: list[dict] = [
+ {"aip_number": 76, "title": "Asset Partitions"},
+ {"aip_number": 99, "title": "Common Data Access Pattern + AI"},
+ {"aip_number": 103, "title": "Task State Management"},
+ {"aip_number": 105, "title": "LLM Retry Policy"},
+ {"aip_number": 108, "title": "Language Coordinator Layer"},
+]
+
+SAMPLE_EVIDENCE: dict[int, dict] = {
+ 76: {
+ "spec_text": (
+ "AIP-76 adds partition awareness to Airflow assets. Instead of "
+ "triggering on any update to an asset, Dags can depend on specific "
+ "partitions (e.g. a date-based slice of a dataset). The scheduler "
+ "tracks which partitions have been produced and only triggers "
+ "downstream Dags when the required partitions are available."
+ ),
+ "prs": [
+ "#62400 -- Asset partition model and metadata schema",
+ "#63900 -- Partition-aware scheduling in the DagRun creator",
+ "#65100 -- UI: partition status badges on Asset views",
+ ],
+ "commits": [
+ "Add AssetPartition model with composite key",
+ "Extend DagScheduleAssetReference for partition filters",
+ "Show partition status in React Asset detail view",
+ ],
+ },
+ 99: {
+ "spec_text": (
+ "AIP-99 adds first-class AI/ML operators to the common.ai provider. "
+ "LLMOperator wraps pydantic-ai for structured LLM calls with retries. "
+ "AgentOperator enables multi-turn ReAct agents with tool use. "
+ "LangChain and LlamaIndex hooks bridge framework models to Airflow "
+ "connections. DocumentLoaderOperator parses files for RAG pipelines."
+ ),
+ "prs": [
+ "#61200 -- LLMOperator and PydanticAIHook",
+ "#62800 -- AgentOperator with tool calling and HITL",
+ "#64100 -- LangChain hook and integration",
+ "#65500 -- LlamaIndex embedding and retrieval operators",
+ "#66300 -- DocumentLoaderOperator for multi-format parsing",
+ ],
+ "commits": [
+ "Add LLMOperator with structured output support",
+ "Add AgentOperator with ReAct loop and durable execution",
+ "Add LangChainHook bridging langchain models to connections",
+ "Add LlamaIndex embedding and retrieval operators",
+ "Add DocumentLoaderOperator with PDF and DOCX support",
+ ],
+ },
+ 103: {
+ "spec_text": (
+ "AIP-103 introduces task-level state persistence via "
+ "context['task_state']. Tasks can checkpoint intermediate results "
+ "that survive retries and restarts. The state backend stores "
+ "key-value pairs scoped to a task instance, enabling long-running "
+ "tasks to resume from the last checkpoint rather than starting "
+ "from scratch."
+ ),
+ "prs": [
+ "#65000 -- Task state storage backend and API",
+ "#65800 -- context['task_state'] integration in Task SDK",
+ "#66700 -- State-aware retry logic for LLM tasks",
+ ],
+ "commits": [
+ "Add TaskState model with key-value storage",
+ "Expose task_state in TaskInstanceContext",
+ "Add state checkpoint and restore in retry path",
+ ],
+ },
+ 105: {
+ "spec_text": (
+ "AIP-105 introduces LLMRetryPolicy, an intelligent retry mechanism "
+ "that uses an LLM to classify task failures before deciding whether "
+ "to retry, fail fast, or back off. Instead of static exception-type "
+ "matching, the policy sends the error context to an LLM that "
+ "determines the appropriate action: rate-limit errors trigger "
+ "exponential backoff, auth errors fail immediately, transient "
+ "network errors retry with a short delay."
+ ),
+ "prs": [
+ "#64800 -- RetryPolicy base class in Task SDK (AIP-105 prerequisite)",
+ "#65600 -- LLMRetryPolicy with pydantic-ai error classification",
+ "#66200 -- Integration tests for LLM-classified retry scenarios",
+ ],
+ "commits": [
+ "Add RetryPolicy protocol and RetryRule dataclass",
+ "Implement LLMRetryPolicy with structured ErrorClassification output",
+ "Wire RetryPolicy into task runner retry loop",
+ ],
+ },
+ 108: {
+ "spec_text": (
+ "AIP-108 defines a language coordinator layer that enables tasks "
+ "written in Java, Go, and TypeScript to run alongside Python tasks. "
+ "A lightweight coordinator process manages the non-Python runtime "
+ "lifecycle, handles serialization between the task and the Execution "
+ "API, and provides the same guarantees (heartbeat, state, XCom) that "
+ "Python tasks get from the Task SDK."
+ ),
+ "prs": [
+ "#66100 -- Coordinator protocol specification and protobuf schema",
+ "#66800 -- Java Task SDK with coordinator bridge",
+ "#67200 -- Go Task SDK initial implementation",
+ ],
+ "commits": [
+ "Define coordinator gRPC protocol for multi-language tasks",
+ "Add Java Task SDK with Maven build and coordinator client",
+ "Scaffold Go Task SDK with coordinator handshake",
+ ],
+ },
+}
+
+# ---------------------------------------------------------------------------
+# HTTP helpers -- used when USE_SAMPLE_DATA is False
+# ---------------------------------------------------------------------------
+
+
+def _confluence_rest_get(path: str) -> dict:
+ """GET a Confluence REST API endpoint (public, no auth required)."""
+ url = f"{CONFLUENCE_BASE_URL}{path}"
+ req = urllib.request.Request(url, headers={"Accept": "application/json"})
+ with urllib.request.urlopen(req, timeout=30) as resp:
+ return json.loads(resp.read())
+
+
+def _github_api_get(path: str) -> dict:
+ """GET a GitHub REST API endpoint (public, rate-limited to 10 req/min)."""
+ url = f"https://api.github.com{path}"
+ req = urllib.request.Request(url, headers={"Accept": "application/vnd.github.v3+json"})
+ with urllib.request.urlopen(req, timeout=30) as resp:
+ return json.loads(resp.read())
+
+
+def _strip_html_tags(html: str) -> str:
+ """Remove HTML/Confluence markup, returning plain text."""
+ text = re.sub(r"<[^>]+>", " ", html)
+ return re.sub(r"\s+", " ", text).strip()
+
+
+def _parse_accepted_aips(listing_html: str) -> list[dict]:
+ """Extract accepted AIPs from the rendered AIP listing page."""
+ match = re.search(r"Accepted AIPs.*?(?=<h[1-3]|$)", listing_html, re.DOTALL | re.IGNORECASE)
+ if not match:
+ return []
+ section = match.group()
+ aips = []
+ for m in re.finditer(r"<a[^>]+>([^<]*AIP-(\d+)[^<]*)</a>", section):
+ title = re.sub(r"\s+", " ", m.group(1)).strip()
+ aip_number = int(m.group(2))
+ aips.append({"aip_number": aip_number, "title": title})
+ return aips
+
+
+# ---------------------------------------------------------------------------
+# System prompts
+# ---------------------------------------------------------------------------
+
+ANALYSIS_SYSTEM_PROMPT = """\
+You are an Airflow project analyst. Given an AIP specification and its \
+GitHub evidence (pull requests and commits), produce a structured status \
+assessment.
+
+Be specific about what has been implemented versus what remains. Rate \
+completion percentage based on the ratio of spec goals that have \
+corresponding PRs or commits."""
+
+SYNTHESIS_SYSTEM_PROMPT = """\
+You are an Airflow release coordinator. Given individual AIP status \
+assessments, produce a concise cross-AIP progress report.
+
+Identify the top priorities, shared blockers across AIPs, and recommend \
+where maintainer attention is most needed. Keep the report actionable \
+and under 500 words."""
+
+
+# ---------------------------------------------------------------------------
+# DAG
+# ---------------------------------------------------------------------------
+
+
+# [START example_aip_progress_tracker]
+@dag(catchup=False, tags=["example", "aip_tracker", "common_ai"])
+def example_aip_progress_tracker():
+ """
+ Track AIP progress by analysing Confluence specs against GitHub evidence.
+
+ Task graph::
+
+ fetch_aip_list (@task)
+ → gather_aip_evidence (@task ×N, via Dynamic Task Mapping)
+ → format_analysis_prompt (@task ×N)
+ → analyze_aip (LLMOperator ×N, structured output)
+ → collect_analyses (@task)
+ → synthesize_report (LLMOperator, with UsageLimits)
+ → review_report (ApprovalOperator)
+ """
+
+ # ------------------------------------------------------------------
+ # Step 1: Fetch the list of active AIPs to investigate.
+ # The length of this list determines how many mapped instances are
+ # created in the downstream steps -- N is decided at runtime.
+ # ------------------------------------------------------------------
+ @task
+ def fetch_aip_list() -> list[dict]:
+ if USE_SAMPLE_DATA:
+ return SAMPLE_AIPS
+ # Fetch the AIP listing page and extract the "Accepted" section.
+ # https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals
+ page = _confluence_rest_get(f"/rest/api/content/{AIP_LISTING_PAGE_ID}?expand=body.view")
+ accepted = _parse_accepted_aips(page["body"]["view"]["value"])
+ return accepted or SAMPLE_AIPS
+
+ aip_list = fetch_aip_list()
+
+ # ------------------------------------------------------------------
+ # Step 2: Gather evidence for each AIP from multiple sources.
+ # Each mapped instance fetches one AIP's spec text from the
+ # Confluence wiki (cwiki.apache.org) and searches GitHub for
+ # related PRs and commits. If the GitHub API is rate-limited
+ # for one AIP, only that instance retries.
+ # ------------------------------------------------------------------
+ @task
+ def gather_aip_evidence(aip: dict) -> dict:
+ aip_number = aip["aip_number"]
+ if USE_SAMPLE_DATA:
+ evidence = SAMPLE_EVIDENCE[aip_number]
+ return {
+ "aip_number": aip_number,
+ "title": aip["title"],
+ "spec_text": evidence["spec_text"],
+ "prs": evidence["prs"],
+ "commits": evidence["commits"],
+ }
+ # Fetch spec text from the AIP's Confluence wiki page via CQL search.
+ # Example: https://cwiki.apache.org/confluence/display/AIRFLOW/AIP-103
+ cql = urllib.parse.quote(
+ f'space="AIRFLOW" AND title~"AIP-{aip_number}" AND ancestor={AIP_LISTING_PAGE_ID}'
+ )
+ results = _confluence_rest_get(f"/rest/api/content/search?cql={cql}&expand=body.view&limit=1")
+ spec_text = ""
+ if results.get("results"):
+ raw_html = results["results"][0]["body"]["view"]["value"]
+ spec_text = _strip_html_tags(raw_html)[:3000]
+ # Search GitHub for related PRs.
+ pr_query = urllib.parse.quote(f"AIP-{aip_number} repo:{GITHUB_REPO} is:pr")
+ pr_data = _github_api_get(f"/search/issues?q={pr_query}&per_page=10")
+ prs = [f"#{it['number']} -- {it['title']}" for it in pr_data.get("items", [])]
+ # Search GitHub for related commits.
+ commit_query = urllib.parse.quote(f"AIP-{aip_number} repo:{GITHUB_REPO}")
+ commit_data = _github_api_get(f"/search/commits?q={commit_query}&per_page=10")
+ commits = [it["commit"]["message"].split("\n")[0] for it in commit_data.get("items", [])]
+ return {
+ "aip_number": aip_number,
+ "title": aip["title"],
+ "spec_text": spec_text,
+ "prs": prs,
+ "commits": commits,
+ }
+
+ evidence = gather_aip_evidence.expand(aip=aip_list)
+
+ # ------------------------------------------------------------------
+ # Step 3: Format the gathered evidence into an LLM analysis prompt.
+ # Separating formatting from data gathering keeps each task focused
+ # and makes prompt iteration independent of API logic.
+ # ------------------------------------------------------------------
+ @task
+ def format_analysis_prompt(evidence: dict) -> str:
+ prs_text = "\n".join(f" - {pr}" for pr in evidence["prs"])
+ commits_text = "\n".join(f" - {c}" for c in evidence["commits"])
+ return (
+ f"Analyze AIP-{evidence['aip_number']}: {evidence['title']}\n\n"
+ f"Specification:\n{evidence['spec_text']}\n\n"
+ f"Pull Requests:\n{prs_text}\n\n"
+ f"Recent Commits:\n{commits_text}"
+ )
+
+ prompts = format_analysis_prompt.expand(evidence=evidence)
+
+ # ------------------------------------------------------------------
+ # Step 4: Analyze each AIP with a structured LLM call.
+ # Dynamic Task Mapping creates one LLMOperator instance per AIP.
+ # output_type=AIPStatus enforces the Pydantic schema on the response.
+ # ------------------------------------------------------------------
+ # [START aip_tracker_dtm_analysis]
+ analyses = LLMOperator.partial(
+ task_id="analyze_aip",
+ llm_conn_id=LLM_CONN_ID,
+ system_prompt=ANALYSIS_SYSTEM_PROMPT,
+ output_type=AIPStatus,
+ ).expand(prompt=prompts)
+ # [END aip_tracker_dtm_analysis]
+
+ # ------------------------------------------------------------------
+ # Step 5: Collect all per-AIP analyses into a single context string
+ # for the synthesis step.
+ # ------------------------------------------------------------------
+ @task
+ def collect_analyses(analyses: list) -> str:
+ sections = []
+ for raw in analyses:
+ a = json.loads(raw) if isinstance(raw, str) else raw
+ blockers = ", ".join(a["blockers"]) if a["blockers"] else "None identified"
+ next_steps = ", ".join(a["next_steps"]) if a["next_steps"] else "N/A"
+ sections.append(
+ f"## AIP-{a['aip_number']}: {a['title']}\n"
+ f"Status: {a['implementation_status']} "
+ f"({a['completion_pct']}% complete)\n"
+ f"Summary: {a['spec_summary']}\n"
+ f"Key PRs: {', '.join(a['key_prs'])}\n"
+ f"Blockers: {blockers}\n"
+ f"Next steps: {next_steps}"
+ )
+ return "\n\n".join(sections)
+
+ collected = collect_analyses(analyses.output)
+
+ # ------------------------------------------------------------------
+ # Step 6: Synthesize a cross-AIP progress report.
+ # UsageLimits caps the token spend so a runaway prompt cannot
+ # exhaust the API budget in a single Dag run.
+ # ------------------------------------------------------------------
+ # [START aip_tracker_synthesis]
+ synthesize = LLMOperator(
+ task_id="synthesize_report",
+ llm_conn_id=LLM_CONN_ID,
+ system_prompt=SYNTHESIS_SYSTEM_PROMPT,
+ prompt="""\
+Create a cross-AIP progress report from these individual assessments.
+Prioritize AIPs that are close to completion or have shared blockers.
+
+{{ ti.xcom_pull(task_ids='collect_analyses') }}""",
+ usage_limits=UsageLimits(
+ request_limit=5,
+ input_tokens_limit=20_000,
+ output_tokens_limit=4_000,
+ ),
+ )
+ # [END aip_tracker_synthesis]
+ collected >> synthesize
+
+ # ------------------------------------------------------------------
+ # Step 7: A maintainer reviews the synthesized report before it is
+ # shared on the dev list. The Dag pauses here until the human
+ # approves, requests changes, or the timeout expires.
+ # ------------------------------------------------------------------
+ # [START aip_tracker_hitl]
+ ApprovalOperator(
+ task_id="review_report",
+ subject="Review AIP Progress Report before sharing",
+ body=synthesize.output,
+ response_timeout=timedelta(hours=24),
+ )
+ # [END aip_tracker_hitl]
+
+
+# [END example_aip_progress_tracker]
+
+example_aip_progress_tracker()