Python: Support knowledge_source_params in AzureAISearchContextProvider

python/packages/azure-ai-search/README.md

@@ -40,6 +40,10 @@ explicit values override the file selected by `env_file_path`, then environment
 `mode="agentic"` and `knowledge_base_name` or `AZURE_SEARCH_KNOWLEDGE_BASE_NAME`.
 Attach the provider through your agent's `context_provider` parameter.
 
+In agentic mode, pass `knowledge_source_params` to set per-source retrieval options,
+such as `SearchIndexKnowledgeSourceParams(knowledge_source_name=..., filter_add_on=...)`.
+Each entry replaces the provider's defaults for the source with the same name.
+
 Injected SDK clients bypass connection settings and remain caller-owned unless
 `managed_client=True`. Use async context managers to close owned store/collection
 clients; credentials and embedding clients remain caller-owned.

python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py

@@ -8,11 +8,12 @@
 
 from __future__ import annotations
 
+import copy
 import importlib.metadata
 import inspect
 import logging
 import sys
-from collections.abc import Awaitable, Callable
+from collections.abc import Awaitable, Callable, Sequence
 from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypedDict, overload
 
 from agent_framework import (
@@ -206,6 +207,7 @@ def __init__(
         retrieval_reasoning_effort: RetrievalReasoningEffortLiteral = "minimal",
         query_source_credential: AzureCredentialTypes | None = None,
         agentic_message_history_count: int = _DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT,
+        knowledge_source_params: None = None,
         env_file_path: str | None = None,
         env_file_encoding: str | None = None,
     ) -> None:
@@ -232,6 +234,7 @@ def __init__(
             retrieval_reasoning_effort: Unused in semantic mode.
             query_source_credential: Unused in semantic mode.
             agentic_message_history_count: Unused in semantic mode.
+            knowledge_source_params: Must be ``None`` in semantic mode.
             env_file_path: Optional ``.env`` file checked before process environment variables.
             env_file_encoding: Encoding for the ``.env`` file.
         """
@@ -261,6 +264,7 @@ def __init__(
         retrieval_reasoning_effort: RetrievalReasoningEffortLiteral = "minimal",
         query_source_credential: AzureCredentialTypes | None = None,
         agentic_message_history_count: int = _DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT,
+        knowledge_source_params: Sequence[KnowledgeSourceParams] | None = None,
         env_file_path: str | None = None,
         env_file_encoding: str | None = None,
     ) -> None:
@@ -288,6 +292,9 @@ def __init__(
             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.
+            knowledge_source_params: Per-source retrieval parameters, such as
+                ``SearchIndexKnowledgeSourceParams(filter_add_on=...)``, that replace the
+                defaults for the matching ``knowledge_source_name``.
             env_file_path: Optional ``.env`` file checked before process environment variables.
             env_file_encoding: Encoding for the ``.env`` file.
         """
@@ -317,6 +324,7 @@ def __init__(
         retrieval_reasoning_effort: RetrievalReasoningEffortLiteral = "minimal",
         query_source_credential: AzureCredentialTypes | None = None,
         agentic_message_history_count: int = _DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT,
+        knowledge_source_params: Sequence[KnowledgeSourceParams] | None = None,
         env_file_path: str | None = None,
         env_file_encoding: str | None = None,
     ) -> None:
@@ -344,6 +352,9 @@ def __init__(
             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.
+            knowledge_source_params: Per-source retrieval parameters, such as
+                ``SearchIndexKnowledgeSourceParams(filter_add_on=...)``, that replace the
+                defaults for the matching ``knowledge_source_name``.
             env_file_path: Optional ``.env`` file checked before process environment variables.
             env_file_encoding: Encoding for the ``.env`` file.
         """
@@ -373,6 +384,7 @@ def __init__(
         retrieval_reasoning_effort: RetrievalReasoningEffortLiteral = "minimal",
         query_source_credential: AzureCredentialTypes | None = None,
         agentic_message_history_count: int = _DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT,
+        knowledge_source_params: Sequence[KnowledgeSourceParams] | None = None,
         env_file_path: str | None = None,
         env_file_encoding: str | None = None,
     ) -> None:
@@ -404,6 +416,9 @@ def __init__(
             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.
+            knowledge_source_params: Per-source retrieval parameters, such as
+                ``SearchIndexKnowledgeSourceParams(filter_add_on=...)``, that replace the
+                defaults for the matching ``knowledge_source_name``.
             env_file_path: Optional ``.env`` file checked before process environment variables.
             env_file_encoding: Encoding for the ``.env`` file.
         """
@@ -432,6 +447,7 @@ def __init__(
         retrieval_reasoning_effort: RetrievalReasoningEffortLiteral = "minimal",
         query_source_credential: AzureCredentialTypes | None = None,
         agentic_message_history_count: int = _DEFAULT_AGENTIC_MESSAGE_HISTORY_COUNT,
+        knowledge_source_params: Sequence[KnowledgeSourceParams] | None = None,
         env_file_path: str | None = None,
         env_file_encoding: str | None = None,
     ) -> None:
@@ -462,6 +478,11 @@ def __init__(
             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.
+            knowledge_source_params: Per-source retrieval parameters for agentic mode, such as
+                ``SearchIndexKnowledgeSourceParams(filter_add_on=...)``. Each entry replaces the
+                default parameters for the knowledge source with the same ``knowledge_source_name``;
+                entries for other sources are appended. ``include_reference_source_data`` defaults
+                to ``True`` when left unset.
             env_file_path: Path to environment file for loading settings.
             env_file_encoding: Encoding of the environment file.
         """
@@ -541,6 +562,7 @@ def __init__(
         self.retrieval_reasoning_effort = retrieval_reasoning_effort
         self.query_source_credential = query_source_credential
         self.agentic_message_history_count = agentic_message_history_count
+        self.knowledge_source_params = list(knowledge_source_params or [])
 
         self._use_existing_knowledge_base = False
         if mode == "agentic":
@@ -555,6 +577,12 @@ def __init__(
         if vector_field_name and not embedding_function:
             raise ValueError("embedding_function is required when vector_field_name is specified")
 
+        if knowledge_source_params and mode != "agentic":
+            raise ValueError("knowledge_source_params is only supported in agentic mode.")
+        user_source_names = [param.knowledge_source_name for param in knowledge_source_params or []]
+        if len(user_source_names) != len(set(user_source_names)):
+            raise ValueError("knowledge_source_params must not contain duplicate knowledge_source_name values.")
+
         if mode == "agentic":
             if not _agentic_retrieval_available:
                 raise ImportError(
@@ -829,7 +857,7 @@ async def _ensure_knowledge_base(self) -> None:
                             kind=knowledge_source.kind,
                         )
                     )
-            self._knowledge_source_params = knowledge_source_params
+            self._knowledge_source_params = self._with_knowledge_source_param_overrides(knowledge_source_params)
             self._knowledge_base_initialized = True
             return
 
@@ -843,13 +871,13 @@ async def _ensure_knowledge_base(self) -> None:
             raise ValueError("index_name is required when creating Knowledge Base from index")
 
         knowledge_source_name = f"{self.index_name}-source"
-        self._knowledge_source_params = [
+        self._knowledge_source_params = self._with_knowledge_source_param_overrides([
             KnowledgeSourceParams(
                 knowledge_source_name=knowledge_source_name,
                 include_reference_source_data=True,
                 kind="searchIndex",
             )
-        ]
+        ])
         try:
             await self._index_client.get_knowledge_source(knowledge_source_name)
         except ResourceNotFoundError:
@@ -905,6 +933,21 @@ async def _ensure_knowledge_base(self) -> None:
                 **self._common_client_kwargs(),
             )
 
+    def _with_knowledge_source_param_overrides(
+        self, defaults: list[KnowledgeSourceParams]
+    ) -> list[KnowledgeSourceParams]:
+        """Replace default per-source parameters with the caller-provided ones, matched by source name."""
+        overrides: dict[str, KnowledgeSourceParams] = {}
+        for param in self.knowledge_source_params:
+            if param.include_reference_source_data is None:
+                # Keep ref.source_data populated unless the caller opts out explicitly (#5095).
+                param = copy.deepcopy(param)
+                param.include_reference_source_data = True
+            overrides[param.knowledge_source_name] = param
+        merged = [overrides.pop(param.knowledge_source_name, param) for param in defaults]
+        merged.extend(overrides.values())
+        return merged
+
     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:

python/packages/azure-ai-search/tests/test_aisearch_context_provider.py

@@ -16,6 +16,7 @@
 from agent_framework.exceptions import SettingNotFoundError
 from azure.core.credentials import AzureKeyCredential
 from azure.core.pipeline.transport import AioHttpTransport
+from azure.search.documents.knowledgebases.models import SearchIndexKnowledgeSourceParams
 
 from agent_framework_azure_ai_search import _context_provider
 from agent_framework_azure_ai_search._context_provider import (
@@ -328,6 +329,32 @@ def test_vector_field_without_embedding_raises(self) -> None:
                 vector_field_name="embedding",
             )
 
+    def test_knowledge_source_params_in_semantic_mode_raises(self) -> None:
+        with pytest.raises(ValueError, match="only supported in agentic mode"):
+            cast(Any, AzureAISearchContextProvider)(
+                source_id="s",
+                endpoint="https://test.search.windows.net",
+                index_name="idx",
+                api_key="key",
+                knowledge_source_params=[
+                    _context_provider.KnowledgeSourceParams(knowledge_source_name="src", kind="searchIndex")
+                ],
+            )
+
+    def test_duplicate_knowledge_source_params_raises(self) -> None:
+        with pytest.raises(ValueError, match="duplicate knowledge_source_name"):
+            AzureAISearchContextProvider(
+                source_id="s",
+                endpoint="https://test.search.windows.net",
+                knowledge_base_name="kb",
+                api_key="key",
+                mode="agentic",
+                knowledge_source_params=[
+                    _context_provider.KnowledgeSourceParams(knowledge_source_name="src", kind="searchIndex"),
+                    _context_provider.KnowledgeSourceParams(knowledge_source_name="src", kind="searchIndex"),
+                ],
+            )
+
     def test_agentic_missing_aoai_url_with_index_raises(self) -> None:
         with pytest.raises(ValueError, match="azure_openai_resource_url"):
             cast(Any, AzureAISearchContextProvider)(
@@ -1675,6 +1702,73 @@ async def test_existing_mixed_source_kb_preserves_source_kinds(self) -> None:
             ("unknown-source", "futureKind"),
         ]
 
+    async def test_existing_kb_applies_knowledge_source_param_overrides(self) -> None:
+        index_override = SearchIndexKnowledgeSourceParams(
+            knowledge_source_name="index-source", filter_add_on="category eq 'docs'"
+        )
+        extra_source = SearchIndexKnowledgeSourceParams(
+            knowledge_source_name="extra-source", include_reference_source_data=False
+        )
+        provider = _make_provider(
+            mode="agentic",
+            index_name=None,
+            knowledge_base_name="kb",
+            knowledge_source_params=[index_override, extra_source],
+        )
+        provider.retrieval_reasoning_effort = "minimal"
+
+        provider._index_client = AsyncMock()
+        provider._index_client.get_knowledge_base.return_value = SimpleNamespace(
+            knowledge_sources=[SimpleNamespace(name="web-source"), SimpleNamespace(name="index-source")]
+        )
+        provider._index_client.get_knowledge_source.side_effect = [
+            SimpleNamespace(kind="web"),
+            SimpleNamespace(kind="searchIndex"),
+        ]
+        provider._retrieval_client = AsyncMock()
+        provider._retrieval_client.retrieve.return_value = Mock(response=[], references=None)
+
+        await provider._agentic_search([Message(role="user", contents=["query"])])
+
+        request = provider._retrieval_client.retrieve.call_args.kwargs["retrieval_request"]
+        assert [param.as_dict() for param in request.knowledge_source_params] == [
+            {"knowledgeSourceName": "web-source", "includeReferenceSourceData": True, "kind": "web"},
+            {
+                "knowledgeSourceName": "index-source",
+                "filterAddOn": "category eq 'docs'",
+                "kind": "searchIndex",
+                "includeReferenceSourceData": True,
+            },
+            {"knowledgeSourceName": "extra-source", "includeReferenceSourceData": False, "kind": "searchIndex"},
+        ]
+        # The caller's objects are not mutated when defaults are filled in.
+        assert index_override.include_reference_source_data is None
+
+    async def test_index_backed_kb_applies_knowledge_source_param_overrides(self) -> None:
+        provider = _make_provider(
+            mode="agentic",
+            model="deploy",
+            azure_openai_resource_url="https://aoai.openai.azure.com",
+            knowledge_source_params=[
+                SearchIndexKnowledgeSourceParams(
+                    knowledge_source_name="test-index-source", filter_add_on="category eq 'docs'"
+                )
+            ],
+        )
+        provider._index_client = AsyncMock()
+
+        with patch("agent_framework_azure_ai_search._context_provider.KnowledgeBaseRetrievalClient"):
+            await provider._ensure_knowledge_base()
+
+        assert [param.as_dict() for param in provider._knowledge_source_params] == [
+            {
+                "knowledgeSourceName": "test-index-source",
+                "filterAddOn": "category eq 'docs'",
+                "kind": "searchIndex",
+                "includeReferenceSourceData": True,
+            }
+        ]
+
     async def test_concurrent_existing_kb_initialization_publishes_complete_source_params(self) -> None:
         provider = _make_provider(mode="agentic", index_name=None, knowledge_base_name="kb")
         provider._retrieval_client = AsyncMock()