Skip to main content

​
Overview

Retrieval augmented generation (RAG) is one of the most widely used patterns for giving an LLM access to information it was not originally trained on. Implementing it needs an indexing pipeline that fetches new content, chunks it, embeds it, and loads it into a vector database. What is in that database can then be queried in several ways: by giving an agent a similarity search tool, or by an upstream task in an LLM pipeline that runs the search deterministically and attaches the result to the user prompt. This architecture uses the second approach against the knowledge base of an eCommerce store. A daily Dag embeds new information about customers and products and writes it to a vector database. A second Dag, triggered when a customer completes a purchase, uses the products in that purchase to retrieve related items, then gives that context to an LLM to generate a personalized recommendation shown on the customer’s account page.
For general information on context engineering with Airflow, see the AI Context Engineering with Apache Airflow® eBook.

​
Architecture

RAG document embedding reference architecture diagram. A daily indexing Dag extracts and chunks customer purchase history, product reviews, and product descriptions, creates vector embeddings with any embedding model, and loads them into a vector database alongside the raw chunks. An event-driven retrieval Dag embeds a customer checkout cart, retrieves similar chunks, passes them to an LLM, and publishes personalized product recommendations to the website.
This architecture consists of two Dags, both with access to the same vector database. The indexing Dag has three steps, structured like an ETL pipeline where the transform step is chunking and embedding:
  1. Product descriptions, catalog updates, and customer interaction data are extracted and chunked.
  2. Each chunk is embedded using an embedding model. The diagram shows the LlamaIndexEmbeddingOperator of the Common AI provider, but you can also use custom code calling an embedding model from a @task decorated task.
  3. Embeddings are written to the vector database along with the source text, and metadata such as the source record ID and a last-updated timestamp.
The retrieval Dag runs on an event-driven schedule, as soon as a customer completes a purchase:
  1. The descriptions of the recently purchased products are embedded by the same model, to be used as the query input.
  2. A task upstream of the LLM runs the similarity search, retrieving information related to the newly purchased products.
  3. The retrieved items are formatted into context and passed to the LLM alongside the customer’s purchase.
  4. The LLM generates a short recommendation message, which is published to the customer’s account dashboard.
Note that the query has to be embedded by the same model that embedded the knowledge base. Vectors from two different models are not comparable, so if you change the model you need to re-embed everything.

​
Airflow features

  • Dynamic task mapping: Chunking, embedding, and loading are all mapped over the list of chunks, creating one parallel task instance per chunk or chunk batch at runtime.
  • Event-driven scheduling: Runs the retrieval Dag on a purchase event.
  • @task.llm: Generates the recommendation text from the retrieved context.
  • Automatic retries: Embedding calls often fail due to rate limits, Airflow automatically retries transient failures according to your provided retry policy.

​
Considerations

  • Retrieve deterministically when you know what to look up. Giving an agent a search tool makes sense when the query depends on reasoning. Here the query is always “items similar to what this customer just bought”, so the search can be one deterministic upstream task.
  • Store metadata you will need for invalidation. A source record ID and a last-updated timestamp help you assess which records might need updating or to be deleted over time.
  • Chunk on record and paragraph boundaries. Use chunking libraries to create sensible chunks for example per subsection in a document or per paragraph.
  • Decide what happens when retrieval returns nothing useful. A similarity search always returns its nearest neighbours, however far away they are. Set a distance threshold to avoid odd recommendations if a customer buys a more unusual product.

​
Next steps