Python: Add include_reference_source_data to agentic search requests

[Matt Van Horn (mvanhorn)]

Summary

_agentic_search() in AzureAISearchContextProvider constructs KnowledgeBaseRetrievalRequest without passing knowledge_source_params, so include_reference_source_data is never set. This causes ref.source_data to always be None on returned references, even when the knowledge source has source_data_fields configured.

Changes

• Pass knowledge_source_params=[KnowledgeSourceParams(..., include_reference_source_data=True)] to both KnowledgeBaseRetrievalRequest constructions in _agentic_search() (the minimal reasoning path at line 828 and the non-minimal path at line 836)
• Add KnowledgeSourceParams to both TYPE_CHECKING and runtime import blocks
• Add two unit tests verifying the parameter is passed for both reasoning effort paths

Testing

• Existing 119 tests pass
• New tests verify knowledge_source_params is present with correct knowledge_source_name and include_reference_source_data=True for both the minimal and non-minimal reasoning paths

Fixes #5095

This contribution was developed with AI assistance (Claude Code).

[Eduard van Valkenburg (eavanvalkenburg)]

Matt Van Horn (@mvanhorn) some tests and check failures, please have a look

[Matt Van Horn (mvanhorn)]

Took a look - the 6 failing tests are all in packages/durabletask/tests/test_executors.py (TestOrchestrationAgentExecutorFireAndForget) with a TypeError: argument of type 'Mock' is not iterable in the durabletask dependency. This PR only touches files in packages/azure-ai-search/, so the failures are unrelated to these changes.

[Eduard van Valkenburg (eavanvalkenburg)]

then ensure you are working against the latest main

[Matt Van Horn (mvanhorn)]

Eduard van Valkenburg (@eavanvalkenburg) rebased onto latest main in d1fd0a8. The durabletask test failures should clear if the fix landed upstream. CI running now.

[Evan Mattson (moonbox3)]

Python Test Coverage Report •

File	Stmts	Miss	Cover	Missing
packages/azure-ai-search/agent_framework_azure_ai_search
_context_provider.py	390	19	95%	96–97, 122–125, 536, 615–616, 743–744, 797–798, 842, 847, 852, 946–947, 995
TOTAL	44393	5338	87%

Python Unit Test Overview

Tests	Skipped	Failures	Errors	Time
8769	33 💤	0 ❌	0 🔥	2m 20s ⏱️

[Copilot]

Pull request overview

This PR updates the Python Azure AI Search AzureAISearchContextProvider agentic retrieval request construction so that reference source_data can be returned by opting in via knowledge_source_params (include_reference_source_data=True). It also adds unit tests to validate the request includes these parameters for both minimal and non-minimal reasoning effort paths.

Changes:

• Add KnowledgeSourceParams(..., include_reference_source_data=True) to both KnowledgeBaseRetrievalRequest code paths in _agentic_search().
• Add KnowledgeSourceParams to the type-checking and runtime import blocks.
• Add unit tests asserting knowledge_source_params is present for minimal and non-minimal reasoning effort requests.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
python/packages/azure-ai-search/agent_framework_azure_ai_search/_context_provider.py	Injects knowledge_source_params into agentic retrieval requests so reference source_data can be returned.
python/packages/azure-ai-search/tests/test_aisearch_context_provider.py	Adds tests verifying knowledge_source_params is set for both reasoning-effort branches.

[Tao Chen (TaoChenOSU)]

Please address or resolve the comments from Copilot if they are valid.

[Matt Van Horn (mvanhorn)]

Good catch from Copilot, both comments were valid. Fixed in cee7874: in agentic mode against an existing Knowledge Base, index_name is forced to None, so f"{self.index_name}-source" was producing "None-source". The provider now resolves the real knowledge source names from the existing KB via get_knowledge_base(...).knowledge_sources (and stores [f"{index_name}-source"] for the create-from-index path), then builds knowledge_source_params from the resolved list. Added a test asserting no "None-source" is emitted when index_name is None. Package tests and ruff are green.

[Evan Mattson (moonbox3)]

Can you resolve the comments by Copilot, if no longer needed/applicable, please?

[Eduard van Valkenburg (eavanvalkenburg)]

Matt Van Horn (@mvanhorn) there are some failing tests and some open comments, please have a look, so we can get this fixed.

[Matt Van Horn (mvanhorn)]

Both Copilot comments were valid and are fixed in cee7874 — in agentic mode against an existing Knowledge Base, index_name is forced to None, so the source name no longer resolves to "None-source". I've resolved those two threads since they're now addressed. On the failing tests: the checks I can see here (license/cla, add_label) are green — I think the Python test workflow needs a maintainer to approve/re-run it on this fork PR. Could you kick that off, or point me at a specific failure and I'll dig in right away?

[Evan Mattson (moonbox3)]

Matt Van Horn (@mvanhorn) can you please have a look at the conflict as well?

[Matt Van Horn (mvanhorn)]

Rebased onto main. The conflict was in _agentic_search: main refactored the retrieval request to a request_kwargs dict with the preview reasoning-effort/output-mode gating, while this branch's two commits added include_reference_source_data and then switched the source name to the resolved _knowledge_source_names (to avoid "None-source"). Merged both: kept main's request_kwargs + preview gating and built knowledge_source_params from _knowledge_source_names.

Ran the azure-ai-search tests: 129 pass with the preview SDK (12.1.0b1); on GA 12.0.0 everything passes except the one preview-only messages test (which needs the preview request shape), as expected. ruff clean.

[Evan Mattson (moonbox3)]

Hi Matt Van Horn (@mvanhorn), still seeing a lot of CI/CD test failures.

[Matt Van Horn (mvanhorn)]

Rebased onto main and reworked this for the azure-search-documents 12.0.0 API that CI now uses — the branch was pinned to an older 11.7 beta, which is why KnowledgeBaseRetrievalRequest(messages=...) was blowing up. It's now a small delta on top of main's current structure: passes knowledge_source_params with include_reference_source_data=True per resolved knowledge source (via SearchIndexKnowledgeSourceParams), and resolves real source names for both created and existing KBs. Ran the package suite locally — 128 passed (incl. 2 new tests covering the params being set and omitted), ruff + pyright clean. Should be green now.

[Evan Mattson (moonbox3)]

Thanks for helping see this through, Matt Van Horn (@mvanhorn)

[Matt Van Horn (mvanhorn)]

Thanks for shepherding this through, Evan Mattson (@moonbox3) - include_reference_source_data rounds out the agentic search response.

[Matt Van Horn (mvanhorn)]

Appreciate the review and merge, Evan Mattson (@moonbox3)! include_reference_source_data on agentic search should help folks trace citations.