The airflow-ai-sdk comes with several example DAGs that demonstrate how to use the various decorators and features. All examples can be found in the examples/dags directory of the repository.
Demonstrates using @task.llm to summarize GitHub commits with a language model.
Features demonstrated:
- Using
@task.llmdecorator with a specific model (gpt-4o-mini) - Specifying a detailed system prompt for commit summarization
- Transforming task input to LLM input (joining commits into a string)
- Weekly scheduling for regular changelog generation
Shows how to use @task.llm with Pydantic models to parse structured data from product feedback.
Features demonstrated:
- Creating a custom Pydantic model for structured LLM output
- PII masking in preprocessing
- Using
output_typeparameter with a Pydantic model - Task mapping with
expand()to process multiple feedback items - Conditional execution with
AirflowSkipException
Illustrates how to use @task.agent with custom tools to perform deep research on topics.
Features demonstrated:
- Creating a custom agent with tools
- Using the
@task.agentdecorator - Creating custom tools (web content fetching)
- Integrating with external APIs (DuckDuckGo search)
- Using runtime parameters for dynamic DAG execution
Demonstrates how to use @task.llm_branch to route support tickets based on priority.
Features demonstrated:
- Using
@task.llm_branchfor conditional workflow routing - Setting up branch tasks based on LLM decisions
- Configuring
allow_multiple_branchesparameter - Processing DAG run configuration parameters
- Detailed prompt engineering for classification tasks
Shows how to use @task.embed to create vector embeddings from text.
Features demonstrated:
- Using
@task.embeddecorator - Specifying embedding model parameters
- Generating vector embeddings from text
- Configuring encoding parameters (normalization)
- Task mapping to process multiple texts in parallel
The examples can be run in a local Airflow environment with the following steps:
-
Clone the examples repository:
git clone https://github.com/astronomer/ai-sdk-examples.git
-
Navigate to the examples directory:
cd ai-sdk-examples -
Start the Airflow environment:
astro dev start
For more information on the individual examples, refer to the code comments in each example file.