kaxil commented on code in PR #70441:
URL: https://github.com/apache/airflow/pull/70441#discussion_r3661375927
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/langchain.py:
##########
@@ -171,3 +170,53 @@ def get_embedding_model(self) -> Embeddings:
kind="embedding",
)
return init_embeddings(model_id, **self._connection_kwargs(conn))
+
+
+class LangChainAzureHook(LangChainHook):
+ """
+ Hook for Azure OpenAI via LangChain.
+
+ ``LangChainHook``'s ``api_key`` + ``base_url`` credential surface doesn't
+ match what Azure OpenAI's LangChain classes (``AzureChatOpenAI``,
+ ``AzureOpenAIEmbeddings``) expect: an ``azure_endpoint`` and an
+ ``api_version`` rather than a generic ``base_url``. This subclass maps
+ the connection to those fields instead, mirroring
+
:class:`~airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIAzureHook`.
+
+ Connection fields:
+ - **password**: Azure API key
+ - **host**: Azure endpoint (e.g.
``https://<resource>.openai.azure.com``)
+ - **extra** JSON::
+
+ {"model": "azure_openai:gpt-4o", "api_version":
"2024-07-01-preview"}
+
+ :param llm_conn_id: Airflow connection ID.
+ """
+
+ conn_type = "langchain-azure"
+ default_conn_name = "langchain_azure_default"
+ hook_name = "LangChain (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_openai:gpt-4o", "api_version":
"2024-07-01-preview"}',
+ },
+ }
+
+ def _connection_kwargs(self, conn: Any) -> dict[str, Any]:
+ """Map connection fields to Azure OpenAI's azure_endpoint/api_version
shape."""
+ kwargs: dict[str, Any] = {}
+ if conn.password:
+ kwargs["api_key"] = conn.password
Review 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.
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/langchain.py:
##########
@@ -171,3 +170,53 @@ def get_embedding_model(self) -> Embeddings:
kind="embedding",
)
return init_embeddings(model_id, **self._connection_kwargs(conn))
+
+
+class LangChainAzureHook(LangChainHook):
Review 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?
##########
providers/common/ai/provider.yaml:
##########
@@ -404,6 +404,41 @@ connection-types:
type:
- string
- 'null'
+ - hook-class-name:
airflow.providers.common.ai.hooks.langchain.LangChainAzureHook
+ hook-name: "LangChain (Azure OpenAI)"
+ connection-type: langchain-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: Chat Model
+ description: "Chat model in azure_openai:name format (e.g.
azure_openai:gpt-4o)."
Review 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?
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/langchain.py:
##########
@@ -171,3 +170,53 @@ def get_embedding_model(self) -> Embeddings:
kind="embedding",
)
return init_embeddings(model_id, **self._connection_kwargs(conn))
+
+
+class LangChainAzureHook(LangChainHook):
+ """
+ Hook for Azure OpenAI via LangChain.
+
+ ``LangChainHook``'s ``api_key`` + ``base_url`` credential surface doesn't
+ match what Azure OpenAI's LangChain classes (``AzureChatOpenAI``,
+ ``AzureOpenAIEmbeddings``) expect: an ``azure_endpoint`` and an
+ ``api_version`` rather than a generic ``base_url``. This subclass maps
+ the connection to those fields instead, mirroring
+
:class:`~airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIAzureHook`.
+
+ Connection fields:
+ - **password**: Azure API key
+ - **host**: Azure endpoint (e.g.
``https://<resource>.openai.azure.com``)
+ - **extra** JSON::
+
+ {"model": "azure_openai:gpt-4o", "api_version":
"2024-07-01-preview"}
+
+ :param llm_conn_id: Airflow connection ID.
+ """
+
+ conn_type = "langchain-azure"
Review 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?
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/langchain.py:
##########
@@ -171,3 +170,53 @@ def get_embedding_model(self) -> Embeddings:
kind="embedding",
)
return init_embeddings(model_id, **self._connection_kwargs(conn))
+
+
+class LangChainAzureHook(LangChainHook):
+ """
+ Hook for Azure OpenAI via LangChain.
+
+ ``LangChainHook``'s ``api_key`` + ``base_url`` credential surface doesn't
+ match what Azure OpenAI's LangChain classes (``AzureChatOpenAI``,
+ ``AzureOpenAIEmbeddings``) expect: an ``azure_endpoint`` and an
+ ``api_version`` rather than a generic ``base_url``. This subclass maps
+ the connection to those fields instead, mirroring
+
:class:`~airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIAzureHook`.
+
+ Connection fields:
+ - **password**: Azure API key
+ - **host**: Azure endpoint (e.g.
``https://<resource>.openai.azure.com``)
+ - **extra** JSON::
+
+ {"model": "azure_openai:gpt-4o", "api_version":
"2024-07-01-preview"}
+
+ :param llm_conn_id: Airflow connection ID.
+ """
+
+ conn_type = "langchain-azure"
+ default_conn_name = "langchain_azure_default"
+ hook_name = "LangChain (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_openai:gpt-4o", "api_version":
"2024-07-01-preview"}',
+ },
+ }
+
+ def _connection_kwargs(self, conn: Any) -> dict[str, Any]:
+ """Map connection fields to Azure OpenAI's azure_endpoint/api_version
shape."""
+ kwargs: dict[str, Any] = {}
+ if conn.password:
+ kwargs["api_key"] = conn.password
+ if conn.host:
+ kwargs["azure_endpoint"] = conn.host
+ api_version = conn.extra_dejson.get("api_version")
+ if api_version:
+ kwargs["api_version"] = api_version
Review Comment:
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?
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/langchain.py:
##########
@@ -171,3 +170,53 @@ def get_embedding_model(self) -> Embeddings:
kind="embedding",
)
return init_embeddings(model_id, **self._connection_kwargs(conn))
+
+
+class LangChainAzureHook(LangChainHook):
+ """
+ Hook for Azure OpenAI via LangChain.
+
+ ``LangChainHook``'s ``api_key`` + ``base_url`` credential surface doesn't
+ match what Azure OpenAI's LangChain classes (``AzureChatOpenAI``,
+ ``AzureOpenAIEmbeddings``) expect: an ``azure_endpoint`` and an
Review 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?
##########
providers/common/ai/src/airflow/providers/common/ai/hooks/langchain.py:
##########
@@ -171,3 +170,53 @@ def get_embedding_model(self) -> Embeddings:
kind="embedding",
)
return init_embeddings(model_id, **self._connection_kwargs(conn))
+
+
+class LangChainAzureHook(LangChainHook):
+ """
+ Hook for Azure OpenAI via LangChain.
+
+ ``LangChainHook``'s ``api_key`` + ``base_url`` credential surface doesn't
+ match what Azure OpenAI's LangChain classes (``AzureChatOpenAI``,
+ ``AzureOpenAIEmbeddings``) expect: an ``azure_endpoint`` and an
+ ``api_version`` rather than a generic ``base_url``. This subclass maps
+ the connection to those fields instead, mirroring
+
:class:`~airflow.providers.common.ai.hooks.pydantic_ai.PydanticAIAzureHook`.
+
+ Connection fields:
+ - **password**: Azure API key
+ - **host**: Azure endpoint (e.g.
``https://<resource>.openai.azure.com``)
+ - **extra** JSON::
+
+ {"model": "azure_openai:gpt-4o", "api_version":
"2024-07-01-preview"}
+
+ :param llm_conn_id: Airflow connection ID.
+ """
+
+ conn_type = "langchain-azure"
+ default_conn_name = "langchain_azure_default"
+ hook_name = "LangChain (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_openai:gpt-4o", "api_version":
"2024-07-01-preview"}',
+ },
+ }
+
+ def _connection_kwargs(self, conn: Any) -> dict[str, Any]:
+ """Map connection fields to Azure OpenAI's azure_endpoint/api_version
shape."""
+ kwargs: dict[str, Any] = {}
+ if conn.password:
+ kwargs["api_key"] = conn.password
+ if conn.host:
+ kwargs["azure_endpoint"] = conn.host
+ api_version = conn.extra_dejson.get("api_version")
+ if api_version:
+ kwargs["api_version"] = api_version
+ return kwargs
Review 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?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
