diff --git a/python/packages/azure-ai-search/README.md b/python/packages/azure-ai-search/README.md index 234ccdf7f7f..7c0567019b3 100644 --- a/python/packages/azure-ai-search/README.md +++ b/python/packages/azure-ai-search/README.md @@ -31,6 +31,23 @@ 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. + +```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 347e03c8186..e274523002e 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 @@ -130,6 +131,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: @@ -196,6 +198,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, @@ -221,6 +224,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. @@ -249,6 +253,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, @@ -274,6 +279,7 @@ 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: Optional async Azure credential for per-query user identity forwarding. 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. @@ -302,6 +308,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, @@ -327,6 +334,7 @@ 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: Optional async Azure credential for per-query user identity forwarding. 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. @@ -355,6 +363,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, @@ -384,6 +393,7 @@ 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: Optional async Azure credential for per-query user identity forwarding. 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. @@ -411,6 +421,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, @@ -439,6 +450,7 @@ 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: Optional async Azure credential for per-query user identity forwarding. 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. @@ -514,6 +526,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 @@ -886,10 +899,26 @@ 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 query_source_authorization := await self._query_source_authorization(): + retrieve_kwargs["query_source_authorization"] = query_source_authorization + retrieval_result = await self._retrieval_client.retrieve(**retrieve_kwargs) return self._parse_messages_from_kb_response(retrieval_result) + async def _query_source_authorization(self) -> str | None: + """Return a per-query Azure AI Search authorization token, when configured.""" + if self.query_source_credential is None: + return 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 + return access_token.token + @staticmethod def _prepare_messages_for_kb_search(messages: list[Message]) -> list[KnowledgeBaseMessage]: """Convert framework Messages to KnowledgeBaseMessages for agentic retrieval. 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 17ae6eb06af..4b22b6401db 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 @@ -1351,6 +1351,60 @@ async def test_minimal_reasoning_returns_results(self) -> None: assert len(results) == 1 assert results[0].text == "Answer text" assert results[0].role == "assistant" + retrieve_call = mock_retrieval.retrieve.await_args + assert retrieve_call is not None + assert "query_source_authorization" not in retrieve_call.kwargs + + async def test_query_source_credential_forwards_authorization_token(self) -> None: + query_source_credential = AsyncMock() + query_source_credential.get_token = AsyncMock(return_value=SimpleNamespace(token="user-token")) + provider = _make_provider(query_source_credential=query_source_credential) + provider._knowledge_base_initialized = True + provider.knowledge_base_name = "kb" + provider.retrieval_reasoning_effort = "minimal" + + mock_content = Mock() + mock_content.text = "Answer text" + mock_message = Mock() + mock_message.role = "assistant" + mock_message.content = [mock_content] + mock_result = Mock() + mock_result.response = [mock_message] + mock_result.references = None + + mock_retrieval = AsyncMock() + mock_retrieval.retrieve = AsyncMock(return_value=mock_result) + provider._retrieval_client = mock_retrieval + + with patch( + "agent_framework_azure_ai_search._context_provider.KnowledgeBaseMessageTextContent", + type(mock_content), + ): + results = await provider._agentic_search([Message(role="user", contents=["test query"])]) + + assert len(results) == 1 + assert results[0].text == "Answer text" + query_source_credential.get_token.assert_awaited_once_with("https://search.azure.com/.default") + retrieve_call = mock_retrieval.retrieve.await_args + assert retrieve_call is not None + assert retrieve_call.kwargs["query_source_authorization"] == "user-token" + + async def test_query_source_credential_requires_async_credential(self) -> None: + query_source_credential = Mock() + query_source_credential.get_token = Mock(return_value=SimpleNamespace(token="user-token")) + provider = _make_provider(query_source_credential=query_source_credential) + provider._knowledge_base_initialized = True + provider.knowledge_base_name = "kb" + provider.retrieval_reasoning_effort = "minimal" + + mock_retrieval = AsyncMock() + provider._retrieval_client = mock_retrieval + + with pytest.raises(TypeError, match="query_source_credential must be an async Azure credential"): + await provider._agentic_search([Message(role="user", contents=["test query"])]) + + query_source_credential.get_token.assert_called_once_with("https://search.azure.com/.default") + mock_retrieval.retrieve.assert_not_awaited() async def test_non_minimal_reasoning_uses_messages(self) -> None: provider = _make_provider()