From 5fea63327b8f27ca601cbe015c9cf3fcc4f5a1ce Mon Sep 17 00:00:00 2001
From: nceccarelli <41022718+nceccarelli@users.noreply.github.com>
Date: Mon, 4 Dec 2023 15:56:35 -0500
Subject: [PATCH] Support Azure gov cloud in Azure Cognitive Search retriever
 (#13695)

<!-- Thank you for contributing to LangChain!

Replace this entire comment with:
- **Description:** The existing version hardcoded search.windows.net in
the base url. This is not compatible with the gov cloud. I am allowing
the user to override the default for gov cloud support.,
  - **Issue:** N/A, did not write up in an issue,
  - **Dependencies:** None

Please make sure your PR is passing linting and testing before
submitting. Run `make format`, `make lint` and `make test` to check this
locally.

See contribution guidelines for more information on how to write/run
tests, lint, etc:

https://github.com/langchain-ai/langchain/blob/master/.github/CONTRIBUTING.md

If you're adding a new integration, please include:
1. a test for the integration, preferably unit tests that do not rely on
network access,
2. an example notebook showing its use. It lives in `docs/extras`
directory.

If no one reviews your PR within a few days, please @-mention one of
@baskaryan, @eyurtsev, @hwchase17.
 -->

---------

Co-authored-by: Nicholas Ceccarelli <nceccarelli2@moog.com>
Co-authored-by: Harrison Chase <hw.chase.17@gmail.com>
---
 .../langchain/retrievers/azure_cognitive_search.py     | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/libs/langchain/langchain/retrievers/azure_cognitive_search.py b/libs/langchain/langchain/retrievers/azure_cognitive_search.py
index 27cdd918157153..8824d986cc872b 100644
--- a/libs/langchain/langchain/retrievers/azure_cognitive_search.py
+++ b/libs/langchain/langchain/retrievers/azure_cognitive_search.py
@@ -13,7 +13,10 @@
     AsyncCallbackManagerForRetrieverRun,
     CallbackManagerForRetrieverRun,
 )
-from langchain.utils import get_from_dict_or_env
+from langchain.utils import get_from_dict_or_env, get_from_env
+
+DEFAULT_URL_SUFFIX = "search.windows.net"
+"""Default URL Suffix for endpoint connection - commercial cloud"""
 
 
 class AzureCognitiveSearchRetriever(BaseRetriever):
@@ -54,7 +57,10 @@ def validate_environment(cls, values: Dict) -> Dict:
         return values
 
     def _build_search_url(self, query: str) -> str:
-        base_url = f"https://{self.service_name}.search.windows.net/"
+        url_suffix = get_from_env(
+            "", "AZURE_COGNITIVE_SEARCH_URL_SUFFIX", DEFAULT_URL_SUFFIX
+        )
+        base_url = f"https://{self.service_name}.{url_suffix}/"
         endpoint_path = f"indexes/{self.index_name}/docs?api-version={self.api_version}"
         top_param = f"&$top={self.top_k}" if self.top_k else ""
         return base_url + endpoint_path + f"&search={query}" + top_param