Python: [Feature]: Pass user identity to the AI Search Provider

[damienaicheh]

Description

Today when we use the AI Search Provider using the agent framework package with the AzureAISearchContextProvider class we can't pass the token of the user to authenticate it to AI Search, this will be great to have it.

https://github.com/microsoft/agent-framework/blob/main/python/samples/02-agents/context_providers/azure_ai_search/search_context_agentic.py

Code Sample

Here is the workaround I did to implement it:


from collections.abc import Sequence

from agent_framework import Message
from agent_framework.azure import AzureAISearchContextProvider
from azure.core.credentials_async import AsyncTokenCredential

# Scope used to request a Microsoft Entra token for Azure AI Search query-time
# permission enforcement (the same token is forwarded as the query source identity).
SEARCH_TOKEN_SCOPE = "https://search.azure.com/.default"


class IdentityAwareAzureAISearchContextProvider(AzureAISearchContextProvider):
    """Azure AI Search provider that enforces document-level permissions at query time.

    The index was populated with ``GroupIds`` / ``UserIds`` permission fields that the
    indexer synchronizes from the source ADLS Gen2 ACLs. To have the service trim the
    retrieval results for the caller, every agentic retrieval request must carry the
    caller's Microsoft Entra token in the ``x-ms-query-source-authorization`` header.

    This subclass forwards that token (acquired from ``query_source_credential``, which
    defaults to the provider credential) on each agentic retrieve call.
    """

    def __init__(
        self,
        *args: object,
        query_source_credential: AsyncTokenCredential | None = None,
        **kwargs: object,
    ) -> None:
        super().__init__(*args, **kwargs)
        self._query_source_credential = query_source_credential or self.credential
        self._retrieve_patched = False

    async def _agentic_search(self, messages: Sequence[Message]) -> list[Message]:
        await self._ensure_knowledge_base()
        self._patch_retrieve_with_identity()
        return await super()._agentic_search(list(messages))

    def _patch_retrieve_with_identity(self) -> None:
        client = self._retrieval_client
        if self._retrieve_patched or client is None:
            return

        original_retrieve = client.retrieve
        query_source_credential = self._query_source_credential

        async def retrieve_with_identity(*args: object, **kwargs: object):
            if not kwargs.get("x_ms_query_source_authorization"):
                access_token = await query_source_credential.get_token(
                    SEARCH_TOKEN_SCOPE
                )
                kwargs["x_ms_query_source_authorization"] = access_token.token
            return await original_retrieve(*args, **kwargs)

        client.retrieve = retrieve_with_identity  # type: ignore[method-assign]
        self._retrieve_patched = True


And used it like this:



        aisearch_context_provider = IdentityAwareAzureAISearchContextProvider(
            source_id="search_provider",
            endpoint=search_endpoint,
            credential=credential,
            mode="agentic",
            knowledge_base_name=knowledge_base_name,
            knowledge_base_output_mode="answer_synthesis",
            retrieval_reasoning_effort="low",
        )

Language/SDK

Both

[kmishra1204]

I opened #6915 for the Python side of this request. The patch adds an explicit optional async query_source_credential on AzureAISearchContextProvider and forwards that credential's bearer token only on agentic Knowledge Base retrieval requests, so existing API key and managed-identity flows remain unchanged unless a caller opts in.