This is an automated email from the ASF dual-hosted git repository.
kaxil 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 2c560aad523 Let a generic pydanticai connection set a default model
provider (#75909)
2c560aad523 is described below
commit 2c560aad523d82a80426b36d30fd906b560599d5
Author: Julian LaNeve <[email protected]>
AuthorDate: Wed Oct 28 18:07:15 2026 -0400
Let a generic pydanticai connection set a default model provider (#75909)
The generic pydanticai connection has no platform of its own, so every
model id on it needs a pydantic-ai prefix such as "openai-chat:". A
connection pointed at an OpenAI-compatible gateway then makes every DAG
author pick the prefix the gateway speaks.
An optional "provider" key in the connection extra (shown as Default
Provider in the form) now qualifies bare model ids on that connection,
including when they are forwarded down a fallback chain. The value is
checked with infer_provider_class, so an unknown provider fails when the
model is resolved. Vendor connection types keep their own platform and
ignore the key. Prefixed ids and connections without the key behave as
before.
closes: #75906
Co-authored-by: Kaxil Naik <[email protected]>
---
.../common/ai/docs/connections/pydantic_ai.rst | 61 ++++++--
providers/common/ai/docs/model_providers.rst | 4 +-
providers/common/ai/docs/operators/llm_batch.rst | 8 +-
providers/common/ai/docs/provider_fallback.rst | 8 +-
providers/common/ai/docs/self_hosted_models.rst | 19 +--
providers/common/ai/docs/troubleshooting.rst | 9 +-
providers/common/ai/provider.yaml | 9 +-
.../src/airflow/providers/common/ai/batch/base.py | 14 +-
.../airflow/providers/common/ai/batch/dispatch.py | 13 +-
.../providers/common/ai/get_provider_info.py | 7 +-
.../providers/common/ai/hooks/pydantic_ai.py | 86 ++++++++---
.../providers/common/ai/operators/llm_batch.py | 31 +++-
.../ai/tests/unit/common/ai/batch/test_base.py | 8 ++
.../ai/tests/unit/common/ai/batch/test_dispatch.py | 4 +
.../tests/unit/common/ai/hooks/test_pydantic_ai.py | 160 ++++++++++++++++++++-
.../unit/common/ai/operators/test_llm_batch.py | 80 +++++++++++
16 files changed, 458 insertions(+), 63 deletions(-)
diff --git a/providers/common/ai/docs/connections/pydantic_ai.rst
b/providers/common/ai/docs/connections/pydantic_ai.rst
index edb26f0fa53..9ff9ac13da2 100644
--- a/providers/common/ai/docs/connections/pydantic_ai.rst
+++ b/providers/common/ai/docs/connections/pydantic_ai.rst
@@ -38,10 +38,10 @@ Model
dedicated input in the connection form (via ``conn-fields``) and stores its
value in ``extra["model"]``.
- The ``provider:`` prefix is required here: this generic connection type has
- no platform of its own (unlike the vendor connection types below), so a
- bare name (e.g. ``gpt-5`` without ``openai:``) raises ``ValueError``
- naming this connection rather than being resolved automatically.
+ Without Default Model Provider (below), the ``provider:`` prefix is
required
+ here: this generic connection type has no platform of its own (unlike the
+ vendor connection types below), so a bare name (e.g. ``gpt-5`` without
+ ``openai:``) raises ``ValueError`` naming this connection.
Examples: ``openai:gpt-5``, ``anthropic:claude-sonnet-5``,
``bedrock:us.anthropic.claude-opus-4-6-v1:0``, ``google:gemini-2.5-flash``
@@ -81,6 +81,16 @@ Extra (JSON, optional)
When using the UI, the "Model" field above writes to this same location
automatically.
+Default Model Provider
+ The pydantic-ai provider to prefix a bare model name with, e.g.
+ ``openai-chat``. Stored in ``extra["model_provider"]``. With it set,
``model_id``
+ and the Model field can name a model the way the endpoint does (e.g.
+ ``claude-sonnet-4-5`` instead of ``openai-chat:claude-sonnet-4-5``). A name
+ that already starts with a provider pydantic-ai recognizes is used as is.
+ An unknown provider raises ``ValueError`` when the model is resolved.
+ ``@task.llm_batch`` applies it the same way, and sends ``openai-chat``
models
+ to its OpenAI adapter.
+
Fallback Connections
Other connection IDs to fail over to, in order, while this provider is
unavailable. Stored in ``extra["fallback_conn_ids"]``. Entries may name any
@@ -129,6 +139,39 @@ Ollama (local)
"extra": "{\"model\": \"openai:llama3\"}"
}
+.. _conn-example-openai-compatible-gateway:
+
+OpenAI-compatible gateway
+^^^^^^^^^^^^^^^^^^^^^^^^^
+
+A gateway or proxy that speaks the OpenAI API (LiteLLM, vLLM, an in-house
+gateway) often serves models from several vendors and implements only Chat
+Completions. Set Default Model Provider to ``openai-chat`` so DAG code can use
the
+gateway's own model names:
+
+.. code-block:: json
+
+ {
+ "conn_type": "pydanticai",
+ "host": "https://llm-gateway.example.com/v1";,
+ "password": "gateway-key",
+ "extra": "{\"model_provider\": \"openai-chat\"}"
+ }
+
+.. code-block:: python
+
+ @task.llm(llm_conn_id="my_gateway", model_id="claude-sonnet-4-5")
+ def summarize(text: str) -> str:
+ return f"Summarize: {text}"
+
+``openai:`` resolves to OpenAI's Responses API and ``openai-chat:`` to Chat
+Completions, so pick the one the gateway implements.
+
+A ``name:tag`` model id such as ``llama3:70b`` gets the default too. The
+exception is a name that is itself a pydantic-ai provider: ``mistral:7b`` goes
+to pydantic-ai's own ``mistral`` provider, so write it as
+``openai-chat:mistral:7b`` to send it through the gateway.
+
.. _conn-example-bedrock:
AWS Bedrock
@@ -191,7 +234,9 @@ The hook reads the model from these sources in priority
order:
``fallback_conn_ids`` but no model raises instead, because a spec-file
model cannot
be wrapped in a fallback chain (see :doc:`/hooks/pydantic_ai`).
-Whichever name is chosen, a name that already pins a recognized platform (its
segment
-before the first ``:`` is itself a pydantic-ai provider) is used verbatim; a
bare name is
-qualified with this connection's platform, and this generic connection type
has none,
-so a bare name reaching this step always raises.
+Whichever of sources 1-3 supplies the name, a name that already pins a
recognized platform
+(its segment before the first ``:`` is itself a pydantic-ai provider) is used
verbatim; a
+bare name is qualified with this connection's platform. This generic
connection type takes
+its platform from Default Model Provider (``extra["model_provider"]``);
without one, a bare name reaching
+this step raises. The spec-file model of source 4 goes to pydantic-ai
unchanged, so write it
+with its provider prefix.
diff --git a/providers/common/ai/docs/model_providers.rst
b/providers/common/ai/docs/model_providers.rst
index 671712dac25..e2b45a14bfb 100644
--- a/providers/common/ai/docs/model_providers.rst
+++ b/providers/common/ai/docs/model_providers.rst
@@ -145,8 +145,8 @@ Where the model name goes
Set the model on the connection's **Model** field in ``provider:model`` form,
for example
``anthropic:claude-sonnet-5``. An operator's ``model_id`` overrides it for
that task, so one
connection can serve several models from the same vendor. The generic
``pydanticai``
-connection has no vendor of its own, so a bare model name without a prefix is
rejected; the
-Azure, Bedrock and Vertex connection types each supply their own prefix and
accept a bare
+connection has no vendor of its own, so a bare model name without a prefix is
rejected unless
+the connection sets **Default Model Provider**; the Azure, Bedrock and Vertex
connection types each supply their own prefix and accept a bare
name. :doc:`connections/pydantic_ai` has every field and the full resolution
order.
Reliability across vendors
diff --git a/providers/common/ai/docs/operators/llm_batch.rst
b/providers/common/ai/docs/operators/llm_batch.rst
index 90b7745b8ba..b41258d4707 100644
--- a/providers/common/ai/docs/operators/llm_batch.rst
+++ b/providers/common/ai/docs/operators/llm_batch.rst
@@ -198,7 +198,8 @@ Per-request overrides
A request dict may carry ``system_prompt``, ``max_tokens``, ``params`` (extra
body parameters
for that request only, layered over ``request_params``) and ``model``.
``model`` is written the
-same way ``model_id`` is, ``"<provider>:<model>"``, and must name the same
provider as the batch.
+same way ``model_id`` is, ``"<provider>:<model>"`` or a bare name the
connection's Default
+Provider qualifies, and must select the same adapter as the batch.
Anthropic allows a different model per request; OpenAI requires every request
in a batch to
resolve to the same model and the operator rejects a mixed batch before
submitting it.
@@ -298,7 +299,10 @@ Reaching the batch API through a gateway
The OpenAI adapter reads ``base_url`` from the connection's host, so a
``pydanticai`` connection
pointed at an OpenAI-compatible gateway that exposes ``/v1/files`` and
``/v1/batches`` (LiteLLM
does, and routes to Azure OpenAI, Vertex and Bedrock batch behind it) runs the
whole batch
-through that gateway with ``model_id="openai:<gateway-model-name>"``. That is
the path to take
+through that gateway with ``model_id="openai:<gateway-model-name>"``.
``openai-chat:`` selects
+the same adapter, and a connection whose Default Model Provider is set (see
+:ref:`conn-example-openai-compatible-gateway`) takes the gateway's bare model
names, both in
+``model_id`` and in a request's ``model``. That is the path to take
today for a provider ``@task.llm_batch`` has no adapter for yet; the error
message for a
``pydanticai_azure``, ``pydanticai_bedrock`` or ``pydanticai_vertex``
connection says so.
diff --git a/providers/common/ai/docs/provider_fallback.rst
b/providers/common/ai/docs/provider_fallback.rst
index 7b2821c1054..3dc2595f64a 100644
--- a/providers/common/ai/docs/provider_fallback.rst
+++ b/providers/common/ai/docs/provider_fallback.rst
@@ -68,9 +68,11 @@ Bedrock primary's native id reaches a Bedrock fallback but
not an Azure one. For
a Bedrock primary with ``fallback_conn_ids: ["bedrock_dr"]`` forwards
``"us.anthropic.claude-opus-4-6-v1:0"`` to ``bedrock_dr`` unchanged; the same
primary with
``fallback_conn_ids: ["azure_dr"]`` does not forward it to ``azure_dr`` at
all, and
-``azure_dr`` raises "no model specified" unless its own extra sets a
``model``. See
-:doc:`connections/pydantic_ai_azure`, :doc:`connections/pydantic_ai_bedrock`
and
-:doc:`connections/pydantic_ai_vertex` for how each vendor connection resolves
a bare name.
+``azure_dr`` raises "no model specified" unless its own extra sets a
``model``. A generic
+``pydanticai`` connection with Default Model Provider set counts as being on
that platform. See
+:doc:`connections/pydantic_ai`, :doc:`connections/pydantic_ai_azure`,
+:doc:`connections/pydantic_ai_bedrock` and
:doc:`connections/pydantic_ai_vertex` for how
+each connection type resolves a bare name.
Configure it on the operator
-----------------------------
diff --git a/providers/common/ai/docs/self_hosted_models.rst
b/providers/common/ai/docs/self_hosted_models.rst
index ce449cd65fa..44fcaedce23 100644
--- a/providers/common/ai/docs/self_hosted_models.rst
+++ b/providers/common/ai/docs/self_hosted_models.rst
@@ -293,14 +293,17 @@ above.
Model identifier format
^^^^^^^^^^^^^^^^^^^^^^^^
-The ``extra`` JSON's ``model`` value keeps the ``provider:model`` format the
-hook always expects, but which ``provider`` prefix to use depends on the
-endpoint:
-
-- **vLLM** has no dedicated provider class in pydantic-ai -- ``openai:<model>``
- is the only option. pydantic-ai's ``openai`` provider class talks to
- whatever ``host`` points at, so it works against vLLM's OpenAI-compatible
- API even though the weights being served aren't OpenAI's.
+The ``extra`` JSON's ``model`` value keeps the ``provider:model`` format (or
+set the prefix once as Default Model Provider, see
+:ref:`conn-example-openai-compatible-gateway`), but which ``provider`` prefix
to
+use depends on the endpoint:
+
+- **vLLM** has no dedicated provider class in pydantic-ai, so use one of the
+ OpenAI prefixes: ``openai:<model>`` for the Responses API or
+ ``openai-chat:<model>`` for Chat Completions. pydantic-ai's ``openai``
+ provider class talks to whatever ``host`` points at, so either works against
+ vLLM's OpenAI-compatible API even though the weights being served aren't
+ OpenAI's.
- **Ollama** has two valid options: the dedicated ``ollama:<model>`` prefix
(recommended) applies Ollama-specific model-profile tuning based on the
model name (``llama``, ``gemma``, ``qwen``, ``deepseek``, ``mistral``,
diff --git a/providers/common/ai/docs/troubleshooting.rst
b/providers/common/ai/docs/troubleshooting.rst
index e2b329e74dc..cc4cb5ecf0a 100644
--- a/providers/common/ai/docs/troubleshooting.rst
+++ b/providers/common/ai/docs/troubleshooting.rst
@@ -33,11 +33,16 @@ Model and connection errors
``Connection '...' has no default model provider, so the bare model name '...'
cannot be resolved``
The generic ``pydanticai`` connection type needs a ``provider:`` prefix on
the model
- name, for example ``openai:gpt-5`` rather than ``gpt-5``. Add the prefix,
or use a
- vendor connection type (:doc:`connections/pydantic_ai_azure`,
+ name, for example ``openai:gpt-5`` rather than ``gpt-5``. Add the prefix,
set
+ **Default Model Provider** (``extra["model_provider"]``, e.g.
``openai-chat``) on the connection,
+ or use a vendor connection type (:doc:`connections/pydantic_ai_azure`,
:doc:`connections/pydantic_ai_bedrock`,
:doc:`connections/pydantic_ai_vertex`), which
supply their own platform.
+``Connection '...' sets 'model_provider' to '...', which is not a provider
pydantic-ai recognizes``
+ The connection's **Default Model Provider** is not a pydantic-ai provider.
Check it for a
+ typo; ``openai``, ``openai-chat`` and ``anthropic`` are examples of valid
values.
+
``'...' is not a provider pydantic-ai recognizes``
The text before the first ``:`` in the model name is not a pydantic-ai
provider.
Check it for a typo. If the vendor's own model id contains a ``:``
(Bedrock-style
diff --git a/providers/common/ai/provider.yaml
b/providers/common/ai/provider.yaml
index 158a36a61bd..80ac9b8c31e 100644
--- a/providers/common/ai/provider.yaml
+++ b/providers/common/ai/provider.yaml
@@ -225,7 +225,14 @@ connection-types:
conn-fields:
model:
label: Model
- description: "Model in provider:name format (e.g.
anthropic:claude-sonnet-5, openai:gpt-5)"
+ description: "Model as provider:name (e.g. openai:gpt-5), or bare with
a Default Model Provider"
+ schema:
+ type:
+ - string
+ - 'null'
+ model_provider:
+ label: Default Model Provider
+ description: "Provider for model names without a provider: prefix
(e.g. openai-chat)"
schema:
type:
- string
diff --git a/providers/common/ai/src/airflow/providers/common/ai/batch/base.py
b/providers/common/ai/src/airflow/providers/common/ai/batch/base.py
index 44ef109137d..d0c64683497 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/batch/base.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/batch/base.py
@@ -32,6 +32,7 @@ from typing import TYPE_CHECKING, Any, ClassVar, Literal,
Required
from typing_extensions import TypedDict
+from airflow.providers.common.ai.batch import dispatch
from airflow.providers.common.ai.exceptions import LLMBatchLimitExceededError,
LLMBatchModelMismatchError
if TYPE_CHECKING:
@@ -53,7 +54,8 @@ class BatchRequest(TypedDict, total=False):
prompt: Required[str]
#: Per-request model override, written as ``"<provider>:<model>"`` (the
same form the
- #: operator's own ``model_id`` uses). ``None`` inherits the batch-level
``model_id``.
+ #: operator's own ``model_id`` uses, so a bare name takes the connection's
Default
+ #: Provider). ``None`` inherits the batch-level ``model_id``.
#: Anthropic allows a different model per request in the same batch;
OpenAI requires
#: every request in a batch to resolve to the same model.
model: str | None
@@ -296,15 +298,15 @@ class BatchAdapter(ABC):
Resolve a per-request ``model`` override to this adapter's bare model
name.
``request_model`` must be ``None`` (inherit the batch-level default,
- already bare) or a ``"<provider>:<model>"`` string whose prefix equals
- :attr:`name`. A request naming a different provider is rejected before
- any network call rather than sent to the wrong provider's model
- namespace.
+ already bare) or a ``"<provider>:<model>"`` string whose prefix selects
+ this adapter (:attr:`name`, or an alias such as ``openai-chat``). A
+ request naming a different provider is rejected before any network call
+ rather than sent to the wrong provider's model namespace.
"""
if request_model is None:
return default_bare_model
prefix, sep, bare = request_model.partition(":")
- if not sep or prefix != self.name or not bare:
+ if not sep or dispatch.adapter_name_for_prefix(prefix) != self.name or
not bare:
raise LLMBatchModelMismatchError(
f"Request {request_index} sets model={request_model!r}, which
does not resolve "
f"to the {self.name!r} adapter this batch is using.
Per-request model overrides "
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/batch/dispatch.py
b/providers/common/ai/src/airflow/providers/common/ai/batch/dispatch.py
index 22615e2f46f..b97a6736796 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/batch/dispatch.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/batch/dispatch.py
@@ -59,10 +59,19 @@ _ADAPTER_MODULES: dict[str, tuple[str, str]] = {
"anthropic": ("airflow.providers.common.ai.batch.anthropic",
"AnthropicBatchAdapter"),
}
+#: model_id prefix -> the adapter that serves it under another name.
``openai-chat`` is
+#: pydantic-ai's Chat Completions prefix, and Chat Completions is what the
``openai`` adapter submits.
+_ADAPTER_ALIASES: dict[str, str] = {"openai-chat": "openai"}
+
#: Adapter classes registered at runtime through :func:`register_adapter`.
_REGISTERED_ADAPTERS: dict[str, type[BatchAdapter]] = {}
+def adapter_name_for_prefix(prefix: str) -> str:
+ """Return the name of the adapter that serves a ``model_id`` prefix
(``openai-chat`` -> ``openai``)."""
+ return _ADAPTER_ALIASES.get(prefix, prefix)
+
+
def register_adapter(adapter_cls: type[BatchAdapter]) -> None:
"""
Register an adapter class for its own
:attr:`~airflow.providers.common.ai.batch.base.BatchAdapter.name` prefix.
@@ -113,7 +122,7 @@ def split_model_id(model_id: str | None) -> tuple[str, str]:
raise UnsupportedBatchProviderError(
"model_id must be written as '<provider>:<model>' (e.g.
'openai:gpt-5') to "
f"select a batch adapter; got {model_id!r}. Set it on the operator
or in the "
- "connection's Model field."
+ "connection's Model field, or set the connection's Default Model
Provider."
)
prefix, _, bare = model_id.strip().partition(":")
if not prefix or not bare:
@@ -133,7 +142,7 @@ def get_adapter_class(conn_type: str | None, model_id: str
| None) -> type[Batch
accept this connection type (e.g. ``pydanticai_azure`` today).
"""
prefix, _ = split_model_id(model_id)
- adapter_cls = import_adapter_class(prefix)
+ adapter_cls = import_adapter_class(adapter_name_for_prefix(prefix))
if conn_type not in adapter_cls.conn_types:
raise BatchProviderNotYetSupportedError(
f"The {prefix!r} batch adapter does not support connection type
{conn_type!r}; it "
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 188a2831d0f..85112a45c33 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
@@ -196,7 +196,12 @@ def get_provider_info():
"conn-fields": {
"model": {
"label": "Model",
- "description": "Model in provider:name format (e.g.
anthropic:claude-sonnet-5, openai:gpt-5)",
+ "description": "Model as provider:name (e.g.
openai:gpt-5), or bare with a Default Model Provider",
+ "schema": {"type": ["string", "null"]},
+ },
+ "model_provider": {
+ "label": "Default Model Provider",
+ "description": "Provider for model names without a
provider: prefix (e.g. openai-chat)",
"schema": {"type": ["string", "null"]},
},
"fallback_conn_ids": {
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 8fee9c1da69..90d3a3736ba 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
@@ -51,6 +51,7 @@ R = TypeVar("R")
_UNSET: Any = object()
FALLBACK_CONN_IDS_EXTRA_KEY = "fallback_conn_ids"
+MODEL_PROVIDER_EXTRA_KEY = "model_provider"
if TYPE_CHECKING:
from collections.abc import Callable, MutableMapping
@@ -70,10 +71,13 @@ def _has_recognized_provider_prefix(model_name: str) ->
bool:
``infer_provider_class`` actually recognizes counts as a platform prefix.
"""
prefix, sep, _ = model_name.partition(":")
- if not sep:
- return False
+ return bool(sep) and _is_recognized_provider(prefix)
+
+
+def _is_recognized_provider(name: str) -> bool:
+ """Return whether *name* is a provider ``infer_provider_class``
recognizes."""
try:
- infer_provider_class(prefix)
+ infer_provider_class(name)
except ImportError:
return True # recognized provider; its optional dependency just isn't
installed
except ValueError:
@@ -144,7 +148,8 @@ class PydanticAIHook(BaseHook):
- **password**: API key
- **host**: Base URL (optional, e.g. ``https://api.openai.com/v1``)
- **extra** JSON: ``{"model": "openai:gpt-5",
- "fallback_conn_ids": ["anthropic_prod", "bedrock_dr"]}``
+ "fallback_conn_ids": ["anthropic_prod", "bedrock_dr"]}``; add
+ ``"model_provider": "openai-chat"`` to resolve bare model names on
this connection
:param llm_conn_id: Airflow connection ID for the LLM provider.
:param model_id: Model identifier. A name whose segment before the first
``:``
@@ -156,8 +161,9 @@ class PydanticAIHook(BaseHook):
platform: vendor subclasses (:class:`PydanticAIAzureHook`,
:class:`PydanticAIBedrockHook`, :class:`PydanticAIVertexHook`) each
default
to their own platform via :attr:`model_provider`; the generic
connection
- type has none, so a bare name there raises ``ValueError`` instead of
- reaching pydantic-ai's own, less actionable ``Unknown model`` error.
+ type uses the ``model_provider`` in its extra, and without one a bare
name there
+ raises ``ValueError`` instead of reaching pydantic-ai's own, less
actionable
+ ``Unknown model`` error.
Overrides the model stored in the connection's extra field. Whichever
of
the two configures the primary's model is forwarded (only while still
bare) down the fallback chain -- see :meth:`_resolve_fallback_models`.
@@ -176,8 +182,9 @@ class PydanticAIHook(BaseHook):
default_conn_name = "pydanticai_default"
conn_type = "pydanticai"
hook_name = "Pydantic AI"
- # Platform to prefix a bare model_id with (e.g. "azure"); None for the
generic
- # connection type, which has no platform of its own. Vendor subclasses
override this.
+ # Platform to prefix a bare model_id with (e.g. "azure"). Vendor
subclasses set it; the
+ # generic connection type leaves it None and reads extra["model_provider"]
instead (see
+ # _get_model_provider).
model_provider: str | None = None
def __init__(
@@ -384,6 +391,30 @@ class PydanticAIHook(BaseHook):
return self._model
return await _run_off_loop(self.get_conn)
+ def _get_model_provider(self) -> str | None:
+ """
+ Return the platform to prefix a bare model name with.
+
+ A vendor subclass's own :attr:`model_provider` wins: its credential
fields are mapped
+ for that one platform, so a ``model_provider`` in its extra is
ignored. The generic
+ connection type takes the ``model_provider`` in its extra, if any.
+
+ :raises ValueError: if the extra's ``model_provider`` is not a
provider pydantic-ai recognizes.
+ """
+ if self.model_provider is not None:
+ return self.model_provider
+ _, extra = self._get_conn_and_extra()
+ provider = extra.get(MODEL_PROVIDER_EXTRA_KEY)
+ if not provider:
+ return None
+ if not isinstance(provider, str) or not
_is_recognized_provider(provider):
+ raise ValueError(
+ f"Connection '{self.llm_conn_id}' sets
'{MODEL_PROVIDER_EXTRA_KEY}' to {provider!r}, "
+ "which is not a provider pydantic-ai recognizes (e.g.
'openai', 'openai-chat', "
+ "'anthropic')."
+ )
+ return provider
+
def _qualify_model_name(self, model_name: str, *, forwarded_from_conn_id:
str | None = None) -> str:
"""
Prefix a bare model name with this connection's platform.
@@ -404,9 +435,12 @@ class PydanticAIHook(BaseHook):
return model_name
if _has_recognized_provider_prefix(model_name):
return model_name
- if self.model_provider is not None:
+ model_provider = self._get_model_provider()
+ if model_provider is not None:
prefix, sep, _ = model_name.partition(":")
- if sep and _looks_like_unrecognized_provider_prefix(prefix):
+ # Only a vendor type warns: behind a generic connection's Default
Model Provider, a
+ # "name:tag" id such as Ollama's "llama3:70b" is the normal
spelling, not a typo.
+ if sep and self.model_provider is not None and
_looks_like_unrecognized_provider_prefix(prefix):
self.log.warning(
"Model name '%s' on connection '%s' contains ':' but its
prefix '%s' is not a "
"provider pydantic-ai recognizes; treating the whole
string as a bare %s model id "
@@ -415,19 +449,20 @@ class PydanticAIHook(BaseHook):
model_name,
self.llm_conn_id,
prefix,
- self.model_provider,
- self.model_provider,
+ model_provider,
+ model_provider,
model_name,
prefix,
)
- return f"{self.model_provider}:{model_name}"
+ return f"{model_provider}:{model_name}"
if forwarded_from_conn_id is not None:
raise ValueError(
f"Connection '{self.llm_conn_id}' has no default model
provider, so the bare model "
f"name '{model_name}' -- forwarded from primary connection
'{forwarded_from_conn_id}' "
f"-- cannot be resolved here. Give '{forwarded_from_conn_id}'
a 'provider:model' "
- f"string, or set an explicit 'model' on '{self.llm_conn_id}'."
+ f"string, or set an explicit 'model' or a
'{MODEL_PROVIDER_EXTRA_KEY}' in the extra "
+ f"of '{self.llm_conn_id}'."
)
prefix, sep, _ = model_name.partition(":")
@@ -436,13 +471,15 @@ class PydanticAIHook(BaseHook):
f"Connection '{self.llm_conn_id}' has no default model
provider, and '{prefix}' is "
f"not a provider pydantic-ai recognizes, so '{model_name}'
cannot be resolved. If "
"this is a vendor's own model id containing a ':' (e.g. a
Bedrock-style "
- "version-suffixed id), use a vendor connection type
(Azure/Bedrock/Vertex) instead; "
- f"if '{prefix}' is meant to be a provider prefix, check it for
a typo."
+ "version-suffixed id), use a vendor connection type
(Azure/Bedrock/Vertex) or set a "
+ f"default '{MODEL_PROVIDER_EXTRA_KEY}' in the connection's
extra; if '{prefix}' is "
+ "meant to be a provider prefix, check it for a typo."
)
raise ValueError(
f"Connection '{self.llm_conn_id}' has no default model provider,
so the bare model name "
- f"'{model_name}' cannot be resolved. Use a vendor connection type
(Azure/Bedrock/Vertex) "
- "or set an explicit 'provider:model' string."
+ f"'{model_name}' cannot be resolved. Use a vendor connection type
(Azure/Bedrock/Vertex), "
+ f"set an explicit 'provider:model' string, or set a default
'{MODEL_PROVIDER_EXTRA_KEY}' "
+ "in the connection's extra."
)
def _get_configured_model_name(self) -> str | KnownModelName | None:
@@ -468,14 +505,16 @@ class PydanticAIHook(BaseHook):
name with an embedded ``:`` of its own (a vendor's own native id,
e.g.
Bedrock's version-suffixed ``us.anthropic.claude-opus-4-6-v1:0``)
is only
forwarded when *forwarded_model_provider* matches this
connection's own
- :attr:`model_provider` -- that spelling is only meaningful on the
- platform that produced it.
+ platform (see :meth:`_get_model_provider`) -- that spelling is only
+ meaningful on the platform that produced it.
:param forwarded_from_conn_id: The primary connection's ID, for error
messages
attributing an unresolvable forwarded name to where it actually
came from.
- :param forwarded_model_provider: The primary connection's
:attr:`model_provider`;
+ :param forwarded_model_provider: The primary connection's platform;
see *forwarded_model_id* above for how it gates forwarding.
"""
conn, extra = self._get_conn_and_extra()
+ # Read up front, so a bad ``model_provider`` in the extra fails even
when the model name pins its own.
+ model_provider = self._get_model_provider()
model_name: str | KnownModelName | None =
self._get_configured_model_name()
forwarded = False
@@ -483,7 +522,7 @@ class PydanticAIHook(BaseHook):
not model_name
and forwarded_model_id
and not _has_recognized_provider_prefix(forwarded_model_id)
- and (":" not in forwarded_model_id or forwarded_model_provider ==
self.model_provider)
+ and (":" not in forwarded_model_id or forwarded_model_provider ==
model_provider)
):
model_name = forwarded_model_id
forwarded = True
@@ -570,6 +609,7 @@ class PydanticAIHook(BaseHook):
return []
forwarded_model_id = self._get_configured_model_name()
+ forwarded_model_provider = self._get_model_provider()
self.log.info("Resolving LLM fallback chain: %s", " ->
".join([self.llm_conn_id, *fallback_conn_ids]))
@@ -614,7 +654,7 @@ class PydanticAIHook(BaseHook):
hook._resolve_own_model(
forwarded_model_id=forwarded_model_id,
forwarded_from_conn_id=self.llm_conn_id,
- forwarded_model_provider=self.model_provider,
+ forwarded_model_provider=forwarded_model_provider,
)
)
diff --git
a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_batch.py
b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_batch.py
index 816318c2bec..b576efecb7d 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/operators/llm_batch.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/operators/llm_batch.py
@@ -40,6 +40,11 @@ from airflow.providers.common.ai.exceptions import (
LLMBatchStaleStateError,
LLMBatchTimeoutError,
)
+from airflow.providers.common.ai.hooks.pydantic_ai import (
+ MODEL_PROVIDER_EXTRA_KEY,
+ PydanticAIHook,
+ _has_recognized_provider_prefix,
+)
from airflow.providers.common.ai.triggers.llm_batch import LLMBatchTrigger
from airflow.providers.common.compat.sdk import BaseHook, BaseOperator,
ObjectStoragePath, conf
@@ -89,7 +94,8 @@ class LLMBatchOperator(BaseOperator):
``pydanticai`` connection today.
:param model_id: Model identifier as ``"<provider>:<model>"`` (e.g.
``"openai:gpt-5"``). Falls back to the connection's Model field, as
- ``LLMOperator`` does.
+ ``LLMOperator`` does. A bare name takes the connection's Default
+ Provider, and ``openai-chat:`` selects the OpenAI adapter.
:param system_prompt: System-level instructions applied to every request
that does not set its own ``system_prompt`` override.
:param output_type: Expected output type. Default ``str``. Set to a
@@ -243,9 +249,23 @@ class LLMBatchOperator(BaseOperator):
def _resolve_model_id(self, conn: Connection) -> str | None:
"""``model_id`` on the operator wins; otherwise the connection's Model
field, as ``LLMOperator`` does."""
- if self.model_id:
- return self.model_id
- return conn.extra_dejson.get("model") or None
+ return self._qualify_model_id(self.model_id or
conn.extra_dejson.get("model") or None, conn)
+
+ @staticmethod
+ def _qualify_model_id(model_id: str | None, conn: Connection) -> str |
None:
+ """
+ Prefix a bare model name with the connection's Default Model Provider,
as ``PydanticAIHook`` does.
+
+ Only the generic ``pydanticai`` type reads
``extra["model_provider"]``, and a name whose
+ segment before the first ``:`` is already a pydantic-ai provider is
left as is, so the
+ same connection and name select the same platform on ``@task.llm`` and
here.
+ """
+ if not model_id or conn.conn_type != PydanticAIHook.conn_type:
+ return model_id
+ default_provider = conn.extra_dejson.get(MODEL_PROVIDER_EXTRA_KEY)
+ if not default_provider or _has_recognized_provider_prefix(model_id):
+ return model_id
+ return f"{default_provider}:{model_id}"
def _connection(self) -> Connection:
return BaseHook.get_connection(self.llm_conn_id)
@@ -277,6 +297,9 @@ class LLMBatchOperator(BaseOperator):
output_schema: dict[str, Any] | str = spec.json_schema if
spec.is_structured else "str"
conn = self._connection()
+ for request in requests:
+ if request.get("model"):
+ request["model"] = self._qualify_model_id(request["model"],
conn)
model_id = self._resolve_model_id(conn)
adapter_name, adapter = self._build_adapter(conn)
_, bare_model = dispatch.split_model_id(model_id)
diff --git a/providers/common/ai/tests/unit/common/ai/batch/test_base.py
b/providers/common/ai/tests/unit/common/ai/batch/test_base.py
index dec6b400550..54e56a7a559 100644
--- a/providers/common/ai/tests/unit/common/ai/batch/test_base.py
+++ b/providers/common/ai/tests/unit/common/ai/batch/test_base.py
@@ -132,6 +132,14 @@ class TestResolveRequestModel:
result = adapter.resolve_request_model("fake:m-2",
default_bare_model="m-1", request_index=0)
assert result == "m-2"
+ def test_alias_prefix_resolves_for_the_adapter_it_selects(self):
+ class _OpenAINamedAdapter(_FakeAdapter):
+ name = "openai"
+
+ adapter = _OpenAINamedAdapter()
+ result = adapter.resolve_request_model("openai-chat:m-2",
default_bare_model="m-1", request_index=0)
+ assert result == "m-2"
+
def test_foreign_provider_prefix_is_rejected(self):
adapter = _FakeAdapter()
with pytest.raises(LLMBatchModelMismatchError, match="Request 3"):
diff --git a/providers/common/ai/tests/unit/common/ai/batch/test_dispatch.py
b/providers/common/ai/tests/unit/common/ai/batch/test_dispatch.py
index 1603f6b2c05..cdedc57bfd9 100644
--- a/providers/common/ai/tests/unit/common/ai/batch/test_dispatch.py
+++ b/providers/common/ai/tests/unit/common/ai/batch/test_dispatch.py
@@ -80,6 +80,10 @@ class TestGetAdapterClass:
with pytest.raises(UnsupportedBatchProviderError, match="'groq' has no
batch adapter"):
dispatch.get_adapter_class("pydanticai", "groq:llama")
+ def test_openai_chat_prefix_selects_the_openai_adapter(self):
+ """pydantic-ai's Chat Completions prefix names the endpoint the OpenAI
adapter submits to."""
+ assert dispatch.get_adapter_class("pydanticai", "openai-chat:gpt-5")
is OpenAIBatchAdapter
+
def test_resolve_adapter_name_returns_the_prefix(self):
assert dispatch.resolve_adapter_name("pydanticai", "openai:gpt-5") ==
"openai"
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 596cbd27a76..04f8ebb0813 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
@@ -37,6 +37,7 @@ from pydantic_ai.messages import ModelResponse, TextPart,
ToolCallPart, ToolRetu
from pydantic_ai.models import Model
from pydantic_ai.models.fallback import FallbackModel
from pydantic_ai.models.function import FunctionModel
+from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.models.test import TestModel
from pydantic_ai.providers import infer_provider_class
@@ -737,6 +738,111 @@ class TestPydanticAIHookModelProviderResolution:
assert "openi" in str(exc_info.value)
assert "openi:gpt-5" in str(exc_info.value)
+ @pytest.mark.parametrize(
+ ("model_id", "extra_model"),
+ [
+ pytest.param(None, "claude-sonnet-4-5", id="model-in-extra"),
+ pytest.param("claude-sonnet-4-5", None, id="model-id-param"),
+ ],
+ )
+ def test_generic_connection_provider_extra_qualifies_bare_name(
+ self, registry, infer_model_stub, model_id, extra_model
+ ):
+ """The extra's ``model_provider`` gives the generic connection type a
platform for bare names."""
+ extra = {"model_provider": "openai-chat"}
+ if extra_model:
+ extra["model"] = extra_model
+ registry.add("gateway", extra=extra)
+ hook = PydanticAIHook(llm_conn_id="gateway", model_id=model_id)
+
+ assert hook.get_conn() is
infer_model_stub.models["openai-chat:claude-sonnet-4-5"]
+
+ def
test_generic_connection_provider_extra_leaves_prefixed_name_alone(self,
registry, infer_model_stub):
+ """A name already pinning a recognized provider wins over the extra's
default."""
+ registry.add(
+ "gateway", extra={"model_provider": "openai-chat", "model":
"anthropic:claude-sonnet-4-5"}
+ )
+ hook = PydanticAIHook(llm_conn_id="gateway")
+
+ assert hook.get_conn() is
infer_model_stub.models["anthropic:claude-sonnet-4-5"]
+
+ def test_generic_connection_provider_extra_name_tag_does_not_warn(
+ self, registry, infer_model_stub, caplog
+ ):
+ """Behind a gateway, Ollama-style ``name:tag`` ids are the normal
spelling, not a typo."""
+ registry.add("gateway", extra={"model_provider": "openai-chat",
"model": "llama3:70b"})
+ hook = PydanticAIHook(llm_conn_id="gateway")
+
+ with caplog.at_level(logging.WARNING):
+ model = hook.get_conn()
+
+ assert model is infer_model_stub.models["openai-chat:llama3:70b"]
+ assert not any("typo" in r.message for r in caplog.records)
+
+ def test_generic_connection_provider_extra_builds_the_gateway_model(self):
+ """The documented gateway connection yields a Chat Completions model
aimed at its host."""
+ conn = Connection(
+ conn_id="gateway",
+ conn_type="pydanticai",
+ password="gateway-key",
+ host="https://llm-gateway.example.com/v1";,
+ extra=json.dumps({"model_provider": "openai-chat"}),
+ )
+ hook = PydanticAIHook(llm_conn_id="gateway",
model_id="claude-sonnet-4-5")
+
+ with patch.object(hook, "get_connection", autospec=True,
return_value=conn):
+ model = hook.get_conn()
+
+ assert isinstance(model, OpenAIChatModel)
+ assert model.model_name == "claude-sonnet-4-5"
+ assert model.base_url.rstrip("/") ==
"https://llm-gateway.example.com/v1";
+
+ @pytest.mark.parametrize(
+ ("provider", "model"),
+ [
+ pytest.param("opneai-chat", "gpt-5", id="typo-bare-model"),
+ pytest.param("opneai-chat", "openai:gpt-5",
id="typo-prefixed-model"),
+ pytest.param(["openai-chat"], "gpt-5", id="not-a-string"),
+ ],
+ )
+ def test_generic_connection_unknown_provider_extra_raises(
+ self, registry, infer_model_stub, provider, model
+ ):
+ """A bad ``model_provider`` fails at resolve time even when the model
name would not need it."""
+ registry.add("gateway", extra={"model_provider": provider, "model":
model})
+ hook = PydanticAIHook(llm_conn_id="gateway")
+
+ with pytest.raises(
+ ValueError,
+ match=re.escape(
+ f"Connection 'gateway' sets 'model_provider' to {provider!r},
which is not a provider"
+ ),
+ ):
+ hook.get_conn()
+
+ @pytest.mark.parametrize("provider", [pytest.param(None, id="null"),
pytest.param("", id="empty")])
+ def test_generic_connection_without_provider_extra_raises(self, registry,
infer_model_stub, provider):
+ """An unset or blank ``model_provider`` (the form can save either)
leaves the error pointing at the fix."""
+ registry.add("primary", extra={"model": "gpt-4", "model_provider":
provider})
+ hook = PydanticAIHook(llm_conn_id="primary")
+
+ with pytest.raises(ValueError, match="has no default model provider")
as exc_info:
+ hook.get_conn()
+
+ assert "set a default 'model_provider' in the connection's extra" in
str(exc_info.value)
+
+ def test_vendor_connection_ignores_provider_extra(self, registry,
infer_model_stub):
+ """A vendor type's credential fields only fit its own platform, so its
class attribute wins."""
+ registry.add(
+ "primary",
+ conn_type="pydanticai_azure",
+ hook_class=PydanticAIAzureHook,
+ extra={"model": "gpt-5", "model_provider": "openai-chat"},
+ )
+ hook = PydanticAIAzureHook(llm_conn_id="primary")
+
+ assert hook.get_conn() is infer_model_stub.models["azure:gpt-5"]
+
class TestPydanticAIHookFallback:
def test_no_fallback_returns_the_bare_model(self, registry,
infer_model_stub):
@@ -938,7 +1044,7 @@ class TestPydanticAIHookFallback:
native vendor id is never valid on another platform -- it must not be
forwarded there.
Mutation canary: dropping the ``":" not in forwarded_model_id or
forwarded_model_provider
- == self.model_provider`` gate (reverting to just ``not
_has_recognized_provider_prefix``)
+ == model_provider`` gate (reverting to just ``not
_has_recognized_provider_prefix``)
makes the Bedrock id forward to the Azure fallback and this raise
never fires --
``pytest.raises`` would report no exception raised.
"""
@@ -994,6 +1100,47 @@ class TestPydanticAIHookFallback:
assert "openai_dr" in str(exc_info.value)
assert "claude-opus-4-5" in str(exc_info.value)
+ def test_bare_name_forwarded_to_generic_fallback_with_provider_extra(self,
registry, infer_model_stub):
+ """A forwarded bare name is qualified with the fallback's own
``model_provider``, not the primary's."""
+ registry.add(
+ "gateway",
+ extra={
+ "model_provider": "openai-chat",
+ "model": "claude-sonnet-4-5",
+ "fallback_conn_ids": ["anthropic"],
+ },
+ )
+ registry.add("anthropic", extra={"model_provider": "anthropic"})
+
+ model = PydanticAIHook(llm_conn_id="gateway").get_conn()
+
+ assert isinstance(model, FallbackModel)
+ assert model.models == [
+ infer_model_stub.models["openai-chat:claude-sonnet-4-5"],
+ infer_model_stub.models["anthropic:claude-sonnet-4-5"],
+ ]
+
+ def
test_colon_bearing_bare_name_forwarded_within_same_provider_extra(self,
registry, infer_model_stub):
+ """The embedded-colon forwarding gate treats a generic connection's
``model_provider`` extra as its platform."""
+ registry.add("gateway", extra={"model_provider": "openai-chat",
"fallback_conn_ids": ["second"]})
+ registry.add("second", extra={"model_provider": "openai-chat"})
+
+ model = PydanticAIHook(llm_conn_id="gateway",
model_id="llama3:70b").get_conn()
+
+ assert isinstance(model, FallbackModel)
+ assert model.models == [
+ infer_model_stub.models["openai-chat:llama3:70b"],
+ infer_model_stub.models["openai-chat:llama3:70b"],
+ ]
+
+ def
test_colon_bearing_bare_name_not_forwarded_across_provider_extras(self,
registry, infer_model_stub):
+ registry.add("gateway", extra={"model_provider": "openai-chat",
"fallback_conn_ids": ["second"]})
+ registry.add("second", extra={"model_provider": "anthropic"})
+
+ hook = PydanticAIHook(llm_conn_id="gateway", model_id="llama3:70b")
+ with pytest.raises(ValueError, match="No model specified for
connection 'second'"):
+ hook.get_conn()
+
def test_non_pydanticai_fallback_raises(self, registry, infer_model_stub):
"""``Connection.get_hook`` dispatches on conn_type alone and can
return anything."""
registry.add("primary", extra={"model": "openai:gpt-5.6-sol"})
@@ -2178,6 +2325,17 @@ class TestPydanticAIVertexHook:
prefix = _extract_google_cloud_prefix(description)
_assert_prefix_is_known_provider(prefix)
+ def
test_conn_fields_default_provider_description_example_is_valid_provider(self):
+ """Drift tripwire for the example in the generic connection's Default
Model Provider conn-field."""
+ connection_types = get_provider_info()["connection-types"]
+ generic_conn_fields = next(
+ c["conn-fields"] for c in connection_types if c["connection-type"]
== "pydanticai"
+ )
+ description = generic_conn_fields["model_provider"]["description"]
+ example = re.search(r"\(e\.g\. ([a-z0-9-]+)\)", description)
+ assert example is not None, description
+ _assert_prefix_is_known_provider(example.group(1))
+
def
test_conn_types_ui_field_behaviour_placeholder_prefix_is_valid_provider(self):
"""
Drift tripwire for the ``provider.yaml``
``ui-field-behaviour.placeholders.extra``,
diff --git
a/providers/common/ai/tests/unit/common/ai/operators/test_llm_batch.py
b/providers/common/ai/tests/unit/common/ai/operators/test_llm_batch.py
index 611035a11e8..7fcce493617 100644
--- a/providers/common/ai/tests/unit/common/ai/operators/test_llm_batch.py
+++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm_batch.py
@@ -46,6 +46,7 @@ from airflow.providers.common.ai.exceptions import (
LLMBatchStaleStateError,
LLMBatchStateReadError,
LLMBatchTimeoutError,
+ UnsupportedBatchProviderError,
)
from airflow.providers.common.ai.operators import llm_batch as llm_batch_module
from airflow.providers.common.ai.operators.llm_batch import LLMBatchOperator
@@ -309,6 +310,85 @@ class TestModelFromConnection:
assert fake_adapter.submit_calls[0]["model"] == "gpt-5-mini"
+ @pytest.mark.parametrize(
+ ("model_id", "extra", "expected_model_id", "expected_bare_model"),
+ [
+ pytest.param(
+ None,
+ {"model_provider": "openai-chat", "model": "gpt-5-mini"},
+ "openai-chat:gpt-5-mini",
+ "gpt-5-mini",
+ id="model-field",
+ ),
+ pytest.param(
+ "gpt-5-mini",
+ {"model_provider": "openai-chat"},
+ "openai-chat:gpt-5-mini",
+ "gpt-5-mini",
+ id="model-id",
+ ),
+ pytest.param(
+ "llama3:70b", {"model_provider": "openai"},
"openai:llama3:70b", "llama3:70b", id="name-tag"
+ ),
+ pytest.param(
+ "openai:gpt-5-mini",
+ {"model_provider": "anthropic"},
+ "openai:gpt-5-mini",
+ "gpt-5-mini",
+ id="prefixed",
+ ),
+ ],
+ )
+ def test_default_provider_qualifies_a_bare_model(
+ self, fake_adapter, result_path, model_id, extra, expected_model_id,
expected_bare_model
+ ):
+ """The connection's Default Model Provider resolves a bare name here
as it does on ``@task.llm``."""
+ ctx = _context()
+ with _running(fake_adapter, conn=_connection(extra=extra)):
+ op = _make_operator(result_path=result_path, deferrable=True,
model_id=model_id)
+ with pytest.raises(TaskDeferred):
+ op.execute(ctx)
+
+ key, _ = LLMBatchOperator._identity(ctx)
+ record = state_module.read_state(ObjectStoragePath(result_path), key)
+ assert record is not None
+ assert record.model_id == expected_model_id
+ assert record.adapter == "openai"
+ assert fake_adapter.submit_calls[0]["model"] == expected_bare_model
+
+ def test_default_provider_qualifies_a_bare_per_request_model(self,
fake_adapter, result_path):
+ requests = [{"prompt": "a", "model": "gpt-5-mini"}, {"prompt": "b"}]
+ with _running(fake_adapter, conn=_connection(extra={"model_provider":
"openai-chat"})):
+ op = _make_operator(
+ result_path=result_path, deferrable=True,
model_id="gpt-5-mini", requests=requests
+ )
+ with pytest.raises(TaskDeferred):
+ op.execute(_context())
+
+ submitted = fake_adapter.submit_calls[0]["requests"]
+ assert [r.get("model") for r in submitted] ==
["openai-chat:gpt-5-mini", None]
+ assert requests[0]["model"] == "gpt-5-mini", "the caller's request
dicts are not mutated"
+
+ def test_bare_model_without_default_provider_points_at_it(self,
fake_adapter, result_path):
+ with _running(fake_adapter, conn=_connection(extra={"model":
"gpt-5-mini"})):
+ op = _make_operator(result_path=result_path, deferrable=True,
model_id=None)
+ with pytest.raises(
+ UnsupportedBatchProviderError, match="set the connection's
Default Model Provider"
+ ):
+ op.execute(_context())
+
+ def test_vendor_connection_ignores_default_provider(self, fake_adapter,
result_path):
+ """As on ``PydanticAIHook``, only the generic type reads
``extra["model_provider"]``."""
+ conn = _connection(
+ conn_type="pydanticai_azure", extra={"model_provider": "openai",
"model": "gpt-5-mini"}
+ )
+ with _running(fake_adapter, conn=conn):
+ op = _make_operator(result_path=result_path, deferrable=True,
model_id=None)
+ with pytest.raises(
+ UnsupportedBatchProviderError, match="must be written as
'<provider>:<model>'"
+ ):
+ op.execute(_context())
+
class TestReattachDoesNotResubmit:
def test_second_execute_with_same_identity_and_input_skips_submit(self,
fake_adapter, result_path):
