Calling out to LLMs and AI Agents from within tasks in Airflow DAGs has never been this easy!
The new Airflow AI SDK is an open-source package that uses PydanticAI to interact with compatible LLMs and Agents, orchestrated with Airflow!
In this Quick Notes document, you’ll learn how to set up a project using the AI SDK and get to know the available decorators, including a full example DAG for each one:
- @task.llm - which lets you post inference calls to LLM models to generate text
- @task.llm_branch - which lets you utilize an LLM to decide how to continue in your DAG based on a provided input
- @task.agent - which lets you orchestrate AI Agents from Airflow tasks
- @task.embed - which provides a handy way to generate vector embeddings with models from the sentence-transformer package
All in a brief, handy reference document, made for the busy data engineer who just wants to get to writing pipelines.