kaxil commented on code in PR #72152:
URL: https://github.com/apache/airflow/pull/72152#discussion_r3951298911
##########
providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py:
##########
@@ -930,3 +951,35 @@ def
test_documented_model_prefix_is_a_valid_pydantic_ai_provider(self):
# environment; failing past provider-name resolution is enough to
# prove "google-cloud" is recognized.
pass
+
+ def test_conn_fields_model_description_prefix_is_valid_provider(self):
+ """
+ Drift tripwire for the ``provider.yaml`` conn-field, the actual UI
source.
+
+ Once a hook's ``provider.yaml`` declares ``conn-fields``, the
connection
+ form renders those and ``get_ui_field_behaviour`` placeholders are
never
+ shown (``providers_manager.py``'s ``ui_metadata_loaded``, deprecated
+ since 3.2.0) — so this description, not the placeholder below, is what
+ a user actually copies the model prefix from.
+ """
+ connection_types = get_provider_info()["connection-types"]
+ vertex_conn_fields = next(
+ c["conn-fields"] for c in connection_types if c["connection-type"]
== "pydanticai-vertex"
+ )
+ description = vertex_conn_fields["model"]["description"]
+ match = _GOOGLE_MODEL_PREFIX_RE.search(description)
+ assert match, f"no google model prefix found in description:
{description!r}"
+ _assert_prefix_is_known_provider(match.group())
+
+ def test_ui_field_behaviour_placeholder_prefix_is_valid_provider(self):
+ """
+ Drift tripwire for the ``get_ui_field_behaviour`` placeholder.
+
+ Superseded at runtime by the ``provider.yaml`` conn-field above, but
+ still source code a developer can read and copy from directly, so it
+ needs to stay accurate too.
+ """
+ placeholder =
PydanticAIVertexHook.get_ui_field_behaviour()["placeholders"]["extra"]
Review Comment:
This reads the Python copy, but by the reasoning in the test right above it,
the live placeholder is `provider.yaml`'s
`ui-field-behaviour.placeholders.extra` (`providers_manager.py:1015` picks it
up, and `ui_metadata_loaded` then skips this method). That is a separate copy
of `google-cloud:gemini-2.0-flash` at `provider.yaml:314` and it stays
untested. Since `get_provider_info()` already hands back that block next to
`conn-fields`, covering the rendered one is about a line.
##########
providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py:
##########
@@ -27,13 +29,32 @@
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 (
PydanticAIAzureHook,
PydanticAIBedrockHook,
PydanticAIHook,
PydanticAIVertexHook,
)
+# Matches the `google...` provider key pydantic-ai expects before the
`:model-name`
+# separator, e.g. "google-cloud" out of "google-cloud:gemini-2.0-flash".
+_GOOGLE_MODEL_PREFIX_RE = re.compile(r"google[\w-]*(?=:)")
Review Comment:
This also matches a bare `google`, which is itself a registered pydantic-ai
provider (`GoogleProvider`, the Generative Language API). I checked against
pydantic-ai 2.36.0 (what `pydantic-ai-slim>=2.0.0` resolves to today):
`infer_provider_class("google")` gets past name resolution fine and only raises
`ImportError` on the optional dep, which the helper suppresses. So if either
surface drifts from `google-cloud:` back to `google:`, both new tests still
pass. That is the mix-up the hook warns about at `pydantic_ai.py:524`, where
the wrong prefix silently selects the wrong API. Could these assert
`match.group() == "google-cloud"`, or that the resolved class is
`GoogleCloudProvider`, instead of just "some known provider"?
##########
providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py:
##########
@@ -27,13 +29,32 @@
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 (
PydanticAIAzureHook,
PydanticAIBedrockHook,
PydanticAIHook,
PydanticAIVertexHook,
)
+# Matches the `google...` provider key pydantic-ai expects before the
`:model-name`
+# separator, e.g. "google-cloud" out of "google-cloud:gemini-2.0-flash".
+_GOOGLE_MODEL_PREFIX_RE = re.compile(r"google[\w-]*(?=:)")
+
+
+def _assert_prefix_is_known_provider(prefix: str) -> None:
Review Comment:
The test just above at L939 does this same check inline, with a
`pytest.fail` that names the offending prefix. Worth folding that into the
helper and having L939 call it, so there is one shape for this rather than two,
and a drift fails with a message instead of a bare `ValueError` traceback.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
