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 f34cae8fa6f Suggest common model providers in the pydanticai
connection form (#76053)
f34cae8fa6f is described below
commit f34cae8fa6f59903d75f0db4d3a358ce4e4ea6f7
Author: Kaxil Naik <[email protected]>
AuthorDate: Wed Nov 4 01:41:02 2026 +0000
Suggest common model providers in the pydanticai connection form (#76053)
The Default Model Provider field was a blank text box, so users had to know
pydantic-ai's provider names (`openai-chat` vs `openai`) by heart. The
field's
schema now lists common ones as `examples`, which the connection form
renders as
suggestions on a text input: any other provider pydantic-ai recognizes can
still
be typed. An `enum` dropdown would reject providers pydantic-ai adds until
a new
release of this provider. A test keeps every suggestion resolvable.
The suggestions ship with the provider, but the installed pydantic-ai can
recognize more. The field's help text and the connection docs now say the
list
is common choices, not every provider that works. A test also keeps each
suggestion among the connection's declared external services, so the form
never
suggests a vendor the docs and registry don't list.
* Keep the model provider help text accurate on the registry and share test
lookups
The registry's connection builder shows a conn-field's description but not
its
examples, so the help text now says any provider the installed pydantic-ai
recognizes works instead of describing suggestions. The docs paragraph no
longer
repeats part of the list by hand. The suggestion checks reuse one
connection-type
lookup and one prefix-to-module mapping, and the two pydantic-ai checks on
this
field are one test.
---
.../common/ai/docs/connections/pydantic_ai.rst | 4 +++
providers/common/ai/provider.yaml | 18 +++++++++++++-
.../providers/common/ai/get_provider_info.py | 19 ++++++++++++--
.../tests/unit/common/ai/hooks/test_pydantic_ai.py | 14 ++++++-----
.../tests/unit/common/ai/test_provider_metadata.py | 29 +++++++++++++++++++---
5 files changed, 71 insertions(+), 13 deletions(-)
diff --git a/providers/common/ai/docs/connections/pydantic_ai.rst
b/providers/common/ai/docs/connections/pydantic_ai.rst
index 9ff9ac13da2..831e0ac2b54 100644
--- a/providers/common/ai/docs/connections/pydantic_ai.rst
+++ b/providers/common/ai/docs/connections/pydantic_ai.rst
@@ -91,6 +91,10 @@ Default Model Provider
``@task.llm_batch`` applies it the same way, and sends ``openai-chat``
models
to its OpenAI adapter.
+ The form suggests common providers as you type. The suggestions are not a
full
+ list: any provider the installed pydantic-ai recognizes is accepted,
including
+ ones a newer pydantic-ai adds.
+
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
diff --git a/providers/common/ai/provider.yaml
b/providers/common/ai/provider.yaml
index eb370d97f32..7196a6955d1 100644
--- a/providers/common/ai/provider.yaml
+++ b/providers/common/ai/provider.yaml
@@ -232,11 +232,27 @@ connection-types:
- 'null'
model_provider:
label: Default Model Provider
- description: "Provider for model names without a provider: prefix
(e.g. openai-chat)"
+ description: >
+ Provider for model names without a provider: prefix (e.g.
openai-chat).
+ Any provider the installed pydantic-ai recognizes works.
schema:
type:
- string
- 'null'
+ # Suggestions on a free-text input, not an enum: an enum would
reject providers a newer
+ # pydantic-ai adds until the next release of this provider.
+ examples:
+ - openai-chat
+ - openai
+ - anthropic
+ - ollama
+ - litellm
+ - openrouter
+ - groq
+ - mistral
+ - deepseek
+ - together
+ - fireworks
fallback_conn_ids:
label: Fallback Connections
description: "Connection IDs to fail over to, in order, while this
provider is unavailable."
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 85112a45c33..ffea5c04250 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
@@ -201,8 +201,23 @@ def get_provider_info():
},
"model_provider": {
"label": "Default Model Provider",
- "description": "Provider for model names without a
provider: prefix (e.g. openai-chat)",
- "schema": {"type": ["string", "null"]},
+ "description": "Provider for model names without a
provider: prefix (e.g. openai-chat). Any provider the installed pydantic-ai
recognizes works.\n",
+ "schema": {
+ "type": ["string", "null"],
+ "examples": [
+ "openai-chat",
+ "openai",
+ "anthropic",
+ "ollama",
+ "litellm",
+ "openrouter",
+ "groq",
+ "mistral",
+ "deepseek",
+ "together",
+ "fireworks",
+ ],
+ },
},
"fallback_conn_ids": {
"label": "Fallback Connections",
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 04f8ebb0813..47af1321456 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
@@ -2325,16 +2325,18 @@ 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."""
+ def
test_conn_fields_default_provider_example_and_suggestions_are_valid_providers(self):
+ """Drift tripwire for the Default Model Provider description example
and form suggestions."""
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))
+ field = generic_conn_fields["model_provider"]
+ example = re.search(r"\(e\.g\. ([a-z0-9-]+)\)", field["description"])
+ assert example is not None, field["description"]
+ assert field["schema"]["examples"]
+ for provider in [example.group(1), *field["schema"]["examples"]]:
+ _assert_prefix_is_known_provider(provider)
def
test_conn_types_ui_field_behaviour_placeholder_prefix_is_valid_provider(self):
"""
diff --git a/providers/common/ai/tests/unit/common/ai/test_provider_metadata.py
b/providers/common/ai/tests/unit/common/ai/test_provider_metadata.py
index df10e78b415..16ae1309350 100644
--- a/providers/common/ai/tests/unit/common/ai/test_provider_metadata.py
+++ b/providers/common/ai/tests/unit/common/ai/test_provider_metadata.py
@@ -267,10 +267,19 @@ def _provider_classes_in_module(mod) -> list[type]:
return provider_classes
-def _find_pydanticai_external_services() -> list[str]:
+def _find_pydanticai_connection_type() -> dict:
connection_types = get_provider_info()["connection-types"]
(pydanticai,) = (c for c in connection_types if c["connection-type"] ==
"pydanticai")
- return pydanticai["external-services"]
+ return pydanticai
+
+
+def _find_pydanticai_external_services() -> list[str]:
+ return _find_pydanticai_connection_type()["external-services"]
+
+
+def _module_for_prefix(prefix: str) -> str:
+ """The pydantic-ai provider module behind a model-name prefix
(``openai-chat`` -> ``openai``)."""
+ return PREFIX_ALIASES.get(prefix, prefix.replace("-", "_"))
class TestConnectionTypeExternalServices:
@@ -281,6 +290,18 @@ class TestConnectionTypeExternalServices:
)
+class TestModelProviderSuggestions:
+ def test_every_suggestion_is_a_declared_external_service(self):
+ """The form must not suggest a vendor the docs and registry don't list
for this connection."""
+ pydanticai = _find_pydanticai_connection_type()
+ suggestions =
pydanticai["conn-fields"]["model_provider"]["schema"]["examples"]
+ declared = set(pydanticai["external-services"])
+
+ undeclared = [s for s in suggestions if
LABELS.get(_module_for_prefix(s)) not in declared]
+
+ assert not undeclared, f"suggested providers missing from
external-services: {undeclared}"
+
+
class TestPydanticAIExternalServicesDrift:
def test_every_installed_vendor_is_declared(self):
"""
@@ -381,7 +402,7 @@ class TestPydanticAIExternalServicesDrift:
if ":" not in name:
continue # the bare "test" entry, pydantic-ai's built-in
keyless TestModel
prefix = name.removeprefix("gateway/").split(":", 1)[0]
- module = PREFIX_ALIASES.get(prefix, prefix.replace("-", "_"))
+ module = _module_for_prefix(prefix)
assert module in labelled, f"model prefix {prefix!r} (from
{name!r}) has no label"
def test_gateway_note_vendors_match_known_model_names(self):
@@ -396,7 +417,7 @@ class TestPydanticAIExternalServicesDrift:
for name in known_model_names()
if name.startswith("gateway/")
}
- modules = {PREFIX_ALIASES.get(p, p.replace("-", "_")) for p in
gateway_prefixes}
+ modules = {_module_for_prefix(p) for p in gateway_prefixes}
derived = {LABELS.get(m, f"<unlabelled:{m}>") for m in modules}
assert derived == GATEWAY_NOTE_VENDOR_LABELS, (
