fix(agents): cite the pages the graph tool read

The graph tool records every page read_entity_pages returns in
retrieved_docs, but agents only ever collected internal_search's. A turn
answered from graph pages therefore emitted no sources: on the multi-hop e2e
run, a question answered from two read_entity_pages calls came back with none.

_search_tool_docs gathers both tools' documents the way the executor caches
them, and both the classic/agentic collector and the research agent's
per-step citations use it. Checked against the real executor and a real
graph: a graph-only turn now cites the pages it read.
This commit is contained in:
Alex committed 2026-09-19 14:33:14 +01:00
1 parent 877609dfb6
commit b32d27c912
4 files changed
+59 -17

No files matched your search

+20
View File
@@ -313,6 +313,26 @@ class TestClassicAgentSearchExposure:
"Tool Doc",
]
def test_collect_internal_sources_includes_graph_pages(
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
):
# Pages the graph tool read carry the answer as much as search hits do,
# so they are cited the same way.
from docsgpt.agents.tools.graph_search import GRAPH_TOOL_ID
retriever_config = {"source": {"active_docs": ["b"]}}
agent = ClassicAgent(retriever_config=retriever_config, **agent_base_params)
search = Mock()
search.retrieved_docs = [{"text": "Found", "title": "Search Doc", "source": "b"}]
graph = Mock()
graph.retrieved_docs = [{"text": "Quill is a store.", "title": "quill.md", "source": "b"}]
user = agent.user or ""
agent.tool_executor._loaded_tools[f"internal_search:{INTERNAL_TOOL_ID}:{user}"] = search
agent.tool_executor._loaded_tools[f"graph_search:{GRAPH_TOOL_ID}:{user}"] = graph
agent._collect_internal_sources()
assert [d["title"] for d in agent.retrieved_docs] == ["Search Doc", "quill.md"]
def test_collect_internal_sources_dedupes(
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
):
+14
View File
@@ -779,6 +779,20 @@ class TestCollectStepSources:
assert len(agent.citations.citations) == 2
def test_collects_pages_the_graph_tool_read(
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
):
from docsgpt.agents.tools.graph_search import GRAPH_TOOL_ID
agent = ResearchAgent(**agent_base_params)
graph = Mock()
graph.retrieved_docs = [{"source": "s3", "title": "quill.md", "text": "Quill"}]
agent.tool_executor._loaded_tools[f"graph_search:{GRAPH_TOOL_ID}:{agent.user or ''}"] = graph
agent._collect_step_sources()
assert len(agent.citations.citations) == 1
def test_no_tool_no_error(
self, agent_base_params, mock_llm_creator, mock_llm_handler_creator
):