This is an automated email from the ASF dual-hosted git repository.
Lee-W pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new dfe9cc6ef3c Add Ollama model readiness check to PydanticAIHook (#75369)
dfe9cc6ef3c is described below
commit dfe9cc6ef3c1ab508ae06e7e945d62d810ab23df
Author: Wei Lee <[email protected]>
AuthorDate: Tue Nov 10 10:28:16 2026 +0900
Add Ollama model readiness check to PydanticAIHook (#75369)
* Add Ollama model readiness check to PydanticAIHook
- Add PydanticAIHook.ensure_model_ready(timeout, poll_interval), an opt-in
method that
waits for the primary connection's Ollama model to be loaded before a
task builds its
agent. Other model prefixes raise NotImplementedError.
- Add LocalModelBackend protocol and OllamaBackend (/api/ps, /api/generate)
in
hooks/local_model_backends.py.
- Every request is bounded by the remaining timeout and timeouts surface as
TimeoutError; after the load request, connection failures and 5xx
responses are
retried until the timeout is used up.
- Declare httpx2 as a common.ai dependency.
- Document the method and its limitations in self_hosted_models.rst.
* Wrap an over-long line in the Ollama readiness docs
* Check readiness once after loading instead of polling
OllamaBackend.load() already blocks until the model is loaded, so polling
afterwards has no purpose. Do one readiness check, one blocking load and one
confirming check, and drop poll_interval. Correct the keep_alive comment in
OllamaBackend.load().
* Let ensure_model_ready take any backend with is_ready and load
Add a ModelReadinessBackend Protocol and a backend= argument so any server
with is_ready() and load() works. Ollama stays the built-in default, picked by
a resolver.
* Load Ollama models with a generate-only request and a backend factory
The /api/ps listing is not a reliable readiness signal while a model is
still loading, so ensure_model_ready now only sends the empty-prompt
/api/generate request and hands the whole timeout to load(). is_ready,
normalize_ollama_model_name and the RuntimeError path are removed.
backend= is replaced by backend_factory=, which the hook calls with the
resolved model_name and the connection's base_url so a custom backend, or
OllamaBackend behind an openai: model, gets the same values the hook used.
The OllamaBackend tests now assert the timeout given to the client and the
full request body.
* Let ensure_model_ready fall back to OLLAMA_BASE_URL
- The built-in Ollama path takes the connection host, then the
OLLAMA_BASE_URL environment variable, as pydantic-ai's OllamaProvider
does, and raises ValueError only when both are empty.
- Cover a model set only in the connection extra, and a bare model name
with the model_provider extra.
* Document keep_alive and upstream loading for Ollama readiness
- Explain how @task.llm and AgentOperator users call ensure_model_ready
from an upstream task, and that the load request sends no keep_alive,
so Ollama may unload the model after its own keep-alive time.
- Say that a custom backend_factory receives the connection host or None
and does not read OLLAMA_BASE_URL.
- Build the backend_factory test double with create_autospec so the
keyword signature is checked.
* Type the readiness backend factory and resolve a missing Ollama URL in
the backend
- OllamaBackend accepts base_url=None and falls back to OLLAMA_BASE_URL,
raising ValueError when neither is set, so backend_factory=OllamaBackend
works for connections without a host
- backend_factory is typed as a ModelReadinessBackendFactory protocol with
keyword-only model_name and base_url
- ensure_model_ready splits the model name once
- test stub for the factory no longer has an empty body
* Name the model and the server error when Ollama fails to load it
- OllamaBackend.load re-raises HTTPStatusError with the model name and the
error field Ollama returned, instead of only the status line and URL
- A body that is not JSON, has no error field, or is blank falls back to the
response text, then to the reason phrase
* Point readiness docstrings at the self-hosted models guide instead of
repeating it
- ModelReadinessBackend, ModelReadinessBackendFactory and ensure_model_ready
keep only their API contract and link to the guide with :ref:
- Add anchors for the Ollama waiting and custom backend sections of the
guide
- Move the base_url contract for backend factories into the guide
* Explain when Ollama warm-up helps and which base URL variable applies
- Say that the warm-up matters when model_settings sets a timeout shorter
than pydantic-ai's default HTTP read timeout of 600 seconds
- Note that a warm-up task and the LLM task can reach different Ollama
servers with per-worker Ollama or replicas behind a load balancer
- Limit the OLLAMA_BASE_URL fallback to ollama: models and point openai:
models at OPENAI_BASE_URL
---
providers/common/ai/README.rst | 1 +
.../common/ai/docs/connections/pydantic_ai.rst | 5 +-
providers/common/ai/docs/index.rst | 1 +
providers/common/ai/docs/self_hosted_models.rst | 103 +++++++++++
providers/common/ai/provider.yaml | 1 +
providers/common/ai/pyproject.toml | 1 +
.../providers/common/ai/get_provider_info.py | 5 +-
.../airflow/providers/common/ai/hooks/ollama.py | 98 ++++++++++
.../providers/common/ai/hooks/pydantic_ai.py | 106 ++++++++++-
.../ai/tests/unit/common/ai/hooks/test_ollama.py | 139 ++++++++++++++
.../tests/unit/common/ai/hooks/test_pydantic_ai.py | 204 ++++++++++++++++++++-
uv.lock | 2 +
12 files changed, 658 insertions(+), 8 deletions(-)
diff --git a/providers/common/ai/README.rst b/providers/common/ai/README.rst
index 2fc139628d4..cd4d21bf72f 100644
--- a/providers/common/ai/README.rst
+++ b/providers/common/ai/README.rst
@@ -56,6 +56,7 @@ PIP package Version required
``apache-airflow`` ``>=2.11.0``
``apache-airflow-providers-common-compat`` ``>=1.22.0``
``apache-airflow-providers-standard`` ``>=1.20.0``
+``httpx2`` ``>=2.7``
``pydantic-ai-slim`` ``>=2.36.0``
``structlog`` ``>=24.2.0``
========================================== ==================
diff --git a/providers/common/ai/docs/connections/pydantic_ai.rst
b/providers/common/ai/docs/connections/pydantic_ai.rst
index 831e0ac2b54..7ae9250a64d 100644
--- a/providers/common/ai/docs/connections/pydantic_ai.rst
+++ b/providers/common/ai/docs/connections/pydantic_ai.rst
@@ -61,7 +61,10 @@ API Key (Password field)
``GOOGLE_APPLICATION_CREDENTIALS``).
Host (optional)
- Base URL for the provider's API. Only needed for custom endpoints:
+ Base URL for the provider's API. Only needed for custom endpoints. For
``ollama:``
+ models it can instead come from the ``OLLAMA_BASE_URL`` environment
variable;
+ ``openai:`` models, including the Ollama example below that uses
+ ``openai:llama3``, read ``OPENAI_BASE_URL`` instead:
- Ollama: ``http://localhost:11434/v1``
- vLLM: ``http://localhost:8000/v1``
diff --git a/providers/common/ai/docs/index.rst
b/providers/common/ai/docs/index.rst
index d94b5b20574..cc627cf2842 100644
--- a/providers/common/ai/docs/index.rst
+++ b/providers/common/ai/docs/index.rst
@@ -188,6 +188,7 @@ PIP package Version required
``apache-airflow`` ``>=2.11.0``
``apache-airflow-providers-common-compat`` ``>=1.22.0``
``apache-airflow-providers-standard`` ``>=1.20.0``
+``httpx2`` ``>=2.7``
``pydantic-ai-slim`` ``>=2.36.0``
``structlog`` ``>=24.2.0``
========================================== ==================
diff --git a/providers/common/ai/docs/self_hosted_models.rst
b/providers/common/ai/docs/self_hosted_models.rst
index 44fcaedce23..72757f7e50b 100644
--- a/providers/common/ai/docs/self_hosted_models.rst
+++ b/providers/common/ai/docs/self_hosted_models.rst
@@ -257,6 +257,109 @@ model doesn't have to guess a path. That matters because
``HttpHook`` raises
on any response outside the 2xx/3xx range: if the model guessed wrong, the
task would fail outright rather than the agent recovering gracefully.
+.. _self_hosted_models_ensure_model_ready:
+
+Waiting for an Ollama model to load
+-------------------------------------
+
+Ollama loads a model into memory on the first request. pydantic-ai's default
+HTTP read timeout is 600 seconds, longer than the 300 seconds used below, so
with
+default settings the agent's first request waits for the load itself. Loading
+ahead of time matters when ``model_settings`` sets a shorter ``timeout``: the
first
+request then times out before the model has loaded. In that case, call
+:meth:`~airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIHook.ensure_model_ready`
+before the task builds its agent. It sends Ollama a load request for this
+connection's model and waits until Ollama answers, which it does once the model
+is available (right away if it is already loaded):
+
+.. code-block:: python
+
+ from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+
+ hook = PydanticAIHook(llm_conn_id="pydanticai_default")
+ hook.ensure_model_ready(timeout=300)
+ agent = hook.create_agent(instructions="Answer briefly.")
+
+Nothing calls it for you; it is opt-in. It has these limits:
+
+- Only Ollama is built in. For any other model prefix pass ``backend_factory=``
+ (see below); without it, ``NotImplementedError`` is raised.
+- Only the hook's own connection is checked. Models behind
``fallback_conn_ids``
+ are not loaded.
+- The call blocks and holds the worker slot for the whole wait.
+- The built-in Ollama backend does not send the connection's password, so an
Ollama
+ server behind authentication is not supported.
+
+``@task.llm`` and ``AgentOperator`` build their own hook from ``llm_conn_id``
and
+their ``model_id``, so call ``ensure_model_ready`` from an upstream task with
+the same connection and, if the downstream task sets ``model_id``, the same
+``model_id``. Otherwise the upstream task loads the connection's own model:
+
+.. code-block:: python
+
+ from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+ from airflow.sdk import task
+
+
+ @task
+ def load_model():
+ hook = PydanticAIHook(llm_conn_id="pydanticai_default",
model_id="ollama:llama3")
+ hook.ensure_model_ready(timeout=300)
+
+The load request does not send ``keep_alive``, so Ollama unloads the model
after
+its own keep-alive time (5 minutes by default) like any other model. If the
+downstream LLM task can wait in the queue longer than that, raise the server's
+default with the ``OLLAMA_KEEP_ALIVE`` environment variable.
+
+The upstream task and the LLM task must reach the same Ollama server for the
+load to help. With one Ollama on each worker (a ``localhost`` connection) or
several
+replicas behind a load balancer, the load request can reach a different server
+than the LLM call, and then it has no effect on that call.
+
+``timeout`` is passed as is to the load request. Running out ends in a
+``TimeoutError`` naming the model, or the server when it cannot be reached.
+Nothing is retried, so a wrong host, a server error or a request timing out is
+raised right away.
+
+.. _self_hosted_models_readiness_backend:
+
+Using your own readiness backend
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+``ensure_model_ready`` accepts ``backend_factory=``, a callable that is called
+with the keyword arguments ``model_name`` (the model name without its provider
+prefix) and ``base_url`` (the connection's host, or ``None``) and returns any
+object with a keyword-only ``load(*, timeout)`` method, described by
+:class:`~airflow.providers.common.ai.hooks.pydantic_ai.ModelReadinessBackend`.
+``load`` should return once the model is available. With ``backend_factory=``,
+the model prefix and the connection host are not checked, and errors are raised
+as is. The ``base_url`` the factory receives is only the connection's host,
never
+``OLLAMA_BASE_URL``; the factory decides what to do when it is ``None``.
+
+.. code-block:: python
+
+ class MyServerBackend:
+ def __init__(self, model_name: str, base_url: str | None) -> None: ...
+
+ def load(self, *, timeout: float) -> None: ...
+
+
+ hook.ensure_model_ready(backend_factory=MyServerBackend)
+
+The prefix is not used to detect the server, because ``openai:<model>`` can
+point at Ollama's ``/v1`` endpoint as well as at any other OpenAI-compatible
+server. When it points at Ollama, pass ``OllamaBackend`` as the factory
+(``base_url`` may keep its ``/v1`` suffix). When the connection has no host,
+``OllamaBackend`` uses the ``OLLAMA_BASE_URL`` environment variable, and raises
+``ValueError`` if that is not set either:
+
+.. code-block:: python
+
+ from airflow.providers.common.ai.hooks.ollama import OllamaBackend
+
+ hook = PydanticAIHook(llm_conn_id="pydanticai_default",
model_id="openai:llama3")
+ hook.ensure_model_ready(backend_factory=OllamaBackend)
+
Routing through an AI gateway
--------------------------------
diff --git a/providers/common/ai/provider.yaml
b/providers/common/ai/provider.yaml
index eaffa676ab8..63c917cb7a1 100644
--- a/providers/common/ai/provider.yaml
+++ b/providers/common/ai/provider.yaml
@@ -92,6 +92,7 @@ hooks:
- integration-name: Pydantic AI
python-modules:
- airflow.providers.common.ai.hooks.pydantic_ai
+ - airflow.providers.common.ai.hooks.ollama
- integration-name: MCP Server
python-modules:
- airflow.providers.common.ai.hooks.mcp
diff --git a/providers/common/ai/pyproject.toml
b/providers/common/ai/pyproject.toml
index ca44e12fdc2..9dcff8706a4 100644
--- a/providers/common/ai/pyproject.toml
+++ b/providers/common/ai/pyproject.toml
@@ -68,6 +68,7 @@ dependencies = [
"apache-airflow>=2.11.0",
"apache-airflow-providers-common-compat>=1.22.0",
"apache-airflow-providers-standard>=1.20.0",
+ "httpx2>=2.7",
# 2.33.0 is the first release that works with anthropic>=1: it moved to
httpx2 alongside
# the SDK and stopped passing temperature/top_p/top_k as messages.create()
kwargs, both
# of which raise TypeError on 2.31 and earlier. The cost API this provider
relies on
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py
b/providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py
index 331957e8030..6c172a2a719 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py
@@ -92,7 +92,10 @@ def get_provider_info():
"hooks": [
{
"integration-name": "Pydantic AI",
- "python-modules":
["airflow.providers.common.ai.hooks.pydantic_ai"],
+ "python-modules": [
+ "airflow.providers.common.ai.hooks.pydantic_ai",
+ "airflow.providers.common.ai.hooks.ollama",
+ ],
},
{"integration-name": "MCP Server", "python-modules":
["airflow.providers.common.ai.hooks.mcp"]},
{
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/hooks/ollama.py
b/providers/common/ai/src/airflow/providers/common/ai/hooks/ollama.py
new file mode 100644
index 00000000000..110443f59f4
--- /dev/null
+++ b/providers/common/ai/src/airflow/providers/common/ai/hooks/ollama.py
@@ -0,0 +1,98 @@
+# 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.
+from __future__ import annotations
+
+import os
+
+import httpx2
+
+
+def build_ollama_native_root(base_url: str) -> str:
+ """Return the server root of an Ollama URL, dropping a trailing ``/v1``
(the OpenAI-compatible path)."""
+ return base_url.rstrip("/").removesuffix("/v1")
+
+
+class OllamaBackend:
+ """
+ Load control for a model served by Ollama, through its native
``/api/generate`` endpoint.
+
+ :param base_url: Ollama URL; a trailing ``/v1`` is dropped to reach the
native API. When ``None``
+ or empty, the ``OLLAMA_BASE_URL`` environment variable is used.
+ :param model_name: Model name without the ``ollama:`` prefix.
+ :param transport: Optional httpx2 transport, mainly to stub the server in
tests.
+
+ :raises ValueError: If *base_url* is not given and ``OLLAMA_BASE_URL`` is
not set.
+
+ Known limitation: the connection password is not sent, so no
``Authorization`` header is
+ added and servers behind authentication are not supported.
+ """
+
+ def __init__(
+ self,
+ base_url: str | None,
+ model_name: str,
+ *,
+ transport: httpx2.BaseTransport | None = None,
+ ) -> None:
+ if not (resolved_url := base_url or os.getenv("OLLAMA_BASE_URL")):
+ raise ValueError(
+ "OllamaBackend needs a base_url or the OLLAMA_BASE_URL
environment variable; neither is set"
+ )
+ self.root = build_ollama_native_root(resolved_url)
+ self.model_name = model_name
+ self._transport = transport
+
+ def load(self, *, timeout: float) -> None:
+ """
+ Load the model, waiting at most *timeout* seconds.
+
+ :raises TimeoutError: If the server cannot be reached or does not
finish loading within *timeout*.
+ :raises httpx2.HTTPStatusError: If the server answers with an error
status; the message names the
+ model and carries the server's ``error`` text (the response body
when it has none).
+ """
+ # A /api/generate request with no prompt and stream=False returns once
the model is
+ # available. keep_alive only controls how long the model stays loaded
afterwards.
+ try:
+ with httpx2.Client(transport=self._transport, timeout=timeout) as
client:
+ response = client.post(
+ f"{self.root}/api/generate", json={"model":
self.model_name, "stream": False}
+ )
+ response.raise_for_status()
+ except httpx2.HTTPStatusError as e:
+ raise httpx2.HTTPStatusError(
+ f"Ollama could not load model '{self.model_name}'
({e.response.status_code}): "
+ f"{self._extract_error_detail(e.response)}",
+ request=e.request,
+ response=e.response,
+ ) from e
+ except httpx2.ConnectTimeout as e:
+ raise TimeoutError(f"Could not reach Ollama at {self.root} within
{timeout} seconds.") from e
+ except httpx2.TimeoutException as e:
+ raise TimeoutError(
+ f"Ollama did not finish loading model '{self.model_name}'
within {timeout} seconds."
+ ) from e
+
+ @staticmethod
+ def _extract_error_detail(response: httpx2.Response) -> str:
+ """Return Ollama's ``error`` field, or the body cut to a short length
when there is none."""
+ try:
+ body = response.json()
+ except ValueError:
+ body = None
+ if isinstance(body, dict) and isinstance(error := body.get("error"),
str) and error:
+ return error
+ return response.text.strip()[:200] or response.reason_phrase
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py
b/providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py
index 90d3a3736ba..645f4a41bbd 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py
@@ -20,9 +20,11 @@ import asyncio
import concurrent.futures
import contextvars
import functools
+import math
import re
+import time
from pathlib import Path
-from typing import TYPE_CHECKING, Any, TypeVar, overload
+from typing import TYPE_CHECKING, Any, Protocol, TypeVar, overload
import pydantic_ai
from pydantic_ai import Agent
@@ -32,6 +34,7 @@ from pydantic_ai.models.fallback import FallbackModel
from pydantic_ai.providers import infer_provider_class
from typing_extensions import ParamSpec
+from airflow.providers.common.ai.hooks.ollama import OllamaBackend
from airflow.providers.common.ai.observability import
genai_instrumentation_settings
from airflow.providers.common.ai.utils.toolset_base import MaskingCapability
from airflow.providers.common.compat.module_loading import import_string
@@ -45,6 +48,32 @@ OutputT = TypeVar("OutputT")
P = ParamSpec("P")
R = TypeVar("R")
+
+class ModelReadinessBackend(Protocol):
+ """
+ Server-specific half of :meth:`PydanticAIHook.ensure_model_ready`.
+
+ Any object with this method works; it does not need to subclass this class.
+
+ ``load`` returns once the model is available, waiting at most *timeout*
seconds, and raises
+ ``TimeoutError`` when it exceeds *timeout*. See
:ref:`self_hosted_models_readiness_backend`.
+ """
+
+ def load(self, *, timeout: float) -> None: ...
+
+
+class ModelReadinessBackendFactory(Protocol):
+ """
+ Builds the :class:`ModelReadinessBackend` for
:meth:`PydanticAIHook.ensure_model_ready`.
+
+ Called with the ``model_name`` without its provider prefix and the
connection host as
+ ``base_url`` (``None`` when the connection has no host). ``OllamaBackend``
is one.
+ See :ref:`self_hosted_models_readiness_backend`.
+ """
+
+ def __call__(self, *, model_name: str, base_url: str | None) ->
ModelReadinessBackend: ...
+
+
# Sentinel distinguishing "caller did not pass ``instrument``" from an explicit
# ``instrument=None`` / ``instrument=False`` (which mean "do not instrument,
and
# do not auto-enable it either").
@@ -527,10 +556,7 @@ class PydanticAIHook(BaseHook):
model_name = forwarded_model_id
forwarded = True
if not model_name:
- raise ValueError(
- f"No model specified for connection '{self.llm_conn_id}'. Set
model_id on the "
- "hook or the Model field on the connection."
- )
+ raise self._build_missing_model_error()
model_name = self._qualify_model_name(
model_name,
forwarded_from_conn_id=forwarded_from_conn_id if forwarded else
None,
@@ -562,6 +588,76 @@ class PydanticAIHook(BaseHook):
return infer_model(model_name)
+ def _build_missing_model_error(self) -> ValueError:
+ return ValueError(
+ f"No model specified for connection '{self.llm_conn_id}'. Set
model_id on the "
+ "hook or the Model field on the connection."
+ )
+
+ def _build_readiness_backend(
+ self, model_name: str, backend_factory: ModelReadinessBackendFactory |
None
+ ) -> ModelReadinessBackend:
+ """
+ Build the readiness backend for the provider-prefixed *model_name*.
+
+ *backend_factory* is called with the unprefixed ``model_name`` and the
connection's
+ ``base_url``; without it only Ollama has a built-in backend.
+ """
+ provider_prefix, sep, rest = model_name.partition(":")
+ if backend_factory is None and provider_prefix != "ollama":
+ raise NotImplementedError(
+ f"ensure_model_ready has no built-in readiness backend for
'{model_name}' on "
+ f"connection '{self.llm_conn_id}'; only Ollama is built in.
Pass backend_factory= a "
+ "callable taking model_name= and base_url= that returns an
object with load(); "
+ "OllamaBackend is available for a model that Ollama serves."
+ )
+ conn, _ = self._get_conn_and_extra()
+ base_url = conn.host or None
+ unprefixed_name = rest if sep else model_name
+ if backend_factory is not None:
+ return backend_factory(model_name=unprefixed_name,
base_url=base_url)
+ return OllamaBackend(base_url, unprefixed_name)
+
+ def ensure_model_ready(
+ self,
+ *,
+ timeout: float = 300.0,
+ backend_factory: ModelReadinessBackendFactory | None = None,
+ ) -> None:
+ """
+ Block until this connection's local model is loaded, asking the server
to load it if needed.
+
+ Opt-in, blocks the worker slot, and loads only this hook's own
connection. Details and
+ limits: :ref:`self_hosted_models_ensure_model_ready`.
+
+ :param timeout: Seconds to wait for the load request, passed to the
backend as is.
+ :param backend_factory: Callable that takes the keyword arguments
``model_name`` (the
+ resolved name without its provider prefix) and ``base_url`` (the
connection host, or
+ ``None``) and returns an object with ``load()`` (see
ModelReadinessBackend). When
+ given, the model prefix and connection host checks are skipped; a
missing model still
+ raises ``ValueError``. See
:ref:`self_hosted_models_readiness_backend`.
+ :raises ValueError: If *timeout* is not a positive finite number or no
model is configured;
+ also from the default ``OllamaBackend`` when neither the
connection host nor the
+ ``OLLAMA_BASE_URL`` environment variable is set.
+ :raises NotImplementedError: If *backend_factory* is not given and the
model's provider
+ prefix has no built-in backend (only ``ollama`` has one).
+ :raises TimeoutError: If a request times out.
+ :raises httpx2.TransportError: If a request cannot reach the server.
+ :raises httpx2.HTTPStatusError: If the server answers with an error
status. The default
+ ``OllamaBackend`` words the message with the model name and the
server's error text.
+ """
+ if not math.isfinite(timeout) or timeout <= 0:
+ raise ValueError(f"timeout must be positive and finite, got
{timeout}.")
+ if not (configured_name := self._get_configured_model_name()):
+ raise self._build_missing_model_error()
+ model_name = self._qualify_model_name(configured_name)
+ backend = self._build_readiness_backend(model_name, backend_factory)
+
+ started = time.monotonic()
+ self.log.info("Loading model '%s' for connection '%s'", model_name,
self.llm_conn_id)
+ backend.load(timeout=timeout)
+ self.log.info("Model '%s' is ready after %.1f seconds", model_name,
time.monotonic() - started)
+
def _get_fallback_conn_ids(self) -> list[str]:
"""
Return the configured fallback connection IDs, hook argument winning
over the extra.
diff --git a/providers/common/ai/tests/unit/common/ai/hooks/test_ollama.py
b/providers/common/ai/tests/unit/common/ai/hooks/test_ollama.py
new file mode 100644
index 00000000000..01a7dbea76d
--- /dev/null
+++ b/providers/common/ai/tests/unit/common/ai/hooks/test_ollama.py
@@ -0,0 +1,139 @@
+# 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.
+from __future__ import annotations
+
+import json
+
+import httpx2
+import pytest
+
+from airflow.providers.common.ai.hooks.ollama import OllamaBackend,
build_ollama_native_root
+
+
+def _build_backend(
+ handler, base_url: str = "http://ollama:11434/v1";, model_name: str =
"llama3"
+) -> OllamaBackend:
+ return OllamaBackend(base_url, model_name,
transport=httpx2.MockTransport(handler))
+
+
+class TestBuildNativeRoot:
+ @pytest.mark.parametrize(
+ "base_url",
+ [
+ "http://h:11434/v1";,
+ "http://h:11434/v1/";,
+ "http://h:11434";,
+ "http://h:11434/";,
+ ],
+ )
+ def test_returns_server_root(self, base_url):
+ assert build_ollama_native_root(base_url) == "http://h:11434";
+
+
+class TestOllamaBackendBaseUrl:
+ def test_none_base_url_falls_back_to_ollama_base_url_env(self,
monkeypatch):
+ monkeypatch.setenv("OLLAMA_BASE_URL", "http://env-ollama:11434/v1";)
+
+ assert OllamaBackend(base_url=None, model_name="llama3").root ==
"http://env-ollama:11434";
+
+ def test_explicit_base_url_wins_over_ollama_base_url_env(self,
monkeypatch):
+ monkeypatch.setenv("OLLAMA_BASE_URL", "http://env-ollama:11434/v1";)
+
+ assert OllamaBackend(base_url="http://h:11434";,
model_name="llama3").root == "http://h:11434";
+
+ @pytest.mark.parametrize("base_url", [None, ""], ids=["base-url-none",
"base-url-empty"])
+ @pytest.mark.parametrize("env_value", [None, ""], ids=["env-unset",
"env-empty"])
+ def test_missing_base_url_and_env_raises_value_error(self, monkeypatch,
base_url, env_value):
+ monkeypatch.delenv("OLLAMA_BASE_URL", raising=False)
+ if env_value is not None:
+ monkeypatch.setenv("OLLAMA_BASE_URL", env_value)
+
+ with pytest.raises(ValueError, match="OLLAMA_BASE_URL environment
variable"):
+ OllamaBackend(base_url=base_url, model_name="llama3")
+
+
+class TestOllamaBackendLoad:
+ def
test_posts_only_model_and_stream_to_generate_with_the_given_timeout(self):
+ requests: list[httpx2.Request] = []
+
+ def handler(request: httpx2.Request) -> httpx2.Response:
+ requests.append(request)
+ return httpx2.Response(200, json={"response": "", "done": True,
"done_reason": "load"})
+
+ _build_backend(handler).load(timeout=12.0)
+
+ (request,) = requests
+ assert (request.method, request.url.path) == ("POST", "/api/generate")
+ assert json.loads(request.content) == {"model": "llama3", "stream":
False}
+ assert request.extensions["timeout"]["read"] == 12.0
+
+ def test_http_error_names_model_and_server_error_field(self):
+ response = httpx2.Response(404, json={"error": "model 'llama3' not
found"})
+ backend = _build_backend(lambda request: response)
+
+ with pytest.raises(httpx2.HTTPStatusError) as exc_info:
+ backend.load(timeout=5.0)
+
+ assert str(exc_info.value) == ("Ollama could not load model 'llama3'
(404): model 'llama3' not found")
+ assert exc_info.value.response.status_code == 404
+ assert isinstance(exc_info.value.__cause__, httpx2.HTTPStatusError)
+
+ @pytest.mark.parametrize(
+ ("response", "expected_detail"),
+ [
+ pytest.param(
+ httpx2.Response(502, text="<html>Bad Gateway</html>"),
+ "<html>Bad Gateway</html>",
+ id="non-json-body",
+ ),
+ pytest.param(
+ httpx2.Response(500, json={"message": "boom"}),
'{"message":"boom"}', id="json-without-error"
+ ),
+ pytest.param(httpx2.Response(500, text="x" * 500), "x" * 200,
id="long-body-truncated"),
+ pytest.param(httpx2.Response(500), "Internal Server Error",
id="empty-body"),
+ pytest.param(httpx2.Response(500, text="\n"), "Internal Server
Error", id="whitespace-body"),
+ ],
+ )
+ def test_http_error_falls_back_to_body_text(self, response,
expected_detail):
+ backend = _build_backend(lambda request: response)
+
+ with pytest.raises(httpx2.HTTPStatusError) as exc_info:
+ backend.load(timeout=5.0)
+
+ assert str(exc_info.value) == (
+ f"Ollama could not load model 'llama3' ({response.status_code}):
{expected_detail}"
+ )
+
+ def test_timeout_is_raised_as_timeout_error_naming_model_and_seconds(self):
+ def handler(request: httpx2.Request) -> httpx2.Response:
+ raise httpx2.ReadTimeout("slow", request=request)
+
+ with pytest.raises(TimeoutError, match=r"'llama3'.*5\.0 seconds") as
exc_info:
+ _build_backend(handler).load(timeout=5.0)
+
+ assert isinstance(exc_info.value.__cause__, httpx2.ReadTimeout)
+
+ def test_connect_timeout_says_server_was_not_reached(self):
+ def handler(request: httpx2.Request) -> httpx2.Response:
+ raise httpx2.ConnectTimeout("no route", request=request)
+
+ with pytest.raises(
+ TimeoutError, match=r"Could not reach Ollama at
http://ollama:11434 within 5\.0"
+ ) as exc_info:
+ _build_backend(handler).load(timeout=5.0)
+
+ assert isinstance(exc_info.value.__cause__, httpx2.ConnectTimeout)
diff --git a/providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py
b/providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py
index 47af1321456..4eea2c0983c 100644
--- a/providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py
+++ b/providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py
@@ -27,8 +27,9 @@ import time
from concurrent.futures import ThreadPoolExecutor
from contextvars import ContextVar
from pathlib import Path
-from unittest.mock import MagicMock, call, patch
+from unittest.mock import MagicMock, call, create_autospec, patch
+import httpx2
import pytest
from pydantic_ai import Agent
from pydantic_ai.capabilities import AbstractCapability
@@ -44,6 +45,7 @@ from pydantic_ai.providers import infer_provider_class
from airflow.models.connection import Connection
from airflow.providers.common.ai.get_provider_info import get_provider_info
from airflow.providers.common.ai.hooks.pydantic_ai import (
+ ModelReadinessBackend,
PydanticAIAzureHook,
PydanticAIBedrockHook,
PydanticAIHook,
@@ -62,6 +64,10 @@ from airflow.providers_manager import HookInfo
_GOOGLE_MODEL_PREFIX_RE = re.compile(r"google[\w-]*(?=:)")
+def build_backend(*, model_name: str, base_url: str | None) ->
ModelReadinessBackend:
+ raise NotImplementedError
+
+
def _assert_prefix_is_known_provider(prefix: str) -> None:
"""
Assert pydantic-ai's provider registry recognizes ``prefix``.
@@ -1830,6 +1836,202 @@ class TestPydanticAIHookTestConnection:
assert "unexpected keyword argument 'api_key'" in message
+class TestPydanticAIHookEnsureModelReady:
+ _HOOK_MODULE = "airflow.providers.common.ai.hooks.pydantic_ai"
+
+ @staticmethod
+ def _build_hook(
+ *,
+ model_id: str | None = "ollama:llama3",
+ host: str | None = "http://ollama:11434/v1";,
+ extra: dict | None = None,
+ **hook_kwargs,
+ ) -> tuple[PydanticAIHook, Connection]:
+ hook = PydanticAIHook(llm_conn_id="local_conn", model_id=model_id,
**hook_kwargs)
+ conn = Connection(
+ conn_id="local_conn",
+ conn_type="pydanticai",
+ host=host,
+ extra=json.dumps(extra) if extra else None,
+ )
+ return hook, conn
+
+ @patch(f"{_HOOK_MODULE}.OllamaBackend", autospec=True)
+ def test_builds_backend_from_host_and_unprefixed_model_name(self,
mock_backend_cls):
+ hook, conn = self._build_hook()
+
+ with patch.object(hook, "get_connection", return_value=conn):
+ hook.ensure_model_ready()
+
+ mock_backend_cls.assert_called_once_with("http://ollama:11434/v1";,
"llama3")
+
+ @pytest.mark.parametrize(
+ ("model_id", "expected_name"),
+ [
+ ("ollama:library/llama3", "library/llama3"),
+ ("ollama:llama3:70b", "llama3:70b"),
+ ("ollama:hf.co/org/model:q4", "hf.co/org/model:q4"),
+ ],
+ )
+ @patch(f"{_HOOK_MODULE}.OllamaBackend", autospec=True)
+ def test_unprefixed_name_splits_on_the_first_colon_only(self,
mock_backend_cls, model_id, expected_name):
+ hook, conn = self._build_hook(model_id=model_id)
+
+ with patch.object(hook, "get_connection", return_value=conn):
+ hook.ensure_model_ready()
+
+ mock_backend_cls.assert_called_once_with("http://ollama:11434/v1";,
expected_name)
+
+ @patch(f"{_HOOK_MODULE}.OllamaBackend", autospec=True)
+ def test_load_is_called_once_with_the_given_timeout(self,
mock_backend_cls):
+ hook, conn = self._build_hook()
+
+ with patch.object(hook, "get_connection", return_value=conn):
+ hook.ensure_model_ready(timeout=60.0)
+
+
mock_backend_cls.return_value.load.assert_called_once_with(timeout=60.0)
+
+ @pytest.mark.parametrize("model_id", ["openai:gpt-5",
"anthropic:claude-opus-4-6"])
+ @patch(f"{_HOOK_MODULE}.OllamaBackend", autospec=True)
+ def
test_non_ollama_prefix_raises_not_implemented_pointing_at_backend_factory(
+ self, mock_backend_cls, model_id
+ ):
+ hook, conn = self._build_hook(model_id=model_id)
+
+ with patch.object(hook, "get_connection", return_value=conn):
+ with pytest.raises(NotImplementedError,
match=r"backend_factory=.*OllamaBackend"):
+ hook.ensure_model_ready()
+
+ mock_backend_cls.assert_not_called()
+
+ @pytest.mark.parametrize(
+ ("model_id", "expected_name"),
+ [
+ ("openai:llama3", "llama3"),
+ ("ollama:llama3:70b", "llama3:70b"),
+ ("openai:hf.co/org/model:q4", "hf.co/org/model:q4"),
+ ],
+ )
+ @patch(f"{_HOOK_MODULE}.OllamaBackend", autospec=True)
+ def test_backend_factory_gets_unprefixed_model_name_and_connection_host(
+ self, mock_ollama_cls, model_id, expected_name
+ ):
+ stub = create_autospec(ModelReadinessBackend, instance=True)
+ factory = create_autospec(build_backend, return_value=stub)
+ hook, conn = self._build_hook(model_id=model_id)
+
+ with patch.object(hook, "get_connection", return_value=conn):
+ hook.ensure_model_ready(timeout=60.0, backend_factory=factory)
+
+ factory.assert_called_once_with(model_name=expected_name,
base_url="http://ollama:11434/v1";)
+ stub.load.assert_called_once_with(timeout=60.0)
+ mock_ollama_cls.assert_not_called()
+
+ @patch(f"{_HOOK_MODULE}.OllamaBackend", autospec=True)
+ def test_backend_factory_skips_the_host_check(self, mock_ollama_cls):
+ stub = create_autospec(ModelReadinessBackend, instance=True)
+ factory = create_autospec(build_backend, return_value=stub)
+ hook, conn = self._build_hook(host=None)
+
+ with patch.object(hook, "get_connection", return_value=conn):
+ hook.ensure_model_ready(backend_factory=factory)
+
+ factory.assert_called_once_with(model_name="llama3", base_url=None)
+ stub.load.assert_called_once()
+ mock_ollama_cls.assert_not_called()
+
+ @pytest.mark.parametrize(
+ "extra",
+ [
+ {"model": "ollama:llama3"},
+ {"model": "llama3", "model_provider": "ollama"},
+ ],
+ ids=["prefixed-model-in-extra",
"bare-model-with-model-provider-in-extra"],
+ )
+ @patch(f"{_HOOK_MODULE}.OllamaBackend", autospec=True)
+ def test_model_configured_only_on_the_connection(self, mock_backend_cls,
extra):
+ hook, conn = self._build_hook(model_id=None, extra=extra)
+
+ with patch.object(hook, "get_connection", return_value=conn):
+ hook.ensure_model_ready()
+
+ mock_backend_cls.assert_called_once_with("http://ollama:11434/v1";,
"llama3")
+
+ @pytest.mark.parametrize("host", [None, ""])
+ @patch(f"{_HOOK_MODULE}.OllamaBackend", autospec=True)
+ def test_missing_host_passes_none_base_url_to_ollama_backend(self,
mock_backend_cls, monkeypatch, host):
+ monkeypatch.setenv("OLLAMA_BASE_URL", "http://env-ollama:11434/v1";)
+ hook, conn = self._build_hook(host=host)
+
+ with patch.object(hook, "get_connection", return_value=conn):
+ hook.ensure_model_ready()
+
+ mock_backend_cls.assert_called_once_with(None, "llama3")
+
+ @patch(f"{_HOOK_MODULE}.OllamaBackend", autospec=True)
+ def test_connection_host_wins_over_ollama_base_url_env(self,
mock_backend_cls, monkeypatch):
+ monkeypatch.setenv("OLLAMA_BASE_URL", "http://env-ollama:11434/v1";)
+ hook, conn = self._build_hook()
+
+ with patch.object(hook, "get_connection", return_value=conn):
+ hook.ensure_model_ready()
+
+ mock_backend_cls.assert_called_once_with("http://ollama:11434/v1";,
"llama3")
+
+ def test_missing_model_raises_value_error(self):
+ hook, conn = self._build_hook(model_id=None)
+
+ with patch.object(hook, "get_connection", return_value=conn):
+ with pytest.raises(ValueError, match="No model specified"):
+ hook.ensure_model_ready()
+
+ @pytest.mark.parametrize(
+ ("kwargs", "match"),
+ [
+ ({"timeout": 0}, r"got 0\."),
+ ({"timeout": -1.0}, r"got -1\.0\."),
+ ({"timeout": float("nan")}, r"got nan\."),
+ ({"timeout": float("inf")}, r"got inf\."),
+ ],
+ )
+ @patch(f"{_HOOK_MODULE}.OllamaBackend", autospec=True)
+ def test_non_positive_arguments_raise_value_error(self, mock_backend_cls,
kwargs, match):
+ hook, conn = self._build_hook()
+
+ with patch.object(hook, "get_connection", return_value=conn):
+ with pytest.raises(ValueError, match=match):
+ hook.ensure_model_ready(**kwargs)
+
+ mock_backend_cls.assert_not_called()
+
+ @patch(f"{_HOOK_MODULE}.OllamaBackend", autospec=True)
+ def
test_only_the_primary_connection_is_read_when_fallbacks_are_configured(self,
mock_backend_cls):
+ hook, conn = self._build_hook(fallback_conn_ids=["fallback_conn"])
+
+ with patch.object(hook, "get_connection", autospec=True,
return_value=conn) as mock_get_connection:
+ hook.ensure_model_ready()
+
+ mock_get_connection.assert_called_once_with("local_conn")
+
+ @pytest.mark.parametrize(
+ "error",
+ [
+ httpx2.ConnectError("refused"),
+ httpx2.HTTPStatusError("busy", request=MagicMock(),
response=httpx2.Response(503)),
+ TimeoutError("slow"),
+ ],
+ ids=["connect-error", "http-503", "request-timeout"],
+ )
+ @patch(f"{_HOOK_MODULE}.OllamaBackend", autospec=True)
+ def test_load_errors_are_raised_as_is(self, mock_backend_cls, error):
+ mock_backend_cls.return_value.load.side_effect = error
+ hook, conn = self._build_hook()
+
+ with patch.object(hook, "get_connection", return_value=conn):
+ with pytest.raises(type(error)):
+ hook.ensure_model_ready()
+
+
# ---------------------------------------------------------------------------
# Subclass hook tests
# ---------------------------------------------------------------------------
diff --git a/uv.lock b/uv.lock
index 76132824ca5..9f28d88a752 100644
--- a/uv.lock
+++ b/uv.lock
@@ -4642,6 +4642,7 @@ dependencies = [
{ name = "apache-airflow" },
{ name = "apache-airflow-providers-common-compat" },
{ name = "apache-airflow-providers-standard" },
+ { name = "httpx2" },
{ name = "pydantic-ai-slim" },
{ name = "structlog" },
]
@@ -4777,6 +4778,7 @@ requires-dist = [
{ name = "fastavro", marker = "python_full_version >= '3.14' and extra ==
'avro'", specifier = ">=1.12.1" },
{ name = "fastavro", marker = "python_full_version < '3.14' and extra ==
'avro'", specifier = ">=1.10.0" },
{ name = "grpcio", marker = "extra == 'openshell'", specifier = ">=1.78.0"
},
+ { name = "httpx2", specifier = ">=2.7" },
{ name = "islo", marker = "extra == 'islo'", specifier = ">=0.3.19" },
{ name = "langchain", marker = "extra == 'langchain'", specifier =
">=1.0.0" },
{ name = "llama-index-core", marker = "extra == 'llamaindex'", specifier =
">=0.14.5" },
