Skip to content

Commit 57fc833

Browse files
vikramkokavikramkoka
committed
Add LlamaIndex operators to common.ai provider
 - Adds LlamaIndexHook to bridge Airflow connections to LlamaIndex's Settings singleton. Reuses the pydanticai connection type, supports separate embedding and LLM connections. - Adds EmbeddingOperator to chunk documents and produce embedding vectors via LlamaIndex's SentenceSplitter. Input is list[dict(text, metadata)] (same shape as DocumentLoaderOperator output), output includes chunks with vectors ready for downstream vector store ingest operators (pgvector, Pinecone, Weaviate). - Adds RetrievalOperator to load a persisted LlamaIndex index and perform similarity search. Output is scored chunks ready for synthesis via LLMOperator. Design notes All LlamaIndex imports are lazy (inside execute() / method bodies), so modules parse without llama-index installed. The hook currently hardcodes OpenAI embedding/LLM providers; a follow-up PR will refactor to use BaseAIHook for provider-agnostic model resolution when it lands. What's included ┌─────────────────────────────────────────┬──────────────────────────────────────────┐ │ File │ Purpose │ ├─────────────────────────────────────────┼──────────────────────────────────────────┤ │ hooks/llamaindex.py │ Hook (~110 lines) │ ├─────────────────────────────────────────┼──────────────────────────────────────────┤ │ operators/llamaindex_embedding.py │ EmbeddingOperator (~110 lines) │ ├─────────────────────────────────────────┼──────────────────────────────────────────┤ │ operators/llamaindex_retrieval.py │ RetrievalOperator (~90 lines) │ ├─────────────────────────────────────────┼──────────────────────────────────────────┤ │ tests/.../test_llamaindex.py │ 12 hook tests │ ├─────────────────────────────────────────┼──────────────────────────────────────────┤ │ tests/.../test_llamaindex_embedding.py │ 10 operator tests │ ├─────────────────────────────────────────┼──────────────────────────────────────────┤ │ tests/.../test_llamaindex_retrieval.py │ 8 operator tests │ ├─────────────────────────────────────────┼──────────────────────────────────────────┤ │ docs/hooks/llamaindex.rst │ Hook docs │ ├─────────────────────────────────────────┼──────────────────────────────────────────┤ │ docs/operators/llamaindex_embedding.rst │ EmbeddingOperator docs │ ├─────────────────────────────────────────┼──────────────────────────────────────────┤ │ docs/operators/llamaindex_retrieval.rst │ RetrievalOperator docs │ ├─────────────────────────────────────────┼──────────────────────────────────────────┤ │ provider.yaml │ Integration, hook, operator registration │ ├─────────────────────────────────────────┼──────────────────────────────────────────┤ │ docs/index.rst │ LlamaIndex Hook in Guides toctree │ ├─────────────────────────────────────────┼──────────────────────────────────────────┤ │ docs/operators/index.rst │ Chooser table rows │ └─────────────────────────────────────────┴──────────────────────────────────────────┘ Test plan - uv run --project providers/common/ai pytest providers/common/ai/tests/unit/common/ai/hooks/test_llamaindex.py -xvs (12 tests) - uv run --project providers/common/ai pytest providers/common/ai/tests/unit/common/ai/operators/test_llamaindex_embedding.py providers/common/ai/tests/unit/common/ai/operators/test_llamaindex_retrieval.py -xvs (18 tests) - Hook: init defaults, separate embed_conn_id, connection kwargs extraction, embedding model, LLM, Settings configuration - EmbeddingOperator: output shape, chunking, index persistence, vector inclusion/omission, splitter params - RetrievalOperator: output shape, chunk keys, top_k forwarding, multiple results, storage context --- Was generative AI tooling used to co-author this PR? - Yes — Claude Code (Opus 4.6) Generated-by: Claude Code (Opus 4.6) following https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions
1 parent 8e51e68 commit 57fc833

12 files changed

Lines changed: 1261 additions & 1 deletion

File tree

‎providers/common/ai/docs/hooks/llamaindex.rst‎

