diff --git a/python/packages/azure-ai-search/README.md b/python/packages/azure-ai-search/README.md index 2324ee3f43b..75f3059a8a0 100644 --- a/python/packages/azure-ai-search/README.md +++ b/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. 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 3bcd83c6619..c454d9262ac 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 @@ -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: 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 776950974a2..ed49d8640cf 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 @@ -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()