用 Airflow 构建 LLM 工作流和 Agent

⚠️ 项目已停止维护

Astronomer 已不再积极维护本项目。目前开发工作已经暂停，我们不再接受新的贡献、bug 修复或版本发布。

我们建议迁移至 apache-airflow-providers-common-ai，这是 Apache Airflow 官方提供的 AI 和 LLM 工作流 provider，也是目前新功能的开发所在。有关分步操作说明，请参阅迁移指南。

代码仍保留在这里，你可以根据其许可证条款进行探索、fork 和修改。请注意，它可能无法兼容最新的依赖项或平台，也可能包含安全漏洞。Astronomer 无法为其使用提供任何保证或担保。

如果你有兴趣接手或维护这个项目，我们很乐意与你交流，请通过 oss@astronomer.io 联系我们。感谢你参与这段开源旅程，并帮助优秀的创意延续下去！

这是一个用于在 Apache Airflow 中使用 LLM 的 Python SDK。用户可以通过基于 decorator 的 task，直接在 Airflow pipeline 中调用 LLM，并编排 Agent 调用。

我们发现，在生产环境中对 LLM 工作流和 Agent 进行编排与观测时，依赖 Airflow 这类成熟的编排工具通常很有帮助。这是因为 LLM 工作流与 ETL pipeline、运维流程和 ML 工作流等更传统的工作流具有相同的组织形态。

pip install airflow-ai-sdk[openai]

如果安装时不包含任何可选依赖，你将获得该 package 的精简版本。所有可用的可选依赖都列在 pyproject.toml 中。

使用 @task.llm 创建 LLM task：定义调用语言模型处理文本的 task

使用 @task.agent 创建 Agent task：通过自定义工具编排多步骤 AI 推理

自动解析输出：使用类型提示自动解析并验证 LLM 输出

使用 @task.llm_branch 实现分支：根据 LLM 输出改变 DAG 的控制流

模型支持：支持 Pydantic AI library 中的所有模型（OpenAI、Anthropic、Gemini 等）

使用 @task.embed 创建 embedding task：根据文本生成向量 embedding

你可以在 Airflow AI SDK Decorators & Code Snippets 快速笔记中，找到这些 decorator 的更多信息，以及每个 decorator 对应的完整 DAG 示例！

from typing import Literal
import pendulum
from airflow.decorators import dag, task
from airflow.models.dagrun import DagRun

@task.llm(
    model="gpt-4o-mini",
    output_type=Literal["positive", "negative", "neutral"],
    system_prompt="Classify the sentiment of the given text.",
)
def process_with_llm(dag_run: DagRun) -> str:
    input_text = dag_run.conf.get("input_text")

    # can do pre-processing here (e.g. PII redaction)
    return input_text

@dag(
    schedule=None,
    start_date=pendulum.datetime(2025, 1, 1),
    catchup=False,
    params={"input_text": "I'm very happy with the product."},
)
def sentiment_classification():
    process_with_llm()

sentiment_classification()

如果想通过一个完整的示例环境快速上手，请查看 examples repository。它提供了一个安装了 AI SDK 的完整本地 Airflow 实例，以及 5 个示例 pipeline：

git clone https://github.com/astronomer/ai-sdk-examples.git
cd ai-sdk-examples
astro dev start

如果你尚未安装 Astro CLI，请运行 brew install astro，或查看此处提供的其他安装方式。

有关详细文档，请参阅 docs 目录：