kaxil commented on code in PR #75909:
URL: https://github.com/apache/airflow/pull/75909#discussion_r4405864355
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -404,7 +435,8 @@ def _qualify_model_name(self, model_name: str, *,
forwarded_from_conn_id: str |
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()
Review Comment:
This branch only ran for the vendor types before, so it now sees
Ollama-style `name:tag` ids on gateway connections. With Default Provider set,
`llama3:70b` logs the "looks like it might be a typo" warning on every
resolution, and a tag whose name happens to be a provider, like `mistral:7b`,
skips the default entirely and goes to pydantic-ai's mistral provider. Maybe
skip the warning when the platform came from `extra["provider"]`, and have the
gateway section say that `openai-chat:mistral:7b` is how to write those?
##########
providers/common/ai/docs/connections/pydantic_ai.rst:
##########
@@ -81,6 +81,14 @@ Extra (JSON, optional)
When using the UI, the "Model" field above writes to this same location
automatically.
+Default Provider
+ The pydantic-ai provider to prefix a bare model name with, e.g.
+ ``openai-chat``. Stored in ``extra["provider"]``. With it set, ``model_id``
+ and the Model field can name a model the way the endpoint does (e.g.
Review Comment:
`@task.llm_batch` reads the same Model field but doesn't apply this.
`LLMBatchOperator._resolve_model_id` takes `extra["model"]` as is and
`dispatch.split_model_id` requires `<provider>:<model>`, so on `{"provider":
"openai", "model": "gpt-5"}` the batch task raises `model_id must be written as
'<provider>:<model>'` and sends the user back to the Model field. `openai-chat`
isn't a batch adapter key either (only `openai` and `anthropic` are), and
`operators/llm_batch.rst` points gateway users at exactly this kind of
connection. Could the batch path qualify bare ids the same way (mapping
`openai-chat` to the `openai` adapter), or at least this section and the Model
description in `provider.yaml` say batch still needs the prefix?
##########
providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py:
##########
@@ -737,6 +737,76 @@ def
test_generic_connection_typo_prefix_names_the_bad_segment(self, registry, in
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(
Review Comment:
All the new tests run under `infer_model_stub`, so they check the string
handed to pydantic-ai but never what it builds from it. Could one test build
the real model from the docs' gateway connection (host, password, `{"provider":
"openai-chat"}`) and assert it's an `OpenAIChatModel` pointed at that base URL?
##########
providers/common/ai/docs/connections/pydantic_ai.rst:
##########
@@ -193,5 +229,6 @@ The hook reads the model from these sources in priority
order:
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.
+qualified with this connection's platform. This generic connection type takes
its platform
+from Default Provider (``extra["provider"]``); without one, a bare name
reaching this step
Review Comment:
This reads as if the spec-file model (source 4) gets Default Provider too,
but `create_agent(spec_file=...)` only passes a model when
`_get_conn_if_model_configured()` returns one, so the spec's `model` goes to
`Agent.from_file` unqualified. Scope this sentence to sources 1-3?
##########
providers/common/ai/docs/self_hosted_models.rst:
##########
@@ -293,9 +293,10 @@ 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:
+The ``extra`` JSON's ``model`` value keeps the ``provider:model`` format (or
+set the prefix once as Default 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 -- ``openai:<model>``
Review Comment:
The new lead-in sends readers to the gateway section, which says to pick
`openai-chat` for a server that only does Chat Completions, but this bullet
still says `openai:<model>` is the only option for vLLM. `openai-chat:` works
there too, so this one probably wants rewording alongside.
--
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]
