Ask questions over a growing PDF corpus¶

A folder of quarterly reports keeps growing and people keep asking questions answered somewhere inside it. One Dag keeps a vector index fresh weekly; the other answers a question on demand from retrieved excerpts, citing them by number. Indexing runs on a schedule so no query pays for embedding, and question is a param, so anyone can trigger it from the UI, CLI or REST API.

What this demonstrates¶

Load documents: DocumentLoaderOperator – DocumentLoaderOperator parses a glob of PDFs into documents.
Embed documents: LlamaIndexEmbeddingOperator – LlamaIndexEmbeddingOperator chunks, embeds and persists the index to disk.
Retrieve context: LlamaIndexRetrievalOperator – LlamaIndexRetrievalOperator pulls the top matching chunks for a question.
Single prompts: LLMOperator and @task.llm – LLMOperator answers from the excerpts only, with the question and the context templated in from params and XCom.

Run it¶

Install the provider with the LlamaIndex and PDF extras:

pip install "apache-airflow-providers-common-ai[openai,llamaindex,pdf]"

Create a llamaindex connection named llamaindex_default for embedding and retrieval (see LlamaIndex connection); pydanticai_default writes the answer.

Put some PDFs under /opt/airflow/data/reports/, run the indexing Dag once, then ask a question:

airflow dags test example_llamaindex_index_pdf
airflow dags test example_llamaindex_query \
    --conf '{"question": "What drove the change in operating margin?"}'

The synthesize XCom holds the answer with [n] references into the excerpts that format_context numbered.

The indexing Dag¶

airflow/providers/common/ai/example_dags/example_llamaindex_rag.py[source]

@dag(schedule="@weekly", tags=["example"])
def example_llamaindex_index_pdf():
    """Weekly indexing DAG -- keep the vector index fresh as PDFs arrive.

    The companion query DAG (below) reads the persisted index on demand.
    """
    load = DocumentLoaderOperator(
        task_id="load_pdfs",
        source_path="/opt/airflow/data/reports/*.pdf",
    )

    build_index = LlamaIndexEmbeddingOperator(
        task_id="build_index",
        documents=load.output,
        embed_model="text-embedding-3-small",
        llm_conn_id="llamaindex_default",
        chunk_size=1024,
        chunk_overlap=100,
        persist_dir="/opt/airflow/data/indexes/reports_index",
    )

    load >> build_index

The query Dag¶

airflow/providers/common/ai/example_dags/example_llamaindex_rag.py[source]

@dag(
    schedule=None,
    params={"question": "Summarize the key findings from the latest quarterly report."},
    tags=["example"],
)
def example_llamaindex_query():
    """On-demand query DAG -- retrieve from a pre-built index and synthesize.

    Trigger manually or via API with a ``question`` parameter.
    """
    retrieve = LlamaIndexRetrievalOperator(
        task_id="retrieve",
        query="{{ params.question }}",
        index_persist_dir="/opt/airflow/data/indexes/reports_index",
        embed_model="text-embedding-3-small",
        llm_conn_id="llamaindex_default",
        top_k=5,
    )

    @task
    def format_context(retrieval_result: dict) -> str:
        chunks = retrieval_result["chunks"]
        numbered = [f"[{i + 1}] {chunk['text']}" for i, chunk in enumerate(chunks)]
        return "\n\n".join(numbered)

    context = format_context(retrieve.output)

    synthesize = LLMOperator(
        task_id="synthesize",
        prompt=(
            "Question: {{ params.question }}\n\n"
            "Relevant excerpts:\n{{ ti.xcom_pull(task_ids='format_context') }}\n\n"
            "Provide a detailed answer with references to the excerpt numbers."
        ),
        llm_conn_id="pydanticai_default",
        system_prompt=(
            "You are a research assistant. Answer the question using only the "
            "provided excerpts. Reference excerpt numbers in square brackets."
        ),
    )

    context >> synthesize

Two more shapes¶

example_llamaindex_rag_pipeline runs the same four operators in one Dag with a fixed question. Run it once to see the chain end to end.

example_llamaindex_multi_source loads from two places, tags each document with where it came from through metadata_fields, and embeds them into one index so retrieval can filter by source later:

airflow/providers/common/ai/example_dags/example_llamaindex_rag.py[source]

@dag(schedule=None, tags=["example"])
def example_llamaindex_multi_source():
    """Combine multiple loaders with source-tagging metadata.

    Shows how ``DocumentLoaderOperator`` handles different file formats and
    how ``metadata_fields`` tags documents by source for filtered retrieval
    downstream.
    """
    load_products = DocumentLoaderOperator(
        task_id="load_products",
        source_path="/opt/airflow/data/products.csv",
        metadata_fields={"source": "product_catalog", "department": "engineering"},
    )

    load_docs = DocumentLoaderOperator(
        task_id="load_docs",
        source_path="/opt/airflow/data/documentation/",
        file_extensions=[".md", ".txt"],
        metadata_fields={"source": "documentation"},
    )

    @task
    def merge_documents(products: list[dict], docs: list[dict]) -> list[dict]:
        return products + docs

    merged = merge_documents(load_products.output, load_docs.output)

    embed_all = LlamaIndexEmbeddingOperator(
        task_id="embed_all",
        documents=merged,
        embed_model="text-embedding-3-small",
        llm_conn_id="llamaindex_default",
        persist_dir="/opt/airflow/data/indexes/multi_source_index",
    )

    embed_all

Adapting it¶

Change source_path to your folder or an object storage URI; the loader also reads DOCX, CSV and JSON.
Change the indexing schedule to match how often documents land, or trigger it from an Asset when the upstream Dag that drops the PDFs emits one.
Raise top_k for broad questions, lower it for precise ones. Keep the excerpts-only instruction in the prompt; it stops the model filling gaps from memory.
Compare companies’ 10-K filings grows this into a fan-out over several indexes with a human at both ends.