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 347e03c818..e93b7e061e 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 @@ -60,6 +60,7 @@ KnowledgeBaseRetrievalResponse, KnowledgeRetrievalIntent, KnowledgeRetrievalSemanticIntent, + SearchIndexKnowledgeSourceParams, ) from azure.search.documents.knowledgebases.models import ( KnowledgeRetrievalMinimalReasoningEffort as KBRetrievalMinimalReasoningEffort, @@ -85,6 +86,7 @@ KnowledgeBaseRetrievalResponse, KnowledgeRetrievalIntent, KnowledgeRetrievalSemanticIntent, + SearchIndexKnowledgeSourceParams, ) from azure.search.documents.knowledgebases.models import ( KnowledgeRetrievalMinimalReasoningEffort as KBRetrievalMinimalReasoningEffort, @@ -576,6 +578,7 @@ def __init__( ) self._knowledge_base_initialized = False + self._knowledge_source_names: list[str] = [] def _common_client_kwargs(self) -> dict[str, Any]: """Build the keyword arguments shared by every Azure AI Search client. @@ -786,6 +789,13 @@ async def _ensure_knowledge_base(self) -> None: credential=self.credential, **self._common_client_kwargs(), ) + # Resolve the existing KB's real knowledge source names so agentic + # retrieval can request reference source data per source. Without + # this, source names were left unset ("None-source"). + self._knowledge_source_names = [] + if self._index_client is not None: + kb = await self._index_client.get_knowledge_base(knowledge_base_name) + self._knowledge_source_names = [ks.name for ks in (kb.knowledge_sources or [])] self._knowledge_base_initialized = True return @@ -799,6 +809,7 @@ 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_names = [knowledge_source_name] try: await self._index_client.get_knowledge_source(knowledge_source_name) except ResourceNotFoundError: @@ -882,6 +893,17 @@ async def _agentic_search(self, messages: list[Message]) -> list[Message]: # rejects on the stable/GA SDK. request_kwargs["messages"] = self._prepare_messages_for_kb_search(messages) + # Request reference source data per knowledge source so ref.source_data + # is populated when the source has source_data_fields configured (#5095). + if self._knowledge_source_names: + request_kwargs["knowledge_source_params"] = [ + SearchIndexKnowledgeSourceParams( + knowledge_source_name=name, + include_reference_source_data=True, + ) + for name in self._knowledge_source_names + ] + retrieval_request = KnowledgeBaseRetrievalRequest(**request_kwargs) 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 17ae6eb06a..22bb4df2a7 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 @@ -1505,6 +1505,57 @@ async def test_none_response_returns_default_message(self) -> None: assert len(results) == 1 assert results[0].text == "No results found from Knowledge Base." + async def test_requests_reference_source_data_per_source(self) -> None: + # Regression for #5095: the retrieval request must carry + # knowledge_source_params with include_reference_source_data=True for + # each resolved knowledge source, otherwise ref.source_data is always + # None even when the source has source_data_fields configured. + provider = _make_provider() + provider._knowledge_base_initialized = True + provider.knowledge_base_name = "kb" + provider.retrieval_reasoning_effort = "minimal" + provider._knowledge_source_names = ["test-index-source"] + + 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=["query"])]) + + mock_retrieval.retrieve.assert_awaited_once() + request = mock_retrieval.retrieve.call_args.kwargs["retrieval_request"] + params = request.knowledge_source_params + assert params is not None + assert len(params) == 1 + assert params[0].knowledge_source_name == "test-index-source" + assert params[0].include_reference_source_data is True + + async def test_no_source_params_when_no_sources(self) -> None: + # When no knowledge source names are resolved, the request must not + # carry an empty/placeholder knowledge_source_params list. + provider = _make_provider() + provider._knowledge_base_initialized = True + provider.knowledge_base_name = "kb" + provider.retrieval_reasoning_effort = "minimal" + provider._knowledge_source_names = [] + + 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=["query"])]) + + request = mock_retrieval.retrieve.call_args.kwargs["retrieval_request"] + assert request.knowledge_source_params is None + # -- before_run: agentic mode --------------------------------------------------