Lines changed: 87 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,87 @@
1+
.. Licensed to the Apache Software Foundation (ASF) under one
2+
or more contributor license agreements. See the NOTICE file
3+
distributed with this work for additional information
4+
regarding copyright ownership. The ASF licenses this file
5+
to you under the Apache License, Version 2.0 (the
6+
"License"); you may not use this file except in compliance
7+
with the License. You may obtain a copy of the License at
8+
9+
.. http://www.apache.org/licenses/LICENSE-2.0
10+
11+
.. Unless required by applicable law or agreed to in writing,
12+
software distributed under the License is distributed on an
13+
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
KIND, either express or implied. See the License for the
15+
specific language governing permissions and limitations
16+
under the License.
17+
18+
.. _howto/hook:llamaindex:
19+
20+
``LlamaIndexHook``
21+
==================
22+
23+
Use :class:`~airflow.providers.common.ai.hooks.llamaindex.LlamaIndexHook` to
24+
bridge Airflow connections to LlamaIndex's ``Settings`` singleton. The hook
25+
reuses the ``pydanticai`` connection type, so users configure a single
26+
connection for both pydantic-ai operators and LlamaIndex operators.
27+
28+
.. seealso::
29+
:ref:`Connection configuration <howto/connection:pydanticai>`
30+
31+
What It Does
32+
------------
33+
34+
The hook resolves API keys and base URLs from Airflow connections and uses
35+
them to configure LlamaIndex's embedding models, LLMs, and global settings.
36+
This eliminates manual ``Settings.embed_model = ...`` boilerplate in every
37+
task that uses LlamaIndex.
38+
39+
Configuration
40+
-------------
41+
42+
``LlamaIndexHook`` reuses the ``pydanticai`` connection type. Set the API key
43+
in the **Password** field and optionally a custom endpoint in the **Host**
44+
field.
45+
46+
Separate Embedding and LLM Connections
47+
--------------------------------------
48+
49+
RAG pipelines often use different providers for embeddings and chat. The hook
50+
supports an optional ``embed_conn_id`` parameter that defaults to the main
51+
``llm_conn_id``:
52+
53+
.. code-block:: python
54+
55+
from airflow.providers.common.ai.hooks.llamaindex import LlamaIndexHook
56+
57+
hook = LlamaIndexHook(
58+
llm_conn_id="openai_default",
59+
embed_conn_id="embedding_provider",
60+
embed_model="text-embedding-3-large",
61+
llm_model="gpt-4o",
62+
)
63+
hook.configure_settings()
64+
65+
Parameters
66+
----------
67+
68+
.. list-table::
69+
:header-rows: 1
70+
:widths: 25 15 60
71+
72+
* - Parameter
73+
- Default
74+
- Description
75+
* - ``llm_conn_id``
76+
- ``pydanticai_default``
77+
- Airflow connection ID for the LLM/embedding provider.
78+
* - ``embed_conn_id``
79+
- Same as ``llm_conn_id``
80+
- Separate connection for embeddings (optional).
81+
* - ``embed_model``
82+
- ``text-embedding-3-small``
83+
- Embedding model name.
84+
* - ``llm_model``
85+
- ``None``
86+
- LLM model name. Required for ``get_llm()`` and ``configure_settings()``
87+
LLM setup.

‎providers/common/ai/docs/index.rst‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
Connection types <connections/pydantic_ai>
3838
MCP connection <connections/mcp>
3939
Hooks <hooks/pydantic_ai>
40+
LlamaIndex Hook <hooks/llamaindex>
4041
Toolsets <toolsets>
4142
Operators <operators/index>
4243
HITL Review <hitl_review>

‎providers/common/ai/docs/operators/index.rst‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ Common AI Operators
2121
Choosing the right operator
2222
---------------------------
2323

24-
The common-ai provider ships five operators (and matching ``@task`` decorators). Use this table
24+
The common-ai provider ships several operators (and matching ``@task`` decorators). Use this table
2525
to pick the one that fits your use case:
2626

2727
.. list-table::
@@ -46,6 +46,12 @@ to pick the one that fits your use case:
4646
* - Multi-turn reasoning with tools (DB queries, API calls, etc.)
4747
- :class:`~airflow.providers.common.ai.operators.agent.AgentOperator`
4848
- ``@task.agent``
49+
* - Chunk documents and produce embedding vectors
50+
- :class:`~airflow.providers.common.ai.operators.llamaindex_embedding.EmbeddingOperator`
51+
- —
52+
* - Retrieve relevant chunks from a vector index
53+
- :class:`~airflow.providers.common.ai.operators.llamaindex_retrieval.RetrievalOperator`
54+
- —
4955

5056
**LLMOperator / @task.llm** — stateless, single-turn calls. Use this for classification,
5157
summarization, extraction, or any prompt that produces one response. Supports structured output

