Skip to content

Commit

Permalink
community(azuresearch): allow to use any valid credential (#28873)
Browse files Browse the repository at this point in the history
Add option to use any valid credential type.
Differentiates async cases needed by Azure Search.

This could replace the use of a static token
  • Loading branch information
ianchi authored Dec 23, 2024
1 parent 4b4d09f commit 8d99070
Showing 1 changed file with 18 additions and 3 deletions.
21 changes: 18 additions & 3 deletions libs/community/langchain_community/vectorstores/azuresearch.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,8 @@
logger = logging.getLogger()

if TYPE_CHECKING:
from azure.core.credentials import TokenCredential
from azure.core.credentials_async import AsyncTokenCredential
from azure.search.documents import SearchClient, SearchItemPaged
from azure.search.documents.aio import (
AsyncSearchItemPaged,
Expand Down Expand Up @@ -96,10 +98,13 @@ def _get_search_client(
cors_options: Optional[CorsOptions] = None,
async_: bool = False,
additional_search_client_options: Optional[Dict[str, Any]] = None,
azure_credential: Optional[TokenCredential] = None,
azure_async_credential: Optional[AsyncTokenCredential] = None,
) -> Union[SearchClient, AsyncSearchClient]:
from azure.core.credentials import AccessToken, AzureKeyCredential, TokenCredential
from azure.core.exceptions import ResourceNotFoundError
from azure.identity import DefaultAzureCredential, InteractiveBrowserCredential
from azure.identity.aio import DefaultAzureCredential as AsyncDefaultAzureCredential
from azure.search.documents import SearchClient
from azure.search.documents.aio import SearchClient as AsyncSearchClient
from azure.search.documents.indexes import SearchIndexClient
Expand Down Expand Up @@ -143,12 +148,17 @@ def get_token(
if key.upper() == "INTERACTIVE":
credential = InteractiveBrowserCredential()
credential.get_token("https://search.azure.com/.default")
async_credential = credential
else:
credential = AzureKeyCredential(key)
async_credential = credential
elif azure_ad_access_token is not None:
credential = AzureBearerTokenCredential(azure_ad_access_token)
async_credential = credential
else:
credential = DefaultAzureCredential()
credential = azure_credential or DefaultAzureCredential()
async_credential = azure_async_credential or AsyncDefaultAzureCredential()

index_client: SearchIndexClient = SearchIndexClient(
endpoint=endpoint,
credential=credential,
Expand Down Expand Up @@ -266,7 +276,7 @@ def fmt_err(x: str) -> str:
return AsyncSearchClient(
endpoint=endpoint,
index_name=index_name,
credential=credential,
credential=async_credential,
user_agent=user_agent,
**additional_search_client_options,
)
Expand All @@ -278,7 +288,7 @@ class AzureSearch(VectorStore):
def __init__(
self,
azure_search_endpoint: str,
azure_search_key: str,
azure_search_key: Optional[str],
index_name: str,
embedding_function: Union[Callable, Embeddings],
search_type: str = "hybrid",
Expand All @@ -295,6 +305,8 @@ def __init__(
vector_search_dimensions: Optional[int] = None,
additional_search_client_options: Optional[Dict[str, Any]] = None,
azure_ad_access_token: Optional[str] = None,
azure_credential: Optional[TokenCredential] = None,
azure_async_credential: Optional[AsyncTokenCredential] = None,
**kwargs: Any,
):
try:
Expand Down Expand Up @@ -361,6 +373,7 @@ def __init__(
user_agent=user_agent,
cors_options=cors_options,
additional_search_client_options=additional_search_client_options,
azure_credential=azure_credential,
)
self.async_client = _get_search_client(
azure_search_endpoint,
Expand All @@ -377,6 +390,8 @@ def __init__(
user_agent=user_agent,
cors_options=cors_options,
async_=True,
azure_credential=azure_credential,
azure_async_credential=azure_async_credential,
)
self.search_type = search_type
self.semantic_configuration_name = semantic_configuration_name
Expand Down

0 comments on commit 8d99070

Please sign in to comment.