Skip to content

Python: forward Azure AI Search query-source identity #7278

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Evan Mattson (moonbox3) merged 3 commits into from
Aug 3, 2026
Merged

Python: forward Azure AI Search query-source identity #7278

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
a37d156
Forward Azure AI Search query-source identity
moonbox3 Jul 23, 2026
6738284
Address query source credential review feedback
moonbox3 Jul 23, 2026
f835318
Merge upstream main into azure-search-query-source-identity-latest
moonbox3 Jul 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions python/packages/azure-ai-search/AGENTS.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,11 @@ 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. Both sync
and async Azure token credentials are supported, starting with `azure-search-documents>=12.1.0b1`.

## Usage

```python
Expand Down
21 changes: 20 additions & 1 deletion python/packages/azure-ai-search/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -31,6 +31,25 @@ 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 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(
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:
Expand Down
42 changes: 41 additions & 1 deletion python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
from __future__ import annotations

import importlib.metadata
import inspect
import logging
import sys
from collections.abc import Awaitable, Callable
Expand Down Expand Up @@ -126,6 +127,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"]
Expand All @@ -134,6 +137,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:
Expand Down Expand Up @@ -200,6 +204,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: 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,
Expand All @@ -225,6 +230,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.
Expand Down Expand Up @@ -253,6 +259,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: 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,
Expand All @@ -278,6 +285,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: 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.
Expand Down Expand Up @@ -306,6 +315,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: 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,
Expand All @@ -331,6 +341,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: 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.
Expand Down Expand Up @@ -359,6 +371,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: 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,
Expand Down Expand Up @@ -388,6 +401,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: 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.
Expand Down Expand Up @@ -415,6 +430,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: 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,
Expand Down Expand Up @@ -443,12 +459,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: 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
Expand Down Expand Up @@ -518,6 +539,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
Expand Down Expand Up @@ -868,6 +890,21 @@ 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 `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()

request_kwargs: dict[str, Any] = {"include_activity": True}
Expand Down Expand Up @@ -911,7 +948,10 @@ 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 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)

Expand Down
Loading
Loading