Python: forward Azure AI Search query-source identity#6915
Conversation
There was a problem hiding this comment.
Pull request overview
Adds an opt-in way for Python agentic Azure AI Search retrieval to forward a caller’s query-time identity (Entra token) to Knowledge Base retrieval requests, enabling document-level permission trimming scenarios without changing existing API key/managed identity defaults.
Changes:
- Added an optional
query_source_credentialparameter toAzureAISearchContextProviderand forwarded its token toKnowledgeBaseRetrievalClient.retrieve(...)viax_ms_query_source_authorization. - Documented query-time identity forwarding in the package README.
- Added regression tests covering both default behavior (no header) and forwarding behavior (header present + correct scope).
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py | Adds query_source_credential, obtains the Search scope token, and forwards it on agentic retrieval calls. |
| python/packages/azure-ai-search/tests/test_aisearch_context_provider.py | Adds assertions and a new test verifying the authorization header forwarding behavior. |
| python/packages/azure-ai-search/README.md | Documents the query-time user identity option and shows how to pass query_source_credential. |
| """Return a per-query Azure AI Search authorization token, when configured.""" | ||
| if self.query_source_credential is None: | ||
| return None | ||
| access_token = await self.query_source_credential.get_token(_AZURE_SEARCH_RESOURCE_SCOPE) | ||
| return access_token.token |
| 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 Azure credential for per-query user identity forwarding. |
|
Updated in bd48f59 to address the async credential feedback: the docs now say query_source_credential must be an async Azure credential, _query_source_authorization() raises a clear TypeError if a sync credential is passed, and there is regression coverage for that misuse case. |
|
@kmishra1204 please read the following Contributor License Agreement(CLA). If you agree with the CLA, please reply with the following information.
Contributor License AgreementContribution License AgreementThis Contribution License Agreement (“Agreement”) is agreed to by the party signing below (“You”),
|
Summary
query_source_credentialtoAzureAISearchContextProviderThis addresses the Python side of #6864. The option is explicit so existing API key and managed identity flows keep their current behavior unless a caller supplies an async query-source credential.
Validation
uv run --project python --dev --no-sync pytest -q python/packages/azure-ai-search/tests/test_aisearch_context_provider.py::TestAgenticSearch::test_minimal_reasoning_returns_results python/packages/azure-ai-search/tests/test_aisearch_context_provider.py::TestAgenticSearch::test_query_source_credential_forwards_authorization_token python/packages/azure-ai-search/tests/test_aisearch_context_provider.py::TestAgenticSearch::test_query_source_credential_requires_async_credentialuv run --project python --dev --no-sync ruff check python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py python/packages/azure-ai-search/tests/test_aisearch_context_provider.pyuv run --project python --dev --no-sync pyright python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.pygit diff --check -- python/CHANGELOG.md python/packages/azure-ai-search/README.md python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py python/packages/azure-ai-search/tests/test_aisearch_context_provider.py