This is an automated email from the ASF dual-hosted git repository.
kaxilnaik 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 0b5b873a4af Use conn-fields for Model ID in PydanticAI connection
(#62818)
0b5b873a4af is described below
commit 0b5b873a4af1ce17d44b6bba63a603c15cda6e1b
Author: Kaxil Naik <[email protected]>
AuthorDate: Tue Mar 3 22:42:15 2026 +0000
Use conn-fields for Model ID in PydanticAI connection (#62818)
Replace the repurposed `schema` field with a proper `conn-fields`
entry in provider.yaml. This renders a dedicated "Model" input in
the connection form that stores to `extra_dejson["model"]` — the
same pattern used by 23+ providers (Google, Slack, Snowflake, etc.).
---
.../common/ai/docs/connections/pydantic_ai.rst | 27 ++++++++++++++++++----
providers/common/ai/provider.yaml | 9 +++++++-
.../providers/common/ai/get_provider_info.py | 12 ++++++----
.../providers/common/ai/hooks/pydantic_ai.py | 13 +++++++----
4 files changed, 47 insertions(+), 14 deletions(-)
diff --git a/providers/common/ai/docs/connections/pydantic_ai.rst
b/providers/common/ai/docs/connections/pydantic_ai.rst
index 4fdd2b16852..6d5363397a0 100644
--- a/providers/common/ai/docs/connections/pydantic_ai.rst
+++ b/providers/common/ai/docs/connections/pydantic_ai.rst
@@ -33,6 +33,17 @@ The ``PydanticAIHook`` uses ``pydanticai_default`` by
default.
Configuring the Connection
--------------------------
+Model
+ The model identifier in ``provider:model`` format. This field appears as a
+ dedicated input in the connection form (via ``conn-fields``) and stores its
+ value in ``extra["model"]``.
+
+ Examples: ``openai:gpt-5.3``, ``anthropic:claude-sonnet-4-20250514``,
+ ``bedrock:us.anthropic.claude-opus-4-6-v1:0``, ``google:gemini-2.0-flash``
+
+ The model can also be overridden at the hook/operator level via the
+ ``model_id`` parameter.
+
API Key (Password field)
The API key for your LLM provider. Required for API-key-based providers
(OpenAI, Anthropic, Groq, Mistral). Leave empty for providers using
@@ -48,15 +59,15 @@ Host (optional)
- Any OpenAI-compatible API: the base URL of that service
Extra (JSON, optional)
- A JSON object with additional configuration. The ``model`` key specifies
- the default model in ``provider:model`` format:
+ A JSON object with additional configuration. Programmatic users can set the
+ model directly in extra:
.. code-block:: json
{"model": "openai:gpt-5.3"}
- The model can also be overridden at the hook/operator level via the
- ``model_id`` parameter.
+ When using the UI, the "Model" field above writes to this same location
+ automatically.
Examples
--------
@@ -112,3 +123,11 @@ Leave password empty and configure
``GOOGLE_APPLICATION_CREDENTIALS`` in the env
"conn_type": "pydanticai",
"extra": "{\"model\": \"google:gemini-2.0-flash\"}"
}
+
+Model Resolution Order
+----------------------
+
+The hook reads the model from these sources in priority order:
+
+1. ``model_id`` parameter on the hook/operator
+2. ``model`` in the connection's extra JSON (set by the "Model" conn-field in
the UI)
diff --git a/providers/common/ai/provider.yaml
b/providers/common/ai/provider.yaml
index 116df5e456f..8f7743003b9 100644
--- a/providers/common/ai/provider.yaml
+++ b/providers/common/ai/provider.yaml
@@ -57,7 +57,14 @@ connection-types:
password: API Key
placeholders:
host: "https://api.openai.com/v1 (optional, for custom endpoints)"
- extra: '{"model": "openai:gpt-5"}'
+ conn-fields:
+ model:
+ label: Model
+ description: "Model in provider:name format (e.g.
anthropic:claude-sonnet-4-20250514, openai:gpt-5)"
+ schema:
+ type:
+ - string
+ - 'null'
operators:
- integration-name: Common AI
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 d69c662e0ad..5faf776f471 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
@@ -56,10 +56,14 @@ def get_provider_info():
"ui-field-behaviour": {
"hidden-fields": ["schema", "port", "login"],
"relabeling": {"password": "API Key"},
- "placeholders": {
- "host": "https://api.openai.com/v1 (optional, for
custom endpoints)",
- "extra": '{"model": "openai:gpt-5"}',
- },
+ "placeholders": {"host": "https://api.openai.com/v1
(optional, for custom endpoints)"},
+ },
+ "conn-fields": {
+ "model": {
+ "label": "Model",
+ "description": "Model in provider:name format (e.g.
anthropic:claude-sonnet-4-20250514, openai:gpt-5)",
+ "schema": {"type": ["string", "null"]},
+ }
},
}
],
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 23983d73dbd..b05f0421e9a 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
@@ -39,9 +39,9 @@ class PydanticAIHook(BaseHook):
Bedrock, Ollama, vLLM, etc.).
Connection fields:
+ - **Model** (conn-field): Model in ``provider:model`` format (e.g.
``"anthropic:claude-sonnet-4-20250514"``)
- **password**: API key (OpenAI, Anthropic, Groq, Mistral, etc.)
- **host**: Base URL (optional — for custom endpoints like Ollama,
vLLM, Azure)
- - **extra** JSON: ``{"model": "openai:gpt-5.3"}``
Cloud providers (Bedrock, Vertex) that use native auth chains should leave
password empty and configure environment-based auth (``AWS_PROFILE``,
@@ -76,7 +76,6 @@ class PydanticAIHook(BaseHook):
"relabeling": {"password": "API Key"},
"placeholders": {
"host": "https://api.openai.com/v1 (optional, for custom
endpoints)",
- "extra": '{"model": "openai:gpt-5.3"}',
},
}
@@ -84,8 +83,12 @@ class PydanticAIHook(BaseHook):
"""
Return a configured pydantic-ai Model.
- Reads API key from connection password, model from connection extra
- or ``model_id`` parameter, and base_url from connection host.
+ Reads API key from connection password, base_url from connection host,
+ and model from (in priority order):
+
+ 1. ``model_id`` parameter on the hook
+ 2. ``extra["model"]`` on the connection (set by the "Model" conn-field
in the UI)
+
The result is cached for the lifetime of this hook instance.
"""
if self._model is not None:
@@ -95,7 +98,7 @@ class PydanticAIHook(BaseHook):
model_name: str | KnownModelName = self.model_id or
conn.extra_dejson.get("model", "")
if not model_name:
raise ValueError(
- "No model specified. Set model_id on the hook or 'model' in
the connection's extra JSON."
+ "No model specified. Set model_id on the hook or the Model
field on the connection."
)
api_key = conn.password
base_url = conn.host or None
