Skip to content

Commit

Permalink
azuresearch: allow to use any valid credential
Browse files Browse the repository at this point in the history
Fix async credentials
  • Loading branch information
ianchi committed Dec 21, 2024
1 parent cb4e6ac commit a32c7af
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 a32c7af

Please sign in to comment.