From 76f1a08535c53f6db45eb7dcf2b60a2c87168fff Mon Sep 17 00:00:00 2001 From: Ayrton Werck Date: Fri, 19 Jun 2026 07:33:25 +0000 Subject: [PATCH 1/2] Python: support knowledge_source_params in AzureAISearchContextProvider Add an optional keyword-only `knowledge_source_params` to the agentic overloads of `AzureAISearchContextProvider`, forwarded verbatim into both `KnowledgeBaseRetrievalRequest` constructions in `_agentic_search`. This lets callers set per-source agentic retrieval options such as `filter_add_on` (OData filtering), `include_reference_source_data`, and similar. Default `None` keeps existing behavior unchanged. The parameter is agentic-only: it is omitted from the semantic overload (a type-level contract) and a runtime ValueError guards untyped callers that pass it in semantic mode. Superset of #5095 / #5100. Refs #5560 --- python/packages/azure-ai-search/README.md | 1 + .../_context_provider.py | 19 +++++++ .../tests/test_aisearch_context_provider.py | 51 +++++++++++++++++++ .../azure_ai_search/search_context_agentic.py | 7 +++ 4 files changed, 78 insertions(+) diff --git a/python/packages/azure-ai-search/README.md b/python/packages/azure-ai-search/README.md index fcd3161f94..893fac9dbf 100644 --- a/python/packages/azure-ai-search/README.md +++ b/python/packages/azure-ai-search/README.md @@ -19,5 +19,6 @@ See the [Azure AI Search context provider examples](../../samples/02-agents/cont - Semantic search with hybrid (vector + keyword) queries - Agentic mode with Knowledge Bases for complex multi-hop reasoning +- Per-source retrieval parameters (e.g. OData filters) in agentic mode - Environment variable configuration with Settings class - API key and managed identity authentication 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 9a7ced525a..68c51a2389 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 @@ -64,6 +64,7 @@ KnowledgeBaseRetrievalResponse, KnowledgeRetrievalIntent, KnowledgeRetrievalSemanticIntent, + KnowledgeSourceParams, ) from azure.search.documents.knowledgebases.models import ( KnowledgeRetrievalLowReasoningEffort as KBRetrievalLowReasoningEffort, @@ -239,6 +240,7 @@ def __init__( azure_openai_api_key: str | None = None, knowledge_base_output_mode: KnowledgeBaseOutputModeLiteral = "extractive_data", retrieval_reasoning_effort: RetrievalReasoningEffortLiteral = "minimal", + knowledge_source_params: list[KnowledgeSourceParams] | None = None, agentic_message_history_count: int = _DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT, env_file_path: str | None = None, env_file_encoding: str | None = None, @@ -264,6 +266,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. + knowledge_source_params: Optional per-source retrieval parameters (for example an + OData ``filter_add_on``) forwarded to the agentic retrieval request. 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. @@ -292,6 +296,7 @@ def __init__( azure_openai_api_key: str | None = None, knowledge_base_output_mode: KnowledgeBaseOutputModeLiteral = "extractive_data", retrieval_reasoning_effort: RetrievalReasoningEffortLiteral = "minimal", + knowledge_source_params: list[KnowledgeSourceParams] | None = None, agentic_message_history_count: int = _DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT, env_file_path: str | None = None, env_file_encoding: str | None = None, @@ -317,6 +322,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. + knowledge_source_params: Optional per-source retrieval parameters (for example an + OData ``filter_add_on``) forwarded to the agentic retrieval request. 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. @@ -345,6 +352,7 @@ def __init__( azure_openai_api_key: str | None = None, knowledge_base_output_mode: KnowledgeBaseOutputModeLiteral = "extractive_data", retrieval_reasoning_effort: RetrievalReasoningEffortLiteral = "minimal", + knowledge_source_params: list[KnowledgeSourceParams] | None = None, agentic_message_history_count: int = _DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT, env_file_path: str | None = None, env_file_encoding: str | None = None, @@ -374,6 +382,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. + knowledge_source_params: Optional per-source retrieval parameters (for example an + OData ``filter_add_on``) forwarded to the agentic retrieval request. 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. @@ -401,6 +411,7 @@ def __init__( azure_openai_api_key: str | None = None, knowledge_base_output_mode: KnowledgeBaseOutputModeLiteral = "extractive_data", retrieval_reasoning_effort: RetrievalReasoningEffortLiteral = "minimal", + knowledge_source_params: list[KnowledgeSourceParams] | None = None, agentic_message_history_count: int = _DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT, env_file_path: str | None = None, env_file_encoding: str | None = None, @@ -429,6 +440,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. + knowledge_source_params: Optional per-source retrieval parameters (for example an + OData ``filter_add_on``) forwarded to the agentic retrieval request. 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. @@ -474,6 +487,9 @@ def __init__( if mode == "agentic" and settings.get("index_name") and not model: raise ValueError("model is required for agentic mode when creating Knowledge Base from index.") + if knowledge_source_params is not None and mode != "agentic": + raise ValueError("knowledge_source_params is only supported in agentic mode.") + resolved_credential: AzureKeyCredential | AsyncTokenCredential if credential: resolved_credential = credential # type: ignore[assignment] @@ -505,6 +521,7 @@ def __init__( self.knowledge_base_output_mode = knowledge_base_output_mode self.retrieval_reasoning_effort = retrieval_reasoning_effort self.agentic_message_history_count = agentic_message_history_count + self._knowledge_source_params = knowledge_source_params self._use_existing_knowledge_base = False if mode == "agentic": @@ -830,6 +847,7 @@ async def _agentic_search(self, messages: list[Message]) -> list[Message]: retrieval_reasoning_effort=reasoning_effort, output_mode=output_mode, include_activity=True, + knowledge_source_params=self._knowledge_source_params, ) else: kb_messages = self._prepare_messages_for_kb_search(messages) @@ -838,6 +856,7 @@ async def _agentic_search(self, messages: list[Message]) -> list[Message]: retrieval_reasoning_effort=reasoning_effort, output_mode=output_mode, include_activity=True, + knowledge_source_params=self._knowledge_source_params, ) if not self._retrieval_client: 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 bf3c48d167..13d4a9be04 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 @@ -11,6 +11,7 @@ from agent_framework._sessions import AgentSession, SessionContext from agent_framework.exceptions import SettingNotFoundError from azure.core.credentials import AzureKeyCredential +from azure.search.documents.knowledgebases.models import SearchIndexKnowledgeSourceParams from agent_framework_azure_ai_search._context_provider import AzureAISearchContextProvider @@ -315,6 +316,17 @@ def test_agentic_explicit_index_ignores_env_kb_name(self) -> None: assert provider.knowledge_base_name == "idx-kb" assert provider._use_existing_knowledge_base is False + def test_knowledge_source_params_in_semantic_mode_raises(self) -> None: + with pytest.raises(ValueError, match="agentic mode"): + cast(Any, AzureAISearchContextProvider)( + source_id="s", + endpoint="https://test.search.windows.net", + index_name="idx", + api_key="key", + mode="semantic", + knowledge_source_params=[SearchIndexKnowledgeSourceParams(knowledge_source_name="src")], + ) + # -- __aenter__ / __aexit__ --------------------------------------------------- @@ -1335,6 +1347,45 @@ async def test_none_response_returns_default_message(self) -> None: assert len(results) == 1 assert results[0].text == "No results found from Knowledge Base." + @pytest.mark.parametrize("effort", ["minimal", "medium"]) + async def test_knowledge_source_params_reach_request(self, effort: str) -> None: + provider = _make_provider() + provider._knowledge_base_initialized = True + provider.knowledge_base_name = "kb" + provider.retrieval_reasoning_effort = effort + params = [SearchIndexKnowledgeSourceParams(knowledge_source_name="src", filter_add_on="category eq 'public'")] + provider._knowledge_source_params = params + + mock_result = Mock() + mock_result.response = [] + mock_result.references = None + mock_retrieval = AsyncMock() + mock_retrieval.retrieve = AsyncMock(return_value=mock_result) + provider._retrieval_client = mock_retrieval + + await provider._agentic_search([Message(role="user", contents=["q"])]) + + request = mock_retrieval.retrieve.call_args.kwargs["retrieval_request"] + assert request.knowledge_source_params is params + + async def test_default_sends_no_knowledge_source_params(self) -> None: + provider = _make_provider() + provider._knowledge_base_initialized = True + provider.knowledge_base_name = "kb" + provider.retrieval_reasoning_effort = "minimal" + + mock_result = Mock() + mock_result.response = [] + mock_result.references = None + mock_retrieval = AsyncMock() + mock_retrieval.retrieve = AsyncMock(return_value=mock_result) + provider._retrieval_client = mock_retrieval + + await provider._agentic_search([Message(role="user", contents=["q"])]) + + request = mock_retrieval.retrieve.call_args.kwargs["retrieval_request"] + assert request.knowledge_source_params is None + # -- before_run: agentic mode -------------------------------------------------- diff --git a/python/samples/02-agents/context_providers/azure_ai_search/search_context_agentic.py b/python/samples/02-agents/context_providers/azure_ai_search/search_context_agentic.py index 07c69ecc3f..39d3d5d149 100644 --- a/python/samples/02-agents/context_providers/azure_ai_search/search_context_agentic.py +++ b/python/samples/02-agents/context_providers/azure_ai_search/search_context_agentic.py @@ -85,6 +85,13 @@ async def main() -> None: # Optional: Configure retrieval behavior knowledge_base_output_mode="extractive_data", # or "answer_synthesis" retrieval_reasoning_effort="minimal", # or "medium", "low" + # Optional: per-source params, e.g. an OData filter for multi-tenant isolation + # (import SearchIndexKnowledgeSourceParams from azure.search.documents.knowledgebases.models): + # knowledge_source_params=[ + # SearchIndexKnowledgeSourceParams( + # knowledge_source_name="my-source", filter_add_on="category eq 'public'" + # ) + # ], ) else: # Auto-create Knowledge Base from index From 7e922cf3ba7f8550dcc5bd45e7f81c3aed36ad2b Mon Sep 17 00:00:00 2001 From: Ayrton Werck Date: Fri, 19 Jun 2026 12:22:46 +0000 Subject: [PATCH 2/2] Python: fix knowledge_source_params test type annotations The agentic-retrieval tests fail the zuban/pyrefly/ty type-checking gate (mypy is lenient and passes): the parametrized `effort` was typed `str` but assigned to a `Literal[...]` attribute, and `params` was inferred as `list[SearchIndexKnowledgeSourceParams]` where a `list[KnowledgeSourceParams] | None` is expected (list invariance). Annotate the parametrized `effort` and the `params` list with their proper types so all five gating type-checkers pass. Test-only; no behavior change. Refs #5560 --- .../tests/test_aisearch_context_provider.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) 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 13d4a9be04..bc57e1fc30 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 @@ -3,7 +3,7 @@ import os from types import SimpleNamespace -from typing import Any, cast +from typing import Any, Literal, cast from unittest.mock import AsyncMock, Mock, patch import pytest @@ -11,7 +11,7 @@ from agent_framework._sessions import AgentSession, SessionContext from agent_framework.exceptions import SettingNotFoundError from azure.core.credentials import AzureKeyCredential -from azure.search.documents.knowledgebases.models import SearchIndexKnowledgeSourceParams +from azure.search.documents.knowledgebases.models import KnowledgeSourceParams, SearchIndexKnowledgeSourceParams from agent_framework_azure_ai_search._context_provider import AzureAISearchContextProvider @@ -1348,12 +1348,14 @@ async def test_none_response_returns_default_message(self) -> None: assert results[0].text == "No results found from Knowledge Base." @pytest.mark.parametrize("effort", ["minimal", "medium"]) - async def test_knowledge_source_params_reach_request(self, effort: str) -> None: + async def test_knowledge_source_params_reach_request(self, effort: Literal["minimal", "medium"]) -> None: provider = _make_provider() provider._knowledge_base_initialized = True provider.knowledge_base_name = "kb" provider.retrieval_reasoning_effort = effort - params = [SearchIndexKnowledgeSourceParams(knowledge_source_name="src", filter_add_on="category eq 'public'")] + params: list[KnowledgeSourceParams] = [ + SearchIndexKnowledgeSourceParams(knowledge_source_name="src", filter_add_on="category eq 'public'") + ] provider._knowledge_source_params = params mock_result = Mock()