Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

community(azuresearch): allow to use any valid credential #28873

Merged
merged 1 commit into from
Dec 23, 2024
Merged

community(azuresearch): allow to use any valid credential #28873

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 18 additions & 3 deletions libs/community/langchain_community/vectorstores/azuresearch.py
View file
Open in desktop
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
Loading