‎providers/common/ai/docs/operators/llamaindex_embedding.rst‎

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,135 @@
1+
.. Licensed to the Apache Software Foundation (ASF) under one
2+
or more contributor license agreements. See the NOTICE file
3+
distributed with this work for additional information
4+
regarding copyright ownership. The ASF licenses this file
5+
to you under the Apache License, Version 2.0 (the
6+
"License"); you may not use this file except in compliance
7+
with the License. You may obtain a copy of the License at
8+
9+
.. http://www.apache.org/licenses/LICENSE-2.0
10+
11+
.. Unless required by applicable law or agreed to in writing,
12+
software distributed under the License is distributed on an
13+
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
KIND, either express or implied. See the License for the
15+
specific language governing permissions and limitations
16+
under the License.
17+
18+
.. _howto/operator:llamaindex_embedding:
19+
20+
``EmbeddingOperator``
21+
=====================
22+
23+
Use :class:`~airflow.providers.common.ai.operators.llamaindex_embedding.EmbeddingOperator`
24+
to chunk documents and produce embedding vectors using LlamaIndex. This operator
25+
bridges document loading (Airflow provider hooks returning text) and vector
26+
storage (pgvector, Pinecone, Weaviate ingest operators).
27+
28+
Basic Usage
29+
-----------
30+
31+
Provide a list of documents with ``text`` and ``metadata`` keys. The operator
32+
chunks the documents, embeds them, and returns the results:
33+
34+
.. code-block:: python
35+
36+
from airflow.providers.common.ai.operators.llamaindex_embedding import EmbeddingOperator
37+
38+
embed = EmbeddingOperator(
39+
task_id="embed_docs",
40+
documents=[
41+
{"text": "Airflow is a workflow orchestration platform.", "metadata": {"source": "docs"}},
42+
{"text": "LlamaIndex is a data framework for LLM applications.", "metadata": {"source": "docs"}},
43+
],
44+
llm_conn_id="openai_default",
45+
)
46+
47+
Connection Configuration
48+
------------------------
49+
50+
The operator uses :class:`~airflow.providers.common.ai.hooks.llamaindex.LlamaIndexHook`
51+
internally. Configure your embedding API credentials via the ``pydanticai``
52+
connection type.
53+
54+
.. seealso::
55+
:ref:`Connection configuration <howto/connection:pydanticai>`
56+
57+
Chunking Parameters
58+
-------------------
59+
60+
Control how documents are split into chunks before embedding:
61+
62+
.. code-block:: python
63+
64+
embed = EmbeddingOperator(
65+
task_id="embed_docs",
66+
documents=documents,
67+
llm_conn_id="openai_default",
68+
chunk_size=256,
69+
chunk_overlap=25,
70+
)
71+
72+
Index Persistence
73+
-----------------
74+
75+
Set ``persist_dir`` to save the LlamaIndex index for later retrieval via
76+
:class:`~airflow.providers.common.ai.operators.llamaindex_retrieval.RetrievalOperator`:
77+
78+
.. code-block:: python
79+
80+
embed = EmbeddingOperator(
81+
task_id="embed_docs",
82+
documents=documents,
83+
llm_conn_id="openai_default",
84+
persist_dir="/opt/airflow/data/my_index",
85+
)
86+
87+
Output Shape
88+
------------
89+
90+
The operator returns a dict:
91+
92+
.. code-block:: python
93+
94+
{
95+
"document_count": 2,
96+
"chunk_count": 5,
97+
"persist_dir": "/opt/airflow/data/my_index",
98+
"chunks": [
99+
{"text": "chunk text", "metadata": {"source": "docs"}, "vector": [0.1, ...]},
100+
...
101+
],
102+
}
103+
104+
Each chunk includes ``text``, ``metadata``, and optionally ``vector`` (the
105+
embedding array). The ``chunks`` list is ready for downstream consumption by
106+
vector store ingest operators.
107+
108+
Parameters
109+
----------
110+
111+
.. list-table::
112+
:header-rows: 1
113+
:widths: 25 15 60
114+
115+
* - Parameter
116+
- Default
117+
- Description
118+
* - ``documents``
119+
- (required)
120+
- List of dicts with ``text`` and ``metadata`` keys.
121+
* - ``llm_conn_id``
122+
- ``pydanticai_default``
123+
- Airflow connection ID for the embedding API.
124+
* - ``embed_model``
125+
- ``text-embedding-3-small``
126+
- Embedding model name.
127+
* - ``chunk_size``
128+
- ``512``
129+
- Chunk size for the sentence splitter.
130+
* - ``chunk_overlap``
131+
- ``50``
132+
- Overlap between chunks.
133+
* - ``persist_dir``
134+
- ``None``
135+
- Directory path to persist the index.

