kaxil commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2901281112
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -157,13 +187,229 @@ def test_connection(self) -> tuple[bool, str]:
"""
Test connection by resolving the model.
- Validates that the model string is valid, the provider package is
- installed, and the provider class can be instantiated. Does NOT make an
- LLM API call — that would be expensive, flaky, and fail for reasons
- unrelated to connectivity (quotas, billing, rate limits).
+ Validates that the model string is valid and the provider class can be
+ instantiated with the supplied credentials. Does NOT make an LLM API
+ call — that would be expensive and fail for reasons unrelated to
+ connectivity (quotas, billing, rate limits).
"""
try:
self.get_conn()
return True, "Model resolved successfully."
except Exception as e:
return False, str(e)
+
+
+class PydanticAIAzureHook(PydanticAIHook):
+ """
+ Hook for Azure OpenAI via pydantic-ai.
+
+ Connection fields:
+ - **password**: Azure API key
+ - **host**: Azure endpoint (e.g.
``https://<resource>.openai.azure.com``)
+ - **extra** JSON::
+
+ {"model": "azure:gpt-4o", "api_version": "2024-07-01-preview"}
+
+ :param llm_conn_id: Airflow connection ID.
+ :param model_id: Model identifier, e.g. ``"azure:gpt-4o"``.
+ """
+
+ conn_type = "pydanticaiazure"
+ default_conn_name = "pydanticai_azure_default"
+ hook_name = "Pydantic AI (Azure OpenAI)"
+
+ @staticmethod
+ def get_ui_field_behaviour() -> dict[str, Any]:
+ """Return custom field behaviour for the Airflow connection form."""
+ return {
+ "hidden_fields": ["schema", "port", "login"],
+ "relabeling": {"password": "API Key", "host": "Azure Endpoint"},
+ "placeholders": {
+ "host": "https://<resource>.openai.azure.com",
+ "extra": '{"model": "azure:gpt-4o", "api_version":
"2024-07-01-preview"}',
+ },
+ }
+
+ def _get_provider_kwargs(
+ self,
+ api_key: str | None,
+ base_url: str | None,
+ extra: dict[str, Any],
+ ) -> dict[str, Any]:
+ kwargs: dict[str, Any] = {}
+ if api_key:
+ kwargs["api_key"] = api_key
+ if base_url:
+ kwargs["azure_endpoint"] = base_url
+ if extra.get("api_version"):
+ kwargs["api_version"] = extra["api_version"]
+ return kwargs
+
+
+class PydanticAIBedrockHook(PydanticAIHook):
+ """
+ Hook for AWS Bedrock via pydantic-ai.
+
+ Credentials are resolved in order:
+
+ 1. IAM keys from ``extra`` (``aws_access_key_id`` +
``aws_secret_access_key``,
+ optionally ``aws_session_token``).
+ 2. Bearer token in ``extra`` (``api_key``, maps to env
``AWS_BEARER_TOKEN_BEDROCK``).
+ 3. Environment-variable / instance-role chain (``AWS_PROFILE``, IAM role,
…)
+ when no explicit keys are provided.
+
+ Connection fields:
+ - **extra** JSON::
+
+ {
+ "model": "bedrock:us.anthropic.claude-opus-4-5",
+ "region_name": "us-east-1",
+ "aws_access_key_id": "AKIA...",
+ "aws_secret_access_key": "...",
+ "aws_session_token": "...",
+ "profile_name": "my-aws-profile",
+ "api_key": "bearer-token",
+ "base_url": "https://custom-bedrock-endpoint";,
+ "aws_read_timeout": 60.0,
+ "aws_connect_timeout": 10.0
+ }
+
+ Leave ``aws_access_key_id`` / ``aws_secret_access_key`` and
``api_key``
+ empty to use the default AWS credential chain.
+
+ :param llm_conn_id: Airflow connection ID.
+ :param model_id: Model identifier, e.g.
``"bedrock:us.anthropic.claude-opus-4-5"``.
+ """
+
+ conn_type = "pydanticaibedrock"
+ default_conn_name = "pydanticai_bedrock_default"
+ hook_name = "Pydantic AI (AWS Bedrock)"
+
+ @staticmethod
+ def get_ui_field_behaviour() -> dict[str, Any]:
+ """Return custom field behaviour for the Airflow connection form."""
+ return {
+ "hidden_fields": ["schema", "port", "login", "host", "password"],
+ "relabeling": {},
+ "placeholders": {
+ "extra": (
+ '{"model": "bedrock:us.anthropic.claude-opus-4-5", '
+ '"region_name": "us-east-1"}'
+ " — leave aws_access_key_id empty for IAM role / env-var
auth"
+ ),
+ },
+ }
Review Comment:
The method signature receives `api_key` (from `conn.password`) and
`base_url` (from `conn.host`) but ignores them — instead reading `api_key` and
`base_url` from `extra`. This is correct since the Bedrock connection hides
`host` and `password` in the UI, but having the same field name in two sources
is confusing.
Consider renaming the extra keys to avoid ambiguity, or add a brief note in
the docstring that `api_key`/`base_url` from the method params are
intentionally unused for this subclass.
##########
providers/common/ai/src/airflow/providers/common/ai/operators/llm.py:
##########
@@ -101,8 +101,15 @@ def __init__(
@cached_property
def llm_hook(self) -> PydanticAIHook:
- """Return PydanticAIHook for the configured LLM connection."""
- return PydanticAIHook(llm_conn_id=self.llm_conn_id,
model_id=self.model_id)
+ """
+ Return the correct PydanticAIHook subclass for the configured
connection.
+
+ Delegates to :meth:`~PydanticAIHook.get_hook` which looks up
+ the connection's ``conn_type`` and instantiates the matching subclass
+ (e.g.
:class:`~airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIAzureHook`
+ for ``pydanticaiazure`` connections).
+ """
+ return PydanticAIHook.get_hook(self.llm_conn_id)
def execute(self, context: Context) -> Any:
agent: Agent[None, Any] = self.llm_hook.create_agent(
Review Comment:
`model_id` is silently dropped here. The old code was:
```python
return PydanticAIHook(llm_conn_id=self.llm_conn_id, model_id=self.model_id)
```
`get_hook()` (inherited from `BaseHook`) doesn't accept `model_id`. So if a
user sets `model_id="openai:gpt-5"` on `LLMOperator`, the hook never sees it —
it falls through to `extra.get("model", "")` and either uses the connection's
model or raises "No model specified."
The test change confirms this:
```python
# Before
mock_hook_cls.for_connection.assert_called_once_with("my_llm",
model_id="openai:gpt-5")
# After
mock_hook_cls.get_hook.assert_called_once_with("my_llm")
```
Fix: set `model_id` on the hook after `get_hook()` returns:
```python
hook = PydanticAIHook.get_hook(self.llm_conn_id)
hook.model_id = self.model_id
return hook
```
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -75,62 +78,89 @@ def get_ui_field_behaviour() -> dict[str, Any]:
"hidden_fields": ["schema", "port", "login"],
"relabeling": {"password": "API Key"},
"placeholders": {
- "host": "https://api.openai.com/v1 (optional, for custom
endpoints)",
+ "host": "https://api.openai.com/v1 (optional, for custom
endpoints / Ollama)",
+ "extra": '{"model": "openai:gpt-5.3"}',
},
}
+ # ------------------------------------------------------------------
+ # Core connection / agent API
+ # ------------------------------------------------------------------
+
+ def _get_provider_kwargs(
+ self,
+ api_key: str | None,
+ base_url: str | None,
+ extra: dict[str, Any],
+ ) -> dict[str, Any]:
+ """
+ Return the kwargs to pass to the provider constructor.
+
+ Subclasses override this method to map their connection fields to the
+ parameters expected by their specific provider class. The base
+ implementation handles the common ``api_key`` / ``base_url`` pattern
+ used by OpenAI, Anthropic, Groq, Mistral, Ollama, and most other
+ providers.
+
+ :param api_key: Value of ``conn.password``.
+ :param base_url: Value of ``conn.host``.
+ :param extra: Deserialized ``conn.extra`` JSON.
+ :return: Kwargs forwarded to ``provider_cls(**kwargs)``. Empty dict
+ signals that no explicit credentials are available and the hook
+ should fall back to environment-variable–based auth.
+ """
+ kwargs: dict[str, Any] = {}
+ if api_key:
+ kwargs["api_key"] = api_key
+ if base_url:
+ kwargs["base_url"] = base_url
+ return kwargs
+
def get_conn(self) -> Model:
"""
- Return a configured pydantic-ai Model.
+ Return a configured pydantic-ai ``Model``.
- Reads API key from connection password, base_url from connection host,
- and model from (in priority order):
+ Resolution order:
- 1. ``model_id`` parameter on the hook
- 2. ``extra["model"]`` on the connection (set by the "Model" conn-field
in the UI)
+ 1. **Explicit credentials** — when :meth:`_get_provider_kwargs` returns
+ a non-empty dict the provider class is instantiated with those
kwargs
+ and wrapped in a ``provider_factory``.
+ 2. **Default resolution** — delegates to pydantic-ai ``infer_model``
+ which reads standard env vars (``OPENAI_API_KEY``, ``AWS_PROFILE``,
…).
- The result is cached for the lifetime of this hook instance.
+ The resolved model is cached for the lifetime of this hook instance.
"""
if self._model is not None:
return self._model
conn = self.get_connection(self.llm_conn_id)
- model_name: str | KnownModelName = self.model_id or
conn.extra_dejson.get("model", "")
+
+ extra: dict[str, Any] = conn.extra_dejson
+ model_name: str | KnownModelName = self.model_id or extra.get("model",
"")
if not model_name:
raise ValueError(
"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
- if not api_key and not base_url:
- # No credentials to inject — use default provider resolution
- # (picks up env vars like OPENAI_API_KEY, AWS_PROFILE, etc.)
- self._model = infer_model(model_name)
+ api_key: str | None = conn.password or None
+ base_url: str | None = conn.host or None
+
+ provider_kwargs = self._get_provider_kwargs(api_key, base_url, extra)
+ if provider_kwargs:
+ _kwargs = provider_kwargs # capture for closure
+ self.log.info(
+ "Using explicit credentials for provider with model '%s': %s",
+ model_name,
+ list(provider_kwargs),
+ )
+
+ def _provider_factory(pname: str) -> Any:
+ return infer_provider_class(pname)(**_kwargs)
+
+ self._model = infer_model(model_name,
provider_factory=_provider_factory)
return self._model
Review Comment:
The `try/except TypeError` fallback was removed here. Previously:
```python
try:
return provider_cls(**kwargs)
except TypeError:
return infer_provider(provider_name)
```
With the subclass architecture this matters less, but if someone uses a
generic `pydanticai` connection with a provider that rejects
`api_key`/`base_url`, it crashes instead of falling back gracefully. The
`infer_provider` import was also removed.
Suggestion: keep the fallback as a safety net — it costs nothing and
prevents surprises.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -157,13 +187,229 @@ def test_connection(self) -> tuple[bool, str]:
"""
Test connection by resolving the model.
- Validates that the model string is valid, the provider package is
- installed, and the provider class can be instantiated. Does NOT make an
- LLM API call — that would be expensive, flaky, and fail for reasons
- unrelated to connectivity (quotas, billing, rate limits).
+ Validates that the model string is valid and the provider class can be
+ instantiated with the supplied credentials. Does NOT make an LLM API
+ call — that would be expensive and fail for reasons unrelated to
+ connectivity (quotas, billing, rate limits).
"""
try:
self.get_conn()
return True, "Model resolved successfully."
except Exception as e:
return False, str(e)
+
+
+class PydanticAIAzureHook(PydanticAIHook):
+ """
+ Hook for Azure OpenAI via pydantic-ai.
+
+ Connection fields:
+ - **password**: Azure API key
+ - **host**: Azure endpoint (e.g.
``https://<resource>.openai.azure.com``)
+ - **extra** JSON::
+
+ {"model": "azure:gpt-4o", "api_version": "2024-07-01-preview"}
+
+ :param llm_conn_id: Airflow connection ID.
+ :param model_id: Model identifier, e.g. ``"azure:gpt-4o"``.
+ """
Review Comment:
`pydanticaiazure` is hard to read compared to `pydanticai_azure`. The
`default_conn_name` uses underscores (`pydanticai_azure_default`) which creates
an odd mismatch:
- `conn_type = "pydanticaiazure"`
- `default_conn_name = "pydanticai_azure_default"`
While `gcpbigquery` exists as a precedent in Airflow,
`google_cloud_platform` uses underscores too — there is no strict convention.
`pydanticai_azure` would be more readable and consistent with the
default_conn_name. Same applies to bedrock and vertex.
--
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]
