Skip to content

Add LangChainAzureHook for Azure OpenAI support - #70441

Open
ColtenOuO wants to merge 1 commit into
apache:mainfrom
ColtenOuO:add-langchain-azure-hook
Open

Add LangChainAzureHook for Azure OpenAI support#70441
ColtenOuO wants to merge 1 commit into
apache:mainfrom
ColtenOuO:add-langchain-azure-hook

Conversation

@ColtenOuO

Copy link
Copy Markdown
Contributor

Summary

LangChainHook's docstring already flagged this gap: "Providers with bespoke auth (AWS Bedrock, Google Vertex AI / GenAI, Azure OpenAI, Cohere, HuggingFace) reject these kwargs; per-vendor subclasses can be added later mirroring the pydantic-ai pattern." PydanticAIHook already has that pattern (PydanticAIAzureHook/PydanticAIBedrockHook/PydanticAIVertexHook, each overriding _get_provider_kwargs) -- this PR adds the first LangChain equivalent.

Without it, LangChainHook always calls init_chat_model(model, api_key=..., base_url=...), but Azure OpenAI's LangChain classes (AzureChatOpenAI, AzureOpenAIEmbeddings) don't accept base_url -- they need azure_endpoint and api_version. Verified against the real langchain-openai package: those are the actual constructor fields/aliases.

Changes

  • LangChainHook._connection_kwargs: @staticmethod -> instance method, so subclasses can override it (matches PydanticAIHook's _get_provider_kwargs). No behavior change for the base class.
  • Add LangChainAzureHook(LangChainHook), overriding _connection_kwargs to map conn.host -> azure_endpoint and read api_version from conn.extra.
  • Register the new langchain-azure connection type in provider.yaml; regenerated get_provider_info.py via breeze release-management prepare-provider-documentation --reapply-templates-only --skip-changelog --skip-readme (not hand-edited).

Test plan

  • Added 7 tests: class attributes, UI field behaviour, host -> azure_endpoint mapping, api_version from extra, embedding model uses the same mapping, empty kwargs when no credentials set.
  • Full test_langchain.py passes (28 passed).
  • ruff format/ruff check, breeze run mypy clean (mypy caught the staticmethod/instance-method override incompatibility, since fixed).
  • scripts/ci/prek/check_provider_yaml_files.py passes: LangChainAzureHook registered correctly, 0 errors.

Was generative AI tooling used to co-author this PR?
  • Yes — Claude Code (Sonnet 5)

@ColtenOuO

ColtenOuO commented Jul 25, 2026

Copy link
Copy Markdown
Contributor Author

Happy to add the remaining vendors (ex. Bedrock, Vertex) as follow-up PRs if that's a direction the maintainers would like to take.

@kaxil kaxil left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for this, and for the offer to follow up with Bedrock and Vertex. Specifics are in line comments. The rest is cross-cutting enough that it does not sit on any single line, and it is mostly about the pydantic-ai pattern we shipped rather than about your code.

What can't a user do today?

The existing LangChainHook already reaches Azure OpenAI as long as host is left empty and the worker has AZURE_OPENAI_ENDPOINT and OPENAI_API_VERSION set:

init_chat_model("azure_openai:gpt-4o", api_key="k")
# resolves fine, endpoint and api_version come from env

So what this PR newly enables is carrying the endpoint and API version in the connection instead of the worker env. That is worth wanting, for a secrets backend or for several Azure resources in one deployment, but it is narrower than "Azure doesn't work". Which of those are you actually hitting? A concrete deployment shape would help, because it changes what the right fix is.

Why a new connection type instead of kwargs passthrough?

init_chat_model forwards arbitrary kwargs to the vendor class, so this already works today:

init_chat_model("azure_openai:gpt-4o", api_key="k",
                azure_endpoint="https://x.openai.azure.com",
                api_version="2024-07-01-preview",
                azure_deployment="gpt4o-prod-eu")

Having _connection_kwargs merge an allow-listed set of extra keys would cover Azure, Bedrock, Vertex, Cohere and HuggingFace in roughly five lines, without a new class and conn type per vendor. What does the subclass buy over that? It would need an allow-list rather than a blind passthrough, since init_chat_model sweeps unrecognised kwargs into model_kwargs behind a warning instead of rejecting them.

I am raising this here rather than against the pydantic-ai hooks because of the trajectory. The provider registers 7 conn types today, your Bedrock and Vertex follow-ups take LangChain to 10, and LlamaIndexHook has the identical api_key/api_base shape and the identical gap, so the same argument adds a third row. Each conn type is permanent public surface, so the axis is worth settling before we extend it.

If we do keep per-vendor subclasses, this may belong in its own provider

common.ai's stated scope is pydantic-ai (provider.yaml: "AI/LLM hooks and operators for Airflow pipelines using pydantic-ai"). A full LangChain vendor matrix means langchain-openai, langchain-aws, langchain-google-vertexai, langchain-cohere, langchain-huggingface and possibly langchain-azure-ai as extras. That is a dependency and CVE surface attached to a provider whose actual product is the pydantic-ai operators and toolsets, and LangChain 1.x moves fast enough that its churn would start gating common.ai releases. Airflow already ships per-vendor providers (anthropic, cohere, openai, pinecone, qdrant, weaviate), so apache-airflow-providers-langchain fits that shape better.

Now is the cheap moment to decide. common.ai is 0.x and lifecycle: incubation, and the langchain conn type has only been out since 0.4.0. The toolset bridge from #67791 would stay here regardless, since it exposes common.ai toolsets as LangChain tools.

None of this should turn into you standing up a new provider before anything can merge. If that is where it lands, it gets scoped separately. The code here is clean and well tested.

return init_embeddings(model_id, **self._connection_kwargs(conn))


class LangChainAzureHook(LangChainHook):

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LangChainHook has no test_connection, unlike PydanticAIHook. This PR registers a new conn type in the UI, so users get a Test button that will not tell them whether the endpoint or the api_version is right. Given how many Azure-specific fields have to line up, worth adding one here?

Comment on lines +180 to +181
match what Azure OpenAI's LangChain classes (``AzureChatOpenAI``,
``AzureOpenAIEmbeddings``) expect: an ``azure_endpoint`` and an

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Both of these classes come from langchain-openai, which nothing installs. The provider's extra is just "langchain" = ["langchain>=1.0.0"], so installing apache-airflow-providers-common-ai[langchain] and creating this connection fails inside init_chat_model with an ImportError. Should the extra grow, or should the hook raise something that names the missing package?

