Add LangChainAzureHook for Azure OpenAI support - #70441
Conversation
|
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
left a comment
There was a problem hiding this comment.
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 envSo 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): |
There was a problem hiding this comment.
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?
| match what Azure OpenAI's LangChain classes (``AzureChatOpenAI``, | ||
| ``AzureOpenAIEmbeddings``) expect: an ``azure_endpoint`` and an |
There was a problem hiding this comment.
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" |
There was a problem hiding this comment.
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?
| if conn.password: | ||
| kwargs["api_key"] = conn.password |
There was a problem hiding this comment.
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.
| api_version = conn.extra_dejson.get("api_version") | ||
| if api_version: | ||
| kwargs["api_version"] = api_version |
There was a problem hiding this comment.
Treating api_version as optional behaves differently for chat and embeddings, and neither case is good (checked against langchain-openai 1.4.1):
AzureChatOpenAIhas no default, so omitting it raises a raw pydanticValidationErrorfrom inside the constructor rather than an Airflow-side message naming the connection field.AzureOpenAIEmbeddingsdefaults to2023-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 |
There was a problem hiding this comment.
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)." |
There was a problem hiding this comment.
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?

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."PydanticAIHookalready has that pattern (PydanticAIAzureHook/PydanticAIBedrockHook/PydanticAIVertexHook, each overriding_get_provider_kwargs) -- this PR adds the first LangChain equivalent.Without it,
LangChainHookalways callsinit_chat_model(model, api_key=..., base_url=...), but Azure OpenAI's LangChain classes (AzureChatOpenAI,AzureOpenAIEmbeddings) don't acceptbase_url-- they needazure_endpointandapi_version. Verified against the reallangchain-openaipackage: those are the actual constructor fields/aliases.Changes
LangChainHook._connection_kwargs:@staticmethod-> instance method, so subclasses can override it (matchesPydanticAIHook's_get_provider_kwargs). No behavior change for the base class.LangChainAzureHook(LangChainHook), overriding_connection_kwargsto mapconn.host->azure_endpointand readapi_versionfromconn.extra.langchain-azureconnection type inprovider.yaml; regeneratedget_provider_info.pyviabreeze release-management prepare-provider-documentation --reapply-templates-only --skip-changelog --skip-readme(not hand-edited).Test plan
host -> azure_endpointmapping,api_versionfrom extra, embedding model uses the same mapping, empty kwargs when no credentials set.test_langchain.pypasses (28 passed).ruff format/ruff check,breeze run mypyclean (mypy caught the staticmethod/instance-method override incompatibility, since fixed).scripts/ci/prek/check_provider_yaml_files.pypasses:LangChainAzureHookregistered correctly, 0 errors.Was generative AI tooling used to co-author this PR?