From a37d156ca1b6bd15bce2d6cc529f190b3d3bbfc2 Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Thu, 23 Jul 2026 16:14:39 +0900 Subject: [PATCH 1/2] Forward Azure AI Search query-source identity --- python/packages/azure-ai-search/AGENTS.md | 4 + python/packages/azure-ai-search/README.md | 18 ++ .../_context_provider.py | 38 +++- .../tests/test_aisearch_context_provider.py | 169 ++++++++++++++++++ 4 files changed, 228 insertions(+), 1 deletion(-) diff --git a/python/packages/azure-ai-search/AGENTS.md b/python/packages/azure-ai-search/AGENTS.md index 1cdbc2df43f..696f2e4cbfe 100644 --- a/python/packages/azure-ai-search/AGENTS.md +++ b/python/packages/azure-ai-search/AGENTS.md @@ -27,6 +27,10 @@ can be imported. Agentic **output mode** (`answer_synthesis`) and **extended rea (extractive + minimal) and raises an actionable `ValueError` (citing the installed version) if they are explicitly requested. Semantic mode is unaffected. +Agentic query-time user identity is also preview-only. It is gated by +`_query_source_authorization_available`; when enabled, `query_source_credential` supplies a +per-request Azure AI Search token through the `x-ms-query-source-authorization` header. + ## Usage ```python diff --git a/python/packages/azure-ai-search/README.md b/python/packages/azure-ai-search/README.md index 234ccdf7f7f..e359456f08e 100644 --- a/python/packages/azure-ai-search/README.md +++ b/python/packages/azure-ai-search/README.md @@ -31,6 +31,24 @@ ship only in the preview build. When a stable build is installed, the provider u output with minimal reasoning effort and raises an actionable error if a preview-only option is explicitly requested. Switching channels is a single change — the install — with no code edits. +### Query-time user identity + +Agentic retrieval can forward a caller-specific Azure AI Search authorization token when the +index uses permission fields for document-level access control. Pass an async credential for the +caller via `query_source_credential`; the provider requests the Azure AI Search resource scope and +forwards the token on each Knowledge Base retrieval request. This capability requires a preview +build of `azure-search-documents`, installed with `pip install --pre azure-search-documents`. + +```python +context_provider = AzureAISearchContextProvider( + endpoint=search_endpoint, + credential=application_credential, + mode="agentic", + knowledge_base_name=knowledge_base_name, + query_source_credential=user_credential, +) +``` + ### Basic Usage Example See the [Azure AI Search context provider examples](../../samples/02-agents/context_providers/azure_ai_search/) which demonstrate: diff --git a/python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py b/python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py index e93b7e061ec..1e15fa43563 100644 --- a/python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py +++ b/python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py @@ -9,6 +9,7 @@ from __future__ import annotations import importlib.metadata +import inspect import logging import sys from collections.abc import Awaitable, Callable @@ -124,6 +125,8 @@ KBRetrievalOutputMode = _preview_symbols["KnowledgeRetrievalOutputMode"] _preview_agentic_features_available = True +_query_source_authorization_available = _preview_agentic_features_available + AzureCredentialTypes = TokenCredential | AsyncTokenCredential EmbeddingFunction = Callable[[str], Awaitable[list[float]]] | SupportsGetEmbeddings[str, list[float], Any] KnowledgeBaseOutputModeLiteral = Literal["extractive_data", "answer_synthesis"] @@ -132,6 +135,7 @@ logger = logging.getLogger("agent_framework.azure_ai_search") _DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT = 10 +_AZURE_SEARCH_RESOURCE_SCOPE = "https://search.azure.com/.default" def _installed_search_documents_version() -> str: @@ -198,6 +202,7 @@ def __init__( azure_openai_api_key: str | None = None, knowledge_base_output_mode: KnowledgeBaseOutputModeLiteral = "extractive_data", retrieval_reasoning_effort: RetrievalReasoningEffortLiteral = "minimal", + query_source_credential: AsyncTokenCredential | None = None, agentic_message_history_count: int = _DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT, env_file_path: str | None = None, env_file_encoding: str | None = None, @@ -223,6 +228,7 @@ def __init__( azure_openai_api_key: Unused in semantic mode. knowledge_base_output_mode: Unused in semantic mode. retrieval_reasoning_effort: Unused in semantic mode. + query_source_credential: Unused in semantic mode. agentic_message_history_count: Unused in semantic mode. env_file_path: Optional ``.env`` file checked before process environment variables. env_file_encoding: Encoding for the ``.env`` file. @@ -251,6 +257,7 @@ def __init__( azure_openai_api_key: str | None = None, knowledge_base_output_mode: KnowledgeBaseOutputModeLiteral = "extractive_data", retrieval_reasoning_effort: RetrievalReasoningEffortLiteral = "minimal", + query_source_credential: AsyncTokenCredential | None = None, agentic_message_history_count: int = _DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT, env_file_path: str | None = None, env_file_encoding: str | None = None, @@ -276,6 +283,8 @@ def __init__( azure_openai_api_key: Optional Azure OpenAI API key for Knowledge Base creation. knowledge_base_output_mode: Output mode for Knowledge Base retrieval. retrieval_reasoning_effort: Reasoning effort for query planning. + query_source_credential: Async Azure credential used to authorize each retrieval query. + Requires a preview build of ``azure-search-documents``. agentic_message_history_count: Number of recent messages included in retrieval. env_file_path: Optional ``.env`` file checked before process environment variables. env_file_encoding: Encoding for the ``.env`` file. @@ -304,6 +313,7 @@ def __init__( azure_openai_api_key: str | None = None, knowledge_base_output_mode: KnowledgeBaseOutputModeLiteral = "extractive_data", retrieval_reasoning_effort: RetrievalReasoningEffortLiteral = "minimal", + query_source_credential: AsyncTokenCredential | None = None, agentic_message_history_count: int = _DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT, env_file_path: str | None = None, env_file_encoding: str | None = None, @@ -329,6 +339,8 @@ def __init__( azure_openai_api_key: Unused when connecting to an existing Knowledge Base. knowledge_base_output_mode: Output mode for Knowledge Base retrieval. retrieval_reasoning_effort: Reasoning effort for query planning. + query_source_credential: Async Azure credential used to authorize each retrieval query. + Requires a preview build of ``azure-search-documents``. agentic_message_history_count: Number of recent messages included in retrieval. env_file_path: Optional ``.env`` file checked before process environment variables. env_file_encoding: Encoding for the ``.env`` file. @@ -357,6 +369,7 @@ def __init__( azure_openai_api_key: str | None = None, knowledge_base_output_mode: KnowledgeBaseOutputModeLiteral = "extractive_data", retrieval_reasoning_effort: RetrievalReasoningEffortLiteral = "minimal", + query_source_credential: AsyncTokenCredential | None = None, agentic_message_history_count: int = _DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT, env_file_path: str | None = None, env_file_encoding: str | None = None, @@ -386,6 +399,8 @@ def __init__( azure_openai_api_key: Optional Azure OpenAI API key for Knowledge Base creation. knowledge_base_output_mode: Output mode for Knowledge Base retrieval. retrieval_reasoning_effort: Reasoning effort for query planning. + query_source_credential: Async Azure credential used to authorize each retrieval query. + Requires a preview build of ``azure-search-documents``. agentic_message_history_count: Number of recent messages included in retrieval. env_file_path: Optional ``.env`` file checked before process environment variables. env_file_encoding: Encoding for the ``.env`` file. @@ -413,6 +428,7 @@ def __init__( azure_openai_api_key: str | None = None, knowledge_base_output_mode: KnowledgeBaseOutputModeLiteral = "extractive_data", retrieval_reasoning_effort: RetrievalReasoningEffortLiteral = "minimal", + query_source_credential: AsyncTokenCredential | None = None, agentic_message_history_count: int = _DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT, env_file_path: str | None = None, env_file_encoding: str | None = None, @@ -441,6 +457,8 @@ def __init__( azure_openai_api_key: Azure OpenAI API key. knowledge_base_output_mode: Output mode for Knowledge Base retrieval. retrieval_reasoning_effort: Reasoning effort for Knowledge Base query planning. + query_source_credential: Async Azure credential used to authorize each agentic retrieval query. + Requires a preview build of ``azure-search-documents``. agentic_message_history_count: Number of recent messages for agentic mode. env_file_path: Path to environment file for loading settings. env_file_encoding: Encoding of the environment file. @@ -516,6 +534,7 @@ def __init__( self.azure_openai_api_key = azure_openai_api_key self.knowledge_base_output_mode = knowledge_base_output_mode self.retrieval_reasoning_effort = retrieval_reasoning_effort + self.query_source_credential = query_source_credential self.agentic_message_history_count = agentic_message_history_count self._use_existing_knowledge_base = False @@ -865,6 +884,13 @@ async def _ensure_knowledge_base(self) -> None: async def _agentic_search(self, messages: list[Message]) -> list[Message]: """Perform agentic retrieval with multi-hop reasoning.""" + if self.query_source_credential is not None and not _query_source_authorization_available: + installed = _installed_search_documents_version() + raise ValueError( + "query_source_credential requires a preview build of azure-search-documents " + f"(installed: {installed}). Install it with `pip install --pre azure-search-documents`." + ) + await self._ensure_knowledge_base() request_kwargs: dict[str, Any] = {"include_activity": True} @@ -908,7 +934,17 @@ async def _agentic_search(self, messages: list[Message]) -> list[Message]: if not self._retrieval_client: raise RuntimeError("Retrieval client not initialized.") - retrieval_result = await self._retrieval_client.retrieve(retrieval_request=retrieval_request) + retrieve_kwargs: dict[str, Any] = {"retrieval_request": retrieval_request} + if self.query_source_credential is not None: + access_token_result = self.query_source_credential.get_token(_AZURE_SEARCH_RESOURCE_SCOPE) + if not inspect.isawaitable(access_token_result): + raise TypeError( + "query_source_credential must be an async Azure credential. " + "Pass an azure.core.credentials_async.AsyncTokenCredential." + ) + access_token = await access_token_result + retrieve_kwargs["headers"] = {"x-ms-query-source-authorization": access_token.token} + retrieval_result = await self._retrieval_client.retrieve(**retrieve_kwargs) return self._parse_messages_from_kb_response(retrieval_result) diff --git a/python/packages/azure-ai-search/tests/test_aisearch_context_provider.py b/python/packages/azure-ai-search/tests/test_aisearch_context_provider.py index 22bb4df2a76..b3287a2a475 100644 --- a/python/packages/azure-ai-search/tests/test_aisearch_context_provider.py +++ b/python/packages/azure-ai-search/tests/test_aisearch_context_provider.py @@ -13,6 +13,7 @@ from agent_framework._sessions import AgentSession, SessionContext from agent_framework.exceptions import SettingNotFoundError from azure.core.credentials import AzureKeyCredential +from azure.core.pipeline.transport import AioHttpTransport from agent_framework_azure_ai_search import _context_provider from agent_framework_azure_ai_search._context_provider import ( @@ -54,6 +55,10 @@ async def __anext__(self): return doc +class _TransportRequestCaptured(Exception): + """Stop a test after the real Azure SDK has built the outgoing HTTP request.""" + + def _make_mock_index( fields: list[SimpleNamespace] | None = None, profiles: list[SimpleNamespace] | None = None, @@ -1898,6 +1903,170 @@ def test_multiple_messages_with_references(self) -> None: class TestBeforeRunAgentic: """Tests for before_run in agentic mode.""" + async def test_query_source_credential_requires_preview_sdk_before_transport(self) -> None: + query_source_credential = AsyncMock() + query_source_credential.get_token = AsyncMock(return_value=SimpleNamespace(token="user-token")) + mock_index_client = AsyncMock() + mock_index_client.get_knowledge_base.return_value = SimpleNamespace(knowledge_sources=[]) + + with ( + patch.object(_context_provider, "_query_source_authorization_available", False), + patch( + "agent_framework_azure_ai_search._context_provider.SearchIndexClient", + return_value=mock_index_client, + ), + patch.object( + AioHttpTransport, + "send", + new_callable=AsyncMock, + side_effect=AssertionError("HTTP transport must not be reached"), + ) as transport_send, + ): + provider = AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + knowledge_base_name="kb", + api_key="key", + mode="agentic", + query_source_credential=query_source_credential, + ) + session = AgentSession(session_id="test-session") + context = SessionContext( + input_messages=[Message(role="user", contents=["agentic question"])], + session_id="test-session", + ) + + with pytest.raises(ValueError, match="query_source_credential requires a preview build"): + await provider.before_run( + agent=cast(Any, None), + session=session, + context=context, + state=session.state.setdefault(provider.source_id, {}), + ) + + transport_send.assert_not_awaited() + + async def test_query_source_credential_is_serialized_as_http_header(self) -> None: + query_source_credential = AsyncMock() + query_source_credential.get_token = AsyncMock(return_value=SimpleNamespace(token="user-token")) + mock_index_client = AsyncMock() + mock_index_client.get_knowledge_base.return_value = SimpleNamespace(knowledge_sources=[]) + captured_headers: dict[str, str] = {} + + async def capture_request(_transport: AioHttpTransport, request: Any, **_kwargs: Any) -> None: + captured_headers.update(request.headers) + raise _TransportRequestCaptured + + with ( + patch.object(_context_provider, "_query_source_authorization_available", True), + patch( + "agent_framework_azure_ai_search._context_provider.SearchIndexClient", + return_value=mock_index_client, + ), + patch.object(AioHttpTransport, "send", new=capture_request), + ): + provider = AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + knowledge_base_name="kb", + api_key="key", + mode="agentic", + query_source_credential=query_source_credential, + ) + session = AgentSession(session_id="test-session") + context = SessionContext( + input_messages=[Message(role="user", contents=["agentic question"])], + session_id="test-session", + ) + + with pytest.raises(_TransportRequestCaptured): + await provider.before_run( + agent=cast(Any, None), + session=session, + context=context, + state=session.state.setdefault(provider.source_id, {}), + ) + + assert captured_headers["x-ms-query-source-authorization"] == "user-token" + + async def test_without_query_source_credential_omits_authorization(self) -> None: + mock_index_client = AsyncMock() + mock_index_client.get_knowledge_base.return_value = SimpleNamespace(knowledge_sources=[]) + mock_retrieval_client = AsyncMock() + mock_retrieval_client.retrieve.return_value = SimpleNamespace(response=[], references=None) + + with ( + patch( + "agent_framework_azure_ai_search._context_provider.SearchIndexClient", + return_value=mock_index_client, + ), + patch( + "agent_framework_azure_ai_search._context_provider.KnowledgeBaseRetrievalClient", + return_value=mock_retrieval_client, + ), + ): + provider = AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + knowledge_base_name="kb", + api_key="key", + mode="agentic", + ) + session = AgentSession(session_id="test-session") + context = SessionContext( + input_messages=[Message(role="user", contents=["agentic question"])], + session_id="test-session", + ) + + await provider.before_run( + agent=cast(Any, None), + session=session, + context=context, + state=session.state.setdefault(provider.source_id, {}), + ) + + assert mock_retrieval_client.retrieve.await_args is not None + assert "headers" not in mock_retrieval_client.retrieve.await_args.kwargs + + async def test_sync_query_source_credential_is_rejected_before_retrieval(self) -> None: + query_source_credential = Mock() + query_source_credential.get_token.return_value = SimpleNamespace(token="user-token") + mock_index_client = AsyncMock() + mock_index_client.get_knowledge_base.return_value = SimpleNamespace(knowledge_sources=[]) + mock_retrieval_client = AsyncMock() + + with ( + patch.object(_context_provider, "_query_source_authorization_available", True), + patch( + "agent_framework_azure_ai_search._context_provider.SearchIndexClient", + return_value=mock_index_client, + ), + patch( + "agent_framework_azure_ai_search._context_provider.KnowledgeBaseRetrievalClient", + return_value=mock_retrieval_client, + ), + ): + provider = cast(Any, AzureAISearchContextProvider)( + endpoint="https://test.search.windows.net", + knowledge_base_name="kb", + api_key="key", + mode="agentic", + query_source_credential=query_source_credential, + ) + session = AgentSession(session_id="test-session") + context = SessionContext( + input_messages=[Message(role="user", contents=["agentic question"])], + session_id="test-session", + ) + + with pytest.raises(TypeError, match="query_source_credential must be an async Azure credential"): + await provider.before_run( + agent=cast(Any, None), + session=session, + context=context, + state=session.state.setdefault(provider.source_id, {}), + ) + + query_source_credential.get_token.assert_called_once_with("https://search.azure.com/.default") + mock_retrieval_client.retrieve.assert_not_awaited() + async def test_agentic_mode_calls_agentic_search(self) -> None: provider = _make_provider() provider.mode = "agentic" From 6738284aefe00cdf9f2cc29f7414fdfaadba9881 Mon Sep 17 00:00:00 2001 From: Evan Mattson Date: Thu, 23 Jul 2026 18:21:26 +0900 Subject: [PATCH 2/2] Address query source credential review feedback --- python/packages/azure-ai-search/AGENTS.md | 3 +- python/packages/azure-ai-search/README.md | 11 +-- .../_context_provider.py | 50 ++++++------ .../tests/test_aisearch_context_provider.py | 76 ++++++++++++++++--- 4 files changed, 101 insertions(+), 39 deletions(-) diff --git a/python/packages/azure-ai-search/AGENTS.md b/python/packages/azure-ai-search/AGENTS.md index 696f2e4cbfe..b8e93143d88 100644 --- a/python/packages/azure-ai-search/AGENTS.md +++ b/python/packages/azure-ai-search/AGENTS.md @@ -29,7 +29,8 @@ they are explicitly requested. Semantic mode is unaffected. Agentic query-time user identity is also preview-only. It is gated by `_query_source_authorization_available`; when enabled, `query_source_credential` supplies a -per-request Azure AI Search token through the `x-ms-query-source-authorization` header. +per-request Azure AI Search token through the `x-ms-query-source-authorization` header. Both sync +and async Azure token credentials are supported, starting with `azure-search-documents>=12.1.0b1`. ## Usage diff --git a/python/packages/azure-ai-search/README.md b/python/packages/azure-ai-search/README.md index e359456f08e..35b35652d68 100644 --- a/python/packages/azure-ai-search/README.md +++ b/python/packages/azure-ai-search/README.md @@ -21,7 +21,7 @@ nothing to configure in code: | Channel | Install | Data-plane `api-version` (chosen by the SDK) | | --- | --- | --- | | **Stable** | `pip install azure-search-documents` (`>=12.0.0`) | `2026-04-01` | -| **Preview** | `pip install --pre azure-search-documents` (e.g. `12.1.0b1`) | `2026-05-01-preview` | +| **Preview** | `pip install --pre "azure-search-documents>=12.1.0b1"` | `2026-05-01-preview` | The provider never pins an `api-version`; the installed build selects its own, so newer releases work without code changes. @@ -34,10 +34,11 @@ explicitly requested. Switching channels is a single change — the install — ### Query-time user identity Agentic retrieval can forward a caller-specific Azure AI Search authorization token when the -index uses permission fields for document-level access control. Pass an async credential for the -caller via `query_source_credential`; the provider requests the Azure AI Search resource scope and -forwards the token on each Knowledge Base retrieval request. This capability requires a preview -build of `azure-search-documents`, installed with `pip install --pre azure-search-documents`. +index uses permission fields for document-level access control. Pass a sync or async Azure token +credential for the caller via `query_source_credential`; the provider requests the Azure AI Search +resource scope and forwards the token on each Knowledge Base retrieval request. This capability +requires `azure-search-documents>=12.1.0b1`, installed with +`pip install --pre "azure-search-documents>=12.1.0b1"`. ```python context_provider = AzureAISearchContextProvider( diff --git a/python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py b/python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py index 1e15fa43563..34f2fb67dbb 100644 --- a/python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py +++ b/python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py @@ -202,7 +202,7 @@ def __init__( azure_openai_api_key: str | None = None, knowledge_base_output_mode: KnowledgeBaseOutputModeLiteral = "extractive_data", retrieval_reasoning_effort: RetrievalReasoningEffortLiteral = "minimal", - query_source_credential: AsyncTokenCredential | None = None, + query_source_credential: AzureCredentialTypes | None = None, agentic_message_history_count: int = _DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT, env_file_path: str | None = None, env_file_encoding: str | None = None, @@ -257,7 +257,7 @@ def __init__( azure_openai_api_key: str | None = None, knowledge_base_output_mode: KnowledgeBaseOutputModeLiteral = "extractive_data", retrieval_reasoning_effort: RetrievalReasoningEffortLiteral = "minimal", - query_source_credential: AsyncTokenCredential | None = None, + query_source_credential: AzureCredentialTypes | None = None, agentic_message_history_count: int = _DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT, env_file_path: str | None = None, env_file_encoding: str | None = None, @@ -283,8 +283,8 @@ def __init__( azure_openai_api_key: Optional Azure OpenAI API key for Knowledge Base creation. knowledge_base_output_mode: Output mode for Knowledge Base retrieval. retrieval_reasoning_effort: Reasoning effort for query planning. - query_source_credential: Async Azure credential used to authorize each retrieval query. - Requires a preview build of ``azure-search-documents``. + query_source_credential: Sync or async Azure credential used to authorize each retrieval query. + Requires ``azure-search-documents>=12.1.0b1``. agentic_message_history_count: Number of recent messages included in retrieval. env_file_path: Optional ``.env`` file checked before process environment variables. env_file_encoding: Encoding for the ``.env`` file. @@ -313,7 +313,7 @@ def __init__( azure_openai_api_key: str | None = None, knowledge_base_output_mode: KnowledgeBaseOutputModeLiteral = "extractive_data", retrieval_reasoning_effort: RetrievalReasoningEffortLiteral = "minimal", - query_source_credential: AsyncTokenCredential | None = None, + query_source_credential: AzureCredentialTypes | None = None, agentic_message_history_count: int = _DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT, env_file_path: str | None = None, env_file_encoding: str | None = None, @@ -339,8 +339,8 @@ def __init__( azure_openai_api_key: Unused when connecting to an existing Knowledge Base. knowledge_base_output_mode: Output mode for Knowledge Base retrieval. retrieval_reasoning_effort: Reasoning effort for query planning. - query_source_credential: Async Azure credential used to authorize each retrieval query. - Requires a preview build of ``azure-search-documents``. + query_source_credential: Sync or async Azure credential used to authorize each retrieval query. + Requires ``azure-search-documents>=12.1.0b1``. agentic_message_history_count: Number of recent messages included in retrieval. env_file_path: Optional ``.env`` file checked before process environment variables. env_file_encoding: Encoding for the ``.env`` file. @@ -369,7 +369,7 @@ def __init__( azure_openai_api_key: str | None = None, knowledge_base_output_mode: KnowledgeBaseOutputModeLiteral = "extractive_data", retrieval_reasoning_effort: RetrievalReasoningEffortLiteral = "minimal", - query_source_credential: AsyncTokenCredential | None = None, + query_source_credential: AzureCredentialTypes | None = None, agentic_message_history_count: int = _DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT, env_file_path: str | None = None, env_file_encoding: str | None = None, @@ -399,8 +399,8 @@ def __init__( azure_openai_api_key: Optional Azure OpenAI API key for Knowledge Base creation. knowledge_base_output_mode: Output mode for Knowledge Base retrieval. retrieval_reasoning_effort: Reasoning effort for query planning. - query_source_credential: Async Azure credential used to authorize each retrieval query. - Requires a preview build of ``azure-search-documents``. + query_source_credential: Sync or async Azure credential used to authorize each retrieval query. + Requires ``azure-search-documents>=12.1.0b1``. agentic_message_history_count: Number of recent messages included in retrieval. env_file_path: Optional ``.env`` file checked before process environment variables. env_file_encoding: Encoding for the ``.env`` file. @@ -428,7 +428,7 @@ def __init__( azure_openai_api_key: str | None = None, knowledge_base_output_mode: KnowledgeBaseOutputModeLiteral = "extractive_data", retrieval_reasoning_effort: RetrievalReasoningEffortLiteral = "minimal", - query_source_credential: AsyncTokenCredential | None = None, + query_source_credential: AzureCredentialTypes | None = None, agentic_message_history_count: int = _DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT, env_file_path: str | None = None, env_file_encoding: str | None = None, @@ -457,14 +457,17 @@ def __init__( azure_openai_api_key: Azure OpenAI API key. knowledge_base_output_mode: Output mode for Knowledge Base retrieval. retrieval_reasoning_effort: Reasoning effort for Knowledge Base query planning. - query_source_credential: Async Azure credential used to authorize each agentic retrieval query. - Requires a preview build of ``azure-search-documents``. + query_source_credential: Sync or async Azure credential used to authorize each agentic retrieval query. + Requires ``azure-search-documents>=12.1.0b1``. agentic_message_history_count: Number of recent messages for agentic mode. env_file_path: Path to environment file for loading settings. env_file_encoding: Encoding of the environment file. """ super().__init__(source_id) + 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.") + required: list[str | tuple[str, ...]] ignored_agentic_field: Literal["index_name", "knowledge_base_name"] | None = None explicit_index_name = index_name is not None @@ -888,8 +891,16 @@ async def _agentic_search(self, messages: list[Message]) -> list[Message]: installed = _installed_search_documents_version() raise ValueError( "query_source_credential requires a preview build of azure-search-documents " - f"(installed: {installed}). Install it with `pip install --pre azure-search-documents`." + f"(installed: {installed}). Install `azure-search-documents>=12.1.0b1`." + ) + + query_source_authorization: str | None = None + if self.query_source_credential is not None: + access_token_result = self.query_source_credential.get_token(_AZURE_SEARCH_RESOURCE_SCOPE) + access_token = ( + await access_token_result if inspect.isawaitable(access_token_result) else access_token_result ) + query_source_authorization = access_token.token await self._ensure_knowledge_base() @@ -935,15 +946,8 @@ async def _agentic_search(self, messages: list[Message]) -> list[Message]: if not self._retrieval_client: raise RuntimeError("Retrieval client not initialized.") retrieve_kwargs: dict[str, Any] = {"retrieval_request": retrieval_request} - if self.query_source_credential is not None: - access_token_result = self.query_source_credential.get_token(_AZURE_SEARCH_RESOURCE_SCOPE) - if not inspect.isawaitable(access_token_result): - raise TypeError( - "query_source_credential must be an async Azure credential. " - "Pass an azure.core.credentials_async.AsyncTokenCredential." - ) - access_token = await access_token_result - retrieve_kwargs["headers"] = {"x-ms-query-source-authorization": access_token.token} + if query_source_authorization is not None: + retrieve_kwargs["headers"] = {"x-ms-query-source-authorization": query_source_authorization} retrieval_result = await self._retrieval_client.retrieve(**retrieve_kwargs) return self._parse_messages_from_kb_response(retrieval_result) diff --git a/python/packages/azure-ai-search/tests/test_aisearch_context_provider.py b/python/packages/azure-ai-search/tests/test_aisearch_context_provider.py index b3287a2a475..afb2687530a 100644 --- a/python/packages/azure-ai-search/tests/test_aisearch_context_provider.py +++ b/python/packages/azure-ai-search/tests/test_aisearch_context_provider.py @@ -334,6 +334,22 @@ def test_agentic_with_index_generates_kb_name(self) -> None: assert provider._use_existing_knowledge_base is False assert provider.knowledge_base_name == "idx-kb" + def test_invalid_query_source_credential_raises_before_client_construction(self) -> None: + with ( + patch("agent_framework_azure_ai_search._context_provider.SearchIndexClient") as index_client_cls, + pytest.raises(TypeError, match="query_source_credential must be an Azure TokenCredential"), + ): + cast(Any, AzureAISearchContextProvider)( + source_id="s", + endpoint="https://test.search.windows.net", + knowledge_base_name="my-kb", + api_key="key", + mode="agentic", + query_source_credential=object(), + ) + + index_client_cls.assert_not_called() + def test_agentic_explicit_kb_ignores_env_index_name(self) -> None: with patch.dict(os.environ, {"AZURE_SEARCH_INDEX_NAME": "env-index"}, clear=False): provider = AzureAISearchContextProvider( @@ -1987,6 +2003,43 @@ async def capture_request(_transport: AioHttpTransport, request: Any, **_kwargs: assert captured_headers["x-ms-query-source-authorization"] == "user-token" + async def test_query_source_token_failure_happens_before_knowledge_base_access(self) -> None: + query_source_credential = AsyncMock() + query_source_credential.get_token = AsyncMock(side_effect=RuntimeError("token acquisition failed")) + mock_index_client = AsyncMock() + mock_index_client.get_knowledge_base.return_value = SimpleNamespace(knowledge_sources=[]) + + with ( + patch.object(_context_provider, "_query_source_authorization_available", True), + patch( + "agent_framework_azure_ai_search._context_provider.SearchIndexClient", + return_value=mock_index_client, + ), + ): + provider = AzureAISearchContextProvider( + endpoint="https://test.search.windows.net", + knowledge_base_name="kb", + api_key="key", + mode="agentic", + query_source_credential=query_source_credential, + ) + session = AgentSession(session_id="test-session") + context = SessionContext( + input_messages=[Message(role="user", contents=["agentic question"])], + session_id="test-session", + ) + + with pytest.raises(RuntimeError, match="token acquisition failed"): + await provider.before_run( + agent=cast(Any, None), + session=session, + context=context, + state=session.state.setdefault(provider.source_id, {}), + ) + + mock_index_client.get_knowledge_base.assert_not_awaited() + mock_index_client.create_or_update_knowledge_base.assert_not_awaited() + async def test_without_query_source_credential_omits_authorization(self) -> None: mock_index_client = AsyncMock() mock_index_client.get_knowledge_base.return_value = SimpleNamespace(knowledge_sources=[]) @@ -2025,12 +2078,13 @@ async def test_without_query_source_credential_omits_authorization(self) -> None assert mock_retrieval_client.retrieve.await_args is not None assert "headers" not in mock_retrieval_client.retrieve.await_args.kwargs - async def test_sync_query_source_credential_is_rejected_before_retrieval(self) -> None: + async def test_sync_query_source_credential_is_forwarded_to_retrieval(self) -> None: query_source_credential = Mock() query_source_credential.get_token.return_value = SimpleNamespace(token="user-token") mock_index_client = AsyncMock() mock_index_client.get_knowledge_base.return_value = SimpleNamespace(knowledge_sources=[]) mock_retrieval_client = AsyncMock() + mock_retrieval_client.retrieve.return_value = SimpleNamespace(response=[], references=None) with ( patch.object(_context_provider, "_query_source_authorization_available", True), @@ -2043,7 +2097,7 @@ async def test_sync_query_source_credential_is_rejected_before_retrieval(self) - return_value=mock_retrieval_client, ), ): - provider = cast(Any, AzureAISearchContextProvider)( + provider = AzureAISearchContextProvider( endpoint="https://test.search.windows.net", knowledge_base_name="kb", api_key="key", @@ -2056,16 +2110,18 @@ async def test_sync_query_source_credential_is_rejected_before_retrieval(self) - session_id="test-session", ) - with pytest.raises(TypeError, match="query_source_credential must be an async Azure credential"): - await provider.before_run( - agent=cast(Any, None), - session=session, - context=context, - state=session.state.setdefault(provider.source_id, {}), - ) + await provider.before_run( + agent=cast(Any, None), + session=session, + context=context, + state=session.state.setdefault(provider.source_id, {}), + ) query_source_credential.get_token.assert_called_once_with("https://search.azure.com/.default") - mock_retrieval_client.retrieve.assert_not_awaited() + assert mock_retrieval_client.retrieve.await_args is not None + assert mock_retrieval_client.retrieve.await_args.kwargs["headers"] == { + "x-ms-query-source-authorization": "user-token" + } async def test_agentic_mode_calls_agentic_search(self) -> None: provider = _make_provider()