:param llm_conn_id: Airflow connection ID.
"""

conn_type = "langchain-azure"

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LangChain registers two Azure providers in _BUILTIN_PROVIDERS: azure_openai (langchain-openai) and azure_ai (langchain-azure-ai, AI Foundry). Different classes, different kwargs.

langchain-azure does not say which one this is, and it is the name you would want if Foundry support is ever added. Worth making it langchain-azure-openai while the conn type is still unreleased?

Comment on lines +215 to +216
if conn.password:
kwargs["api_key"] = conn.password

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

conn.password is the only auth path here, but AzureChatOpenAI also accepts azure_ad_token and azure_ad_token_provider, and managed identity is how most enterprise Azure OpenAI deployments authenticate (AKS workload identity, no key stored in the connection at all).

A hook that exists specifically to handle Azure credentials, but only reads a static key, misses those deployments. Deliberate scope cut for a follow-up? PydanticAIAzureHook has the same limitation, so this is not something you introduced.

Comment on lines +219 to +221
api_version = conn.extra_dejson.get("api_version")
if api_version:
kwargs["api_version"] = api_version

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Treating api_version as optional behaves differently for chat and embeddings, and neither case is good (checked against langchain-openai 1.4.1):

  • AzureChatOpenAI has no default, so omitting it raises a raw pydantic ValidationError from inside the constructor rather than an Airflow-side message naming the connection field.
  • AzureOpenAIEmbeddings defaults to 2023-05-15, so the same omission silently pins a three-year-old API version instead of failing.

Given that carrying this field is much of the reason the subclass exists, should it be required with a clear error, the way _resolve_model_id already does for the model id?

api_version = conn.extra_dejson.get("api_version")
if api_version:
kwargs["api_version"] = api_version
return kwargs

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nothing here can set azure_deployment. Azure deployment names are user-chosen, so they routinely differ from the model name, and the SDK routes on it:

azure_deployment=None            -> .../openai/deployments/gpt-4o/chat/completions
azure_deployment='gpt4o-prod-eu' -> .../openai/deployments/gpt4o-prod-eu/chat/completions

So anyone whose deployment is not named exactly after the model cannot use this hook. Deliberate?

conn-fields:
model:
label: Chat Model
description: "Chat model in azure_openai:name format (e.g. azure_openai:gpt-4o)."

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nothing enforces this format. Putting openai:gpt-4o on a langchain-azure connection does not raise: init_chat_model transfers azure_endpoint and api_version into model_kwargs behind a UserWarning and hands back a plain ChatOpenAI, so it fails later at request time with an error that points nowhere useful. Worth defaulting the prefix, or validating it in the hook?

@ColtenOuO

Copy link
Copy Markdown
Contributor Author

Thanks a lot for the review, and for putting this much thought into the feedback!

To be honest, what got me started on this PR was this part of the docstring:

Providers with bespoke auth (AWS Bedrock, Google Vertex AI / GenAI, Azure OpenAI, Cohere, HuggingFace) reject these kwargs; per-vendor subclasses can be added later mirroring the pydantic-ai pattern.

So I followed the pydantic-ai pattern and put together an Azure version, hoping to help extend it.


What can't a user do today?

Honestly — nothing is broken. Users really can use Azure OpenAI today without an Azure hook.

As I mentioned above, because the docstring brings this up, I took it to mean that adding per-vendor subclasses was a planned direction for the future. (Or maybe I misread what that comment meant — please let me know if that's the case ><)

But what I can confirm is that there's no problem with the current usage at all.


Why a new connection type instead of kwargs passthrough?

You're right that a plain allow-list of per-vendor keys would be considerably easier to maintain.

The advantage I see in the subclasses is on the UI side: when a user picks a vendor, we can show the fields that vendor actually needs, which makes the experience friendlier and gives them an easy way to see what LangChain vendors are supported right now.

image

As in the screenshot above — if we pick Azure as the vendor, both the Standard Fields and the Extra Fields immediately adapt to it and hint at the keys the user should fill in.

Maybe there's a way to ask the user to choose the vendor first when they select LangChain, and then list the corresponding Standard Fields and Extra Fields? If that were possible, I think the user experience would be even friendlier ><

That said, it would probably mean touching things inside airflow-core, so it may not be a good approach either.

I may well be overthinking this (sorry, I do that a lot xD) — if we set the UI experience aside, I agree that a plain allow-list is the better way to maintain this.


If we do keep per-vendor subclasses, this may belong in its own provider

I agree on all three points.


I've also gone through all the code-level suggestions and I'm happy to handle them.

Before that, though, I'd like to wait until the direction on subclasses is settled before moving on to the next step.

Thanks again for raising these questions and suggestions!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants