`LlamaIndexHook`¶

Use LlamaIndexHook to bridge an Airflow connection to LlamaIndex chat and embedding models. The hook reads credentials (API key, optional base URL) from a connection of type llamaindex and returns native LlamaIndex objects ready to pass to VectorStoreIndex(..., embed_model=...), load_index_from_storage(..., embed_model=...), or index.as_retriever(..., llm=...).

The hook deliberately does not mutate LlamaIndex’s global Settings singleton. Operators pass the resolved model directly to LlamaIndex constructors, so concurrent tasks in the same worker don’t race on shared state.

get_llm() and get_embedding_model() return LlamaIndex’s OpenAI / OpenAIEmbedding classes, which validate model= client-side against LlamaIndex’s OpenAI-only model-name allowlists before any request is sent. Pointing host at an Ollama or vLLM endpoint does not add support for those backends: their model names (e.g. llama3.2) are never in the OpenAI allowlist, so the call fails on the model name, not on connectivity. get_embedding_model() raises immediately at construction; get_llm() defers the error until the first call that reads .metadata (.chat() / .complete()).

OpenAI by default, BYO for other vendors¶

LlamaIndex does not ship a universal init_chat_model / init_embedding_model equivalent (each vendor is a separate package under llama-index-llms-* / llama-index-embeddings-* with its own constructor kwargs). The hook therefore covers the OpenAI-compatible surface that matches LlamaIndex’s own resolve_embed_model("default") behaviour:

hook.get_embedding_model() returns an OpenAIEmbedding configured from the connection.
hook.get_llm() returns an OpenAI LLM configured from the connection.

For other vendors (Cohere, Bedrock, Vertex AI, HuggingFace, …), instantiate the LlamaIndex class directly in a @task and pass it to the operator’s embed_model= / llm= parameter – both LlamaIndexEmbeddingOperator and LlamaIndexRetrievalOperator accept a pre-built BaseEmbedding / LLM instance and bypass the hook:

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

@dag(schedule=None, tags=["example"])
def example_llamaindex_byo_embed_model():
    """Use a non-OpenAI embedding by instantiating the LlamaIndex class directly.

    LlamaIndex doesn't ship a universal init helper, so the operator accepts
    a pre-built ``BaseEmbedding`` instance and bypasses the hook entirely.
    Install the matching extra:
    ``pip install llama-index-embeddings-cohere``.
    """

    @task
    def build_cohere_embedder():
        from llama_index.embeddings.cohere import CohereEmbedding

        from airflow.providers.common.compat.sdk import BaseHook

        conn = BaseHook.get_connection("cohere_default")
        return CohereEmbedding(model_name="embed-english-v3.0", cohere_api_key=conn.password)

    @task
    def empty_doc_list() -> list[dict]:
        return [{"text": "Cohere demo content", "metadata": {}}]

    embed = LlamaIndexEmbeddingOperator(
        task_id="embed",
        documents=empty_doc_list(),
        embed_model=build_cohere_embedder(),
        persist_dir="/opt/airflow/data/cohere_index",
    )

    embed

Install the per-vendor LlamaIndex integration package separately: pip install llama-index-embeddings-cohere, ...-bedrock, ...-huggingface, llama-index-llms-anthropic, etc.

Connection Configuration¶

The hook reads credentials from the Airflow connection of type llamaindex:

password – API key (passed as api_key to OpenAIEmbedding / OpenAI).
host – Optional base URL (passed as api_base). Only useful for an OpenAI-compatible proxy that accepts OpenAI’s exact model names (e.g. an internal gateway) – not Ollama or vLLM (see above).
extra JSON – {"embed_model": "text-embedding-3-small", "llm_model": "gpt-4o"} – default model identifiers stored on the connection.

Parameters¶

Parameter	Default	Description
`llm_conn_id`	`llamaindex_default`	Airflow connection ID for the LLM/embedding provider.
`embed_conn_id`	`None` (falls back to `llm_conn_id`)	Optional separate Airflow connection ID for the embedding provider.
`embed_model`	`None` (falls back to `extra["embed_model"]`)	Embedding model name, e.g. `text-embedding-3-small`.
`llm_model`	`None` (falls back to `extra["llm_model"]`)	LLM model name, e.g. `gpt-4o`. Required when calling `get_llm()`.

Dependencies¶

Install the llamaindex extra:

pip install apache-airflow-providers-common-ai[llamaindex]

That extra installs llama-index-core, llama-index-embeddings-openai, and llama-index-llms-openai – enough to back the hook’s default OpenAI return values. For other LlamaIndex vendor packages, install their integration package separately.

LlamaIndexHook¶

OpenAI by default, BYO for other vendors¶

Connection Configuration¶

Parameters¶

Dependencies¶

`LlamaIndexHook`¶