Editor’s note: Kenten Danas is a speaker for ODSC AI West this October 28th-30th. Be sure to check out her talk, Orchestrating Workflows for GenAI Applications with Apache Airflow, there!

Productionizing GenAI applications often relies on multi-stage, complex pipelines that manage data ingestion, creation of vector embeddings, vector storage, querying, and response construction. Apache Airflow, a leading open-source orchestration framework, provides the structure and flexibility required to implement complex GenAI workflows with relatively basic Python code.

In this tutorial, you’ll learn how to use Airflow to orchestrate a basic RAG pipeline that includes embedding book descriptions with the OpenAI API, ingesting the embeddings into a PostgreSQL database with pgvector installed, and querying the database for books that match a user-provided mood. To learn about a more advanced implementation of this topic, join my workshop at ODSC West where I’ll cover:

• More on Airflow and orchestration basics
• Writing GenAI pipelines in Airflow, including how to use the Airflow AI SDK for interacting with LLMs
• Airflow features for running GenAI pipelines in production that optimize reliability and performance

Why use Airflow for GenAI Applications?

The pipelines that help power GenAI applications need to run reliably, scale efficiently, and recover from failure without manual intervention. When using Airflow for orchestration, actions in your AI pipelines become aware of each other, and you have a central location to monitor, edit, and troubleshoot all your workflows, from the ETL/ELT pipelines that provide data for your models, to RAG, to prompt engineering, to fine tuning. 

Airflow provides many benefits for GenAI workflows, including:

• Reliable task execution: APIs of AI models often come with transient errors and rate limits. Airflow can automatically retry individual tasks, send notifications in case of issues and handle timeouts.
• Tool agnosticism: Airflow can connect to any application in your data ecosystem that allows connections through an API. This means you are never locked into one model vendor and can quickly switch individual tasks to the latest and best model and AI tooling. Prebuilt provider packages exist to connect to many common data tools. You can find a comprehensive catalogue of provider packages in the Astronomer registry.
• High extensibility: Since Airflow pipelines are written in Python, you can build on top of the existing codebase and extend the functionality of Airflow to meet your needs. Almost anything you can do in Python, you can do in Airflow.Infin
• ite scalability: Given enough computing power, you can orchestrate as many processes as you need, no matter the complexity of your pipelines.
• Observability: The Airflow UI provides an immediate overview of all your data pipelines and can provide the source of truth for workflows in your whole data ecosystem. 

Prerequisites

This tutorial requires the following tools:

• The Astro CLI. This is a free tool provided by Astronomer used to run Airflow locally. If you have a different way you are already running Airflow, you can use that as well.
• An OpenAI API key of at least tier 1 if you want to use OpenAI for vectorization. If you do not want to use OpenAI, you can adapt the create_embeddings function at the start of the DAG to use a different vectorizer.

This tutorial uses a local PostgreSQL database created as a Docker container. The image comes with pgvector preinstalled.

Step 1: Configure your Astro project

1. Create a new Astro project:

$ mkdir astro-pgvector-tutorial && cd astro-pgvector-tutorial
$ astro dev init

2. Add the following two packages to your requirements.txt file to install the pgvector Airflow provider and the OpenAI Python client in your Astro project:

apache-airflow-providers-pgvector==1.0.0
openai==1.3.2

3. This tutorial uses a local PostgreSQL database running in a Docker container. To add a second PostgreSQL container to your Astro project, create a new file in your project’s root directory called docker-compose.override.yml and add the following. The ankane/pgvector image builds a PostgreSQL database with pgvector preinstalled.

services:
  postgres_pgvector:
    image: ankane/pgvector
    volumes:
      - ${PWD}/include/postgres:/var/lib/postgresql/data
      - ${PWD}/include:/include
    networks:
      - airflow
    ports:
      - 5433:5432
    environment:
      - POSTGRES_USER=postgres
      - POSTGRES_PASSWORD=postgres
# Airflow containers
  scheduler:
    networks:
      - airflow
  api-server:
    networks:
      - airflow
  triggerer:
    networks:
      - airflow
  postgres:
    networks:
      - airflow

4. To create an Airflow connection to the PostgreSQL database, add the following to your .env file. If you are using the OpenAI API for embeddings you will need to update the OPENAI_API_KEY environment variable.

AIRFLOW_CONN_POSTGRES_DEFAULT='{
    "conn_type": "postgres",
    "login": "postgres",
    "password": "postgres",
    "host": "host.docker.internal",
    "port": 5433,
    "schema": "postgres"
}'
OPENAI_API_KEY=""

Step 2: Add your data

The dag in this tutorial runs a query on vectorized book descriptions from Goodreads, but you can adjust the dag to use any data you want.

1. Create a new file called book_data.txt in the include directory.

2. Copy the book description from the book_data.txt file in Astronomer’s GitHub for a list of great books.

Note: If you want to add your own books make sure the data is in the following format:

 integer> :::