LLM Workflows then Agents: Getting Started with Apache Airflow

This repository contains an SDK for working with LLMs from Apache Airflow, based on Pydantic AI. It allows users to call LLMs and orchestrate agent calls directly within their Airflow pipelines using decorator-based tasks. The SDK leverages the familiar Airflow @task syntax with extensions like @task.llm, @task.llm_branch, and @task.agent.

To get started, check out the examples repository here, which offers a full local Airflow instance with the AI SDK installed and 5 example pipelines. To run this locally, run:

If you don't have the Astro CLI installed, run brew install astro (or see other options here).

If you already have Airflow running, you can also install the package with any optional dependencies you need:

Note that installing the package with no optional dependencies will install the slim version of the package, which does not include any LLM models or tools. The available optional packages are listed here. While this SDK offers the optional dependencies for convenience sake, you can also install the optional dependencies from Pydantic AI directly.

Table of Contents:

• Features
• Motivation
• Examples
  • LLM calls from a DAG (summarize Airflow's commits)
  • LLM calls with structured outputs using @task.llm (user feedback -> sentiment and feature requests)
  • Agent calls with @task.agent (deep research agent)
  • Changing dag control flow with @task.llm_branch (support ticket routing)
• Future Work

• LLM tasks with @task.llm: Define tasks that call language models (e.g. GPT-3.5-turbo) to process text.
• Agent tasks with @task.agent: Orchestrate multi-step AI reasoning by leveraging custom tools.
• Automatic output parsing: Use function type hints (including Pydantic models) to automatically parse and validate LLM outputs.
• Branching with @task.llm_branch: Change the control flow of a DAG based on the output of an LLM.
• Model support: Support for all models in the Pydantic AI library (OpenAI, Anthropic, Gemini, Ollama, Groq, Mistral, Cohere, Bedrock)

We follow the taskflow pattern of Airflow with three decorators:

• @task.llm: Define a task that calls an LLM. Under the hood, this creates a Pydantic AI Agent with no tools.
• @task.agent: Define a task that calls an agent. You can pass in a Pydantic AI Agent directly.
• @task.llm_branch: Define a task that branches the control flow of a DAG based on the output of an LLM. Enforces that the LLM output is one of the downstream task_ids.

The function supplied to each decorator is a translation function that converts the Airflow task's input into the LLM's input. If you don't want to do any translation, you can just return the input unchanged.

AI workflows are becoming increasingly common as organizations look for pragmatic ways to get value out of LLMs. As with any workflow, it's important to have a flexible and scalable way to orchestrate them.

Airflow is a popular choice for orchestrating data pipelines. It's a powerful tool for managing the dependencies between tasks and for scheduling and monitoring them, and has been trusted by data teams everywhere for 10+ years. It comes "batteries included" with a rich set of capabilities:

• Flexible scheduling: run tasks on a fixed schedule, on-demand, or based on external events
• Dynamic task mapping: easily process multiple inputs in parallel with full error handling and observability
• Branching and conditional logic: change the control flow of a DAG based on the output of certain tasks
• Error handling: built-in support for retries, exponential backoff, and timeouts
• Resource management: limit the concurrency of tasks with Airflow Pools
• Monitoring: detailed logs and monitoring capabilities
• Scalability: designed for production workflows

This SDK is designed to make it easy to integrate LLM workflows into your Airflow pipelines. It allows you to do anything from simple LLM calls to complex agentic workflows.

See the full set of examples in the examples/dags directory.

This example shows how to use the @task.llm decorator as part of an Airflow DAG. In the @task.llm decorator, we can specify a model and system prompt. The decorator allows you to transform the Airflow task's input into the LLM's input.

See full example: github_changelog.py

This example demonstrates how to use the @task.llm decorator to call an LLM and return a structured output. In this case, we're using a Pydantic model to validate the output of the LLM. We recommend using the airflow_ai_sdk.BaseModel class to define your Pydantic models in case we add more functionality in the future.

See full example: product_feedback_summarization.py

This example shows how to build an AI agent that can autonomously invoke external tools (e.g., a knowledge base search) when answering a user question.

See full example: deep_research.py

This example demonstrates how to use the @task.llm_branch decorator to change the control flow of a DAG based on the output of an LLM. In this case, we're routing support tickets based on the severity of the ticket.

See full example: support_ticket_routing.py