You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "cetingokhan (via GitHub)" <gi...@apache.org> on 2026/03/03 20:23:57 UTC
[PR] added azureopenai support & extandable architecture for pydanticai [airflow]
cetingokhan opened a new pull request, #62816:
URL: https://github.com/apache/airflow/pull/62816
This pull request refactors the model resolution logic in the `PydanticAIHook` to use a modular builder pattern, improving extensibility and clarity for handling different AI providers, especially Azure OpenAI. It introduces dedicated builder classes for Azure OpenAI, custom endpoints, and default resolution, and updates documentation and tests to match the new structure.
**Builder pattern for model resolution:**
* Added new builder classes: `AzureOpenAIBuilder`, `CustomEndpointBuilder`, `DefaultBuilder`, and a `ProviderBuilder` protocol in the `builders` package to modularize how models are constructed from Airflow connection details. [[1]](diffhunk://#diff-d1854c84b6c6196fef61024285e7bb3f1424f1a434a3f5507930158567cba3e5R1-R31) [[2]](diffhunk://#diff-0a37576ea742b0fc055960de587a234457f0c7385de3c4ac0b3ff2180b058f50R1-R99) [[3]](diffhunk://#diff-1592522b00129d036aebc50dcfbe3654e3ab2f61288cfde48944438911cc1630R1-R58) [[4]](diffhunk://#diff-6c40bb2d6dde6e5186a35f360a251ae28d3ad997fc1859b7f2da85c9ebc98aebR1-R92)
**Refactoring in `PydanticAIHook`:**
* Replaced direct calls to `infer_model` and provider factory logic with a prioritized builder selection process (`AzureOpenAIBuilder` → `CustomEndpointBuilder` → `DefaultBuilder`), improving support for Azure OpenAI and custom endpoints.
**Documentation and UI improvements:**
* Updated docstrings and UI field behavior in `PydanticAIHook` to clarify connection fields for Azure OpenAI and provide examples for required extras like `api_version` and `azure_deployment`. [[1]](diffhunk://#diff-dd850d1a210dfcbdb51ec2fcb2ef47a967449aa4b683d6e7256cbc7a2b68055cL39-R58) [[2]](diffhunk://#diff-dd850d1a210dfcbdb51ec2fcb2ef47a967449aa4b683d6e7256cbc7a2b68055cL78-R148)
**Testing updates:**
* Modified unit tests to patch `infer_model` and `infer_provider_class` from the new builder modules instead of the hook, ensuring tests match the refactored code structure. [[1]](diffhunk://#diff-d0cd7707473a4d994023a0ccdc500b47fec345067d38e9351604282d54d571e8L41-R43) [[2]](diffhunk://#diff-d0cd7707473a4d994023a0ccdc500b47fec345067d38e9351604282d54d571e8L75-R77) [[3]](diffhunk://#diff-d0cd7707473a4d994023a0ccdc500b47fec345067d38e9351604282d54d571e8L95-R97) [[4]](diffhunk://#diff-d0cd7707473a4d994023a0ccdc500b47fec345067d38e9351604282d54d571e8L125-R126) [[5]](diffhunk://#diff-d0cd7707473a4d994023a0ccdc500b47fec345067d38e9351604282d54d571e8L142-R144)
**Connection validation enhancements:**
* Improved connection testing to validate Azure-specific requirements (presence of `api_version` and `host`) and clarified error handling for missing model configuration. <!-- SPDX-License-Identifier: Apache-2.0
https://www.apache.org/licenses/LICENSE-2.0 -->
<!--
Thank you for contributing!
Please provide above a brief description of the changes made in this pull request.
Write a good git commit message following this guide: http://chris.beams.io/posts/git-commit/
Please make sure that your code changes are covered with tests.
And in case of new features or big changes remember to adjust the documentation.
Feel free to ping (in general) for the review if you do not see reaction for a few days
(72 Hours is the minimum reaction time you can expect from volunteers) - we sometimes miss notifications.
In case of an existing issue, reference it using one of the following:
* closes: #ISSUE
* related: #ISSUE
-->
---
##### Was generative AI tooling used to co-author this PR?
<!--
If generative AI tooling has been used in the process of authoring this PR, please
change below checkbox to `[X]` followed by the name of the tool, uncomment the "Generated-by".
-->
- [x] Yes (please specify the tool below)
Cloude Sonnet 4.6 & Gemini 3.1 Pro
Filled some of methods scope and tests created via copilot
<!--
Generated-by: [Tool Name] following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
-->
---
* Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
* For fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
* When adding dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
* For significant user-facing changes create newsfragment: `{pr_number}.significant.rst` or `{issue_number}.significant.rst`, in [airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments).
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on PR #62816:
URL: https://github.com/apache/airflow/pull/62816#issuecomment-4034906559
Hey @cetingokhan — don't be embarrassed, this is exactly how the process is supposed to work\! The PR has genuinely improved a lot through each round — going from the builder pattern to the subclass hooks was a big architectural shift and you landed it well. That's the whole point of code review.
For #62963 — no rush at all, take your time with it. Happy to review when you're ready 👍
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.
cetingokhan commented on PR #62816:
URL: https://github.com/apache/airflow/pull/62816#issuecomment-4006915805
Hi @kaxil, @gopidesupavan
I've refactored the PR to fully implement your suggested approach using dedicated hook subclasses and conn_types. I also removed the builder and dynamic parsing logic.
Could you take a look when you have a chance? Please let me know if there is anything missing or if you'd like me to make any further tweaks/additions.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2908425953
##########
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:
Let's do `pydanticai-bedrock`, `pydanticai-azure` etc, please
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil merged PR #62816:
URL: https://github.com/apache/airflow/pull/62816
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.
cetingokhan commented on PR #62816:
URL: https://github.com/apache/airflow/pull/62816#issuecomment-4000681149
Thanks for the guidance! I completely understand the "thin hook" approach and using a simpler method for Azure and LiteLLM.
I just have one concern and wanted to get your thoughts on it. When I look at other providers, they seem to require entirely different parameters from the connection. For example:
BedrockProvider uses aws_access_key_id, region_name, etc.
GoogleVertexProvider uses service_account_file, project_id, etc.
I'm a bit worried that if we only map api_key and base_url right now, we might eventually have to write endless if/elif blocks in the hook for every single provider's unique parameters in the future.
But of course, if you wish, we can start as you suggested and then revise it as needed.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2892283022
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -75,59 +74,104 @@ 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, model from connection extra
- or ``model_id`` parameter, and base_url from connection host.
- The result is cached for the lifetime of this hook instance.
+ Resolution order:
+
+ 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 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 'model' in the connection's extra JSON."
)
- 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
+
+ # Auto-dispatch: if using base hook with a subclass connection type, borrow
+ # that subclass's _get_provider_kwargs so the correct field mapping is used.
+ if type(self) is PydanticAIHook:
+ hook_cls = _CONN_TYPE_TO_HOOK.get(conn.conn_type or "", PydanticAIHook)
+ provider_kwargs = hook_cls._get_provider_kwargs(self, api_key, base_url, extra)
+ else:
+ 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:
+ try:
+ return infer_provider_class(pname)(**_kwargs)
+ except TypeError:
+ self.log.warning(
+ "Provider '%s' does not accept the supplied kwargs %s; "
Review Comment:
`TypeError` here catches both "provider doesn't accept these kwargs" and genuine bugs (wrong types, missing required args). The warning only logs kwarg *names*, not the actual error. This makes debugging auth failures harder — a user who passes the wrong kwarg name gets silently downgraded to env-var auth with a vague warning.
Suggestion: log the exception too:
```python
except TypeError as exc:
self.log.warning(
"Provider '%s' does not accept kwargs %s (%s); "
"falling back to env-var auth.",
pname, list(_kwargs), exc,
)
```
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -154,13 +198,213 @@ 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)
+
+ @classmethod
+ def for_connection(cls, conn_id: str, model_id: str | None = None) -> PydanticAIHook:
+ """
+ Return the correct :class:`PydanticAIHook` subclass for *conn_id*.
+
+ Looks up the connection's ``conn_type`` in the registered hook map and
+ instantiates the matching subclass. Falls back to
+ :class:`PydanticAIHook` for unknown types.
+
+ :param conn_id: Airflow connection ID.
+ :param model_id: Optional model override forwarded to the hook.
+ """
+ conn = cls.get_connection(conn_id)
+ hook_cls = _CONN_TYPE_TO_HOOK.get(conn.conn_type or "", cls)
+ return hook_cls(llm_conn_id=conn_id, model_id=model_id)
+
+
+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_name_attr = "llm_conn_id"
+ default_conn_name = "pydantic_ai_azure_default"
+ conn_type = "pydanticai_azure"
+ 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"]
+
+ self.log.info(
+ "Using explicit credentials for Azure provider with model '%s': %s",
+ extra.get("model", ""),
Review Comment:
This `self.log.info(...)` duplicates the log in `get_conn()` which also logs when `provider_kwargs` is non-empty. Azure connections will produce two "Using explicit credentials..." log lines; Bedrock and Vertex produce one. Remove this — let the base class handle logging uniformly.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -154,13 +198,213 @@ 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)
+
+ @classmethod
+ def for_connection(cls, conn_id: str, model_id: str | None = None) -> PydanticAIHook:
+ """
+ Return the correct :class:`PydanticAIHook` subclass for *conn_id*.
+
+ Looks up the connection's ``conn_type`` in the registered hook map and
+ instantiates the matching subclass. Falls back to
+ :class:`PydanticAIHook` for unknown types.
+
+ :param conn_id: Airflow connection ID.
+ :param model_id: Optional model override forwarded to the hook.
+ """
+ conn = cls.get_connection(conn_id)
+ hook_cls = _CONN_TYPE_TO_HOOK.get(conn.conn_type or "", cls)
+ return hook_cls(llm_conn_id=conn_id, model_id=model_id)
+
+
+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_name_attr = "llm_conn_id"
+ default_conn_name = "pydantic_ai_azure_default"
+ conn_type = "pydanticai_azure"
+ hook_name = "Pydantic AI (Azure OpenAI)"
Review Comment:
Python evaluates default argument values at function/method definition time. `PydanticAIHook.__init__` binds `llm_conn_id`'s default to `"pydanticai_default"` when the class is created. Subclasses inherit this `__init__` with the already-resolved default, so:
```python
hook = PydanticAIAzureHook() # llm_conn_id="pydanticai_default", not "pydantic_ai_azure_default"
```
The `default_conn_name` attribute on the subclass is never consulted during `__init__`. Same issue for Bedrock and Vertex subclasses.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -75,59 +74,104 @@ 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, model from connection extra
- or ``model_id`` parameter, and base_url from connection host.
- The result is cached for the lifetime of this hook instance.
+ Resolution order:
+
+ 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 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 'model' in the connection's extra JSON."
)
- 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
+
+ # Auto-dispatch: if using base hook with a subclass connection type, borrow
+ # that subclass's _get_provider_kwargs so the correct field mapping is used.
+ if type(self) is PydanticAIHook:
+ hook_cls = _CONN_TYPE_TO_HOOK.get(conn.conn_type or "", PydanticAIHook)
+ provider_kwargs = hook_cls._get_provider_kwargs(self, api_key, base_url, extra)
+ else:
+ provider_kwargs = self._get_provider_kwargs(api_key, base_url, extra)
Review Comment:
This calls a subclass's unbound method with a `PydanticAIHook` instance as `self`:
```python
hook_cls._get_provider_kwargs(self, api_key, base_url, extra)
```
It works today because `_get_provider_kwargs` only touches `self.log` (inherited from BaseHook). But if any subclass later accesses subclass-specific attributes, this breaks silently.
A cleaner alternative: have the operator use `for_connection()` instead of constructing `PydanticAIHook` directly. That returns the right subclass, and no auto-dispatch is needed:
```python
@cached_property
def llm_hook(self) -> PydanticAIHook:
return PydanticAIHook.for_connection(self.llm_conn_id, model_id=self.model_id)
```
This eliminates the `type(self) is PydanticAIHook` check, the `_CONN_TYPE_TO_HOOK` registry, and the unbound method call.
##########
providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py:
##########
@@ -60,7 +60,41 @@ def get_provider_info():
"extra": '{"model": "openai:gpt-5"}',
},
},
- }
+ },
Review Comment:
This file is auto-generated from `provider.yaml` — manual edits here will be overwritten by `prek run update-providers-build-files`. The new connection types need to be added to `provider.yaml` instead.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -154,13 +198,213 @@ 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)
+
+ @classmethod
+ def for_connection(cls, conn_id: str, model_id: str | None = None) -> PydanticAIHook:
+ """
+ Return the correct :class:`PydanticAIHook` subclass for *conn_id*.
+
+ Looks up the connection's ``conn_type`` in the registered hook map and
+ instantiates the matching subclass. Falls back to
+ :class:`PydanticAIHook` for unknown types.
+
+ :param conn_id: Airflow connection ID.
+ :param model_id: Optional model override forwarded to the hook.
+ """
+ conn = cls.get_connection(conn_id)
+ hook_cls = _CONN_TYPE_TO_HOOK.get(conn.conn_type or "", cls)
+ return hook_cls(llm_conn_id=conn_id, model_id=model_id)
+
+
+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_name_attr = "llm_conn_id"
+ default_conn_name = "pydantic_ai_azure_default"
+ conn_type = "pydanticai_azure"
+ 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"]
+
+ self.log.info(
+ "Using explicit credentials for Azure provider with model '%s': %s",
+ extra.get("model", ""),
+ list(kwargs),
+ )
+ return kwargs
+
+
+class PydanticAIBedrockHook(PydanticAIHook):
+ """
+ Hook for AWS Bedrock via pydantic-ai.
+
+ Credentials are resolved in order:
+
+ 1. Explicit keys in ``extra`` (``aws_access_key_id``,
+ ``aws_secret_access_key``, ``aws_session_token``).
+ 2. 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": "..."
+ }
+
+ Leave ``aws_access_key_id`` / ``aws_secret_access_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_name_attr = "llm_conn_id"
+ default_conn_name = "pydantic_ai_bedrock_default"
+ conn_type = "pydanticai_bedrock"
+ 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"
+ ),
+ },
+ }
+
+ def _get_provider_kwargs(
+ self,
+ api_key: str | None,
+ base_url: str | None,
+ extra: dict[str, Any],
+ ) -> dict[str, Any]:
+ _bedrock_keys = (
+ "aws_access_key_id",
+ "aws_secret_access_key",
+ "aws_session_token",
+ "region_name",
+ "profile_name",
+ )
+ return {k: extra[k] for k in _bedrock_keys if extra.get(k) is not None}
+
+
+class PydanticAIVertexHook(PydanticAIHook):
+ """
+ Hook for Google Vertex AI via pydantic-ai.
+
+ Credentials are resolved in order:
+
+ 1. ``service_account_file`` / ``service_account_info`` in ``extra``.
+ 2. Application Default Credentials (``GOOGLE_APPLICATION_CREDENTIALS``,
+ ``gcloud auth application-default login``, …) when no explicit keys
+ are provided.
+
+ Connection fields:
+ - **extra** JSON::
+
+ {
+ "model": "google-vertex:gemini-2.0-flash",
+ "project_id": "my-gcp-project",
+ "location": "us-central1",
+ "service_account_file": "/path/to/sa.json",
+ }
+
+ :param llm_conn_id: Airflow connection ID.
+ :param model_id: Model identifier, e.g. ``"google-vertex:gemini-2.0-flash"``.
+ """
+
+ conn_name_attr = "llm_conn_id"
+ default_conn_name = "pydantic_ai_vertex_default"
+ conn_type = "pydanticai_vertex"
+ hook_name = "Pydantic AI (Google Vertex AI)"
+
+ @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": "google-vertex:gemini-2.0-flash", '
+ '"project_id": "my-project", "location": "us-central1"}'
+ " — leave service_account_file empty for ADC auth"
+ ),
+ },
+ }
+
+ def _get_provider_kwargs(
+ self,
+ api_key: str | None,
+ base_url: str | None,
+ extra: dict[str, Any],
+ ) -> dict[str, Any]:
+ _vertex_keys = ("project_id", "location", "service_account_file", "service_account_info")
+ return {k: extra[k] for k in _vertex_keys if extra.get(k) is not None}
+
Review Comment:
These kwarg names don't match what pydantic-ai's `GoogleProvider` accepts (see [docs](https://ai.pydantic.dev/models/google/)):
| Connection extra | pydantic-ai param | Match? |
|---|---|---|
| `project_id` | `project` | **No** |
| `location` | `location` | Yes |
| `service_account_file` | `credentials` (object) | **No** |
| `service_account_info` | `credentials` (object) | **No** |
`project_id` needs to be mapped to `project`. `service_account_file`/`service_account_info` need to be loaded into a `google.auth.credentials.Credentials` object before passing — the provider doesn't accept file paths directly.
Because of the `TypeError` catch in `_provider_factory`, these mismatches fail silently — the credentials get ignored and it falls back to ADC. Users who explicitly configure credentials in the connection would have no idea they're not being used.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "gopidesupavan (via GitHub)" <gi...@apache.org>.
gopidesupavan commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2902459420
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -75,62 +78,92 @@ 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"}',
Review Comment:
same here.. model_id
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.
cetingokhan commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2914409766
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -157,13 +190,238 @@ 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 = "pydanticai_azure"
+ 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 = "pydanticai_bedrock"
+ 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"
+ ),
+ },
+ }
+
+ def _get_provider_kwargs(
+ self,
+ api_key: str | None,
+ base_url: str | None,
+ extra: dict[str, Any],
+ ) -> dict[str, Any]:
+ """
+ Return kwargs for ``BedrockProvider``.
+
+ .. note::
+ The ``api_key`` and ``base_url`` parameters (sourced from
+ ``conn.password`` and ``conn.host``) are intentionally ignored here.
+ Bedrock connections hide those fields in the UI; all config is
+ stored in ``extra`` instead. The ``api_key`` and ``base_url``
+ keys below refer to *extra* fields, not the method parameters.
+ """
+ _str_keys = (
+ "aws_access_key_id",
+ "aws_secret_access_key",
+ "aws_session_token",
+ "region_name",
+ "profile_name",
+ # Bearer-token auth (alternative to IAM key/secret).
+ # Maps to AWS_BEARER_TOKEN_BEDROCK env var.
+ "api_key",
+ # Custom Bedrock runtime endpoint.
+ "base_url",
+ )
+ kwargs: dict[str, Any] = {k: extra[k] for k in _str_keys if extra.get(k) is not None}
Review Comment:
removed "is not None" conditions
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on PR #62816:
URL: https://github.com/apache/airflow/pull/62816#issuecomment-4001132383
Good question — and I think we should go further than just extras. The right approach for Airflow is **separate connection types** for the cloud providers that need distinct credentials.
Airflow's connection form doesn't support conditional fields (no way to show/hide fields based on a dropdown). But it does have an established pattern for this: the Google provider ships 5 connection types (`google_cloud_platform`, `gcpcloudsql`, `gcpcloudsqldb`, `gcpbigquery`, `dataprep`). Amazon ships 5 (`aws`, `chime`, `emr`, `redshift`, `athena`). Each has its own hook subclass with tailored form fields.
We should follow the same pattern:
| Connection Type | Covers | Form Fields |
|---|---|---|
| `pydanticai` | OpenAI, Anthropic, Groq, Mistral, Deepseek, Ollama, etc. | API Key, Base URL (optional), Model |
| `pydanticai_azure` | Azure OpenAI | API Key, Azure Endpoint, API Version, Model |
| `pydanticai_bedrock` | AWS Bedrock | Region, Profile Name, Model |
| `pydanticai_vertex` | Google Vertex | Project, Location, Model |
Each connection type is a hook subclass that overrides `conn_type`, `get_ui_field_behaviour()`, and the credential-to-provider mapping. Users pick the right connection type, see only the relevant fields, and everything just works.
The generic `pydanticai` type covers most providers — anything that accepts `api_key` + `base_url`. Only the 3 cloud providers with non-standard auth get their own types.
Implementation-wise, each subclass is small:
```python
class PydanticAIAzureHook(PydanticAIHook):
conn_type = "pydanticai_azure"
hook_name = "Pydantic AI (Azure OpenAI)"
@staticmethod
def get_ui_field_behaviour() -> dict[str, Any]:
return {
"hidden_fields": ["schema", "port", "login"],
"relabeling": {"password": "API Key", "host": "Azure Endpoint"},
"placeholders": {"host": "https://myapp.openai.azure.com/"},
}
```
The provider-specific credential mapping lives in an overridable method on the base class. No if/elif chains, no `inspect.signature`, no generic extras dumping.
This is a clean follow-up PR after the base hook lands — it doesn't need to block Azure support in this PR. For this PR, the explicit Azure branch in `_provider_factory` is fine as a first step. The separate connection types are the proper long-term solution.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2908398829
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -157,13 +190,238 @@ 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 = "pydanticai_azure"
+ 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 = "pydanticai_bedrock"
+ 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"
+ ),
+ },
+ }
+
+ def _get_provider_kwargs(
+ self,
+ api_key: str | None,
+ base_url: str | None,
+ extra: dict[str, Any],
+ ) -> dict[str, Any]:
+ """
+ Return kwargs for ``BedrockProvider``.
+
+ .. note::
+ The ``api_key`` and ``base_url`` parameters (sourced from
+ ``conn.password`` and ``conn.host``) are intentionally ignored here.
+ Bedrock connections hide those fields in the UI; all config is
+ stored in ``extra`` instead. The ``api_key`` and ``base_url``
+ keys below refer to *extra* fields, not the method parameters.
+ """
+ _str_keys = (
+ "aws_access_key_id",
+ "aws_secret_access_key",
+ "aws_session_token",
+ "region_name",
+ "profile_name",
+ # Bearer-token auth (alternative to IAM key/secret).
+ # Maps to AWS_BEARER_TOKEN_BEDROCK env var.
+ "api_key",
+ # Custom Bedrock runtime endpoint.
+ "base_url",
+ )
+ kwargs: dict[str, Any] = {k: extra[k] for k in _str_keys if extra.get(k) is not None}
Review Comment:
The base class uses truthiness checks to filter empty values (`if api_key:`, `if base_url:`), but here `is not None` lets empty strings through. If a user clears a conn-field in the UI, `aws_access_key_id=""` would be passed to `BedrockProvider`, causing confusing auth failures instead of falling through to environment-based auth.
Same pattern appears in `PydanticAIVertexHook._get_provider_kwargs` (line ~409).
Suggestion: use truthiness to match the base class:
```python
kwargs = {k: extra[k] for k in _str_keys if extra.get(k)}
```
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -75,62 +78,92 @@ 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:
+ try:
+ return infer_provider_class(pname)(**_kwargs)
+ except TypeError:
Review Comment:
This silently discards the explicit credentials from `_get_provider_kwargs` and falls back to env-var auth. If a subclass carefully builds provider kwargs but the provider class's constructor signature changes, the user would silently get env-var auth instead of the credentials they configured — hard to debug.
Consider logging a warning here:
```python
except TypeError:
self.log.warning(
"Provider '%s' rejected kwargs %s; falling back to env-var auth",
pname, list(_kwargs),
)
return infer_provider(pname)
```
(The closure can capture `self` — just define `_provider_factory` as a regular closure instead of standalone, or pass a logger reference.)
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -157,13 +190,238 @@ 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 = "pydanticai_azure"
+ 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 = "pydanticai_bedrock"
+ 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"
+ ),
+ },
+ }
+
+ def _get_provider_kwargs(
+ self,
+ api_key: str | None,
+ base_url: str | None,
+ extra: dict[str, Any],
+ ) -> dict[str, Any]:
+ """
+ Return kwargs for ``BedrockProvider``.
+
+ .. note::
+ The ``api_key`` and ``base_url`` parameters (sourced from
+ ``conn.password`` and ``conn.host``) are intentionally ignored here.
+ Bedrock connections hide those fields in the UI; all config is
+ stored in ``extra`` instead. The ``api_key`` and ``base_url``
+ keys below refer to *extra* fields, not the method parameters.
+ """
+ _str_keys = (
+ "aws_access_key_id",
+ "aws_secret_access_key",
+ "aws_session_token",
+ "region_name",
+ "profile_name",
+ # Bearer-token auth (alternative to IAM key/secret).
+ # Maps to AWS_BEARER_TOKEN_BEDROCK env var.
+ "api_key",
+ # Custom Bedrock runtime endpoint.
+ "base_url",
+ )
+ kwargs: dict[str, Any] = {k: extra[k] for k in _str_keys if extra.get(k) is not None}
+ # BedrockProvider expects float for timeout values; JSON integers must be coerced.
+ for _timeout_key in ("aws_read_timeout", "aws_connect_timeout"):
+ if extra.get(_timeout_key) is not None:
+ kwargs[_timeout_key] = float(extra[_timeout_key])
+ return kwargs
+
+
+class PydanticAIVertexHook(PydanticAIHook):
+ """
+ Hook for Google Vertex AI (or Generative Language API) via pydantic-ai.
+
+ Credentials are resolved in order:
+
+ 1. ``service_account_info`` (JSON object) in ``extra``
+ — loaded into a ``google.auth.credentials.Credentials``
+ object and passed as ``credentials`` to ``GoogleProvider``.
+ 2. ``api_key`` in ``extra`` — for Generative Language API (non-Vertex) or
+ Vertex API-key auth.
+ 3. Application Default Credentials (``GOOGLE_APPLICATION_CREDENTIALS``,
+ ``gcloud auth application-default login``, Workload Identity, …) when
+ no explicit credentials are provided.
+
+ Connection fields:
+ - **extra** JSON::
+
+ {
+ "model": "google-vertex:gemini-2.0-flash",
+ "project": "my-gcp-project",
+ "location": "us-central1",
+ "service_account_info": {...},
+ "vertexai": true,
+ }
+
+ Use ``"service_account_info"`` to embed the service-account JSON directly
+ (as an object, not a string path).
+
+ Set ``"vertexai": true`` to force Vertex AI mode when only ``api_key`` is
+ provided. Omit ``vertexai`` for the Generative Language API (GLA).
+
+ :param llm_conn_id: Airflow connection ID.
+ :param model_id: Model identifier, e.g. ``"google-vertex:gemini-2.0-flash"``.
+ """
+
+ conn_type = "pydanticai_vertex"
+ default_conn_name = "pydanticai_vertex_default"
+ hook_name = "Pydantic AI (Google Vertex AI)"
+
+ @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": "google-vertex:gemini-2.0-flash", '
+ '"project": "my-project", "location": "us-central1", "vertexai": true}'
+ " — add service_account_info (object) for SA auth;"
+ " omit both to use Application Default Credentials"
+ ),
+ },
+ }
+
+ def _get_provider_kwargs(
+ self,
+ api_key: str | None,
+ base_url: str | None,
+ extra: dict[str, Any],
+ ) -> dict[str, Any]:
+ sa_info = extra.get("service_account_info")
+ kwargs: dict[str, Any] = {}
+
+ # Direct GoogleProvider scalar kwargs.
+ for _key in ("api_key", "project", "location", "base_url"):
+ if extra.get(_key) is not None:
Review Comment:
Same empty-string concern as the Bedrock hook — `extra.get(_key) is not None` lets `""` through. Using `if extra.get(_key):` would be consistent with the base class's truthiness pattern.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.
cetingokhan commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2900089165
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -19,42 +19,41 @@
from typing import TYPE_CHECKING, Any, TypeVar, overload
from pydantic_ai import Agent
-from pydantic_ai.models import Model, infer_model
-from pydantic_ai.providers import Provider, infer_provider, infer_provider_class
+from pydantic_ai.models import infer_model
+from pydantic_ai.providers import infer_provider, infer_provider_class
from airflow.providers.common.compat.sdk import BaseHook
OutputT = TypeVar("OutputT")
if TYPE_CHECKING:
- from pydantic_ai.models import KnownModelName
+ from pydantic_ai.models import KnownModelName, Model
class PydanticAIHook(BaseHook):
"""
Hook for LLM access via pydantic-ai.
- Manages connection credentials and model creation. Uses pydantic-ai's
- model inference to support any provider (OpenAI, Anthropic, Google,
- Bedrock, Ollama, vLLM, etc.).
+ Covers providers that use a standard ``api_key`` + optional ``base_url``
+ (OpenAI, Anthropic, Groq, Mistral, DeepSeek, Ollama, vLLM, …).
- 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)
+ For cloud providers with non-standard auth use the dedicated subclasses:
+ :class:`PydanticAIAzureHook`, :class:`PydanticAIBedrockHook`,
+ :class:`PydanticAIVertexHook`.
- Cloud providers (Bedrock, Vertex) that use native auth chains should leave
- password empty and configure environment-based auth (``AWS_PROFILE``,
- ``GOOGLE_APPLICATION_CREDENTIALS``).
+ Connection fields:
+ - **password**: API key
+ - **host**: Base URL (optional, e.g. ``https://api.openai.com/v1``)
+ - **extra** JSON: ``{"model": "openai:gpt-5.3"}``
:param llm_conn_id: Airflow connection ID for the LLM provider.
:param model_id: Model identifier in ``provider:model`` format (e.g. ``"openai:gpt-5.3"``).
Overrides the model stored in the connection's extra field.
"""
conn_name_attr = "llm_conn_id"
- default_conn_name = "pydanticai_default"
- conn_type = "pydanticai"
+ default_conn_name = "pydantic_ai_default"
+ conn_type = "pydantic_ai"
Review Comment:
I fixed them! Also I change new hook names becouse of following warning
Connection schemes (type: pydanticai_azure) shall not contain '_' according to RFC3986
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.
cetingokhan commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2902540196
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -19,33 +19,32 @@
from typing import TYPE_CHECKING, Any, TypeVar, overload
from pydantic_ai import Agent
-from pydantic_ai.models import Model, infer_model
-from pydantic_ai.providers import Provider, infer_provider, infer_provider_class
+from pydantic_ai.models import infer_model
+from pydantic_ai.providers import infer_provider, infer_provider_class
from airflow.providers.common.compat.sdk import BaseHook
OutputT = TypeVar("OutputT")
if TYPE_CHECKING:
- from pydantic_ai.models import KnownModelName
+ from pydantic_ai.models import KnownModelName, Model
class PydanticAIHook(BaseHook):
"""
Hook for LLM access via pydantic-ai.
- Manages connection credentials and model creation. Uses pydantic-ai's
- model inference to support any provider (OpenAI, Anthropic, Google,
- Bedrock, Ollama, vLLM, etc.).
+ Covers providers that use a standard ``api_key`` + optional ``base_url``
+ (OpenAI, Anthropic, Groq, Mistral, DeepSeek, Ollama, vLLM, …).
- 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)
+ For cloud providers with non-standard auth use the dedicated subclasses:
+ :class:`PydanticAIAzureHook`, :class:`PydanticAIBedrockHook`,
+ :class:`PydanticAIVertexHook`.
- Cloud providers (Bedrock, Vertex) that use native auth chains should leave
- password empty and configure environment-based auth (``AWS_PROFILE``,
- ``GOOGLE_APPLICATION_CREDENTIALS``).
+ Connection fields:
+ - **password**: API key
+ - **host**: Base URL (optional, e.g. ``https://api.openai.com/v1``)
+ - **extra** JSON: ``{"model": "openai:gpt-5.3"}``
Review Comment:
should be like following?
```
def __init__(
self,
llm_conn_id: str | None = None,
model_id: str | None = None, -----> model: str | None = None
**kwargs,
) -> None:
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "gopidesupavan (via GitHub)" <gi...@apache.org>.
gopidesupavan commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2914097527
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -19,33 +19,32 @@
from typing import TYPE_CHECKING, Any, TypeVar, overload
from pydantic_ai import Agent
-from pydantic_ai.models import Model, infer_model
-from pydantic_ai.providers import Provider, infer_provider, infer_provider_class
+from pydantic_ai.models import infer_model
+from pydantic_ai.providers import infer_provider, infer_provider_class
from airflow.providers.common.compat.sdk import BaseHook
OutputT = TypeVar("OutputT")
if TYPE_CHECKING:
- from pydantic_ai.models import KnownModelName
+ from pydantic_ai.models import KnownModelName, Model
class PydanticAIHook(BaseHook):
"""
Hook for LLM access via pydantic-ai.
- Manages connection credentials and model creation. Uses pydantic-ai's
- model inference to support any provider (OpenAI, Anthropic, Google,
- Bedrock, Ollama, vLLM, etc.).
+ Covers providers that use a standard ``api_key`` + optional ``base_url``
+ (OpenAI, Anthropic, Groq, Mistral, DeepSeek, Ollama, vLLM, …).
- 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)
+ For cloud providers with non-standard auth use the dedicated subclasses:
+ :class:`PydanticAIAzureHook`, :class:`PydanticAIBedrockHook`,
+ :class:`PydanticAIVertexHook`.
- Cloud providers (Bedrock, Vertex) that use native auth chains should leave
- password empty and configure environment-based auth (``AWS_PROFILE``,
- ``GOOGLE_APPLICATION_CREDENTIALS``).
+ Connection fields:
+ - **password**: API key
+ - **host**: Base URL (optional, e.g. ``https://api.openai.com/v1``)
+ - **extra** JSON: ``{"model": "openai:gpt-5.3"}``
Review Comment:
agree make sense :)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.
cetingokhan commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2914455654
##########
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:
I see! I think I misread your earlier feedback about removing the try/except as a general rule for the whole refactor. I've just pushed the update with the self.log.warning to avoid any 'head-scratching' for users. ;)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.
cetingokhan commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2914412694
##########
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:
done
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.
cetingokhan commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2916134167
##########
providers/common/ai/provider.yaml:
##########
@@ -73,6 +73,176 @@ connection-types:
type:
- string
- 'null'
+ - hook-class-name: airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIAzureHook
+ connection-type: pydanticai-azure
+ ui-field-behaviour:
+ hidden-fields:
+ - schema
+ - port
+ - login
+ relabeling:
+ password: API Key
+ host: Azure Endpoint
+ placeholders:
+ host: "https://<resource>.openai.azure.com"
+ conn-fields:
+ model:
+ label: Model
+ description: "Azure model identifier (e.g. azure:gpt-4o)"
+ schema:
+ type:
+ - string
+ - 'null'
+ api_version:
+ label: API Version
+ description: "Azure OpenAI API version (e.g. 2024-07-01-preview). Falls back to OPENAI_API_VERSION."
+ schema:
+ type:
+ - string
+ - 'null'
+ - hook-class-name: airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIBedrockHook
+ connection-type: pydanticai-bedrock
+ ui-field-behaviour:
+ hidden-fields:
+ - schema
+ - port
+ - login
+ - host
+ - password
+ relabeling: {}
+ placeholders: {}
+ conn-fields:
+ model:
+ label: Model
+ description: "Bedrock model identifier (e.g. bedrock:us.anthropic.claude-opus-4-5)"
+ schema:
+ type:
+ - string
+ - 'null'
+ region_name:
+ label: AWS Region
+ description: "AWS region (e.g. us-east-1). Falls back to AWS_DEFAULT_REGION env var."
+ schema:
+ type:
+ - string
+ - 'null'
+ aws_access_key_id:
+ label: AWS Access Key ID
+ description: "IAM access key. Leave empty to use instance role / environment credential chain."
+ schema:
+ type:
+ - string
+ - 'null'
+ aws_secret_access_key:
+ label: AWS Secret Access Key
+ description: "IAM secret key."
+ schema:
+ type:
+ - string
+ - 'null'
+ aws_session_token:
+ label: AWS Session Token
+ description: "Temporary session token (optional)."
+ schema:
+ type:
+ - string
+ - 'null'
+ profile_name:
+ label: AWS Profile Name
+ description: "Named AWS credentials profile (optional)."
+ schema:
+ type:
+ - string
+ - 'null'
+ api_key:
+ label: Bearer Token
+ description: "AWS bearer token (alt. to IAM key/secret). Falls back to AWS_BEARER_TOKEN_BEDROCK."
+ schema:
+ type:
+ - string
+ - 'null'
+ base_url:
+ label: Custom Endpoint URL
+ description: "Override the Bedrock runtime endpoint URL (optional)."
+ schema:
+ type:
+ - string
+ - 'null'
+ aws_read_timeout:
+ label: Read Timeout (s)
+ description: "boto3 read timeout in seconds (float, optional)."
+ schema:
+ type:
+ - number
+ - 'null'
+ aws_connect_timeout:
+ label: Connect Timeout (s)
+ description: "boto3 connect timeout in seconds (float, optional)."
+ schema:
+ type:
+ - number
+ - 'null'
+ - hook-class-name: airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIVertexHook
+ connection-type: pydanticai-vertex
+ ui-field-behaviour:
+ hidden-fields:
+ - schema
+ - port
+ - login
+ - host
+ - password
+ relabeling: {}
+ placeholders: {}
+ conn-fields:
+ model:
+ label: Model
+ description: "Google model identifier (e.g. google-vertex:gemini-2.0-flash)"
+ schema:
+ type:
+ - string
+ - 'null'
+ project:
+ label: GCP Project
+ description: "Google Cloud project ID. Falls back to GOOGLE_CLOUD_PROJECT env var."
+ schema:
+ type:
+ - string
+ - 'null'
+ location:
+ label: Location / Region
+ description: "Vertex AI region (e.g. us-central1). Falls back to GOOGLE_CLOUD_LOCATION env var."
+ schema:
+ type:
+ - string
+ - 'null'
+ vertexai:
+ label: Force Vertex AI Mode
+ description: "Force Vertex AI mode. Auto-detected when project/location/credentials are set."
+ schema:
+ type:
+ - boolean
+ - 'null'
+ api_key:
+ label: API Key
+ description: "Google API key for Gen Language API or Vertex AI. Falls back to GOOGLE_API_KEY."
+ schema:
+ type:
+ - string
+ - 'null'
+ service_account_info:
+ label: Service Account Info
+ description: "Service account key as inline dict. Mutually exclusive with service_account_file."
Review Comment:
Done!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2892259289
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -75,62 +74,104 @@ 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."
+ "No model specified. Set model_id on the hook or 'model' in the connection's extra JSON."
+ )
Review Comment:
This calls a subclass's unbound method with a `PydanticAIHook` instance as `self`:
```python
hook_cls._get_provider_kwargs(self, api_key, base_url, extra)
```
It works today because `_get_provider_kwargs` only touches `self.log` (inherited from BaseHook). But if any subclass later accesses subclass-specific attributes, this breaks silently.
A cleaner alternative: have the operator use `for_connection()` instead of constructing `PydanticAIHook` directly. That returns the right subclass, and no auto-dispatch is needed:
```python
@cached_property
def llm_hook(self) -> PydanticAIHook:
return PydanticAIHook.for_connection(self.llm_conn_id, model_id=self.model_id)
```
This eliminates the `type(self) is PydanticAIHook` check, the `_CONN_TYPE_TO_HOOK` registry, and the unbound method call.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -75,62 +74,104 @@ 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."
+ "No model specified. Set model_id on the hook or 'model' in the connection's extra JSON."
+ )
+
+ api_key: str | None = conn.password or None
+ base_url: str | None = conn.host or None
+
+ # Auto-dispatch: if using base hook with a subclass connection type, borrow
+ # that subclass's _get_provider_kwargs so the correct field mapping is used.
+ if type(self) is PydanticAIHook:
+ hook_cls = _CONN_TYPE_TO_HOOK.get(conn.conn_type or "", PydanticAIHook)
+ provider_kwargs = hook_cls._get_provider_kwargs(self, api_key, base_url, extra)
+ else:
+ provider_kwargs = self._get_provider_kwargs(api_key, base_url, extra)
+ if provider_kwargs:
+ _kwargs = provider_kwargs # capture for closure
+ self.log.info(
Review Comment:
`TypeError` here catches both "provider doesn't accept these kwargs" and genuine bugs (wrong types, missing required args). The warning only logs kwarg *names*, not the actual error. This makes debugging auth failures harder — a user who passes the wrong kwarg name gets silently downgraded to env-var auth with a vague warning.
Suggestion: log the exception too:
```python
except TypeError as exc:
self.log.warning(
"Provider '%s' does not accept kwargs %s (%s); "
"falling back to env-var auth.",
pname, list(_kwargs), exc,
)
```
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -157,13 +198,213 @@ 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)
+
+ @classmethod
+ def for_connection(cls, conn_id: str, model_id: str | None = None) -> PydanticAIHook:
+ """
+ Return the correct :class:`PydanticAIHook` subclass for *conn_id*.
+
+ Looks up the connection's ``conn_type`` in the registered hook map and
+ instantiates the matching subclass. Falls back to
+ :class:`PydanticAIHook` for unknown types.
+
+ :param conn_id: Airflow connection ID.
+ :param model_id: Optional model override forwarded to the hook.
+ """
+ conn = cls.get_connection(conn_id)
+ hook_cls = _CONN_TYPE_TO_HOOK.get(conn.conn_type or "", cls)
+ return hook_cls(llm_conn_id=conn_id, model_id=model_id)
+
+
+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_name_attr = "llm_conn_id"
+ default_conn_name = "pydantic_ai_azure_default"
+ conn_type = "pydanticai_azure"
+ 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"):
Review Comment:
This `self.log.info(...)` duplicates the log in `get_conn()` (line ~144) which also logs when `provider_kwargs` is non-empty. Azure connections will produce two "Using explicit credentials..." log lines; Bedrock and Vertex produce one. Remove this — let the base class handle logging uniformly.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -157,13 +198,213 @@ 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)
+
+ @classmethod
+ def for_connection(cls, conn_id: str, model_id: str | None = None) -> PydanticAIHook:
+ """
+ Return the correct :class:`PydanticAIHook` subclass for *conn_id*.
+
+ Looks up the connection's ``conn_type`` in the registered hook map and
+ instantiates the matching subclass. Falls back to
+ :class:`PydanticAIHook` for unknown types.
+
+ :param conn_id: Airflow connection ID.
+ :param model_id: Optional model override forwarded to the hook.
+ """
+ conn = cls.get_connection(conn_id)
+ hook_cls = _CONN_TYPE_TO_HOOK.get(conn.conn_type or "", cls)
+ return hook_cls(llm_conn_id=conn_id, model_id=model_id)
+
+
+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_name_attr = "llm_conn_id"
+ default_conn_name = "pydantic_ai_azure_default"
+ conn_type = "pydanticai_azure"
+ 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"]
+
+ self.log.info(
+ "Using explicit credentials for Azure provider with model '%s': %s",
+ extra.get("model", ""),
+ list(kwargs),
+ )
+ return kwargs
+
+
+class PydanticAIBedrockHook(PydanticAIHook):
+ """
+ Hook for AWS Bedrock via pydantic-ai.
+
+ Credentials are resolved in order:
+
+ 1. Explicit keys in ``extra`` (``aws_access_key_id``,
+ ``aws_secret_access_key``, ``aws_session_token``).
+ 2. 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": "..."
+ }
+
+ Leave ``aws_access_key_id`` / ``aws_secret_access_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_name_attr = "llm_conn_id"
+ default_conn_name = "pydantic_ai_bedrock_default"
+ conn_type = "pydanticai_bedrock"
+ 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"
+ ),
+ },
+ }
+
+ def _get_provider_kwargs(
+ self,
+ api_key: str | None,
+ base_url: str | None,
+ extra: dict[str, Any],
+ ) -> dict[str, Any]:
+ _bedrock_keys = (
+ "aws_access_key_id",
+ "aws_secret_access_key",
+ "aws_session_token",
+ "region_name",
+ "profile_name",
+ )
+ return {k: extra[k] for k in _bedrock_keys if extra.get(k) is not None}
+
+
+class PydanticAIVertexHook(PydanticAIHook):
+ """
+ Hook for Google Vertex AI via pydantic-ai.
+
+ Credentials are resolved in order:
+
+ 1. ``service_account_file`` / ``service_account_info`` in ``extra``.
+ 2. Application Default Credentials (``GOOGLE_APPLICATION_CREDENTIALS``,
+ ``gcloud auth application-default login``, …) when no explicit keys
+ are provided.
+
+ Connection fields:
+ - **extra** JSON::
+
+ {
+ "model": "google-vertex:gemini-2.0-flash",
+ "project_id": "my-gcp-project",
+ "location": "us-central1",
+ "service_account_file": "/path/to/sa.json",
+ }
+
Review Comment:
These kwarg names don't match what pydantic-ai's `GoogleProvider` accepts (see [docs](https://ai.pydantic.dev/models/google/)):
| Connection extra | pydantic-ai param | Match? |
|---|---|---|
| `project_id` | `project` | **No** |
| `location` | `location` | Yes |
| `service_account_file` | `credentials` (object) | **No** |
| `service_account_info` | `credentials` (object) | **No** |
`project_id` needs to be mapped to `project`. `service_account_file`/`service_account_info` need to be loaded into a `google.auth.credentials.Credentials` object before passing — the provider doesn't accept file paths directly.
Because of the `TypeError` catch in `_provider_factory`, these mismatches fail silently — the credentials get ignored and it falls back to ADC. Users who explicitly configure credentials in the connection would have no idea they're not being used.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -157,13 +198,213 @@ 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)
+
+ @classmethod
+ def for_connection(cls, conn_id: str, model_id: str | None = None) -> PydanticAIHook:
+ """
+ Return the correct :class:`PydanticAIHook` subclass for *conn_id*.
+
+ Looks up the connection's ``conn_type`` in the registered hook map and
+ instantiates the matching subclass. Falls back to
+ :class:`PydanticAIHook` for unknown types.
+
+ :param conn_id: Airflow connection ID.
+ :param model_id: Optional model override forwarded to the hook.
+ """
+ conn = cls.get_connection(conn_id)
+ hook_cls = _CONN_TYPE_TO_HOOK.get(conn.conn_type or "", cls)
+ return hook_cls(llm_conn_id=conn_id, model_id=model_id)
+
+
+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:
Python evaluates default argument values at function/method definition time. `PydanticAIHook.__init__` binds `llm_conn_id`'s default to `"pydanticai_default"` when the class is created. Subclasses inherit this `__init__` with the already-resolved default, so:
```python
hook = PydanticAIAzureHook() # llm_conn_id="pydanticai_default", not "pydantic_ai_azure_default"
```
The `default_conn_name` attribute on the subclass is never consulted during `__init__`. Same issue for Bedrock and Vertex subclasses.
##########
providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py:
##########
@@ -105,25 +100,13 @@ def get_provider_info():
{
"integration-name": "Common AI",
"python-modules": [
- "airflow.providers.common.ai.operators.agent",
"airflow.providers.common.ai.operators.llm",
Review Comment:
This file is auto-generated from `provider.yaml` — manual edits here will be overwritten by `prek run update-providers-build-files`. The new connection types need to be added to `provider.yaml` instead.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.
cetingokhan commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2900083136
##########
providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py:
##########
@@ -74,6 +74,139 @@ def get_provider_info():
}
},
},
Review Comment:
You're right! I actually thought I had already updated it based on your previous feedback, but something must have gone wrong during the merge/rebase process, and the manual edits probably snuck back in.
To fix this, I've now synced with the latest main, pulled the original get_provider_info.py, and let prek regenerate it properly. It should be correct now!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on PR #62816:
URL: https://github.com/apache/airflow/pull/62816#issuecomment-4018124877
## Code Review (updated PR)
Big improvement from the earlier iterations. The separate connection types with `_get_provider_kwargs()` as the extension point is the right pattern.
### [HIGH] `model_id` no longer forwarded from operator to hook
`LLMOperator.llm_hook` changed from:
```python
return PydanticAIHook.for_connection(self.llm_conn_id, model_id=self.model_id)
```
to:
```python
return PydanticAIHook.get_hook(self.llm_conn_id)
```
`get_hook()` (inherited from `BaseHook`) doesn't accept `model_id`. So if a user sets `model_id="openai:gpt-5"` on `LLMOperator`, it's silently dropped — the hook never sees it. The hook falls through to `extra.get("model", "")` and either uses the connection's model or raises "No model specified."
The test confirms the regression:
```python
# Before: model_id forwarded
mock_hook_cls.for_connection.assert_called_once_with("my_llm", model_id="openai:gpt-5")
# After: model_id dropped
mock_hook_cls.get_hook.assert_called_once_with("my_llm")
```
**Fix:** Either keep `for_connection()` (it was a clean API), or set `model_id` on the hook after `get_hook()` returns:
```python
@cached_property
def llm_hook(self) -> PydanticAIHook:
hook = PydanticAIHook.get_hook(self.llm_conn_id)
hook.model_id = self.model_id
return hook
```
---
### [MEDIUM] TypeError fallback removed — defensive safety net gone
The base hook's `_provider_factory` lost the `try/except TypeError` fallback:
```python
# Before (defensive):
try:
return provider_cls(**kwargs)
except TypeError:
return infer_provider(provider_name)
# After (brittle):
return infer_provider_class(pname)(**_kwargs)
```
With the subclass architecture, this matters less because the base hook should only handle standard providers. But if someone uses a generic `pydanticai` connection with an exotic provider that rejects `api_key`/`base_url`, it crashes instead of falling back gracefully. The `infer_provider` import was also removed, so restoring this later requires adding it back.
**Suggestion:** Keep the fallback as a safety net. It costs nothing and prevents surprises.
---
### [MEDIUM] `conn_type` readability — `pydanticaiazure` vs `pydanticai_azure`
The conn_types changed from `pydanticai_azure` to `pydanticaiazure` (same for bedrock, vertex). This is hard to read. While `gcpbigquery` exists as a precedent, `google_cloud_platform` also uses underscores. There's no strict convention.
`pydanticai_azure` is more readable and consistent with `default_conn_name = "pydanticai_azure_default"`. The mismatch between conn_type (`pydanticaiazure`) and default_conn_name (`pydanticai_azure_default`) is confusing.
---
### [LOW] Bedrock/Vertex `get_ui_field_behaviour` placeholder text still shows raw JSON
The Bedrock and Vertex hooks have `conn-fields` defined in provider.yaml (proper form widgets), but their `get_ui_field_behaviour().placeholders` still show raw JSON in the `extra` placeholder. Since the fields now have dedicated form widgets, the `extra` placeholder is misleading — users should fill in the individual fields, not paste JSON into extra.
---
### [LOW] Bedrock `_get_provider_kwargs` has `api_key` and `base_url` in both signature and extra keys
The method receives `api_key` (from `conn.password`) and `base_url` (from `conn.host`), but ignores them. Instead, `api_key` and `base_url` are read from `extra`. The comment explains this, but the same logical field has two possible sources. Since the Bedrock connection hides `host` and `password` in the UI this is consistent — just worth documenting.
---
**Summary**: 1 High (model_id dropped), 2 Medium (TypeError fallback, conn_type naming), 2 Low. The architecture is solid — the `model_id` regression is the main blocker.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.
cetingokhan commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2900129133
##########
providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py:
##########
@@ -16,13 +16,20 @@
# under the License.
from __future__ import annotations
+import json
+import sys
from unittest.mock import MagicMock, patch
import pytest
from pydantic_ai.models import Model
from airflow.models.connection import Connection
-from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+from airflow.providers.common.ai.hooks.pydantic_ai import (
+ PydanticAIAzureHook,
+ PydanticAIBedrockHook,
+ PydanticAIHook,
Review Comment:
You're absolutely right. Since the previous review suggested that subclasses might not need their own default_conn_name, I had revised them to share the base class's default. I added those tests specifically to verify that sharing was working as intended.
However, I see your point regarding the conn_type mismatch and the confusion it creates. I've now separated the default_conn_name for each subclass again and updated the tests to match. Thanks for catching that! :)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.
cetingokhan commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2901762480
##########
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:
```
def _provider_factory(pname: str) -> Any:
try:
return infer_provider_class(pname)(**_kwargs)
except TypeError:
return infer_provider(pname)
```
I've added the try/except TypeError fallback and the infer_provider import back as suggested.
Regarding the warning log: I haven't added it yet because we previously discussed keeping the logs clean and removed it in the earlier rounds. However, I’m a bit hesitant about a "silent fallback."
Don't you think it would be more helpful for the user to know that their connection credentials were rejected and the hook is falling back to environment variables? If it’s silent, they might struggle to understand why their Airflow connection settings aren't being applied.
What do you think—should we keep it silent or add a simple warning log?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2880383612
##########
providers/common/ai/src/airflow/providers/common/ai/builders/base.py:
##########
@@ -0,0 +1,58 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Protocol
+
+if TYPE_CHECKING:
+ from pydantic_ai.models import KnownModelName, Model
+
+
+class ProviderBuilder(Protocol):
Review Comment:
Do we need this abstraction? The current `get_conn()` is a straightforward if/else that delegates to pydantic-ai's own `infer_model()`. Adding a Protocol + 3 builder classes + a dispatch loop for what's essentially a single new code path (Azure) feels like premature abstraction for a problem that doesn't exist yet. If/when we genuinely need pluggable resolution, we can introduce it then.
Also — `ProviderBuilder` is declared as a `Protocol` (structural typing), but the concrete classes inherit from it (nominal typing). These are two different patterns — if you want an inheritance hierarchy, use `ABC`; if you want duck typing, don't inherit from the `Protocol` in the concrete classes. Mixing both is confusing for contributors.
##########
providers/common/ai/src/airflow/providers/common/ai/builders/azure_openai.py:
##########
@@ -0,0 +1,99 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+ from pydantic_ai.models import KnownModelName, Model
+
+from airflow.providers.common.ai.builders.base import ProviderBuilder
+
+
+class AzureOpenAIBuilder(ProviderBuilder):
+ """Builds a pydantic-ai Model backed by an AsyncAzureOpenAI client."""
+
+ def supports(self, extra: dict[str, Any], api_key: str | None, base_url: str | None) -> bool:
+ """Return True when the connection extra contains 'api_version' for Azure."""
+ return bool(extra.get("api_version"))
+
+ def build(
+ self,
+ model_name: str | KnownModelName,
+ extra: dict[str, Any],
+ api_key: str | None,
+ base_url: str | None,
+ ) -> Model:
+ try:
+ from openai import AsyncAzureOpenAI
+ from pydantic_ai.models.openai import OpenAIChatModel
+ from pydantic_ai.providers.openai import OpenAIProvider
+ except ImportError as exc:
+ raise ImportError(
+ "The 'openai' and 'pydantic-ai[openai]' packages are required for Azure OpenAI connections. "
+ "Install them with: pip install 'pydantic-ai[openai]'"
+ ) from exc
+
+ api_version: str | None = extra.get("api_version")
+ if not api_version:
+ raise ValueError(
+ "Connection extra must contain 'api_version' for Azure OpenAI. "
+ 'Example: {"api_version": "2024-07-01-preview"}'
+ )
+
+ if not base_url:
+ raise ValueError(
+ "Connection 'host' must be set to the Azure endpoint. "
+ "Example: https://<resource>.openai.azure.com"
+ )
+
+ client_kwargs: dict[str, Any] = {
+ "api_version": api_version,
+ "azure_endpoint": base_url,
+ }
+ if api_key:
+ client_kwargs["api_key"] = api_key
+
+ azure_deployment: str | None = extra.get("azure_deployment")
+ if azure_deployment:
+ client_kwargs["azure_deployment"] = azure_deployment
+
+ azure_ad_token: str | None = extra.get("azure_ad_token")
+ if azure_ad_token:
+ client_kwargs["azure_ad_token"] = azure_ad_token
+
+ azure_ad_token_provider_path: str | None = extra.get("azure_ad_token_provider")
+ if azure_ad_token_provider_path:
+ client_kwargs["azure_ad_token_provider"] = self._import_callable(azure_ad_token_provider_path)
+
+ azure_client = AsyncAzureOpenAI(**client_kwargs)
+
+ # Strip provider prefix if present ("openai:gpt-5.2" → "gpt-5.2")
+ slug = model_name.split(":", 1)[-1] if ":" in model_name else model_name
+ return OpenAIChatModel(slug, provider=OpenAIProvider(openai_client=azure_client))
+
+ @staticmethod
+ def _import_callable(dotted_path: str) -> Any:
Review Comment:
This calls `importlib.import_module()` on a user-provided string from connection extras — module imports can run arbitrary code at import time. Connection extras are editable by any user with connection-edit permissions, which is a lower-privilege surface than DAG deployment.
If we end up needing a custom token provider path, this should at least be documented as a security-sensitive field. But since pydantic-ai's `AzureProvider` accepts `api_key` directly, we may not need this at all for the initial implementation.
##########
providers/common/ai/src/airflow/providers/common/ai/builders/azure_openai.py:
##########
@@ -0,0 +1,99 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+ from pydantic_ai.models import KnownModelName, Model
+
+from airflow.providers.common.ai.builders.base import ProviderBuilder
+
+
+class AzureOpenAIBuilder(ProviderBuilder):
+ """Builds a pydantic-ai Model backed by an AsyncAzureOpenAI client."""
+
+ def supports(self, extra: dict[str, Any], api_key: str | None, base_url: str | None) -> bool:
+ """Return True when the connection extra contains 'api_version' for Azure."""
+ return bool(extra.get("api_version"))
+
+ def build(
+ self,
+ model_name: str | KnownModelName,
+ extra: dict[str, Any],
+ api_key: str | None,
+ base_url: str | None,
+ ) -> Model:
+ try:
+ from openai import AsyncAzureOpenAI
Review Comment:
pydantic-ai already has native Azure support via `AzureProvider` — no need to drop down to the raw `openai` SDK:
```python
from pydantic_ai.providers.azure import AzureProvider
from pydantic_ai.models.openai import OpenAIChatModel
provider = AzureProvider(
azure_endpoint=base_url,
api_version=api_version,
api_key=api_key,
)
model = OpenAIChatModel(model_name, provider=provider)
```
See https://ai.pydantic.dev/models/openai/#azure
By constructing `AsyncAzureOpenAI` directly we're bypassing pydantic-ai's own provider abstraction (which may handle retries, error mapping, etc.) and coupling ourselves to `openai` SDK internals.
Using `AzureProvider` would also make a separate builder class unnecessary — it could just be a few lines in `get_conn()`.
##########
providers/common/ai/src/airflow/providers/common/ai/builders/custom_endpoint.py:
##########
@@ -0,0 +1,92 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from pydantic_ai.models import infer_model
+from pydantic_ai.providers import Provider, infer_provider, infer_provider_class
+
+if TYPE_CHECKING:
+ from pydantic_ai.models import KnownModelName, Model
+
+from airflow.providers.common.ai.builders.base import ProviderBuilder
+
+
+class CustomEndpointBuilder(ProviderBuilder):
+ """
+ Builds a pydantic-ai Model with explicitly overridden credentials or endpoints.
+
+ Used for general connections (OpenAI, Gemini, Groq, Ollama) where password or host is set.
+ """
+
+ def supports(self, extra: dict[str, Any], api_key: str | None, base_url: str | None) -> bool:
+ """Return True if api_key or base_url is explicitly provided."""
+ return bool(api_key or base_url)
+
+ def build(
+ self,
+ model_name: str | KnownModelName,
+ extra: dict[str, Any],
+ api_key: str | None,
+ base_url: str | None,
+ ) -> Model:
+ return infer_model(model_name, provider_factory=self._make_provider_factory(api_key, base_url))
+
+ @staticmethod
+ def _make_provider_factory(api_key: str | None, base_url: str | None):
+ """
+ Return a provider factory closure for non-Azure custom endpoints.
+
+ Falls back to default provider resolution when the provider's
+ constructor does not accept `api_key` / `base_url`
+ """
+
+ def _factory(provider_name: str) -> Provider[Any]:
+ provider_cls = infer_provider_class(provider_name)
+ kwargs: dict[str, Any] = {}
+ if api_key:
+ kwargs["api_key"] = api_key
+ if base_url:
+ kwargs["base_url"] = base_url
+ try:
+ return provider_cls(**kwargs)
+ except TypeError:
+ return infer_provider(provider_name)
+
+ return _factory
+
+
+class DefaultBuilder(ProviderBuilder):
Review Comment:
`DefaultBuilder.build()` is `return infer_model(model_name)` — a single call. The existing hook does this in two lines: `if not api_key and not base_url: return infer_model(model_name)`. Does this one-liner need its own class with Protocol conformance and a `supports()` method?
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -75,60 +89,62 @@ 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)",
- "extra": '{"model": "openai:gpt-5.3"}',
+ "host": (
+ "https://api.openai.com/v1 — or Azure endpoint: https://<resource>.openai.azure.com"
+ ),
+ "extra": ('{"model": "openai:gpt-5.3"} — Azure: also add "api_version", "azure_deployment"'),
},
}
+ # ------------------------------------------------------------------
+ # Core connection / agent API
+ # ------------------------------------------------------------------
+
def get_conn(self) -> Model:
"""
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.
- The result is cached for the lifetime of this hook instance.
+ Resolution is delegated to builders based on connection characteristics:
+
+ 1. **Azure OpenAI** — when ``api_version`` is present in connection extra.
+ 2. **Custom endpoint** — when ``password`` or ``host`` are set.
+ 3. **Default resolution** — delegates to pydantic-ai ``infer_model``.
+
+ The resolved model is cached for the lifetime of this hook instance.
"""
if self._model is not None:
return self._model
+ from airflow.providers.common.ai.builders import (
+ AzureOpenAIBuilder,
+ CustomEndpointBuilder,
+ DefaultBuilder,
+ ProviderBuilder,
+ )
+
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 'model' in the connection's extra JSON."
)
- 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)
- return self._model
+ api_key: str | None = conn.password or None
+ base_url: str | None = conn.host or None
+
+ builders: list[ProviderBuilder] = [
+ AzureOpenAIBuilder(),
+ CustomEndpointBuilder(),
+ DefaultBuilder(),
+ ]
+
+ for builder in builders:
+ if builder.supports(extra, api_key, base_url):
+ self._model = builder.build(model_name, extra, api_key, base_url)
+ return self._model
- def _provider_factory(provider_name: str) -> Provider[Any]:
- """
- Create a provider with credentials from the Airflow connection.
-
- Falls back to default provider resolution if the provider's constructor
- doesn't accept api_key/base_url (e.g. Google Vertex, Bedrock).
- """
- provider_cls = infer_provider_class(provider_name)
- kwargs: dict[str, Any] = {}
- if api_key:
- kwargs["api_key"] = api_key
- if base_url:
- kwargs["base_url"] = base_url
- try:
- return provider_cls(**kwargs)
- except TypeError:
- # Provider doesn't accept these kwargs (e.g. Google Vertex/GLA
- # use ADC, Bedrock uses boto session). Fall back to default
- # provider resolution which reads credentials from the environment.
- return infer_provider(provider_name)
-
- self._model = infer_model(model_name, provider_factory=_provider_factory)
- return self._model
+ raise RuntimeError("No suitable ProviderBuilder found to construct the model.")
Review Comment:
This line is unreachable — `DefaultBuilder.supports()` always returns `True`, so the loop will always match on the third iteration. Dead code that implies the loop might not match, which could mislead future readers.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "gopidesupavan (via GitHub)" <gi...@apache.org>.
gopidesupavan commented on PR #62816:
URL: https://github.com/apache/airflow/pull/62816#issuecomment-4000582807
> Hi @kaxil
>
> I have taken your "thin hook" suggestion into consideration. However, I believe that relying solely on environment variables might not be sufficient for Airflow users; being able to fetch configurations directly from connection extras is highly valuable.
>
> I've reviewed all the current pydanticai providers and noticed that only AzureProvider and LiteLLMProvider require their own specific parameters instead of the standard base_url.
>
> To handle this distinction cleanly without over-engineering or building registries, I am planning to introduce a simple _build_provider_kwargs(conn) helper method. This method will parse the connection and return the appropriate kwargs dictionary to feed directly into the resolved provider_cls.
>
> What are your thoughts on proceeding with this approach? Does this align better with the design goals?
>
> ```
> def _build_provider_kwargs(
> provider_cls: type,
> api_key: str | None,
> base_url: str | None,
> extra: dict[str, Any],
> ) -> dict[str, Any]:
> sig = inspect.signature(provider_cls.__init__)
> accepted: set[str] = {
> name
> for name, param in sig.parameters.items()
> if name != "self"
> and param.kind
> not in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD)
> }
>
> candidates: dict[str, Any] = {
> "api_key": api_key,
> "base_url": base_url,
> # Alias used by AzureProvider instead of base_url
> "azure_endpoint": base_url,
> # Alias used by LiteLLMProvider instead of base_url
> "api_base": base_url,
> **{k: v for k, v in extra.items() if k != "model"},
> }
>
> return {k: v for k, v in candidates.items() if k in accepted and v is not None}
>
> def get_conn(self) -> Model:
> .....
> if api_key or base_url:
> provider_name = model_name.split(":")[0] if ":" in model_name else "openai"
> provider_cls = infer_provider_class(provider_name)
> kwargs: dict[str, Any] = self._build_provider_kwargs(provider_cls, api_key, base_url, extra)
> try:
> provider = provider_cls(**kwargs)
> except TypeError:
> provider = infer_provider(provider_name)
> ....
> ```
What if we simply read key-value pairs from conn_extras and pass them directly to the provider? We could define a schema in the connection extras like {provider_extra_args: {...}}, then document that this field and point users to the pydantic-ai docs to configure it correctly themselves ?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.
cetingokhan commented on PR #62816:
URL: https://github.com/apache/airflow/pull/62816#issuecomment-4000702696
> > Hi @kaxil
> > I have taken your "thin hook" suggestion into consideration. However, I believe that relying solely on environment variables might not be sufficient for Airflow users; being able to fetch configurations directly from connection extras is highly valuable.
> > I've reviewed all the current pydanticai providers and noticed that only AzureProvider and LiteLLMProvider require their own specific parameters instead of the standard base_url.
> > To handle this distinction cleanly without over-engineering or building registries, I am planning to introduce a simple _build_provider_kwargs(conn) helper method. This method will parse the connection and return the appropriate kwargs dictionary to feed directly into the resolved provider_cls.
> > What are your thoughts on proceeding with this approach? Does this align better with the design goals?
> > ```
> > def _build_provider_kwargs(
> > provider_cls: type,
> > api_key: str | None,
> > base_url: str | None,
> > extra: dict[str, Any],
> > ) -> dict[str, Any]:
> > sig = inspect.signature(provider_cls.__init__)
> > accepted: set[str] = {
> > name
> > for name, param in sig.parameters.items()
> > if name != "self"
> > and param.kind
> > not in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD)
> > }
> >
> > candidates: dict[str, Any] = {
> > "api_key": api_key,
> > "base_url": base_url,
> > # Alias used by AzureProvider instead of base_url
> > "azure_endpoint": base_url,
> > # Alias used by LiteLLMProvider instead of base_url
> > "api_base": base_url,
> > **{k: v for k, v in extra.items() if k != "model"},
> > }
> >
> > return {k: v for k, v in candidates.items() if k in accepted and v is not None}
> >
> > def get_conn(self) -> Model:
> > .....
> > if api_key or base_url:
> > provider_name = model_name.split(":")[0] if ":" in model_name else "openai"
> > provider_cls = infer_provider_class(provider_name)
> > kwargs: dict[str, Any] = self._build_provider_kwargs(provider_cls, api_key, base_url, extra)
> > try:
> > provider = provider_cls(**kwargs)
> > except TypeError:
> > provider = infer_provider(provider_name)
> > ....
> > ```
>
> What if we simply read key-value pairs from conn_extras and pass them directly to the provider? We could define a schema in the connection extras like {provider_extra_args: {...}}, then document that this field and point users to the pydantic-ai docs to configure it correctly themselves ?
Yes we can follow KISS principe so then user could manage what need to pass extra properties for thier provider
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2895755196
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -19,42 +19,41 @@
from typing import TYPE_CHECKING, Any, TypeVar, overload
from pydantic_ai import Agent
-from pydantic_ai.models import Model, infer_model
-from pydantic_ai.providers import Provider, infer_provider, infer_provider_class
+from pydantic_ai.models import infer_model
+from pydantic_ai.providers import infer_provider, infer_provider_class
from airflow.providers.common.compat.sdk import BaseHook
OutputT = TypeVar("OutputT")
if TYPE_CHECKING:
- from pydantic_ai.models import KnownModelName
+ from pydantic_ai.models import KnownModelName, Model
class PydanticAIHook(BaseHook):
"""
Hook for LLM access via pydantic-ai.
- Manages connection credentials and model creation. Uses pydantic-ai's
- model inference to support any provider (OpenAI, Anthropic, Google,
- Bedrock, Ollama, vLLM, etc.).
+ Covers providers that use a standard ``api_key`` + optional ``base_url``
+ (OpenAI, Anthropic, Groq, Mistral, DeepSeek, Ollama, vLLM, …).
- 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)
+ For cloud providers with non-standard auth use the dedicated subclasses:
+ :class:`PydanticAIAzureHook`, :class:`PydanticAIBedrockHook`,
+ :class:`PydanticAIVertexHook`.
- Cloud providers (Bedrock, Vertex) that use native auth chains should leave
- password empty and configure environment-based auth (``AWS_PROFILE``,
- ``GOOGLE_APPLICATION_CREDENTIALS``).
+ Connection fields:
+ - **password**: API key
+ - **host**: Base URL (optional, e.g. ``https://api.openai.com/v1``)
+ - **extra** JSON: ``{"model": "openai:gpt-5.3"}``
:param llm_conn_id: Airflow connection ID for the LLM provider.
:param model_id: Model identifier in ``provider:model`` format (e.g. ``"openai:gpt-5.3"``).
Overrides the model stored in the connection's extra field.
"""
conn_name_attr = "llm_conn_id"
- default_conn_name = "pydanticai_default"
- conn_type = "pydanticai"
+ default_conn_name = "pydantic_ai_default"
+ conn_type = "pydantic_ai"
Review Comment:
Why change this :)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.
cetingokhan commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2901762480
##########
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:
I replaced with following;
```
def _provider_factory(pname: str) -> Any:
try:
return infer_provider_class(pname)(**_kwargs)
except TypeError:
return infer_provider(pname)
```
?
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.
cetingokhan commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2900084261
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -157,13 +194,266 @@ 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)
+
+ @classmethod
+ def for_connection(cls, conn_id: str, model_id: str | None = None) -> PydanticAIHook:
+ """
+ Return the correct :class:`PydanticAIHook` subclass for *conn_id*.
+
+ Looks up the connection's ``conn_type`` in the registered hook map and
+ instantiates the matching subclass. Falls back to
+ :class:`PydanticAIHook` for unknown types.
+
+ :param conn_id: Airflow connection ID.
+ :param model_id: Optional model override forwarded to the hook.
+ """
+ conn = cls.get_connection(conn_id)
+ hook_cls = _CONN_TYPE_TO_HOOK.get(conn.conn_type or "", cls)
+ return hook_cls(llm_conn_id=conn_id, model_id=model_id)
+
+
+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 = "pydanticai_azure"
+ 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 = "pydanticai_bedrock"
+ 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"
+ ),
+ },
+ }
+
+ def _get_provider_kwargs(
+ self,
+ api_key: str | None,
+ base_url: str | None,
+ extra: dict[str, Any],
+ ) -> dict[str, Any]:
+ _str_keys = (
+ "aws_access_key_id",
+ "aws_secret_access_key",
+ "aws_session_token",
+ "region_name",
+ "profile_name",
+ # Bearer-token auth (alternative to IAM key/secret).
+ # Maps to AWS_BEARER_TOKEN_BEDROCK env var.
+ "api_key",
+ # Custom Bedrock runtime endpoint.
+ "base_url",
+ )
+ kwargs: dict[str, Any] = {k: extra[k] for k in _str_keys if extra.get(k) is not None}
+ # BedrockProvider expects float for timeout values; JSON integers must be coerced.
+ for _timeout_key in ("aws_read_timeout", "aws_connect_timeout"):
+ if extra.get(_timeout_key) is not None:
+ kwargs[_timeout_key] = float(extra[_timeout_key])
+ return kwargs
+
+
+class PydanticAIVertexHook(PydanticAIHook):
+ """
+ Hook for Google Vertex AI (or Generative Language API) via pydantic-ai.
+
+ Credentials are resolved in order:
+
+ 1. ``service_account_file`` (path string) or ``service_account_info`` (JSON
+ object) in ``extra`` — loaded into a ``google.auth.credentials.Credentials``
+ object and passed as ``credentials`` to ``GoogleProvider``.
+ 2. ``api_key`` in ``extra`` — for Generative Language API (non-Vertex) or
+ Vertex API-key auth.
+ 3. Application Default Credentials (``GOOGLE_APPLICATION_CREDENTIALS``,
+ ``gcloud auth application-default login``, Workload Identity, …) when
+ no explicit credentials are provided.
+
+ Connection fields:
+ - **extra** JSON::
+
+ {
+ "model": "google-vertex:gemini-2.0-flash",
+ "project": "my-gcp-project",
+ "location": "us-central1",
+ "service_account_file": "/path/to/sa.json",
+ "vertexai": true,
+ }
+
+ Use ``"service_account_info"`` instead of ``"service_account_file"`` to
+ embed the service-account JSON directly (as an object, not a string path).
+ Setting both at the same time raises ``ValueError``.
+
+ Set ``"vertexai": true`` to force Vertex AI mode when only ``api_key`` is
+ provided. Omit ``vertexai`` for the Generative Language API (GLA).
+
+ :param llm_conn_id: Airflow connection ID.
+ :param model_id: Model identifier, e.g. ``"google-vertex:gemini-2.0-flash"``.
+ """
+
+ conn_type = "pydanticai_vertex"
+ hook_name = "Pydantic AI (Google Vertex AI)"
+
+ @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": "google-vertex:gemini-2.0-flash", '
+ '"project": "my-project", "location": "us-central1", "vertexai": true}'
+ " — add service_account_file (path) or service_account_info (object) for SA auth;"
+ " omit both to use Application Default Credentials"
+ ),
+ },
+ }
+
+ def _get_provider_kwargs(
+ self,
+ api_key: str | None,
+ base_url: str | None,
+ extra: dict[str, Any],
+ ) -> dict[str, Any]:
+ sa_file = extra.get("service_account_file")
+ sa_info = extra.get("service_account_info")
+ if sa_file and sa_info:
+ raise ValueError(
+ "Specify 'service_account_file' or 'service_account_info' in the connection extra, not both."
+ )
+
+ kwargs: dict[str, Any] = {}
+
+ # Direct GoogleProvider scalar kwargs.
+ for _key in ("api_key", "project", "location", "base_url"):
+ if extra.get(_key) is not None:
+ kwargs[_key] = extra[_key]
+
+ # Optional vertexai bool flag (force Vertex AI mode for API-key auth).
+ _vertexai = extra.get("vertexai")
+ if _vertexai is not None:
+ kwargs["vertexai"] = bool(_vertexai)
+
+ # Service-account credentials — loaded lazily to avoid importing
+ # google-auth on non-Vertex code paths (optional heavy dependency).
+ if sa_file:
+ from google.oauth2 import service_account # lazy: optional dep
+
+ kwargs["credentials"] = service_account.Credentials.from_service_account_file(
Review Comment:
refactored with option1 ;)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.
cetingokhan commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2900085520
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -157,13 +194,266 @@ 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)
+
+ @classmethod
+ def for_connection(cls, conn_id: str, model_id: str | None = None) -> PydanticAIHook:
+ """
+ Return the correct :class:`PydanticAIHook` subclass for *conn_id*.
+
+ Looks up the connection's ``conn_type`` in the registered hook map and
+ instantiates the matching subclass. Falls back to
+ :class:`PydanticAIHook` for unknown types.
+
+ :param conn_id: Airflow connection ID.
+ :param model_id: Optional model override forwarded to the hook.
+ """
+ conn = cls.get_connection(conn_id)
+ hook_cls = _CONN_TYPE_TO_HOOK.get(conn.conn_type or "", cls)
+ return hook_cls(llm_conn_id=conn_id, model_id=model_id)
+
+
+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 = "pydanticai_azure"
+ 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 = "pydanticai_bedrock"
+ 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"
+ ),
+ },
+ }
+
+ def _get_provider_kwargs(
Review Comment:
Added comment which you point in option1
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.
cetingokhan commented on PR #62816:
URL: https://github.com/apache/airflow/pull/62816#issuecomment-4034536032
Hi @kaxil and @gopidesupavan ,
I would like to sincerely thank you for the time and patience you've invested in reviewing my PRs. I feel a bit embarrassed knowing that I’ve taken up so much of your time with my numerous mistakes, especially when I know you could have implemented these features much faster yourselves. I truly appreciate the guidance you've provided despite the extra work I've caused.
These reviews are proving to be incredibly valuable for me in understanding the "Airflow Gurus" perspective on architecture. I believe that as I adopt this mindset, my future PRs will require much less back-and-forth, though I know this takes time and experience.
Regarding my other PR (LLMDataQualityOperator #62963 ), I am fully committed to it. However, if you have a strict deadline or if you feel the review process will be too taxing for your current schedule, I would completely understand if you prefer to put it on hold. That said, if you are willing to bear with me, I am more than eager to make all the necessary improvements and see it through to completion!
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.
cetingokhan commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2916134167
##########
providers/common/ai/provider.yaml:
##########
@@ -73,6 +73,176 @@ connection-types:
type:
- string
- 'null'
+ - hook-class-name: airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIAzureHook
+ connection-type: pydanticai-azure
+ ui-field-behaviour:
+ hidden-fields:
+ - schema
+ - port
+ - login
+ relabeling:
+ password: API Key
+ host: Azure Endpoint
+ placeholders:
+ host: "https://<resource>.openai.azure.com"
+ conn-fields:
+ model:
+ label: Model
+ description: "Azure model identifier (e.g. azure:gpt-4o)"
+ schema:
+ type:
+ - string
+ - 'null'
+ api_version:
+ label: API Version
+ description: "Azure OpenAI API version (e.g. 2024-07-01-preview). Falls back to OPENAI_API_VERSION."
+ schema:
+ type:
+ - string
+ - 'null'
+ - hook-class-name: airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIBedrockHook
+ connection-type: pydanticai-bedrock
+ ui-field-behaviour:
+ hidden-fields:
+ - schema
+ - port
+ - login
+ - host
+ - password
+ relabeling: {}
+ placeholders: {}
+ conn-fields:
+ model:
+ label: Model
+ description: "Bedrock model identifier (e.g. bedrock:us.anthropic.claude-opus-4-5)"
+ schema:
+ type:
+ - string
+ - 'null'
+ region_name:
+ label: AWS Region
+ description: "AWS region (e.g. us-east-1). Falls back to AWS_DEFAULT_REGION env var."
+ schema:
+ type:
+ - string
+ - 'null'
+ aws_access_key_id:
+ label: AWS Access Key ID
+ description: "IAM access key. Leave empty to use instance role / environment credential chain."
+ schema:
+ type:
+ - string
+ - 'null'
+ aws_secret_access_key:
+ label: AWS Secret Access Key
+ description: "IAM secret key."
+ schema:
+ type:
+ - string
+ - 'null'
+ aws_session_token:
+ label: AWS Session Token
+ description: "Temporary session token (optional)."
+ schema:
+ type:
+ - string
+ - 'null'
+ profile_name:
+ label: AWS Profile Name
+ description: "Named AWS credentials profile (optional)."
+ schema:
+ type:
+ - string
+ - 'null'
+ api_key:
+ label: Bearer Token
+ description: "AWS bearer token (alt. to IAM key/secret). Falls back to AWS_BEARER_TOKEN_BEDROCK."
+ schema:
+ type:
+ - string
+ - 'null'
+ base_url:
+ label: Custom Endpoint URL
+ description: "Override the Bedrock runtime endpoint URL (optional)."
+ schema:
+ type:
+ - string
+ - 'null'
+ aws_read_timeout:
+ label: Read Timeout (s)
+ description: "boto3 read timeout in seconds (float, optional)."
+ schema:
+ type:
+ - number
+ - 'null'
+ aws_connect_timeout:
+ label: Connect Timeout (s)
+ description: "boto3 connect timeout in seconds (float, optional)."
+ schema:
+ type:
+ - number
+ - 'null'
+ - hook-class-name: airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIVertexHook
+ connection-type: pydanticai-vertex
+ ui-field-behaviour:
+ hidden-fields:
+ - schema
+ - port
+ - login
+ - host
+ - password
+ relabeling: {}
+ placeholders: {}
+ conn-fields:
+ model:
+ label: Model
+ description: "Google model identifier (e.g. google-vertex:gemini-2.0-flash)"
+ schema:
+ type:
+ - string
+ - 'null'
+ project:
+ label: GCP Project
+ description: "Google Cloud project ID. Falls back to GOOGLE_CLOUD_PROJECT env var."
+ schema:
+ type:
+ - string
+ - 'null'
+ location:
+ label: Location / Region
+ description: "Vertex AI region (e.g. us-central1). Falls back to GOOGLE_CLOUD_LOCATION env var."
+ schema:
+ type:
+ - string
+ - 'null'
+ vertexai:
+ label: Force Vertex AI Mode
+ description: "Force Vertex AI mode. Auto-detected when project/location/credentials are set."
+ schema:
+ type:
+ - boolean
+ - 'null'
+ api_key:
+ label: API Key
+ description: "Google API key for Gen Language API or Vertex AI. Falls back to GOOGLE_API_KEY."
+ schema:
+ type:
+ - string
+ - 'null'
+ service_account_info:
+ label: Service Account Info
+ description: "Service account key as inline dict. Mutually exclusive with service_account_file."
Review Comment:
changed with following description;
"Service account key as inline dict (JSON with type, project_id, private_key, etc.)."
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on PR #62816:
URL: https://github.com/apache/airflow/pull/62816#issuecomment-4038706188
@cetingokhan looks like there is a conflict!
Could you resolve that and we are good to go
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2914866165
##########
providers/common/ai/provider.yaml:
##########
@@ -73,6 +73,176 @@ connection-types:
type:
- string
- 'null'
+ - hook-class-name: airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIAzureHook
+ connection-type: pydanticai-azure
+ ui-field-behaviour:
+ hidden-fields:
+ - schema
+ - port
+ - login
+ relabeling:
+ password: API Key
+ host: Azure Endpoint
+ placeholders:
+ host: "https://<resource>.openai.azure.com"
+ conn-fields:
+ model:
+ label: Model
+ description: "Azure model identifier (e.g. azure:gpt-4o)"
+ schema:
+ type:
+ - string
+ - 'null'
+ api_version:
+ label: API Version
+ description: "Azure OpenAI API version (e.g. 2024-07-01-preview). Falls back to OPENAI_API_VERSION."
+ schema:
+ type:
+ - string
+ - 'null'
+ - hook-class-name: airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIBedrockHook
+ connection-type: pydanticai-bedrock
+ ui-field-behaviour:
+ hidden-fields:
+ - schema
+ - port
+ - login
+ - host
+ - password
+ relabeling: {}
+ placeholders: {}
+ conn-fields:
+ model:
+ label: Model
+ description: "Bedrock model identifier (e.g. bedrock:us.anthropic.claude-opus-4-5)"
+ schema:
+ type:
+ - string
+ - 'null'
+ region_name:
+ label: AWS Region
+ description: "AWS region (e.g. us-east-1). Falls back to AWS_DEFAULT_REGION env var."
+ schema:
+ type:
+ - string
+ - 'null'
+ aws_access_key_id:
+ label: AWS Access Key ID
+ description: "IAM access key. Leave empty to use instance role / environment credential chain."
+ schema:
+ type:
+ - string
+ - 'null'
+ aws_secret_access_key:
+ label: AWS Secret Access Key
+ description: "IAM secret key."
+ schema:
+ type:
+ - string
+ - 'null'
+ aws_session_token:
+ label: AWS Session Token
+ description: "Temporary session token (optional)."
+ schema:
+ type:
+ - string
+ - 'null'
+ profile_name:
+ label: AWS Profile Name
+ description: "Named AWS credentials profile (optional)."
+ schema:
+ type:
+ - string
+ - 'null'
+ api_key:
+ label: Bearer Token
+ description: "AWS bearer token (alt. to IAM key/secret). Falls back to AWS_BEARER_TOKEN_BEDROCK."
+ schema:
+ type:
+ - string
+ - 'null'
+ base_url:
+ label: Custom Endpoint URL
+ description: "Override the Bedrock runtime endpoint URL (optional)."
+ schema:
+ type:
+ - string
+ - 'null'
+ aws_read_timeout:
+ label: Read Timeout (s)
+ description: "boto3 read timeout in seconds (float, optional)."
+ schema:
+ type:
+ - number
+ - 'null'
+ aws_connect_timeout:
+ label: Connect Timeout (s)
+ description: "boto3 connect timeout in seconds (float, optional)."
+ schema:
+ type:
+ - number
+ - 'null'
+ - hook-class-name: airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIVertexHook
+ connection-type: pydanticai-vertex
+ ui-field-behaviour:
+ hidden-fields:
+ - schema
+ - port
+ - login
+ - host
+ - password
+ relabeling: {}
+ placeholders: {}
+ conn-fields:
+ model:
+ label: Model
+ description: "Google model identifier (e.g. google-vertex:gemini-2.0-flash)"
+ schema:
+ type:
+ - string
+ - 'null'
+ project:
+ label: GCP Project
+ description: "Google Cloud project ID. Falls back to GOOGLE_CLOUD_PROJECT env var."
+ schema:
+ type:
+ - string
+ - 'null'
+ location:
+ label: Location / Region
+ description: "Vertex AI region (e.g. us-central1). Falls back to GOOGLE_CLOUD_LOCATION env var."
+ schema:
+ type:
+ - string
+ - 'null'
+ vertexai:
+ label: Force Vertex AI Mode
+ description: "Force Vertex AI mode. Auto-detected when project/location/credentials are set."
+ schema:
+ type:
+ - boolean
+ - 'null'
+ api_key:
+ label: API Key
+ description: "Google API key for Gen Language API or Vertex AI. Falls back to GOOGLE_API_KEY."
+ schema:
+ type:
+ - string
+ - 'null'
+ service_account_info:
+ label: Service Account Info
+ description: "Service account key as inline dict. Mutually exclusive with service_account_file."
Review Comment:
Nit: "Mutually exclusive with service_account_file" — we removed `service_account_file` a few rounds ago, so this reference is stale. Same in `get_provider_info.py`.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.
cetingokhan commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2901765798
##########
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:
I thought it would be better to pass this information to the class using `hook_params`, and I didn't want to add it later because any exceptions that might arise could be reflected within the class itself.
```
hook_params = {
"model_id": self.model_id,
}
return PydanticAIHook.get_hook(self.llm_conn_id, hook_params=hook_params)
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.
cetingokhan commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2901759000
##########
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:
I added a note because I thought changing the parameter names that Bedrock users are used to might cause confusion.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2908416181
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -19,33 +19,32 @@
from typing import TYPE_CHECKING, Any, TypeVar, overload
from pydantic_ai import Agent
-from pydantic_ai.models import Model, infer_model
-from pydantic_ai.providers import Provider, infer_provider, infer_provider_class
+from pydantic_ai.models import infer_model
+from pydantic_ai.providers import infer_provider, infer_provider_class
from airflow.providers.common.compat.sdk import BaseHook
OutputT = TypeVar("OutputT")
if TYPE_CHECKING:
- from pydantic_ai.models import KnownModelName
+ from pydantic_ai.models import KnownModelName, Model
class PydanticAIHook(BaseHook):
"""
Hook for LLM access via pydantic-ai.
- Manages connection credentials and model creation. Uses pydantic-ai's
- model inference to support any provider (OpenAI, Anthropic, Google,
- Bedrock, Ollama, vLLM, etc.).
+ Covers providers that use a standard ``api_key`` + optional ``base_url``
+ (OpenAI, Anthropic, Groq, Mistral, DeepSeek, Ollama, vLLM, …).
- 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)
+ For cloud providers with non-standard auth use the dedicated subclasses:
+ :class:`PydanticAIAzureHook`, :class:`PydanticAIBedrockHook`,
+ :class:`PydanticAIVertexHook`.
- Cloud providers (Bedrock, Vertex) that use native auth chains should leave
- password empty and configure environment-based auth (``AWS_PROFILE``,
- ``GOOGLE_APPLICATION_CREDENTIALS``).
+ Connection fields:
+ - **password**: API key
+ - **host**: Base URL (optional, e.g. ``https://api.openai.com/v1``)
+ - **extra** JSON: ``{"model": "openai:gpt-5.3"}``
Review Comment:
@gopidesupavan @cetingokhan Good catch on the naming question\! `model_id` is the right name for the Python parameter — it matches the operator interface (`LLMOperator(model_id="openai:gpt-5.3")`) and signals that this is a `provider:model` identifier string, not the model object itself.
The `model` key in `extra` JSON is the conn-field name (what shows up in the Airflow connection form). Having them slightly different actually helps — `model_id` is the programmatic override, `model` is the UI field. Both are well-established in the codebase already, so no rename needed here.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.
cetingokhan commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2914402659
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -157,13 +190,238 @@ 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 = "pydanticai_azure"
+ 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 = "pydanticai_bedrock"
+ 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"
+ ),
+ },
+ }
+
+ def _get_provider_kwargs(
+ self,
+ api_key: str | None,
+ base_url: str | None,
+ extra: dict[str, Any],
+ ) -> dict[str, Any]:
+ """
+ Return kwargs for ``BedrockProvider``.
+
+ .. note::
+ The ``api_key`` and ``base_url`` parameters (sourced from
+ ``conn.password`` and ``conn.host``) are intentionally ignored here.
+ Bedrock connections hide those fields in the UI; all config is
+ stored in ``extra`` instead. The ``api_key`` and ``base_url``
+ keys below refer to *extra* fields, not the method parameters.
+ """
+ _str_keys = (
+ "aws_access_key_id",
+ "aws_secret_access_key",
+ "aws_session_token",
+ "region_name",
+ "profile_name",
+ # Bearer-token auth (alternative to IAM key/secret).
+ # Maps to AWS_BEARER_TOKEN_BEDROCK env var.
+ "api_key",
+ # Custom Bedrock runtime endpoint.
+ "base_url",
+ )
+ kwargs: dict[str, Any] = {k: extra[k] for k in _str_keys if extra.get(k) is not None}
+ # BedrockProvider expects float for timeout values; JSON integers must be coerced.
+ for _timeout_key in ("aws_read_timeout", "aws_connect_timeout"):
+ if extra.get(_timeout_key) is not None:
+ kwargs[_timeout_key] = float(extra[_timeout_key])
+ return kwargs
+
+
+class PydanticAIVertexHook(PydanticAIHook):
+ """
+ Hook for Google Vertex AI (or Generative Language API) via pydantic-ai.
+
+ Credentials are resolved in order:
+
+ 1. ``service_account_info`` (JSON object) in ``extra``
+ — loaded into a ``google.auth.credentials.Credentials``
+ object and passed as ``credentials`` to ``GoogleProvider``.
+ 2. ``api_key`` in ``extra`` — for Generative Language API (non-Vertex) or
+ Vertex API-key auth.
+ 3. Application Default Credentials (``GOOGLE_APPLICATION_CREDENTIALS``,
+ ``gcloud auth application-default login``, Workload Identity, …) when
+ no explicit credentials are provided.
+
+ Connection fields:
+ - **extra** JSON::
+
+ {
+ "model": "google-vertex:gemini-2.0-flash",
+ "project": "my-gcp-project",
+ "location": "us-central1",
+ "service_account_info": {...},
+ "vertexai": true,
+ }
+
+ Use ``"service_account_info"`` to embed the service-account JSON directly
+ (as an object, not a string path).
+
+ Set ``"vertexai": true`` to force Vertex AI mode when only ``api_key`` is
+ provided. Omit ``vertexai`` for the Generative Language API (GLA).
+
+ :param llm_conn_id: Airflow connection ID.
+ :param model_id: Model identifier, e.g. ``"google-vertex:gemini-2.0-flash"``.
+ """
+
+ conn_type = "pydanticai_vertex"
+ default_conn_name = "pydanticai_vertex_default"
+ hook_name = "Pydantic AI (Google Vertex AI)"
+
+ @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": "google-vertex:gemini-2.0-flash", '
+ '"project": "my-project", "location": "us-central1", "vertexai": true}'
+ " — add service_account_info (object) for SA auth;"
+ " omit both to use Application Default Credentials"
+ ),
+ },
+ }
+
+ def _get_provider_kwargs(
+ self,
+ api_key: str | None,
+ base_url: str | None,
+ extra: dict[str, Any],
+ ) -> dict[str, Any]:
+ sa_info = extra.get("service_account_info")
+ kwargs: dict[str, Any] = {}
+
+ # Direct GoogleProvider scalar kwargs.
+ for _key in ("api_key", "project", "location", "base_url"):
+ if extra.get(_key) is not None:
Review Comment:
done
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "gopidesupavan (via GitHub)" <gi...@apache.org>.
gopidesupavan commented on PR #62816:
URL: https://github.com/apache/airflow/pull/62816#issuecomment-4042787373
> Hi @kaxil and @gopidesupavan ,
>
> I would like to sincerely thank you for the time and patience you've invested in reviewing my PRs. I feel a bit embarrassed knowing that I’ve taken up so much of your time with my numerous mistakes, especially when I know you could have implemented these features much faster yourselves. I truly appreciate the guidance you've provided despite the extra work I've caused.
>
> These reviews are proving to be incredibly valuable for me in understanding the "Airflow Gurus" perspective on architecture. I believe that as I adopt this mindset, my future PRs will require much less back-and-forth, though I know this takes time and experience.
>
> Regarding my other PR (LLMDataQualityOperator #62963 ), I am fully committed to it. However, if you have a strict deadline or if you feel the review process will be too taxing for your current schedule, I would completely understand if you prefer to put it on hold. That said, if you are willing to bear with me, I am more than eager to make all the necessary improvements and see it through to completion!
no worries @cetingokhan :) we are all here learning from each other ..
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.
cetingokhan commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2895825191
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -19,42 +19,41 @@
from typing import TYPE_CHECKING, Any, TypeVar, overload
from pydantic_ai import Agent
-from pydantic_ai.models import Model, infer_model
-from pydantic_ai.providers import Provider, infer_provider, infer_provider_class
+from pydantic_ai.models import infer_model
+from pydantic_ai.providers import infer_provider, infer_provider_class
from airflow.providers.common.compat.sdk import BaseHook
OutputT = TypeVar("OutputT")
if TYPE_CHECKING:
- from pydantic_ai.models import KnownModelName
+ from pydantic_ai.models import KnownModelName, Model
class PydanticAIHook(BaseHook):
"""
Hook for LLM access via pydantic-ai.
- Manages connection credentials and model creation. Uses pydantic-ai's
- model inference to support any provider (OpenAI, Anthropic, Google,
- Bedrock, Ollama, vLLM, etc.).
+ Covers providers that use a standard ``api_key`` + optional ``base_url``
+ (OpenAI, Anthropic, Groq, Mistral, DeepSeek, Ollama, vLLM, …).
- 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)
+ For cloud providers with non-standard auth use the dedicated subclasses:
+ :class:`PydanticAIAzureHook`, :class:`PydanticAIBedrockHook`,
+ :class:`PydanticAIVertexHook`.
- Cloud providers (Bedrock, Vertex) that use native auth chains should leave
- password empty and configure environment-based auth (``AWS_PROFILE``,
- ``GOOGLE_APPLICATION_CREDENTIALS``).
+ Connection fields:
+ - **password**: API key
+ - **host**: Base URL (optional, e.g. ``https://api.openai.com/v1``)
+ - **extra** JSON: ``{"model": "openai:gpt-5.3"}``
:param llm_conn_id: Airflow connection ID for the LLM provider.
:param model_id: Model identifier in ``provider:model`` format (e.g. ``"openai:gpt-5.3"``).
Overrides the model stored in the connection's extra field.
"""
conn_name_attr = "llm_conn_id"
- default_conn_name = "pydanticai_default"
- conn_type = "pydanticai"
+ default_conn_name = "pydantic_ai_default"
+ conn_type = "pydantic_ai"
Review Comment:
It mixed my previous commits which before your #62817 pr(I think I didn't merge correctly). I changed other hooks but I forgot this :) I'm rolling back right now ;)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2892259289
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -75,62 +74,104 @@ 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."
+ "No model specified. Set model_id on the hook or 'model' in the connection's extra JSON."
+ )
Review Comment:
This calls a subclass's unbound method with a `PydanticAIHook` instance as `self`:
```python
hook_cls._get_provider_kwargs(self, api_key, base_url, extra)
```
It works today because `_get_provider_kwargs` only touches `self.log` (inherited from BaseHook). But if any subclass later accesses subclass-specific attributes, this breaks silently.
A cleaner alternative: have the operator use `for_connection()` instead of constructing `PydanticAIHook` directly. That returns the right subclass, and no auto-dispatch is needed:
```python
@cached_property
def llm_hook(self) -> PydanticAIHook:
return PydanticAIHook.for_connection(self.llm_conn_id, model_id=self.model_id)
```
This eliminates the `type(self) is PydanticAIHook` check, the `_CONN_TYPE_TO_HOOK` registry, and the unbound method call.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -157,13 +198,213 @@ 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)
+
+ @classmethod
+ def for_connection(cls, conn_id: str, model_id: str | None = None) -> PydanticAIHook:
+ """
+ Return the correct :class:`PydanticAIHook` subclass for *conn_id*.
+
+ Looks up the connection's ``conn_type`` in the registered hook map and
+ instantiates the matching subclass. Falls back to
+ :class:`PydanticAIHook` for unknown types.
+
+ :param conn_id: Airflow connection ID.
+ :param model_id: Optional model override forwarded to the hook.
+ """
+ conn = cls.get_connection(conn_id)
+ hook_cls = _CONN_TYPE_TO_HOOK.get(conn.conn_type or "", cls)
+ return hook_cls(llm_conn_id=conn_id, model_id=model_id)
+
+
+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_name_attr = "llm_conn_id"
+ default_conn_name = "pydantic_ai_azure_default"
+ conn_type = "pydanticai_azure"
+ 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"]
+
+ self.log.info(
+ "Using explicit credentials for Azure provider with model '%s': %s",
+ extra.get("model", ""),
+ list(kwargs),
+ )
+ return kwargs
+
+
+class PydanticAIBedrockHook(PydanticAIHook):
+ """
+ Hook for AWS Bedrock via pydantic-ai.
+
+ Credentials are resolved in order:
+
+ 1. Explicit keys in ``extra`` (``aws_access_key_id``,
+ ``aws_secret_access_key``, ``aws_session_token``).
+ 2. 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": "..."
+ }
+
+ Leave ``aws_access_key_id`` / ``aws_secret_access_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_name_attr = "llm_conn_id"
+ default_conn_name = "pydantic_ai_bedrock_default"
+ conn_type = "pydanticai_bedrock"
+ 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"
+ ),
+ },
+ }
+
+ def _get_provider_kwargs(
+ self,
+ api_key: str | None,
+ base_url: str | None,
+ extra: dict[str, Any],
+ ) -> dict[str, Any]:
+ _bedrock_keys = (
+ "aws_access_key_id",
+ "aws_secret_access_key",
+ "aws_session_token",
+ "region_name",
+ "profile_name",
+ )
+ return {k: extra[k] for k in _bedrock_keys if extra.get(k) is not None}
+
+
+class PydanticAIVertexHook(PydanticAIHook):
+ """
+ Hook for Google Vertex AI via pydantic-ai.
+
+ Credentials are resolved in order:
+
+ 1. ``service_account_file`` / ``service_account_info`` in ``extra``.
+ 2. Application Default Credentials (``GOOGLE_APPLICATION_CREDENTIALS``,
+ ``gcloud auth application-default login``, …) when no explicit keys
+ are provided.
+
+ Connection fields:
+ - **extra** JSON::
+
+ {
+ "model": "google-vertex:gemini-2.0-flash",
+ "project_id": "my-gcp-project",
+ "location": "us-central1",
+ "service_account_file": "/path/to/sa.json",
+ }
+
Review Comment:
These kwarg names don't match what pydantic-ai's `GoogleProvider` accepts (see [docs](https://ai.pydantic.dev/models/google/)):
| Connection extra | pydantic-ai param | Match? |
|---|---|---|
| `project_id` | `project` | **No** |
| `location` | `location` | Yes |
| `service_account_file` | `credentials` (object) | **No** |
| `service_account_info` | `credentials` (object) | **No** |
`project_id` needs to be mapped to `project`. `service_account_file`/`service_account_info` need to be loaded into a `google.auth.credentials.Credentials` object before passing — the provider doesn't accept file paths directly.
Because of the `TypeError` catch in `_provider_factory`, these mismatches fail silently — the credentials get ignored and it falls back to ADC. Users who explicitly configure credentials in the connection would have no idea they're not being used.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -157,13 +198,213 @@ 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)
+
+ @classmethod
+ def for_connection(cls, conn_id: str, model_id: str | None = None) -> PydanticAIHook:
+ """
+ Return the correct :class:`PydanticAIHook` subclass for *conn_id*.
+
+ Looks up the connection's ``conn_type`` in the registered hook map and
+ instantiates the matching subclass. Falls back to
+ :class:`PydanticAIHook` for unknown types.
+
+ :param conn_id: Airflow connection ID.
+ :param model_id: Optional model override forwarded to the hook.
+ """
+ conn = cls.get_connection(conn_id)
+ hook_cls = _CONN_TYPE_TO_HOOK.get(conn.conn_type or "", cls)
+ return hook_cls(llm_conn_id=conn_id, model_id=model_id)
+
+
+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:
Python evaluates default argument values at function/method definition time. `PydanticAIHook.__init__` binds `llm_conn_id`'s default to `"pydanticai_default"` when the class is created. Subclasses inherit this `__init__` with the already-resolved default, so:
```python
hook = PydanticAIAzureHook() # llm_conn_id="pydanticai_default", not "pydantic_ai_azure_default"
```
The `default_conn_name` attribute on the subclass is never consulted during `__init__`. Same issue for Bedrock and Vertex subclasses.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -75,62 +74,104 @@ 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."
+ "No model specified. Set model_id on the hook or 'model' in the connection's extra JSON."
+ )
+
+ api_key: str | None = conn.password or None
+ base_url: str | None = conn.host or None
+
+ # Auto-dispatch: if using base hook with a subclass connection type, borrow
+ # that subclass's _get_provider_kwargs so the correct field mapping is used.
+ if type(self) is PydanticAIHook:
+ hook_cls = _CONN_TYPE_TO_HOOK.get(conn.conn_type or "", PydanticAIHook)
+ provider_kwargs = hook_cls._get_provider_kwargs(self, api_key, base_url, extra)
+ else:
+ provider_kwargs = self._get_provider_kwargs(api_key, base_url, extra)
+ if provider_kwargs:
+ _kwargs = provider_kwargs # capture for closure
+ self.log.info(
Review Comment:
`TypeError` here catches both "provider doesn't accept these kwargs" and genuine bugs (wrong types, missing required args). The warning only logs kwarg *names*, not the actual error. This makes debugging auth failures harder — a user who passes the wrong kwarg name gets silently downgraded to env-var auth with a vague warning.
Suggestion: log the exception too:
```python
except TypeError as exc:
self.log.warning(
"Provider '%s' does not accept kwargs %s (%s); "
"falling back to env-var auth.",
pname, list(_kwargs), exc,
)
```
##########
providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py:
##########
@@ -105,25 +100,13 @@ def get_provider_info():
{
"integration-name": "Common AI",
"python-modules": [
- "airflow.providers.common.ai.operators.agent",
"airflow.providers.common.ai.operators.llm",
Review Comment:
This file is auto-generated from `provider.yaml` — manual edits here will be overwritten by `prek run update-providers-build-files`. The new connection types need to be added to `provider.yaml` instead.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -157,13 +198,213 @@ 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)
+
+ @classmethod
+ def for_connection(cls, conn_id: str, model_id: str | None = None) -> PydanticAIHook:
+ """
+ Return the correct :class:`PydanticAIHook` subclass for *conn_id*.
+
+ Looks up the connection's ``conn_type`` in the registered hook map and
+ instantiates the matching subclass. Falls back to
+ :class:`PydanticAIHook` for unknown types.
+
+ :param conn_id: Airflow connection ID.
+ :param model_id: Optional model override forwarded to the hook.
+ """
+ conn = cls.get_connection(conn_id)
+ hook_cls = _CONN_TYPE_TO_HOOK.get(conn.conn_type or "", cls)
+ return hook_cls(llm_conn_id=conn_id, model_id=model_id)
+
+
+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_name_attr = "llm_conn_id"
+ default_conn_name = "pydantic_ai_azure_default"
+ conn_type = "pydanticai_azure"
+ 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"):
Review Comment:
This `self.log.info(...)` duplicates the log in `get_conn()` (line ~144) which also logs when `provider_kwargs` is non-empty. Azure connections will produce two "Using explicit credentials..." log lines; Bedrock and Vertex produce one. Remove this — let the base class handle logging uniformly.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.
cetingokhan commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2900080636
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -157,13 +194,266 @@ 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)
+
+ @classmethod
+ def for_connection(cls, conn_id: str, model_id: str | None = None) -> PydanticAIHook:
+ """
+ Return the correct :class:`PydanticAIHook` subclass for *conn_id*.
+
+ Looks up the connection's ``conn_type`` in the registered hook map and
+ instantiates the matching subclass. Falls back to
+ :class:`PydanticAIHook` for unknown types.
+
+ :param conn_id: Airflow connection ID.
+ :param model_id: Optional model override forwarded to the hook.
+ """
+ conn = cls.get_connection(conn_id)
+ hook_cls = _CONN_TYPE_TO_HOOK.get(conn.conn_type or "", cls)
+ return hook_cls(llm_conn_id=conn_id, model_id=model_id)
+
+
+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 = "pydanticai_azure"
+ 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 = "pydanticai_bedrock"
+ 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"
+ ),
+ },
+ }
+
+ def _get_provider_kwargs(
+ self,
+ api_key: str | None,
+ base_url: str | None,
+ extra: dict[str, Any],
+ ) -> dict[str, Any]:
+ _str_keys = (
+ "aws_access_key_id",
+ "aws_secret_access_key",
+ "aws_session_token",
+ "region_name",
+ "profile_name",
+ # Bearer-token auth (alternative to IAM key/secret).
+ # Maps to AWS_BEARER_TOKEN_BEDROCK env var.
+ "api_key",
+ # Custom Bedrock runtime endpoint.
+ "base_url",
+ )
+ kwargs: dict[str, Any] = {k: extra[k] for k in _str_keys if extra.get(k) is not None}
+ # BedrockProvider expects float for timeout values; JSON integers must be coerced.
+ for _timeout_key in ("aws_read_timeout", "aws_connect_timeout"):
+ if extra.get(_timeout_key) is not None:
+ kwargs[_timeout_key] = float(extra[_timeout_key])
+ return kwargs
+
+
+class PydanticAIVertexHook(PydanticAIHook):
+ """
+ Hook for Google Vertex AI (or Generative Language API) via pydantic-ai.
+
+ Credentials are resolved in order:
+
+ 1. ``service_account_file`` (path string) or ``service_account_info`` (JSON
+ object) in ``extra`` — loaded into a ``google.auth.credentials.Credentials``
+ object and passed as ``credentials`` to ``GoogleProvider``.
+ 2. ``api_key`` in ``extra`` — for Generative Language API (non-Vertex) or
+ Vertex API-key auth.
+ 3. Application Default Credentials (``GOOGLE_APPLICATION_CREDENTIALS``,
+ ``gcloud auth application-default login``, Workload Identity, …) when
+ no explicit credentials are provided.
+
+ Connection fields:
+ - **extra** JSON::
+
+ {
+ "model": "google-vertex:gemini-2.0-flash",
+ "project": "my-gcp-project",
+ "location": "us-central1",
+ "service_account_file": "/path/to/sa.json",
+ "vertexai": true,
+ }
+
+ Use ``"service_account_info"`` instead of ``"service_account_file"`` to
+ embed the service-account JSON directly (as an object, not a string path).
+ Setting both at the same time raises ``ValueError``.
+
+ Set ``"vertexai": true`` to force Vertex AI mode when only ``api_key`` is
+ provided. Omit ``vertexai`` for the Generative Language API (GLA).
+
+ :param llm_conn_id: Airflow connection ID.
+ :param model_id: Model identifier, e.g. ``"google-vertex:gemini-2.0-flash"``.
+ """
+
+ conn_type = "pydanticai_vertex"
+ hook_name = "Pydantic AI (Google Vertex AI)"
+
+ @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": "google-vertex:gemini-2.0-flash", '
+ '"project": "my-project", "location": "us-central1", "vertexai": true}'
+ " — add service_account_file (path) or service_account_info (object) for SA auth;"
+ " omit both to use Application Default Credentials"
+ ),
+ },
+ }
+
+ def _get_provider_kwargs(
+ self,
+ api_key: str | None,
+ base_url: str | None,
+ extra: dict[str, Any],
+ ) -> dict[str, Any]:
+ sa_file = extra.get("service_account_file")
+ sa_info = extra.get("service_account_info")
+ if sa_file and sa_info:
+ raise ValueError(
+ "Specify 'service_account_file' or 'service_account_info' in the connection extra, not both."
+ )
+
+ kwargs: dict[str, Any] = {}
+
+ # Direct GoogleProvider scalar kwargs.
+ for _key in ("api_key", "project", "location", "base_url"):
+ if extra.get(_key) is not None:
+ kwargs[_key] = extra[_key]
+
+ # Optional vertexai bool flag (force Vertex AI mode for API-key auth).
+ _vertexai = extra.get("vertexai")
+ if _vertexai is not None:
+ kwargs["vertexai"] = bool(_vertexai)
+
+ # Service-account credentials — loaded lazily to avoid importing
+ # google-auth on non-Vertex code paths (optional heavy dependency).
+ if sa_file:
+ from google.oauth2 import service_account # lazy: optional dep
+
+ kwargs["credentials"] = service_account.Credentials.from_service_account_file(
+ sa_file,
+ scopes=["https://www.googleapis.com/auth/cloud-platform"],
+ )
+ elif sa_info:
+ from google.oauth2 import service_account # lazy: optional dep
+
+ kwargs["credentials"] = service_account.Credentials.from_service_account_info(
+ sa_info,
+ scopes=["https://www.googleapis.com/auth/cloud-platform"],
+ )
+
+ return kwargs
+
+
+# ---------------------------------------------------------------------------
+# Hook registry — maps conn_type → hook class for use by for_connection()
+# ---------------------------------------------------------------------------
+_CONN_TYPE_TO_HOOK: dict[str, type[PydanticAIHook]] = {
Review Comment:
During development, I hadn't defined the hooks in the provider.yaml file, so the tests couldn't discover them. That's why I ended up adding the for_connection() method as a workaround.
I didn't know ProviderManager handled it automatically under the hood like that—thanks for the heads-up! :)
I've updated it as you suggested: adding them to the yaml file made _CONN_TYPE_TO_HOOK and for_connection() completely unnecessary, so I've removed them. ;)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "gopidesupavan (via GitHub)" <gi...@apache.org>.
gopidesupavan commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2902457819
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -19,33 +19,32 @@
from typing import TYPE_CHECKING, Any, TypeVar, overload
from pydantic_ai import Agent
-from pydantic_ai.models import Model, infer_model
-from pydantic_ai.providers import Provider, infer_provider, infer_provider_class
+from pydantic_ai.models import infer_model
+from pydantic_ai.providers import infer_provider, infer_provider_class
from airflow.providers.common.compat.sdk import BaseHook
OutputT = TypeVar("OutputT")
if TYPE_CHECKING:
- from pydantic_ai.models import KnownModelName
+ from pydantic_ai.models import KnownModelName, Model
class PydanticAIHook(BaseHook):
"""
Hook for LLM access via pydantic-ai.
- Manages connection credentials and model creation. Uses pydantic-ai's
- model inference to support any provider (OpenAI, Anthropic, Google,
- Bedrock, Ollama, vLLM, etc.).
+ Covers providers that use a standard ``api_key`` + optional ``base_url``
+ (OpenAI, Anthropic, Groq, Mistral, DeepSeek, Ollama, vLLM, …).
- 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)
+ For cloud providers with non-standard auth use the dedicated subclasses:
+ :class:`PydanticAIAzureHook`, :class:`PydanticAIBedrockHook`,
+ :class:`PydanticAIVertexHook`.
- Cloud providers (Bedrock, Vertex) that use native auth chains should leave
- password empty and configure environment-based auth (``AWS_PROFILE``,
- ``GOOGLE_APPLICATION_CREDENTIALS``).
+ Connection fields:
+ - **password**: API key
+ - **host**: Base URL (optional, e.g. ``https://api.openai.com/v1``)
+ - **extra** JSON: ``{"model": "openai:gpt-5.3"}``
Review Comment:
please make this consistent , operators uses `model_id` as param
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.
cetingokhan commented on PR #62816:
URL: https://github.com/apache/airflow/pull/62816#issuecomment-4000464220
Hi @kaxil
I have taken your "thin hook" suggestion into consideration. However, I believe that relying solely on environment variables might not be sufficient for Airflow users; being able to fetch configurations directly from connection extras is highly valuable.
I've reviewed all the current pydanticai providers and noticed that only AzureProvider and LiteLLMProvider require their own specific parameters instead of the standard base_url.
To handle this distinction cleanly without over-engineering or building registries, I am planning to introduce a simple _build_provider_kwargs(conn) helper method. This method will parse the connection and return the appropriate kwargs dictionary to feed directly into the resolved provider_cls.
What are your thoughts on proceeding with this approach? Does this align better with the design goals?
```@staticmethod
def _build_provider_kwargs(
provider_cls: type,
api_key: str | None,
base_url: str | None,
extra: dict[str, Any],
) -> dict[str, Any]:
sig = inspect.signature(provider_cls.__init__)
accepted: set[str] = {
name
for name, param in sig.parameters.items()
if name != "self"
and param.kind
not in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD)
}
candidates: dict[str, Any] = {
"api_key": api_key,
"base_url": base_url,
# Alias used by AzureProvider instead of base_url
"azure_endpoint": base_url,
# Alias used by LiteLLMProvider instead of base_url
"api_base": base_url,
**{k: v for k, v in extra.items() if k != "model"},
}
return {k: v for k, v in candidates.items() if k in accepted and v is not None}
def get_conn(self) -> Model:
.....
if api_key or base_url:
provider_name = model_name.split(":")[0] if ":" in model_name else "openai"
provider_cls = infer_provider_class(provider_name)
kwargs: dict[str, Any] = self._build_provider_kwargs(provider_cls, api_key, base_url, extra)
try:
provider = provider_cls(**kwargs)
except TypeError:
provider = infer_provider(provider_name)
....
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on PR #62816:
URL: https://github.com/apache/airflow/pull/62816#issuecomment-4000519830
Thanks for the follow-up! I think there's a misunderstanding — I'm not suggesting env-vars-only. The hook already reads credentials from the Airflow Connection (`conn.password` → `api_key`, `conn.host` → `base_url`). The env-var path only kicks in when no credentials are set on the connection (for Bedrock/Vertex users who rely on IAM auth).
What I'm suggesting is: instead of introspecting provider constructors with `inspect.signature`, just add explicit branches for the providers that need special handling. Right now that's only Azure and LiteLLM. For example, something along these lines in `_provider_factory`:
```python
provider_cls = infer_provider_class(provider_name)
if provider_name == "azure":
return provider_cls(azure_endpoint=base_url, api_key=api_key)
if provider_name == "litellm":
return provider_cls(api_base=base_url, api_key=api_key)
# Default path — works for OpenAI, Anthropic, Groq, Mistral, etc.
kwargs = {}
if api_key:
kwargs["api_key"] = api_key
if base_url:
kwargs["base_url"] = base_url
return provider_cls(**kwargs)
```
This is explicit, easy to read, and doesn't break if pydantic-ai renames a constructor parameter. The `inspect.signature` approach is clever but fragile — if a provider adds or renames a param, the introspection silently changes behavior. Explicit branches make the mapping obvious and reviewable.
For connection extras: I'm fine with passing extra fields through to the provider constructor for the specific branches that need them (e.g. `api_version` for Azure). But that should be per-provider, not a generic "dump all extras into kwargs" pattern.
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2896204369
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -157,13 +194,266 @@ 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)
+
+ @classmethod
+ def for_connection(cls, conn_id: str, model_id: str | None = None) -> PydanticAIHook:
+ """
+ Return the correct :class:`PydanticAIHook` subclass for *conn_id*.
+
+ Looks up the connection's ``conn_type`` in the registered hook map and
+ instantiates the matching subclass. Falls back to
+ :class:`PydanticAIHook` for unknown types.
+
+ :param conn_id: Airflow connection ID.
+ :param model_id: Optional model override forwarded to the hook.
+ """
+ conn = cls.get_connection(conn_id)
+ hook_cls = _CONN_TYPE_TO_HOOK.get(conn.conn_type or "", cls)
+ return hook_cls(llm_conn_id=conn_id, model_id=model_id)
+
+
+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 = "pydanticai_azure"
+ 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 = "pydanticai_bedrock"
+ 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"
+ ),
+ },
+ }
+
+ def _get_provider_kwargs(
+ self,
+ api_key: str | None,
+ base_url: str | None,
+ extra: dict[str, Any],
+ ) -> dict[str, Any]:
+ _str_keys = (
+ "aws_access_key_id",
+ "aws_secret_access_key",
+ "aws_session_token",
+ "region_name",
+ "profile_name",
+ # Bearer-token auth (alternative to IAM key/secret).
+ # Maps to AWS_BEARER_TOKEN_BEDROCK env var.
+ "api_key",
+ # Custom Bedrock runtime endpoint.
+ "base_url",
+ )
+ kwargs: dict[str, Any] = {k: extra[k] for k in _str_keys if extra.get(k) is not None}
+ # BedrockProvider expects float for timeout values; JSON integers must be coerced.
+ for _timeout_key in ("aws_read_timeout", "aws_connect_timeout"):
+ if extra.get(_timeout_key) is not None:
+ kwargs[_timeout_key] = float(extra[_timeout_key])
+ return kwargs
+
+
+class PydanticAIVertexHook(PydanticAIHook):
+ """
+ Hook for Google Vertex AI (or Generative Language API) via pydantic-ai.
+
+ Credentials are resolved in order:
+
+ 1. ``service_account_file`` (path string) or ``service_account_info`` (JSON
+ object) in ``extra`` — loaded into a ``google.auth.credentials.Credentials``
+ object and passed as ``credentials`` to ``GoogleProvider``.
+ 2. ``api_key`` in ``extra`` — for Generative Language API (non-Vertex) or
+ Vertex API-key auth.
+ 3. Application Default Credentials (``GOOGLE_APPLICATION_CREDENTIALS``,
+ ``gcloud auth application-default login``, Workload Identity, …) when
+ no explicit credentials are provided.
+
+ Connection fields:
+ - **extra** JSON::
+
+ {
+ "model": "google-vertex:gemini-2.0-flash",
+ "project": "my-gcp-project",
+ "location": "us-central1",
+ "service_account_file": "/path/to/sa.json",
+ "vertexai": true,
+ }
+
+ Use ``"service_account_info"`` instead of ``"service_account_file"`` to
+ embed the service-account JSON directly (as an object, not a string path).
+ Setting both at the same time raises ``ValueError``.
+
+ Set ``"vertexai": true`` to force Vertex AI mode when only ``api_key`` is
+ provided. Omit ``vertexai`` for the Generative Language API (GLA).
+
+ :param llm_conn_id: Airflow connection ID.
+ :param model_id: Model identifier, e.g. ``"google-vertex:gemini-2.0-flash"``.
+ """
+
+ conn_type = "pydanticai_vertex"
+ hook_name = "Pydantic AI (Google Vertex AI)"
+
+ @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": "google-vertex:gemini-2.0-flash", '
+ '"project": "my-project", "location": "us-central1", "vertexai": true}'
+ " — add service_account_file (path) or service_account_info (object) for SA auth;"
+ " omit both to use Application Default Credentials"
+ ),
+ },
+ }
+
+ def _get_provider_kwargs(
+ self,
+ api_key: str | None,
+ base_url: str | None,
+ extra: dict[str, Any],
+ ) -> dict[str, Any]:
+ sa_file = extra.get("service_account_file")
+ sa_info = extra.get("service_account_info")
+ if sa_file and sa_info:
+ raise ValueError(
+ "Specify 'service_account_file' or 'service_account_info' in the connection extra, not both."
+ )
+
+ kwargs: dict[str, Any] = {}
+
+ # Direct GoogleProvider scalar kwargs.
+ for _key in ("api_key", "project", "location", "base_url"):
+ if extra.get(_key) is not None:
+ kwargs[_key] = extra[_key]
+
+ # Optional vertexai bool flag (force Vertex AI mode for API-key auth).
+ _vertexai = extra.get("vertexai")
+ if _vertexai is not None:
+ kwargs["vertexai"] = bool(_vertexai)
+
+ # Service-account credentials — loaded lazily to avoid importing
+ # google-auth on non-Vertex code paths (optional heavy dependency).
+ if sa_file:
+ from google.oauth2 import service_account # lazy: optional dep
+
+ kwargs["credentials"] = service_account.Credentials.from_service_account_file(
+ sa_file,
+ scopes=["https://www.googleapis.com/auth/cloud-platform"],
+ )
+ elif sa_info:
+ from google.oauth2 import service_account # lazy: optional dep
+
+ kwargs["credentials"] = service_account.Credentials.from_service_account_info(
+ sa_info,
+ scopes=["https://www.googleapis.com/auth/cloud-platform"],
+ )
+
+ return kwargs
+
+
+# ---------------------------------------------------------------------------
+# Hook registry — maps conn_type → hook class for use by for_connection()
+# ---------------------------------------------------------------------------
+_CONN_TYPE_TO_HOOK: dict[str, type[PydanticAIHook]] = {
Review Comment:
`_CONN_TYPE_TO_HOOK` + `for_connection()` duplicates what Airflow's `ProvidersManager` already does. The provider manager reads `provider.yaml` and maps `conn_type -> hook class` automatically. The operator can just do:
```python
@cached_property
def llm_hook(self) -> PydanticAIHook:
return PydanticAIHook.get_hook(conn_id=self.llm_conn_id)
```
Or use `ProvidersManager().hooks` if you need more control. The manual registry will drift when someone adds a new subclass and forgets to update this dict.
Remove `_CONN_TYPE_TO_HOOK`, `for_connection()`, and let the framework handle the dispatch.
##########
providers/common/ai/tests/unit/common/ai/hooks/test_pydantic_ai.py:
##########
@@ -16,13 +16,20 @@
# under the License.
from __future__ import annotations
+import json
+import sys
from unittest.mock import MagicMock, patch
import pytest
from pydantic_ai.models import Model
from airflow.models.connection import Connection
-from airflow.providers.common.ai.hooks.pydantic_ai import PydanticAIHook
+from airflow.providers.common.ai.hooks.pydantic_ai import (
+ PydanticAIAzureHook,
+ PydanticAIBedrockHook,
+ PydanticAIHook,
Review Comment:
The `test_azure_hook_uses_own_default_conn_name` test asserts `hook.llm_conn_id == "pydanticai_default"`. That means `PydanticAIAzureHook()` with no args will try to fetch a connection of type `pydanticai` (the generic one), not `pydanticai_azure`. The conn_type won't match.
Consider giving each subclass its own `default_conn_name` (e.g. `pydanticai_azure_default`), or at least make the test name/comment reflect that this is intentional sharing, not "uses own" default.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -75,62 +75,99 @@ 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:
+ try:
+ return infer_provider_class(pname)(**_kwargs)
+ except TypeError as exc:
Review Comment:
This `except TypeError` still swallows real bugs (flagged in the previous round). It catches both "provider doesn't accept these kwargs" and genuine errors like wrong types or missing required args. The warning only logs kwarg *names*, not the actual error.
With the subclass approach, each hook maps kwargs explicitly for its provider. A `TypeError` from the provider constructor is almost certainly a real bug now, not an expected mismatch. Silently falling back to env-var auth hides misconfiguration from users.
For the base class (OpenAI/Anthropic/Groq), `api_key` and `base_url` are standard params that all these providers accept. The fallback was a workaround for the old generic approach where you didn't know what the provider accepted. That's no longer the case.
Remove the try/except and let the TypeError propagate. If a user misconfigures their connection, they should see a clear error, not a silent downgrade.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -157,13 +194,266 @@ 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)
+
+ @classmethod
+ def for_connection(cls, conn_id: str, model_id: str | None = None) -> PydanticAIHook:
+ """
+ Return the correct :class:`PydanticAIHook` subclass for *conn_id*.
+
+ Looks up the connection's ``conn_type`` in the registered hook map and
+ instantiates the matching subclass. Falls back to
+ :class:`PydanticAIHook` for unknown types.
+
+ :param conn_id: Airflow connection ID.
+ :param model_id: Optional model override forwarded to the hook.
+ """
+ conn = cls.get_connection(conn_id)
+ hook_cls = _CONN_TYPE_TO_HOOK.get(conn.conn_type or "", cls)
+ return hook_cls(llm_conn_id=conn_id, model_id=model_id)
+
+
+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 = "pydanticai_azure"
+ 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 = "pydanticai_bedrock"
+ 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"
+ ),
+ },
+ }
+
+ def _get_provider_kwargs(
Review Comment:
Bedrock's `_get_provider_kwargs` reads `api_key` and `base_url` from `extra` (lines 338-340), but the base class `get_conn()` also passes `conn.password` as `api_key` and `conn.host` as `base_url` to this method. Those two arguments are silently ignored here.
This works correctly since Bedrock hides `host`/`password` in the UI, but it's confusing for anyone reading the code. The method signature says it accepts `api_key` and `base_url`, then throws them away.
Two options:
1. Add a brief comment at the top: `# Bedrock reads all config from extra; api_key/base_url from conn.password/host are unused.`
2. Or read `api_key`/`base_url` from the positional args too (mapping them to the Bedrock equivalents), so the standard connection fields still work even though the UI hides them.
Option 1 is fine for now.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -157,13 +194,266 @@ 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)
+
+ @classmethod
+ def for_connection(cls, conn_id: str, model_id: str | None = None) -> PydanticAIHook:
+ """
+ Return the correct :class:`PydanticAIHook` subclass for *conn_id*.
+
+ Looks up the connection's ``conn_type`` in the registered hook map and
+ instantiates the matching subclass. Falls back to
+ :class:`PydanticAIHook` for unknown types.
+
+ :param conn_id: Airflow connection ID.
+ :param model_id: Optional model override forwarded to the hook.
+ """
+ conn = cls.get_connection(conn_id)
+ hook_cls = _CONN_TYPE_TO_HOOK.get(conn.conn_type or "", cls)
+ return hook_cls(llm_conn_id=conn_id, model_id=model_id)
+
+
+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 = "pydanticai_azure"
+ 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 = "pydanticai_bedrock"
+ 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"
+ ),
+ },
+ }
+
+ def _get_provider_kwargs(
+ self,
+ api_key: str | None,
+ base_url: str | None,
+ extra: dict[str, Any],
+ ) -> dict[str, Any]:
+ _str_keys = (
+ "aws_access_key_id",
+ "aws_secret_access_key",
+ "aws_session_token",
+ "region_name",
+ "profile_name",
+ # Bearer-token auth (alternative to IAM key/secret).
+ # Maps to AWS_BEARER_TOKEN_BEDROCK env var.
+ "api_key",
+ # Custom Bedrock runtime endpoint.
+ "base_url",
+ )
+ kwargs: dict[str, Any] = {k: extra[k] for k in _str_keys if extra.get(k) is not None}
+ # BedrockProvider expects float for timeout values; JSON integers must be coerced.
+ for _timeout_key in ("aws_read_timeout", "aws_connect_timeout"):
+ if extra.get(_timeout_key) is not None:
+ kwargs[_timeout_key] = float(extra[_timeout_key])
+ return kwargs
+
+
+class PydanticAIVertexHook(PydanticAIHook):
+ """
+ Hook for Google Vertex AI (or Generative Language API) via pydantic-ai.
+
+ Credentials are resolved in order:
+
+ 1. ``service_account_file`` (path string) or ``service_account_info`` (JSON
+ object) in ``extra`` — loaded into a ``google.auth.credentials.Credentials``
+ object and passed as ``credentials`` to ``GoogleProvider``.
+ 2. ``api_key`` in ``extra`` — for Generative Language API (non-Vertex) or
+ Vertex API-key auth.
+ 3. Application Default Credentials (``GOOGLE_APPLICATION_CREDENTIALS``,
+ ``gcloud auth application-default login``, Workload Identity, …) when
+ no explicit credentials are provided.
+
+ Connection fields:
+ - **extra** JSON::
+
+ {
+ "model": "google-vertex:gemini-2.0-flash",
+ "project": "my-gcp-project",
+ "location": "us-central1",
+ "service_account_file": "/path/to/sa.json",
+ "vertexai": true,
+ }
+
+ Use ``"service_account_info"`` instead of ``"service_account_file"`` to
+ embed the service-account JSON directly (as an object, not a string path).
+ Setting both at the same time raises ``ValueError``.
+
+ Set ``"vertexai": true`` to force Vertex AI mode when only ``api_key`` is
+ provided. Omit ``vertexai`` for the Generative Language API (GLA).
+
+ :param llm_conn_id: Airflow connection ID.
+ :param model_id: Model identifier, e.g. ``"google-vertex:gemini-2.0-flash"``.
+ """
+
+ conn_type = "pydanticai_vertex"
+ hook_name = "Pydantic AI (Google Vertex AI)"
+
+ @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": "google-vertex:gemini-2.0-flash", '
+ '"project": "my-project", "location": "us-central1", "vertexai": true}'
+ " — add service_account_file (path) or service_account_info (object) for SA auth;"
+ " omit both to use Application Default Credentials"
+ ),
+ },
+ }
+
+ def _get_provider_kwargs(
+ self,
+ api_key: str | None,
+ base_url: str | None,
+ extra: dict[str, Any],
+ ) -> dict[str, Any]:
+ sa_file = extra.get("service_account_file")
+ sa_info = extra.get("service_account_info")
+ if sa_file and sa_info:
+ raise ValueError(
+ "Specify 'service_account_file' or 'service_account_info' in the connection extra, not both."
+ )
+
+ kwargs: dict[str, Any] = {}
+
+ # Direct GoogleProvider scalar kwargs.
+ for _key in ("api_key", "project", "location", "base_url"):
+ if extra.get(_key) is not None:
+ kwargs[_key] = extra[_key]
+
+ # Optional vertexai bool flag (force Vertex AI mode for API-key auth).
+ _vertexai = extra.get("vertexai")
+ if _vertexai is not None:
+ kwargs["vertexai"] = bool(_vertexai)
+
+ # Service-account credentials — loaded lazily to avoid importing
+ # google-auth on non-Vertex code paths (optional heavy dependency).
+ if sa_file:
+ from google.oauth2 import service_account # lazy: optional dep
+
+ kwargs["credentials"] = service_account.Credentials.from_service_account_file(
Review Comment:
`service_account_file` takes a raw file path from connection extras and passes it to `Credentials.from_service_account_file()`. Anyone with connection-edit permissions (lower privilege than DAG deployment) can point this at arbitrary files on the worker filesystem.
This is the same class of concern as the `importlib.import_module` issue from the previous round.
Options:
1. Drop `service_account_file` and only support `service_account_info` (inline JSON). Users embed the key content directly in the connection. No path traversal possible.
2. If you keep it, add a note in the docstring that this field requires the same trust level as DAG deployment access.
I'd go with option 1 for the initial PR. `service_account_info` covers the same use case without the risk.
##########
providers/common/ai/src/airflow/providers/common/ai/get_provider_info.py:
##########
@@ -74,6 +74,139 @@ def get_provider_info():
}
},
},
Review Comment:
This file is auto-generated from `provider.yaml`. Manual edits here get overwritten when someone runs `prek run update-providers-build-files`. (Flagged this in the previous round too.)
Your `provider.yaml` changes look correct, so just delete these manual edits and regenerate:
```bash
prek run update-providers-build-files
```
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2908415260
##########
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:
Yes, definitely add the warning\! A silent fallback is confusing — if someone carefully sets up their Airflow connection with explicit credentials and those get quietly ignored, they'd have a hard time figuring out why. A simple `self.log.warning("Provider '%s' rejected kwargs %s; falling back to env-var auth", pname, list(_kwargs))` would save a lot of head-scratching. (I just left the same suggestion in my latest review too — great minds think alike\!)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.
cetingokhan commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2914406140
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/pydantic_ai.py:
##########
@@ -75,62 +78,92 @@ 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:
+ try:
+ return infer_provider_class(pname)(**_kwargs)
+ except TypeError:
Review Comment:
done
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.
cetingokhan commented on code in PR #62816:
URL: https://github.com/apache/airflow/pull/62816#discussion_r2901761000
##########
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:
I added the underscore symbol, but I want you to know that you will still see the warning like this :)
`[2026-03-08 15:05:07] WARNING - Connection schemes (type: pydanticai_azure) shall not contain '_' according to RFC3986.`
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] added azureopenai support & extandable architecture for pydanticai [airflow]
Posted by "cetingokhan (via GitHub)" <gi...@apache.org>.
cetingokhan commented on PR #62816:
URL: https://github.com/apache/airflow/pull/62816#issuecomment-3993581554
Thanks for the comments.
Actually, I initially only added Azure support to the get_conn method in the pydantic_ai.py file for needs like api_version, but considering that similar needs might arise in other LLM providers and anticipating a significant increase in LLM providers, I thought creating a protocol would provide more flexibility. However, from what you wrote, I understand the necessary logic; following a restrictive approach rather than focusing on flexibility is a more appropriate perspective.
So, I will try to proceed as you suggested ;)
--
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: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org