diff --git a/python/packages/azure-ai-search/AGENTS.md b/python/packages/azure-ai-search/AGENTS.md index fe2c296652b..3d946b6cae6 100644 --- a/python/packages/azure-ai-search/AGENTS.md +++ b/python/packages/azure-ai-search/AGENTS.md @@ -75,8 +75,9 @@ validate query credentials before resolving connection settings or creating owne Both constructors create owned clients through `_create_index_client`, reusing `AzureAISearchSettings` and the public AF `load_settings` / `SecretString` API with the `AZURE_SEARCH_` prefix. Explicit settings override the selected `.env` file, which overrides the process environment. `env_file_path` and -`env_file_encoding` are supported on both. Injected clients bypass settings entirely and reject explicit -connection/file overrides; store-created collections reuse its resolved client without another load. +`env_file_encoding` are supported on both. Injected clients take precedence over connection/file +overrides, bypass settings entirely, and remain caller-owned; store-created collections reuse its +resolved client without another load. The store tracks open collections in insertion order with constant-time removal. A collection's private close callback removes its registration and is cleared even when client cleanup fails; store shutdown still attempts all remaining owned clients without closing borrowed credentials. diff --git a/python/packages/azure-ai-search/README.md b/python/packages/azure-ai-search/README.md index 99075ce8f9a..88f15497daf 100644 --- a/python/packages/azure-ai-search/README.md +++ b/python/packages/azure-ai-search/README.md @@ -44,9 +44,10 @@ In agentic mode, pass `knowledge_source_params` to set per-source retrieval opti such as `SearchIndexKnowledgeSourceParams(knowledge_source_name=..., filter_add_on=...)`. Each entry replaces the provider's defaults for the source with the same name. -Injected SDK clients bypass connection settings and remain caller-owned unless -`managed_client=True`. Use async context managers to close owned store/collection -clients; credentials and embedding clients remain caller-owned. +Injected SDK clients take precedence over connection settings and remain caller-owned, +including when `managed_client=True` is passed for compatibility. Use async context +managers to close owned store/collection clients; credentials and embedding clients +remain caller-owned. ## Vector collections and stores (experimental) diff --git a/python/packages/azure-ai-search/agent_framework_azure_ai_search/_vector_store.py b/python/packages/azure-ai-search/agent_framework_azure_ai_search/_vector_store.py index 4cfd141ad4f..be010c805db 100644 --- a/python/packages/azure-ai-search/agent_framework_azure_ai_search/_vector_store.py +++ b/python/packages/azure-ai-search/agent_framework_azure_ai_search/_vector_store.py @@ -249,7 +249,7 @@ class AzureAISearchCollection(BaseVectorCollection[str, ModelT], BaseVectorSearc """An Azure Search index with batch CRUD and vector/keyword-hybrid retrieval. Credentials are always caller-owned. Clients created here are closed by ``close()`` - or the async context manager; injected clients are borrowed unless ``managed_client=True``. + or the async context manager; injected clients are always borrowed. Existing indexes are never updated by ``ensure_collection_exists``. """ @@ -289,9 +289,9 @@ def __init__( credential: Caller-owned async Azure token or key credential. api_key: API key, falling back to AZURE_SEARCH_API_KEY. search_client: Borrowed document client; may be used without an index client. - Injected clients bypass settings and cannot be combined with connection or file overrides. - index_client: Borrowed index administration client, mutually exclusive with connection/file overrides. - managed_client: Explicitly take ownership of injected clients when True. + Injected clients take precedence over connection and file settings. + index_client: Borrowed index administration client; takes precedence over connection and file settings. + managed_client: Retained for compatibility; injected clients remain caller-owned. vector_search: Native SDK profiles, algorithms, vectorizers and compression. index_options: Native SDK SearchIndex configuration, excluding name/fields/vector_search. is_alias: Use an existing index alias; disallow index creation/deletion. @@ -302,18 +302,12 @@ def __init__( """ if query_source_credential is not None and not callable(getattr(query_source_credential, "get_token", None)): raise TypeError("query_source_credential must be an Azure TokenCredential or AsyncTokenCredential.") - if (index_client is not None or search_client is not None) and any( - value is not None for value in (endpoint, credential, api_key, env_file_path, env_file_encoding) - ): - raise ValueError( - "Injected clients cannot be combined with endpoint, credential, api_key, or env_file options." - ) super().__init__( record_type, definition=definition, collection_name=collection_name, embedding_generator=embedding_generator, - managed_client=managed_client if managed_client is not None else False, + managed_client=False, ) if self.definition.key_field.is_auto_generated: raise NotImplementedError("Azure Search requires application-provided string keys.") @@ -344,16 +338,12 @@ def __init__( env_file_encoding=env_file_encoding, ) self._owned_clients.append(index_client) - elif index_client is not None and managed_client: - self._owned_clients.append(index_client) self.index_client = index_client if search_client is None: if index_client is None: raise ValueError("An index_client or search_client is required.") search_client = index_client.get_search_client(self.collection_name, user_agent=get_user_agent()) self._owned_clients.append(search_client) - elif managed_client: - self._owned_clients.append(search_client) self.search_client = search_client self.managed_client = bool(self._owned_clients) @@ -908,7 +898,7 @@ class AzureAISearchStore(BaseVectorStore): """Factory and index administration for an Azure AI Search service. Closing a store closes its collection clients and its owned index client. - Injected index clients and credentials remain caller-owned by default. + Injected index clients and credentials remain caller-owned. """ def __init__( @@ -931,10 +921,10 @@ def __init__( endpoint: Service URL, or AZURE_SEARCH_ENDPOINT. credential: Caller-owned async token or key credential. api_key: Plain or masked key, or AZURE_SEARCH_API_KEY. - index_client: Borrowed SDK index client, mutually exclusive with connection/file overrides. + index_client: Borrowed SDK index client; takes precedence over connection and file settings. Bypasses environment and file settings. embedding_generator: Default local embedding client for collections. - managed_client: Take ownership of an injected index client when True. + managed_client: Retained for compatibility; injected index clients remain caller-owned. allow_preview: Permit explicitly requested, SDK-supported preview capabilities. query_source_credential: Caller identity forwarded on every collection read. env_file_path: Optional settings file. @@ -942,13 +932,7 @@ def __init__( """ if query_source_credential is not None and not callable(getattr(query_source_credential, "get_token", None)): raise TypeError("query_source_credential must be an Azure TokenCredential or AsyncTokenCredential.") - if index_client is not None and any( - value is not None for value in (endpoint, credential, api_key, env_file_path, env_file_encoding) - ): - raise ValueError( - "Injected clients cannot be combined with endpoint, credential, api_key, or env_file options." - ) - owned = index_client is None or managed_client is True + owned = index_client is None super().__init__(embedding_generator=embedding_generator, managed_client=owned) self.index_client = ( index_client diff --git a/python/packages/azure-ai-search/tests/test_vector_store.py b/python/packages/azure-ai-search/tests/test_vector_store.py index 3efd49c9c0e..94e5ff34ebf 100644 --- a/python/packages/azure-ai-search/tests/test_vector_store.py +++ b/python/packages/azure-ai-search/tests/test_vector_store.py @@ -888,7 +888,7 @@ async def test_owned_and_borrowed_clients(client: Mock) -> None: client.close.assert_awaited_once() async with AzureAISearchCollection(dict, definition=definition(), search_client=client, managed_client=True): pass - assert client.close.await_count == 2 + client.close.assert_awaited_once() @pytest.mark.parametrize("managed", [False, True]) @@ -926,7 +926,7 @@ async def test_store_releases_closed_collections_but_retains_open_clients(manage await store.close() await store.close() active_client.close.assert_awaited_once() - assert index_client.close.await_count == int(managed) + index_client.close.assert_not_awaited() assert not store._collections gc.collect() assert active_ref() is None @@ -945,7 +945,7 @@ async def test_store_releases_collection_after_close_error(client: Mock, error: await collection.close() await store.close() client.close.assert_awaited_once() - index_client.close.assert_awaited_once() + index_client.close.assert_not_awaited() async def test_close_remaining_clients_if_one_close_fails(client: Mock) -> None: @@ -960,7 +960,7 @@ async def test_close_remaining_clients_if_one_close_fails(client: Mock) -> None: await store.close() client.close.assert_awaited_once() other_client.close.assert_awaited_once() - index_client.close.assert_awaited_once() + index_client.close.assert_not_awaited() assert not store._collections @@ -1218,18 +1218,19 @@ async def test_injected_clients_bypass_settings_and_preserve_ownership( await connector.close() await connector.close() settings_loader.assert_not_called() - assert index_client.close.await_count == int(managed and kind != "collection-search") - assert client.close.await_count == int(kind == "collection-index" or (managed and kind != "store")) + index_client.close.assert_not_awaited() + assert client.close.await_count == int(kind == "collection-index") @pytest.mark.parametrize("kind", ["store", "collection-index", "collection-search", "collection-both"]) @pytest.mark.parametrize("override", ["endpoint", "credential", "api_key", "env_file_path", "env_file_encoding"]) -def test_injected_clients_reject_explicit_settings_before_resolution( +async def test_injected_clients_ignore_explicit_settings_before_resolution( client: Mock, kind: str, override: str, ) -> None: index_client = Mock(spec=SearchIndexClient) + index_client.get_search_client.return_value = client connection_options = { "endpoint": "https://explicit.search.windows.net", "credential": AzureKeyCredential("credential-key"), @@ -1244,16 +1245,24 @@ def test_injected_clients_reject_explicit_settings_before_resolution( options["search_client"] = client with ( patch("agent_framework_azure_ai_search._vector_store.load_settings") as settings_loader, - pytest.raises(ValueError, match="Injected clients cannot be combined"), + patch("agent_framework_azure_ai_search._vector_store.SearchIndexClient") as factory, ): if kind == "store": - AzureAISearchStore(**options) + store = AzureAISearchStore(**options, managed_client=True) + assert store.index_client is index_client + await store.close() else: - AzureAISearchCollection(dict, definition=definition(), **options) + collection = AzureAISearchCollection(dict, definition=definition(), **options, managed_client=True) + assert collection.index_client is (None if kind == "collection-search" else index_client) + if kind == "collection-search": + assert collection.search_client is client + await collection.close() settings_loader.assert_not_called() - index_client.get_search_client.assert_not_called() - index_client.close.assert_not_called() - client.close.assert_not_called() + factory.assert_not_called() + assert index_client.get_search_client.call_count == int(kind == "collection-index") + index_client.close.assert_not_awaited() + if kind != "collection-index": + client.close.assert_not_awaited() async def test_store_collections_reuse_resolved_clients_without_settings(