# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
"""Hook for LlamaIndex integration with Airflow connections."""
from __future__ import annotations
import inspect
from typing import TYPE_CHECKING, Any
from airflow.providers.common.compat.sdk import (
AirflowOptionalProviderFeatureException,
BaseHook,
)
if TYPE_CHECKING:
from llama_index.core.base.embeddings.base import BaseEmbedding
from llama_index.core.llms.llm import LLM
[docs]
class LlamaIndexHook(BaseHook):
"""
Bridge an Airflow connection to LlamaIndex chat and embedding models.
.. note::
Experimental: this can change or be removed in a minor release of this provider.
See :ref:`howto/stability`.
The hook resolves credentials (API key, optional API base URL) from the
Airflow connection and returns native LlamaIndex objects ready to pass
to ``VectorStoreIndex(..., embed_model=...)``,
``load_index_from_storage(..., embed_model=...)``, or
``index.as_retriever(..., llm=...)``.
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. For other vendors (Cohere, Bedrock, Vertex, HuggingFace, ...)
instantiate the LlamaIndex class directly in your ``@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 in that case.
.. note::
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.
.. note::
``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
still 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()``).
Connection fields:
* **password**: API key passed as ``api_key=``.
* **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, whose model
catalogs are rejected regardless of ``host`` (see note above).
* **extra** JSON: ``{"embed_model": "text-embedding-3-small",
"llm_model": "gpt-5"}`` -- default model identifiers stored on the
connection.
:param llm_conn_id: Airflow connection ID for the LLM provider. Falls
back to :attr:`default_conn_name` (``"llamaindex_default"``) when
not provided.
:param embed_conn_id: Optional separate Airflow connection ID for the
embedding provider. Falls back to ``llm_conn_id`` when not set.
:param embed_model: Embedding model name (e.g.
``"text-embedding-3-small"``). Overrides ``extra["embed_model"]``
on the connection.
:param llm_model: LLM model name (e.g. ``"gpt-5"``). Overrides
``extra["llm_model"]`` on the connection. Required when calling
:meth:`get_llm`.
:param embedding_kwargs: Additional keyword arguments to pass to the embedding
model constructor without filtering. Connection ``api_key`` and ``api_base``
values take precedence at the top level, but nested options supported by the
underlying library can override hook-provided request values, including
credentials, the model, and the input. Only pass trusted values.
"""
[docs]
conn_name_attr = "llm_conn_id"
[docs]
default_conn_name = "llamaindex_default"
[docs]
conn_type = "llamaindex"
[docs]
hook_name = "LlamaIndex"
def __init__(
self,
llm_conn_id: str | None = None,
embed_conn_id: str | None = None,
embed_model: str | None = None,
llm_model: str | None = None,
*,
embedding_kwargs: dict[str, Any] | None = None,
**kwargs: Any,
) -> None:
super().__init__(**kwargs)
# Resolve at runtime so a future per-vendor subclass with its own
# ``default_conn_name`` is honoured.
[docs]
self.llm_conn_id = llm_conn_id if llm_conn_id is not None else self.default_conn_name
[docs]
self.embed_conn_id = embed_conn_id if embed_conn_id is not None else self.llm_conn_id
[docs]
self.embed_model = embed_model
[docs]
self.embedding_kwargs = embedding_kwargs or {}
[docs]
self.llm_model = llm_model
@staticmethod
[docs]
def get_ui_field_behaviour() -> dict[str, Any]:
"""Return custom field behaviour for the Airflow connection form."""
return {
"hidden_fields": ["schema", "port", "login"],
"relabeling": {"password": "API Key"},
"placeholders": {
"host": "https://api.openai.com/v1 (optional, for an OpenAI-compatible proxy)",
"extra": '{"embed_model": "text-embedding-3-small", "llm_model": "gpt-5"}',
},
}
@staticmethod
def _resolve_model(
conn_extra: dict[str, Any],
*,
constructor_value: str | None,
extra_key: str,
kind: str,
) -> str:
"""Resolve a model identifier from the constructor arg or connection extra."""
model_id = constructor_value or conn_extra.get(extra_key)
if not model_id:
raise ValueError(
f"No {kind} model identifier set. Pass {extra_key}= to the hook "
f'constructor or set extra={{"{extra_key}": "model-name"}} on '
"the connection."
)
return model_id
@staticmethod
def _connection_kwargs(conn: Any) -> dict[str, Any]:
"""Return shared OpenAI-style kwargs (api_key, api_base) from the connection."""
kwargs: dict[str, Any] = {}
if conn.password:
kwargs["api_key"] = conn.password
if conn.host:
kwargs["api_base"] = conn.host
return kwargs
[docs]
def get_embedding_model(self) -> BaseEmbedding:
"""
Return a LlamaIndex embedding model configured from the Airflow connection.
Uses ``embed_conn_id`` (falls back to ``llm_conn_id``) for credentials.
Returns an ``OpenAIEmbedding`` instance; for other vendors,
instantiate the LlamaIndex class directly and pass it to the
operator's ``embed_model=`` parameter.
"""
# Lazy: llama-index is an optional extra; importing at module level
# would break common.ai for users who haven't installed ``[llamaindex]``.
try:
from llama_index.embeddings.openai import OpenAIEmbedding
except ImportError as e:
raise AirflowOptionalProviderFeatureException(e)
conn = self.get_connection(self.embed_conn_id)
model_id = self._resolve_model(
conn.extra_dejson,
constructor_value=self.embed_model,
extra_key="embed_model",
kind="embedding",
)
connection_kwargs = self._connection_kwargs(conn)
overridden_keys = sorted(self.embedding_kwargs.keys() & connection_kwargs.keys())
if overridden_keys:
self.log.warning("Connection parameters override embedding_kwargs values: %s", overridden_keys)
kwargs = {**self.embedding_kwargs, **connection_kwargs}
supported_kwargs = {
name
for name, parameter in inspect.signature(OpenAIEmbedding.__init__).parameters.items()
if name != "self"
and parameter.kind not in {inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD}
} | set(OpenAIEmbedding.model_fields)
unsupported_keys = sorted(self.embedding_kwargs.keys() - supported_kwargs)
if unsupported_keys:
self.log.warning("OpenAIEmbedding ignores unsupported embedding_kwargs: %s", unsupported_keys)
return OpenAIEmbedding(model=model_id, **kwargs)
[docs]
def get_llm(self) -> LLM:
"""
Return a LlamaIndex LLM configured from the Airflow connection.
Returns an ``OpenAI`` LLM instance; for other vendors, instantiate
the LlamaIndex class directly and pass it to the operator's ``llm=``
parameter.
"""
try:
from llama_index.llms.openai import OpenAI
except ImportError as e:
raise AirflowOptionalProviderFeatureException(e)
conn = self.get_connection(self.llm_conn_id)
model_id = self._resolve_model(
conn.extra_dejson,
constructor_value=self.llm_model,
extra_key="llm_model",
kind="llm",
)
return OpenAI(model=model_id, **self._connection_kwargs(conn))
[docs]
def test_connection(self) -> tuple[bool, str]:
"""
Test connection by resolving the LLM.
Validates that the model identifier is valid and the provider can be
instantiated with the supplied credentials. Does NOT make an LLM API
call -- that would be expensive and fail for reasons unrelated to
connectivity (quotas, billing, rate limits).
"""
try:
self.get_llm()
return True, "Model resolved successfully."
except Exception as e:
return False, str(e)