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"") for m in modules} assert derived == GATEWAY_NOTE_VENDOR_LABELS, (