77from collections .abc import Sequence
88from typing import Any , ClassVar , Final , Generic , TypeVar
99
10- from azure .core .credentials import AzureKeyCredential , TokenCredential
10+ from azure .core .credentials import AzureKeyCredential
1111from azure .core .credentials_async import AsyncTokenCredential
1212from azure .search .documents .aio import SearchClient
1313from azure .search .documents .indexes .aio import SearchIndexClient
@@ -149,23 +149,47 @@ class AzureAISearchSettings(KernelBaseSettings):
149149
150150
151151def _get_search_client (
152- search_index_client : SearchIndexClient , collection_name : str | None , ** kwargs : Any
152+ endpoint : str ,
153+ collection_name : str | None ,
154+ credential : "AzureKeyCredential | AsyncTokenCredential" ,
155+ ** kwargs : Any ,
153156) -> SearchClient :
154157 """Create a search client for a collection."""
155158 if not collection_name :
156159 raise VectorStoreInitializationException ("Collection name is required to create a search client." )
157160 try :
158- return SearchClient (search_index_client . _endpoint , collection_name , search_index_client . _credential , ** kwargs )
161+ return SearchClient (endpoint , collection_name , credential , ** kwargs )
159162 except ValueError as exc :
160163 raise VectorStoreInitializationException (
161164 f"Failed to create Azure Cognitive Search client for collection { collection_name } ."
162165 ) from exc
163166
164167
168+ def _resolve_credential (
169+ azure_ai_search_settings : AzureAISearchSettings ,
170+ azure_credential : AzureKeyCredential | None = None ,
171+ token_credential : "AsyncTokenCredential | None" = None ,
172+ ) -> "AzureKeyCredential | AsyncTokenCredential" :
173+ """Resolve the credential to use for Azure AI Search.
174+
175+ Args:
176+ azure_ai_search_settings: Azure AI Search settings.
177+ azure_credential: Optional Azure credentials (default: {None}).
178+ token_credential: Optional Token credential (default: {None}).
179+ """
180+ if azure_credential :
181+ return azure_credential
182+ if token_credential :
183+ return token_credential
184+ if azure_ai_search_settings .api_key :
185+ return AzureKeyCredential (azure_ai_search_settings .api_key .get_secret_value ())
186+ raise ServiceInitializationError ("Error: missing Azure AI Search client credentials." )
187+
188+
165189def _get_search_index_client (
166190 azure_ai_search_settings : AzureAISearchSettings ,
167191 azure_credential : AzureKeyCredential | None = None ,
168- token_credential : "AsyncTokenCredential | TokenCredential | None" = None ,
192+ token_credential : "AsyncTokenCredential | None" = None ,
169193) -> SearchIndexClient :
170194 """Return a client for Azure AI Search.
171195
@@ -174,20 +198,11 @@ def _get_search_index_client(
174198 azure_credential: Optional Azure credentials (default: {None}).
175199 token_credential: Optional Token credential (default: {None}).
176200 """
177- # Credentials
178- credential : "AzureKeyCredential | AsyncTokenCredential | TokenCredential | None" = None
179- if azure_credential :
180- credential = azure_credential
181- elif token_credential :
182- credential = token_credential
183- elif azure_ai_search_settings .api_key :
184- credential = AzureKeyCredential (azure_ai_search_settings .api_key .get_secret_value ())
185- else :
186- raise ServiceInitializationError ("Error: missing Azure AI Search client credentials." )
201+ credential = _resolve_credential (azure_ai_search_settings , azure_credential , token_credential )
187202
188203 return SearchIndexClient (
189204 endpoint = str (azure_ai_search_settings .endpoint ),
190- credential = credential , # type: ignore
205+ credential = credential ,
191206 headers = prepend_semantic_kernel_to_user_agent ({}) if APP_INFO else None ,
192207 )
193208
@@ -286,6 +301,8 @@ class AzureAISearchCollection(
286301
287302 search_client : SearchClient
288303 search_index_client : SearchIndexClient
304+ search_endpoint : str | None = None
305+ search_credential : Any = None
289306 supported_key_types : ClassVar [set [str ] | None ] = {"str" }
290307 supported_vector_types : ClassVar [set [str ] | None ] = {"float" , "int" }
291308 supported_search_types : ClassVar [set [SearchType ]] = {SearchType .VECTOR , SearchType .KEYWORD_HYBRID }
@@ -299,6 +316,7 @@ def __init__(
299316 search_index_client : SearchIndexClient | None = None ,
300317 search_client : SearchClient | None = None ,
301318 embedding_generator : "EmbeddingGeneratorBase | None" = None ,
319+ search_credential : "AzureKeyCredential | AsyncTokenCredential | None" = None ,
302320 ** kwargs : Any ,
303321 ) -> None :
304322 """Initializes a new instance of the AzureAISearchCollection class.
@@ -319,13 +337,16 @@ def __init__(
319337 used for creating and deleting indexes.
320338 search_client: The search client for interacting with Azure AI Search,
321339 used for record operations.
340+ search_credential: The credential used to authenticate with Azure AI Search.
341+ If not provided, it will be resolved from azure_credentials, token_credentials,
342+ or api_key in kwargs/environment.
322343 embedding_generator: The embedding generator, optional.
323344 **kwargs: Additional keyword arguments, including:
324345 The same keyword arguments used for AzureAISearchVectorStore:
325- search_endpoint: str | None = None,
346+ search_endpoint: The endpoint of the Azure AI Search service, optional.
326347 api_key: str | None = None,
327348 azure_credentials: AzureKeyCredential | None = None,
328- token_credentials: AsyncTokenCredential | TokenCredential | None = None,
349+ token_credentials: AsyncTokenCredential | None = None,
329350 env_file_path: str | None = None,
330351 env_file_encoding: str | None = None
331352
@@ -343,6 +364,8 @@ def __init__(
343364 collection_name = collection_name ,
344365 search_client = search_client ,
345366 search_index_client = search_index_client ,
367+ search_endpoint = kwargs .get ("search_endpoint" ),
368+ search_credential = search_credential ,
346369 managed_search_index_client = False ,
347370 managed_client = False ,
348371 embedding_generator = embedding_generator ,
@@ -360,14 +383,24 @@ def __init__(
360383 )
361384 except ValidationError as exc :
362385 raise VectorStoreInitializationException ("Failed to create Azure Cognitive Search settings." ) from exc
386+ endpoint = str (azure_ai_search_settings .endpoint )
387+ credential = search_credential or _resolve_credential (
388+ azure_ai_search_settings ,
389+ azure_credential = kwargs .get ("azure_credentials" ),
390+ token_credential = kwargs .get ("token_credentials" ),
391+ )
363392 super ().__init__ (
364393 record_type = record_type ,
365394 definition = definition ,
366395 collection_name = azure_ai_search_settings .index_name ,
367396 search_client = _get_search_client (
368- search_index_client = search_index_client , collection_name = azure_ai_search_settings .index_name
397+ endpoint = endpoint ,
398+ collection_name = azure_ai_search_settings .index_name ,
399+ credential = credential ,
369400 ),
370401 search_index_client = search_index_client ,
402+ search_endpoint = endpoint ,
403+ search_credential = credential ,
371404 managed_search_index_client = False ,
372405 embedding_generator = embedding_generator ,
373406 )
@@ -383,6 +416,12 @@ def __init__(
383416 )
384417 except ValidationError as exc :
385418 raise VectorStoreInitializationException ("Failed to create Azure Cognitive Search settings." ) from exc
419+ endpoint = str (azure_ai_search_settings .endpoint )
420+ credential = search_credential or _resolve_credential (
421+ azure_ai_search_settings ,
422+ azure_credential = kwargs .get ("azure_credentials" ),
423+ token_credential = kwargs .get ("token_credentials" ),
424+ )
386425 search_index_client = _get_search_index_client (
387426 azure_ai_search_settings = azure_ai_search_settings ,
388427 azure_credential = kwargs .get ("azure_credentials" ),
@@ -393,10 +432,13 @@ def __init__(
393432 definition = definition ,
394433 collection_name = azure_ai_search_settings .index_name ,
395434 search_client = _get_search_client (
396- search_index_client = search_index_client ,
397- collection_name = azure_ai_search_settings .index_name , # type: ignore
435+ endpoint = endpoint ,
436+ collection_name = azure_ai_search_settings .index_name ,
437+ credential = credential ,
398438 ),
399439 search_index_client = search_index_client ,
440+ search_endpoint = endpoint ,
441+ search_credential = credential ,
400442 embedding_generator = embedding_generator ,
401443 )
402444
@@ -711,20 +753,24 @@ class AzureAISearchStore(VectorStore):
711753 """Azure AI Search store implementation."""
712754
713755 search_index_client : SearchIndexClient
756+ search_endpoint : str | None = None
757+ search_credential : Any = None
714758
715759 def __init__ (
716760 self ,
717761 search_endpoint : str | None = None ,
718762 api_key : str | None = None ,
719763 azure_credentials : "AzureKeyCredential | None" = None ,
720- token_credentials : "AsyncTokenCredential | TokenCredential | None" = None ,
764+ token_credentials : "AsyncTokenCredential | None" = None ,
721765 search_index_client : SearchIndexClient | None = None ,
722766 embedding_generator : "EmbeddingGeneratorBase | None" = None ,
723767 env_file_path : str | None = None ,
724768 env_file_encoding : str | None = None ,
725769 ) -> None :
726770 """Initializes a new instance of the AzureAISearchStore class."""
727771 managed_client : bool = False
772+ endpoint : str | None = None
773+ credential : AzureKeyCredential | AsyncTokenCredential | None = None
728774 if not search_index_client :
729775 try :
730776 azure_ai_search_settings = AzureAISearchSettings (
@@ -735,15 +781,26 @@ def __init__(
735781 )
736782 except ValidationError as exc :
737783 raise VectorStoreInitializationException ("Failed to create Azure AI Search settings." ) from exc
784+ endpoint = str (azure_ai_search_settings .endpoint )
785+ credential = _resolve_credential (
786+ azure_ai_search_settings ,
787+ azure_credential = azure_credentials ,
788+ token_credential = token_credentials ,
789+ )
738790 search_index_client = _get_search_index_client (
739791 azure_ai_search_settings = azure_ai_search_settings ,
740792 azure_credential = azure_credentials ,
741793 token_credential = token_credentials ,
742794 )
743795 managed_client = True
796+ else :
797+ endpoint = search_endpoint
798+ credential = azure_credentials or token_credentials or (AzureKeyCredential (api_key ) if api_key else None )
744799
745800 super ().__init__ (
746801 search_index_client = search_index_client ,
802+ search_endpoint = endpoint ,
803+ search_credential = credential ,
747804 managed_client = managed_client ,
748805 embedding_generator = embedding_generator ,
749806 )
@@ -777,6 +834,8 @@ def get_collection(
777834 search_index_client = self .search_index_client ,
778835 search_client = search_client ,
779836 embedding_generator = embedding_generator or self .embedding_generator ,
837+ search_credential = self .search_credential ,
838+ search_endpoint = self .search_endpoint ,
780839 ** kwargs ,
781840 )
782841
0 commit comments