added azureopenai support & extandable architecture for pydanticai

[cetingokhan]

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] [2] [3] [4]

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] [2]

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] [2] [3] [4] [5]

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.

---

Was generative AI tooling used to co-author this PR?

• Yes (please specify the tool below)
  Cloude Sonnet 4.6 & Gemini 3.1 Pro
  Filled some of methods scope and tests created via copilot

---

• Read the 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) is needed.
• When adding dependency, check compliance with the ASF 3rd Party License Policy.
• For significant user-facing changes create newsfragment: {pr_number}.significant.rst or {issue_number}.significant.rst, in airflow-core/newsfragments.

[kaxil]

Thanks for the contribution! Adding Azure OpenAI support is a good idea.

My main concern: pydantic-ai already ships with native AzureProvider support — see the docs. You can do:

from pydantic_ai.providers.azure import AzureProvider

model = OpenAIChatModel(
    'gpt-5.2',
    provider=AzureProvider(
        azure_endpoint='https://myresource.openai.azure.com',
        api_version='2024-07-01-preview',
        api_key='...',
    ),
)

Or even just use "azure:gpt-5.2" as the model string (with env vars set). So we don't need to manually construct AsyncAzureOpenAI clients from the openai SDK — pydantic-ai handles Azure natively.

The current hook's get_conn() is ~25 lines and delegates to pydantic-ai's infer_model() + provider_factory. Azure support could be a simple conditional branch using AzureProvider, without the 4-file builder pattern. The whole point of using pydantic-ai is that it abstracts provider differences for us — we should lean on that rather than wrapping it in another layer.

Specific comments inline.

[kaxil]

#62816 (review)

[cetingokhan]

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 ;)

[cetingokhan]

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)
    ....

[kaxil]

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:

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.

[gopidesupavan]

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 ?

[cetingokhan]

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.

[cetingokhan]

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

[kaxil]

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 (already merged)	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 already covers most providers — anything that accepts api_key + base_url. Only the 3 cloud providers with non-standard auth need their own types.

Implementation-wise, each subclass is small:

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.

So to answer your concern directly: we won't end up with endless if/elif. Each cloud provider gets its own connection type + hook subclass as separate PRs. No factory or registry needed.

[cetingokhan]

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.

[kaxil]

Thanks for the quick turnaround — removing the builder pattern and using subclass hooks instead is a much better direction.

A few things still need attention before this can merge. The most important:

1. provider.yaml not updated — get_provider_info.py is auto-generated from provider.yaml. Your manual edits to get_provider_info.py will be overwritten on the next prek run update-providers-build-files. The new connection types (pydanticai_azure, pydanticai_bedrock, pydanticai_vertex) need to be added to provider.yaml under connection-types:, then regenerated.

2. Vertex kwargs don't match pydantic-ai's GoogleProvider — see inline comment.

3. Subclass default_conn_name doesn't propagate — PydanticAIHook.__init__ binds llm_conn_id's default at class-definition time to "pydanticai_default". Subclasses inherit this bound default, so PydanticAIAzureHook() (no args) looks up "pydanticai_default", not "pydantic_ai_azure_default". Either override __init__ in each subclass or use a sentinel that resolves at runtime.

Inline comments below.

[kaxil]

(Reposting inline comments with correct line positions — previous batch was off. PR-level feedback from the earlier review still applies.)

[cetingokhan]

Fixed all comments ;)

[kaxil]

Big improvement over the builder/protocol version. The separate connection types + subclass approach is clean and follows existing Airflow patterns.

A few things still need fixing before this is ready.

[kaxil]

Architecture looks good — subclass hooks with _get_provider_kwargs() as the extension point is the right pattern. A few issues inline.

[kaxil]

Left 3 inline comments on this round of review. The subclass architecture and _get_provider_kwargs() extension point are clean. The model_id regression from last round is fixed. Main feedback is about empty-string handling consistency and silent fallback logging.

[cetingokhan]

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!

[kaxil]

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 👍

[kaxil]

Looks good! All feedback from previous rounds addressed. One tiny nit inline.

[kaxil]

@cetingokhan looks like there is a conflict!

Could you resolve that and we are good to go

[cetingokhan]

All done!

[gopidesupavan]

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 ..