‎providers/common/ai/docs/operators/llamaindex_retrieval.rst‎

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
.. Licensed to the Apache Software Foundation (ASF) under one
2+
or more contributor license agreements. See the NOTICE file
3+
distributed with this work for additional information
4+
regarding copyright ownership. The ASF licenses this file
5+
to you under the Apache License, Version 2.0 (the
6+
"License"); you may not use this file except in compliance
7+
with the License. You may obtain a copy of the License at
8+
9+
.. http://www.apache.org/licenses/LICENSE-2.0
10+
11+
.. Unless required by applicable law or agreed to in writing,
12+
software distributed under the License is distributed on an
13+
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
KIND, either express or implied. See the License for the
15+
specific language governing permissions and limitations
16+
under the License.
17+
18+
.. _howto/operator:llamaindex_retrieval:
19+
20+
``RetrievalOperator``
21+
=====================
22+
23+
Use :class:`~airflow.providers.common.ai.operators.llamaindex_retrieval.RetrievalOperator`
24+
to retrieve relevant document chunks from a persisted LlamaIndex index. The
25+
operator performs similarity search against the provided query and returns
26+
results ready for downstream synthesis via ``LLMOperator``.
27+
28+
Basic Usage
29+
-----------
30+
31+
Provide a query string and the path to a previously persisted index:
32+
33+
.. code-block:: python
34+
35+
from airflow.providers.common.ai.operators.llamaindex_retrieval import RetrievalOperator
36+
37+
retrieve = RetrievalOperator(
38+
task_id="retrieve_context",
39+
query="What are Airflow's key features?",
40+
index_persist_dir="/opt/airflow/data/my_index",
41+
llm_conn_id="openai_default",
42+
)
43+
44+
Query Templating
45+
----------------
46+
47+
The ``query`` field supports Jinja templating, so it can be set dynamically
48+
from upstream task output or Dag run configuration:
49+
50+
.. code-block:: python
51+
52+
retrieve = RetrievalOperator(
53+
task_id="retrieve_context",
54+
query="{{ dag_run.conf['question'] }}",
55+
index_persist_dir="/opt/airflow/data/my_index",
56+
llm_conn_id="openai_default",
57+
top_k=10,
58+
)
59+
60+
Output Shape
61+
------------
62+
63+
The operator returns a dict:
64+
65+
.. code-block:: python
66+
67+
{
68+
"question": "What are Airflow's key features?",
69+
"chunks": [
70+
{
71+
"text": "Airflow provides ...",
72+
"score": 0.95,
73+
"metadata": {"source": "overview.txt"},
74+
"source": "node-abc123",
75+
},
76+
...
77+
],
78+
}
79+
80+
Each chunk includes ``text``, ``score`` (similarity), ``metadata``, and
81+
``source`` (the LlamaIndex node ID). This output pairs naturally with
82+
``LLMOperator`` for RAG synthesis using Jinja templates.
83+
84+
Parameters
85+
----------
86+
87+
.. list-table::
88+
:header-rows: 1
89+
:widths: 25 15 60
90+
91+
* - Parameter
92+
- Default
93+
- Description
94+
* - ``query``
95+
- (required)
96+
- The query string to search for. Supports Jinja templating.
97+
* - ``index_persist_dir``
98+
- (required)
99+
- Path to the persisted LlamaIndex index directory.
100+
* - ``llm_conn_id``
101+
- ``pydanticai_default``
102+
- Airflow connection ID for the embedding API.
103+
* - ``embed_model``
104+
- ``text-embedding-3-small``
105+
- Embedding model name.
106+
* - ``top_k``
107+
- ``5``
108+
- Number of top results to retrieve.

‎providers/common/ai/provider.yaml‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,12 @@ integrations:
4848
- integration-name: MCP Server
4949
external-doc-url: https://modelcontextprotocol.io/
5050
tags: [ai]
51+
- integration-name: LlamaIndex
52+
external-doc-url: https://docs.llamaindex.ai/
53+
how-to-guide:
54+
- /docs/apache-airflow-providers-common-ai/operators/llamaindex_embedding.rst
55+
- /docs/apache-airflow-providers-common-ai/operators/llamaindex_retrieval.rst
56+
tags: [ai]
5157

5258
hooks:
5359
- integration-name: Pydantic AI
@@ -56,6 +62,9 @@ hooks:
5662
- integration-name: MCP Server
5763
python-modules:
5864
- airflow.providers.common.ai.hooks.mcp
65+
- integration-name: LlamaIndex
66+
python-modules:
67+
- airflow.providers.common.ai.hooks.llamaindex
5968

6069
plugins:
6170
- name: hitl_review
@@ -323,6 +332,8 @@ operators:
323332
- airflow.providers.common.ai.operators.llm_branch
324333
- airflow.providers.common.ai.operators.llm_sql
325334
- airflow.providers.common.ai.operators.llm_schema_compare
335+
- airflow.providers.common.ai.operators.llamaindex_embedding
336+
- airflow.providers.common.ai.operators.llamaindex_retrieval
326337

327338
task-decorators:
328339
- class-name: airflow.providers.common.ai.decorators.agent.agent_task

0 commit comments

Comments
 (0)