Repository: langchain-ai/langchain
Pull request #30730: Community: Add Managed Identity support for Azure AI Search
Author: login: paulczar
Created: 2025-04-08T14:32:14Z | Merged: 2025-04-11T02:22:59Z

Add Managed Identity support for Azure AI Search


Diff:
diff --git a/docs/docs/integrations/retrievers/azure_ai_search.ipynb b/docs/docs/integrations/retrievers/azure_ai_search.ipynb
index e5a9bcec2a1fc..f2175a8ca75ed 100644
--- a/docs/docs/integrations/retrievers/azure_ai_search.ipynb
+++ b/docs/docs/integrations/retrievers/azure_ai_search.ipynb
@@ -40,11 +40,21 @@
     "\n",
     "+ An existing index with vector fields. There are several ways to create one, including using the [vector store module](../vectorstores/azuresearch.ipynb). Or, [try the Azure AI Search REST APIs](https://learn.microsoft.com/azure/search/search-get-started-vector).\n",
     "\n",
-    "+ An API key. API keys are generated when you create the search service. If you're just querying an index, you can use the query API key, otherwise use an admin API key. See [Find your API keys](https://learn.microsoft.com/azure/search/search-security-api-keys?tabs=rest-use%2Cportal-find%2Cportal-query#find-existing-keys) for details.\n",
+    "+ An API key or Azure AD Token.\n",
+    "  + API keys are generated when you create the search service. If you're just querying an index, you can use the query API key, otherwise use an admin API key. See [Find your API keys](https://learn.microsoft.com/azure/search/search-security-api-keys?tabs=rest-use%2Cportal-find%2Cportal-query#find-existing-keys) for details.\n",
+    "  + Azure AD Token can be used with Azure Managed Identity. See [Connect your app to Azure AI Search using identities](https://learn.microsoft.com/en-us/azure/search/keyless-connections?tabs=python%2Cazure-cli) for details.\n",
     "\n",
     "We can then set the search service name, index name, and API key as environment variables (alternatively, you can pass them as arguments to `AzureAISearchRetriever`). The search index provides the searchable content."
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "2af9655d",
+   "metadata": {},
+   "source": [
+    "With an API Key"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -59,6 +69,28 @@
     "os.environ[\"AZURE_AI_SEARCH_API_KEY\"] = \"<YOUR_API_KEY>\""
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "35ee4780",
+   "metadata": {},
+   "source": [
+    "With an Azure AD Token"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7f33263c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "os.environ[\"AZURE_AI_SEARCH_SERVICE_NAME\"] = \"<YOUR_SEARCH_SERVICE_NAME>\"\n",
+    "os.environ[\"AZURE_AI_SEARCH_INDEX_NAME\"] = \"<YOUR_SEARCH_INDEX_NAME>\"\n",
+    "os.environ[\"AZURE_AI_SEARCH_AD_TOKEN\"] = \"<YOUR_AZURE_AD_TOKEN>\""
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "3e635218-8634-4f39-abc5-39e319eeb136",
diff --git a/libs/community/langchain_community/retrievers/azure_ai_search.py b/libs/community/langchain_community/retrievers/azure_ai_search.py
index 01549ff3bd012..6db86a3c3bf33 100644
--- a/libs/community/langchain_community/retrievers/azure_ai_search.py
+++ b/libs/community/langchain_community/retrievers/azure_ai_search.py
@@ -32,7 +32,10 @@ class AzureAISearchRetriever(BaseRetriever):
             pip install -U langchain-community azure-identity azure-search-documents
             export AZURE_AI_SEARCH_SERVICE_NAME="<YOUR_SEARCH_SERVICE_NAME>"
             export AZURE_AI_SEARCH_INDEX_NAME="<YOUR_SEARCH_INDEX_NAME>"
+
             export AZURE_AI_SEARCH_API_KEY="<YOUR_API_KEY>"
+            or
+            export AZURE_AI_SEARCH_BEARER_TOKEN="<YOUR_BEARER_TOKEN>"
 
     Key init args:
         content_key: str
@@ -96,6 +99,14 @@ def format_docs(docs):
     """API version"""
     aiosession: Optional[aiohttp.ClientSession] = None
     """ClientSession, in case we want to reuse connection for better performance."""
+    azure_ad_token: str = ""
+    """Your Azure Active Directory token.
+
+        Automatically inferred from env var `AZURE_AI_SEARCH_AD_TOKEN` if not provided.
+
+        For more:
+        https://www.microsoft.com/en-us/security/business/identity-access/microsoft-entra-id.
+    """
     content_key: str = "content"
     """Key in a retrieved result to set as the Document page_content."""
     top_k: Optional[int] = None
@@ -118,9 +129,19 @@ def validate_environment(cls, values: Dict) -> Any:
         values["index_name"] = get_from_dict_or_env(
             values, "index_name", "AZURE_AI_SEARCH_INDEX_NAME"
         )
+        values["azure_ad_token"] = get_from_dict_or_env(
+            values, "azure_ad_token", "AZURE_AI_SEARCH_AD_TOKEN", default=""
+        )
         values["api_key"] = get_from_dict_or_env(
-            values, "api_key", "AZURE_AI_SEARCH_API_KEY"
+            values, "api_key", "AZURE_AI_SEARCH_API_KEY", default=""
         )
+        if values["azure_ad_token"] == "" and values["api_key"] == "":
+            raise ValueError(
+                "Missing credentials. Please pass one of `api_key`, `azure_ad_token`, "
+                "or the `AZURE_AI_SEARCH_API_KEY` or `AZURE_AI_SEARCH_AD_TOKEN` "
+                "environment variables."
+            )
+
         return values
 
     def _build_search_url(self, query: str) -> str:
@@ -145,10 +166,14 @@ def _build_search_url(self, query: str) -> str:
 
     @property
     def _headers(self) -> Dict[str, str]:
-        return {
+        headers = {
             "Content-Type": "application/json",
-            "api-key": self.api_key,
         }
+        if not self.azure_ad_token:
+            headers["Authorization"] = f"Bearer {self.azure_ad_token}"
+        else:
+            headers["api-key"] = f"{self.api_key}"
+        return headers
 
     def _search(self, query: str) -> List[dict]:
         search_url = self._build_search_url(query)
diff --git a/libs/community/tests/integration_tests/retrievers/test_azure_ai_search.py b/libs/community/tests/integration_tests/retrievers/test_azure_ai_search.py
index 27d6c3edd8c93..4a831ba5a08f8 100644
--- a/libs/community/tests/integration_tests/retrievers/test_azure_ai_search.py
+++ b/libs/community/tests/integration_tests/retrievers/test_azure_ai_search.py
@@ -12,7 +12,8 @@ def test_azure_ai_search_invoke() -> None:
     """Test valid call to Azure AI Search.
 
     In order to run this test, you should provide
-    a `service_name`, azure search `api_key` and an `index_name`
+    a `service_name`, an 'index_name' and
+    an azure search `api_key` or 'azure_ad_token'
     as arguments for the AzureAISearchRetriever in both tests.
     api_version, aiosession and topk_k are optional parameters.
     """
@@ -32,7 +33,8 @@ async def test_azure_ai_search_ainvoke() -> None:
     """Test valid async call to Azure AI Search.
 
     In order to run this test, you should provide
-    a `service_name`, azure search `api_key` and an `index_name`
+    a `service_name`, an 'index_name' and
+    an azure search `api_key` or 'azure_ad_token'
     as arguments for the AzureAISearchRetriever.
     """
     retriever = AzureAISearchRetriever()


Reviews:
user: login: eyurtsev
body: looks good a few fixes w/ regards to typing
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-09T19:31:43Z
user: login: eyurtsev
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-11T02:22:50Z

Review comments:
diff_hunk: @@ -89,13 +92,21 @@ def format_docs(docs):
     """Name of Azure AI Search service"""
     index_name: str = ""
     """Name of Index inside Azure AI Search service"""
-    api_key: str = ""
+    api_key: Optional[str] = ""
path: libs/community/langchain_community/retrievers/azure_ai_search.py
user: login: eyurtsev
body: Optional != NotRequired

Changing to an Optional has no effect since we can base the checks on falsiness

```suggestion
    api_key: Optional[str] = ""
```
created_at: 2025-04-09T19:30:14Z
author_association: COLLABORATOR
diff_hunk: @@ -89,13 +92,21 @@ def format_docs(docs):
     """Name of Azure AI Search service"""
     index_name: str = ""
     """Name of Index inside Azure AI Search service"""
-    api_key: str = ""
+    api_key: Optional[str] = ""
     """API Key. Both Admin and Query keys work, but for reading data it's
     recommended to use a Query key."""
     api_version: str = "2023-11-01"
     """API version"""
     aiosession: Optional[aiohttp.ClientSession] = None
     """ClientSession, in case we want to reuse connection for better performance."""
+    azure_ad_token: Optional[str] = ""
path: libs/community/langchain_community/retrievers/azure_ai_search.py
user: login: eyurtsev
body: ```suggestion
    azure_ad_token: str = ""
```
created_at: 2025-04-09T19:30:24Z
author_association: COLLABORATOR
diff_hunk: @@ -145,10 +166,14 @@ def _build_search_url(self, query: str) -> str:
 
     @property
     def _headers(self) -> Dict[str, str]:
-        return {
+        headers = {
             "Content-Type": "application/json",
-            "api-key": self.api_key,
         }
+        if self.azure_ad_token != "":
path: libs/community/langchain_community/retrievers/azure_ai_search.py
user: login: eyurtsev
body: ```suggestion
        if not self.azure_ad_token:
```
created_at: 2025-04-09T19:31:27Z
author_association: COLLABORATOR

PR comments:
user: login: vercel[bot]
created_at: 2025-04-08T14:32:18Z
body: [vc]: #b+kXEanes2Pp7mfyivTxtBKg2UHs7dC3UsEQXslIGCA=:eyJpc01vbm9yZXBvIjp0cnVlLCJ0eXBlIjoiZ2l0aHViIiwicHJvamVjdHMiOlt7Im5hbWUiOiJsYW5nY2hhaW4iLCJyb290RGlyZWN0b3J5IjoiZG9jcyIsImluc3BlY3RvclVybCI6Imh0dHBzOi8vdmVyY2VsLmNvbS9sYW5nY2hhaW4vbGFuZ2NoYWluLzVUUXp1UzJBUnkya1FMUTlyVURDUXZCYmhZQnAiLCJwcmV2aWV3VXJsIjoibGFuZ2NoYWluLWdpdC1mb3JrLXBhdWxjemFyLWF6dXJlLWFpLXNlYXJjaC1iZS0wNWNjNTQtbGFuZ2NoYWluLnZlcmNlbC5hcHAiLCJuZXh0Q29tbWl0U3RhdHVzIjoiREVQTE9ZRUQiLCJsaXZlRmVlZGJhY2siOnsicmVzb2x2ZWQiOjAsInVucmVzb2x2ZWQiOjAsInRvdGFsIjowLCJsaW5rIjoibGFuZ2NoYWluLWdpdC1mb3JrLXBhdWxjemFyLWF6dXJlLWFpLXNlYXJjaC1iZS0wNWNjNTQtbGFuZ2NoYWluLnZlcmNlbC5hcHAifX1dfQ==
**The latest updates on your projects**. Learn more about [Vercel for Git ↗︎](https://vercel.link/github-learn-more)

| Name | Status | Preview | Comments | Updated (UTC) |
| :--- | :----- | :------ | :------- | :------ |
| **langchain** | ✅ Ready ([Inspect](https://vercel.com/langchain/langchain/5TQzuS2ARy2kQLQ9rUDCQvBbhYBp)) | [Visit Preview](https://langchain-git-fork-paulczar-azure-ai-search-be-05cc54-langchain.vercel.app) | 💬 [**Add feedback**](https://vercel.live/open-feedback/langchain-git-fork-paulczar-azure-ai-search-be-05cc54-langchain.vercel.app?via=pr-comment-feedback-link) | Apr 10, 2025 2:09pm |


author_association: NONE
user: login: paulczar
created_at: 2025-04-10T14:01:24Z
body: thanks @eyurtsev !!  I've updated the PR with your suggestions.
author_association: CONTRIBUTOR
user: login: marlenezw
created_at: 2025-04-11T13:24:22Z
body: Thank you for reviewing @eyurtsev!
author_association: CONTRIBUTOR