diff --git a/src/aiq_agent/agents/shallow_researcher/agent.py b/src/aiq_agent/agents/shallow_researcher/agent.py index f1fa7ca4a..3b296ed1f 100644 --- a/src/aiq_agent/agents/shallow_researcher/agent.py +++ b/src/aiq_agent/agents/shallow_researcher/agent.py @@ -17,6 +17,7 @@ from __future__ import annotations +import asyncio import logging import os import re @@ -39,6 +40,7 @@ from aiq_agent.common import load_prompt from aiq_agent.common import render_prompt_template from aiq_agent.common.callbacks import SUPPRESS_OUTPUT_ARTIFACT_TAG +from aiq_agent.common.citation_verification import CitationIntegrityError from aiq_agent.common.citation_verification import EmptySourceRegistryError from aiq_agent.common.citation_verification import SourceEntry from aiq_agent.common.citation_verification import SourceRegistry @@ -58,6 +60,105 @@ # Path to this agent's directory (for loading prompts) AGENT_DIR = Path(__file__).parent +_SOURCE_SECTION_HEADING_RE = re.compile( + r"^[^\S\n]*(?:" + r"#{1,6}[^\S\n]+(?:Sources|References):?" + r"|\*\*(?:Sources|References):?\*\*:?" + r"|(?:Sources|References):?" + r")[^\S\n]*$", + re.IGNORECASE, +) +_INLINE_CITATION_RE = re.compile(r"\[(\d+)\]") +_REFERENCE_ENTRY_LINE_RE = re.compile(r"^[^\S\n]*(?P[-*][^\S\n]*)?\[(?P\d+)\][^\S\n]+(?P.+)$") +_REFERENCE_TARGET_RE = re.compile( + r"(?:https?://\S+|[^\s,]+\.\w{2,5}(?:,[^\n]+)?|[A-Za-z0-9]+(?:_+[A-Za-z0-9]+)+)$", + re.IGNORECASE, +) + + +def _reference_entry_number(line: str, previous_number: int | None) -> int | None: + """Return a reference number only for an unambiguous definition line.""" + match = _REFERENCE_ENTRY_LINE_RE.fullmatch(line) + if match is None: + return None + + number = int(match.group("number")) + if previous_number is not None and number <= previous_number: + # Reference definitions are unique and ordered. A repeated/reset + # marker begins answer prose, even when it immediately follows them. + return None + + target = match.group("target").strip() + if match.group("bullet") is None and _REFERENCE_TARGET_RE.search(target) is None: + # An unbulleted marker-first sentence is answer text. Verified + # unbulleted definitions carry a URL, file/page key, or tool key. + return None + return number + + +def _source_section_spans(report_text: str) -> list[tuple[int, int]]: + """Locate canonical source headings with definitions, or empty trailing headings.""" + lines = report_text.splitlines(keepends=True) + offsets: list[int] = [] + offset = 0 + for line in lines: + offsets.append(offset) + offset += len(line) + + spans: list[tuple[int, int]] = [] + line_index = 0 + while line_index < len(lines): + heading = lines[line_index].rstrip("\r\n") + if _SOURCE_SECTION_HEADING_RE.fullmatch(heading) is None: + line_index += 1 + continue + + first_definition = line_index + 1 + while first_definition < len(lines) and not lines[first_definition].strip(): + first_definition += 1 + + after_definitions = first_definition + previous_number: int | None = None + while after_definitions < len(lines): + number = _reference_entry_number(lines[after_definitions].rstrip("\r\n"), previous_number) + if number is None: + break + previous_number = number + after_definitions += 1 + + if after_definitions > first_definition: + # Blank lines after a definition block are section separators, not + # answer content. Consume them so preserved prose keeps its spacing. + next_content = after_definitions + while next_content < len(lines) and not lines[next_content].strip(): + next_content += 1 + end = offsets[next_content] if next_content < len(lines) else len(report_text) + spans.append((offsets[line_index], end)) + line_index = next_content + elif first_definition == len(lines): + spans.append((offsets[line_index], len(report_text))) + break + else: + # An exact heading followed by prose is answer content, not a + # reference section, and must not be removed. + line_index += 1 + + return spans + + +def _remove_source_sections(report_text: str, spans: Sequence[tuple[int, int]]) -> str: + """Remove only recognized source-section spans, preserving surrounding prose.""" + if not spans: + return report_text + + pieces: list[str] = [] + previous_end = 0 + for start, end in spans: + pieces.append(report_text[previous_end:start]) + previous_end = end + pieces.append(report_text[previous_end:]) + return "".join(pieces) + def _append_minimal_citation(report_text: str, source: SourceEntry) -> str: """Append one verified citation when the model omitted references.""" @@ -65,23 +166,9 @@ def _append_minimal_citation(report_text: str, source: SourceEntry) -> str: if not citation_target: return report_text - # verify_citations may strip every citation line under a **References:** - # (or ## References / ## Sources) header and leave the empty header - # behind. Drop that trailing header before we append our own so the final - # output has exactly one references section. - content = report_text.rstrip() - content = re.sub( - r"\n{1,2}\*\*References:?\*\*\s*$", - "", - content, - flags=re.IGNORECASE, - ).rstrip() - content = re.sub( - r"\n{1,2}#{2,3}\s+(?:References|Sources)\s*$", - "", - content, - flags=re.IGNORECASE, - ).rstrip() + # Replace only canonical source sections. Similar-looking answer headings + # and prose (for example, "Sources of renewable energy") are content. + content = _remove_source_sections(report_text, _source_section_spans(report_text)).rstrip() if content.endswith((".", "!", "?")): content = f"{content[:-1]} [1]{content[-1]}" else: @@ -96,6 +183,36 @@ def _append_minimal_citation(report_text: str, source: SourceEntry) -> str: return f"{content}\n\n**References:**\n{reference}" +def _has_citation_integrity(report_text: str, valid_citations: Sequence[dict[str, Any]]) -> bool: + """Return whether a verified report has both a source and an inline marker.""" + valid_numbers = { + int(number) + for citation in valid_citations + if (number := citation.get("number")) is not None and str(number).isdigit() + } + if not valid_numbers: + return False + + source_sections = _source_section_spans(report_text) + if not source_sections: + return False + # Definition labels are not inline citations. Remove every recognized + # source block so a later duplicate block cannot satisfy the invariant. + prose = _remove_source_sections(report_text, source_sections) + return any(int(number) in valid_numbers for number in _INLINE_CITATION_RE.findall(prose)) + + +def _format_citation_repair_sources(sources: Sequence[SourceEntry]) -> str: + """Render numbered source lines that a repair pass can copy verbatim.""" + lines: list[str] = [] + for number, source in enumerate(sources, 1): + if source.url: + lines.append(f"- [{number}] Source {number} - {source.url}") + elif source.citation_key: + lines.append(f"- [{number}] {source.citation_key}") + return "\n".join(lines) + + class ShallowResearcherAgent: """ Shallow research agent for fast, bounded research with tool-calling. @@ -129,6 +246,7 @@ def __init__( system_prompt: str | None = None, max_llm_turns: int = 10, max_tool_iterations: int = 5, + citation_repair_timeout: float = 60.0, callbacks: list[Any] | None = None, ) -> None: """ @@ -142,12 +260,15 @@ def __init__( max_llm_turns: Maximum LLM interaction turns (default 10). max_tool_iterations: Maximum tool-calling iterations before forcing synthesis (default 5). + citation_repair_timeout: Maximum seconds for the one-shot citation + repair call (default 60). callbacks: Optional list of LangGraph callbacks. """ self.llm_provider = llm_provider self.tools = list(tools) self.max_llm_turns = max_llm_turns self.max_tool_iterations = max_tool_iterations + self.citation_repair_timeout = citation_repair_timeout self.callbacks = callbacks or [] # Load prompts @@ -188,6 +309,59 @@ def _get_llm(self) -> BaseChatModel: """Get the LLM for shallow research.""" return self.llm_provider.get(LLMRole.RESEARCHER) + async def _repair_missing_citations( + self, + messages: Sequence[Any], + sources: Sequence[SourceEntry], + ) -> str: + """Run one bounded, tool-free repair against captured source identities.""" + source_catalog = _format_citation_repair_sources(sources) + if not source_catalog: + raise CitationIntegrityError() + + repair_system = SystemMessage( + content=( + "You are a deterministic citation-repair editor. Do not answer the original question again from " + "memory and do not call tools. Rewrite only the immediately preceding draft. Keep only claims " + "supported by prior tool results. Your response is invalid unless it contains at least one inline " + "[N] marker and a final **References:** section copied from the allowed reference lines." + ) + ) + repair_request = HumanMessage( + content=( + "The immediately preceding draft failed the citation contract. Rewrite it once using only claims " + "supported by the prior tool results. Preserve the answer's meaning, remove unsupported claims, " + "and do not call tools. Add an inline [N] marker after each externally verified claim and finish " + "with a `**References:**` section. Copy the corresponding allowed reference lines verbatim; never " + "invent or reconstruct a URL. Return only the repaired report.\n\n" + f"Allowed reference lines:\n{source_catalog}" + ) + ) + repair_config: dict[str, Any] = {"tags": [SUPPRESS_OUTPUT_ARTIFACT_TAG]} + if self.callbacks: + repair_config["callbacks"] = self.callbacks + + try: + response = await asyncio.wait_for( + self._get_llm().ainvoke( + [repair_system, *messages, repair_request], + config=repair_config, + ), + timeout=self.citation_repair_timeout, + ) + except Exception as ex: + logger.warning( + "Shallow citation repair failed (error_type=%s detail_%s)", + type(ex).__name__, + log_content_metadata(ex), + ) + raise CitationIntegrityError() from ex + + repaired_content = getattr(response, "content", None) + if not isinstance(repaired_content, str) or not repaired_content.strip(): + raise CitationIntegrityError() + return repaired_content + def _build_graph(self) -> CompiledStateGraph: """Build the LangGraph StateGraph.""" @@ -230,6 +404,7 @@ async def agent_node(state: ShallowResearchAgentState) -> dict[str, Any]: processed_history = list(messages) try: + draft_config = {"tags": [SUPPRESS_OUTPUT_ARTIFACT_TAG]} if iterations >= self.max_tool_iterations: logger.warning("Max iterations (%d) reached. Forcing synthesis.", iterations) @@ -243,14 +418,13 @@ async def agent_node(state: ShallowResearchAgentState) -> dict[str, Any]: ) full_messages = [system_message] + processed_history + [synthesis_anchor] - response = await self._get_llm().ainvoke(full_messages) + response = await self._get_llm().ainvoke(full_messages, config=draft_config) return {"messages": [response], "tool_iterations": iterations} llm = self._get_llm() llm_with_tools = llm.bind_tools(self.tools) if self.tools else llm full_messages = [system_message] + processed_history - pre_evidence_config = {"tags": [SUPPRESS_OUTPUT_ARTIFACT_TAG]} if iterations == 0 else None - response = await llm_with_tools.ainvoke(full_messages, config=pre_evidence_config) + response = await llm_with_tools.ainvoke(full_messages, config=draft_config) if self.tools and iterations == 0 and not getattr(response, "tool_calls", None): logger.warning("Shallow researcher returned an answer before collecting evidence; retrying once") @@ -263,7 +437,7 @@ async def agent_node(state: ShallowResearchAgentState) -> dict[str, Any]: retry_llm = llm.bind_tools(self.tools, parallel_tool_calls=False) response = await retry_llm.ainvoke( full_messages + [response, tool_required], - config=pre_evidence_config, + config=draft_config, ) retry_tool_calls = getattr(response, "tool_calls", None) or [] if len(retry_tool_calls) != 1 or retry_tool_calls[0].get("name") not in source_tool_names: @@ -422,12 +596,40 @@ async def run(self, state: ShallowResearchAgentState) -> ShallowResearchAgentSta ) content = verification.verified_report sources = registry.all_sources() - if not verification.valid_citations and len(sources) == 1: + citation_integrity = _has_citation_integrity(content, verification.valid_citations) + if not citation_integrity and len(sources) == 1: content = _append_minimal_citation(content, sources[0]) + elif not citation_integrity: + logger.info( + "Shallow report is missing citation integrity; attempting one bounded repair " + "(registered_sources=%d)", + len(sources), + ) + content = await self._repair_missing_citations(validated_result["messages"], sources) + repair_verification = verify_citations(content, registry, reference_sources=sources) + content = repair_verification.verified_report + if not _has_citation_integrity(content, repair_verification.valid_citations): + logger.warning( + "Shallow citation repair did not restore integrity " + "(registered_sources=%d verified_sources=%d)", + len(sources), + len(repair_verification.valid_citations), + ) # Step 2: sanitize report (strip body URLs, shortened URLs, unsafe URLs) sanitization = sanitize_report(content) content = sanitization.sanitized_report final_verification = verify_citations(content, registry) + if not _has_citation_integrity( + final_verification.verified_report, + final_verification.valid_citations, + ): + logger.warning( + "Shallow report failed final citation integrity check " + "(registered_sources=%d verified_sources=%d)", + len(registry.all_sources()), + len(final_verification.valid_citations), + ) + raise CitationIntegrityError() content = final_verification.verified_report final_cited_urls = list( dict.fromkeys( diff --git a/src/aiq_agent/agents/shallow_researcher/prompts/researcher.j2 b/src/aiq_agent/agents/shallow_researcher/prompts/researcher.j2 index eba305767..efd5770ad 100644 --- a/src/aiq_agent/agents/shallow_researcher/prompts/researcher.j2 +++ b/src/aiq_agent/agents/shallow_researcher/prompts/researcher.j2 @@ -33,6 +33,7 @@ Cite sources inline with [1], [2], etc. Include a `**References:**` section at t - Use only URLs that appeared in tool results. Do not add URLs from memory. - If you use a tool result to answer, include at least one inline citation and a `**References:**` section. - For tool results without URLs or document citation keys, cite the exact tool name from the tool call: `- [1] mcp_time__get_current_time`. +- Before returning an answer based on tool results, verify that it contains both an inline [N] marker and a non-empty `**References:**` section. - Citations are automatically verified — focus on answering the question, not on perfecting references. **Example**: diff --git a/tests/aiq_agent/agents/shallow_researcher/test_agent.py b/tests/aiq_agent/agents/shallow_researcher/test_agent.py index 00e9d96f6..eeacd9531 100644 --- a/tests/aiq_agent/agents/shallow_researcher/test_agent.py +++ b/tests/aiq_agent/agents/shallow_researcher/test_agent.py @@ -15,6 +15,8 @@ """Tests for the ShallowResearcherAgent.""" +import asyncio +import re from unittest.mock import AsyncMock from unittest.mock import MagicMock from unittest.mock import call @@ -23,14 +25,17 @@ import pytest from langchain_core.messages import AIMessage from langchain_core.messages import HumanMessage +from langchain_core.messages import SystemMessage from langchain_core.tools import tool from aiq_agent.agents.shallow_researcher.agent import ShallowResearcherAgent from aiq_agent.agents.shallow_researcher.agent import _append_minimal_citation +from aiq_agent.agents.shallow_researcher.agent import _has_citation_integrity from aiq_agent.agents.shallow_researcher.models import ShallowResearchAgentState from aiq_agent.common import LLMProvider from aiq_agent.common import LLMRole from aiq_agent.common.callbacks import SUPPRESS_OUTPUT_ARTIFACT_TAG +from aiq_agent.common.citation_verification import CitationIntegrityError from aiq_agent.common.citation_verification import EmptySourceRegistryError from aiq_agent.common.citation_verification import EmptySourceRegistryReason from aiq_agent.common.citation_verification import SourceEntry @@ -66,6 +71,7 @@ def _bypass_citation_pipeline(self): patch.object(SourceRegistry, "all_sources", return_value=[SourceEntry(url="https://example.com")]), patch("aiq_agent.agents.shallow_researcher.agent.verify_citations") as mock_verify, patch("aiq_agent.agents.shallow_researcher.agent.sanitize_report") as mock_sanitize, + patch("aiq_agent.agents.shallow_researcher.agent._has_citation_integrity", return_value=True), ): mock_verify.side_effect = lambda content, reg: MagicMock(verified_report=content, removed_citations=[]) mock_sanitize.side_effect = lambda content: MagicMock(sanitized_report=content) @@ -102,6 +108,7 @@ def test_init_with_defaults(self, mock_llm_provider, real_tool): assert len(agent.tools) == 1 assert agent.max_llm_turns == 10 assert agent.max_tool_iterations == 5 + assert agent.citation_repair_timeout == 60.0 assert agent.callbacks == [] assert agent.system_prompt is not None @@ -290,17 +297,38 @@ def test_load_system_prompt_fallback(self, mock_llm_provider, real_tool): ) assert "research" in agent.system_prompt.lower() - def test_default_prompt_requires_tool_result_references(self, mock_llm_provider, real_tool): - """Default prompt tells the model to cite non-URL tool results by exact tool name.""" + def test_default_prompt_has_structural_citation_contract(self, mock_llm_provider, real_tool): + """Default prompt preserves the shared research and citation contract.""" agent = ShallowResearcherAgent( llm_provider=mock_llm_provider, tools=[real_tool], ) - assert "If you use a tool result to answer" in agent.system_prompt - assert "exact tool name" in agent.system_prompt + assert re.search(r"\[\d+\]", agent.system_prompt) + assert "**References:**" in agent.system_prompt assert "- [1] mcp_time__get_current_time" in agent.system_prompt + @pytest.mark.asyncio + async def test_citation_repair_timeout_fails_closed(self, mock_llm_provider, mock_llm): + """A stalled provider cannot extend a completed shallow run indefinitely.""" + + async def stalled_repair(*_args, **_kwargs): + await asyncio.sleep(1) + + mock_llm.ainvoke = AsyncMock(side_effect=stalled_repair) + agent = ShallowResearcherAgent( + llm_provider=mock_llm_provider, + tools=[], + citation_repair_timeout=0.001, + ) + + async with asyncio.timeout(0.1): + with pytest.raises(CitationIntegrityError, match="citation_integrity_lost"): + await agent._repair_missing_citations( + [HumanMessage(content="Draft")], + [SourceEntry(url="https://example.com/source")], + ) + @pytest.mark.asyncio async def test_initial_answer_without_tool_call_is_retried(self, mock_llm_provider, mock_llm, real_tool): """An initial memory-only answer is retried and replaced by a tool call.""" @@ -321,7 +349,7 @@ async def test_initial_answer_without_tool_call_is_retried(self, mock_llm_provid call([real_tool]), call([real_tool], parallel_tool_calls=False), ] - for invocation in mock_llm.ainvoke.await_args_list[:2]: + for invocation in mock_llm.ainvoke.await_args_list: assert invocation.kwargs["config"] == {"tags": [SUPPRESS_OUTPUT_ARTIFACT_TAG]} @pytest.mark.asyncio @@ -467,9 +495,11 @@ async def test_run_returns_updated_tool_iterations(self, mock_llm_provider, mock async def test_forced_synthesis_adds_instruction_message(self, mock_llm_provider, mock_llm, real_tool): """Test that forced synthesis adds instruction to synthesize.""" captured_messages = [] + captured_configs = [] - async def capture_messages(messages): + async def capture_messages(messages, *, config): captured_messages.append(messages) + captured_configs.append(config) return AIMessage(content="Synthesized response") mock_llm.ainvoke = AsyncMock(side_effect=capture_messages) @@ -494,6 +524,7 @@ async def capture_messages(messages): "synthesize" in str(msg.content).lower() for msg in last_call_messages if hasattr(msg, "content") ) assert synthesis_instruction_found + assert captured_configs == [{"tags": [SUPPRESS_OUTPUT_ARTIFACT_TAG]}] # --------------------------------------------------------------------------- @@ -524,6 +555,24 @@ def weather_observation_tool(location: str) -> str: return f"Current conditions for {location}: clear, 68F" +@tool +def knowledge_search(query: str) -> str: + """Search uploaded documents and return file citation metadata.""" + return ( + "Found 2 relevant document(s):\n\n" + "--- Result 1 ---\n" + "Source: report.pdf\n" + "Page: 15\n" + "Citation: report.pdf, p.15\n\n" + f"Relevant content for {query}.\n\n" + "--- Result 2 ---\n" + "Source: appendix.pdf\n" + "Page: 4\n" + "Citation: appendix.pdf, p.4\n\n" + "Supporting appendix content." + ) + + class TestShallowResearcherSourceRegistryGating: """Tests that shallow source capture is gated by data_source_registry.""" @@ -639,6 +688,53 @@ async def test_missing_tool_result_citation_is_appended(self, mock_llm_provider, "It's currently 4:54 AM in Tokyo [1].\n\n## Sources\n- [1] mcp_time__get_current_time" ) + @pytest.mark.asyncio + async def test_missing_uploaded_document_citation_is_repaired_before_publication(self, mock_llm_provider, mock_llm): + """Knowledge-search summaries enforce file citation keys without requiring public URLs.""" + populate_from_config( + [ + { + "id": "knowledge_layer", + "name": "Knowledge Layer", + "description": "Search uploaded documents.", + "tools": ["knowledge_search"], + } + ] + ) + tool_call_response = AIMessage( + content="", + tool_calls=[{"name": "knowledge_search", "args": {"query": "revenue summary"}, "id": "1"}], + ) + uncited_summary = AIMessage(content="The uploaded report says revenue increased.") + repaired_summary = AIMessage( + content=("The uploaded report says revenue increased [1].\n\n**References:**\n- [1] report.pdf, p.15") + ) + bound_llm = MagicMock() + bound_llm.ainvoke = AsyncMock(side_effect=[tool_call_response, uncited_summary]) + mock_llm.bind_tools = MagicMock(return_value=bound_llm) + mock_llm.ainvoke = AsyncMock(return_value=repaired_summary) + callback = MagicMock() + agent = ShallowResearcherAgent( + llm_provider=mock_llm_provider, + tools=[knowledge_search], + callbacks=[callback], + ) + + result = await agent.run( + ShallowResearchAgentState(messages=[HumanMessage(content="Summarize the uploaded revenue report.")]) + ) + + sources = agent.source_registry.all_sources() + assert len(sources) == 2 + assert sources[0].citation_key == "report.pdf, p.15" + assert sources[0].source_type == "knowledge_layer" + assert sources[1].citation_key == "appendix.pdf, p.4" + assert "The uploaded report says revenue increased [1]." in result.messages[-1].content + assert "[1] report.pdf, p.15" in result.messages[-1].content + assert bound_llm.ainvoke.await_count == 2 + mock_llm.ainvoke.assert_awaited_once() + callback.emit_final_report.assert_called_once_with(result.messages[-1].content, cited_urls=[]) + @pytest.mark.asyncio async def test_missing_url_citation_fallback_emits_authoritative_metadata(self, mock_llm_provider, mock_llm): """A single verified URL appended by fallback must be persisted as cited.""" @@ -675,8 +771,8 @@ async def test_missing_url_citation_fallback_emits_authoritative_metadata(self, ) @pytest.mark.asyncio - async def test_missing_citation_fallback_skips_ambiguous_multi_source_registry(self, mock_llm_provider, mock_llm): - """Do not inject the first captured source when multiple sources exist.""" + async def test_missing_citations_are_repaired_once_for_multiple_sources(self, mock_llm_provider, mock_llm): + """Ambiguous multi-source drafts get one bounded repair instead of a guessed citation.""" populate_from_config( [ { @@ -702,7 +798,18 @@ async def test_missing_citation_fallback_skips_ambiguous_multi_source_registry(s ], ) final_response = AIMessage(content="CUDA is a parallel computing platform.") - mock_llm.ainvoke = AsyncMock(side_effect=[tool_call_response, final_response]) + repaired_response = AIMessage( + content=( + "CUDA is a parallel computing platform [2]. The current time came from the time tool [1].\n\n" + "**References:**\n" + "- [1] mcp_time__get_current_time\n" + "- [2] Source 2 - https://docs.nvidia.com/cuda/" + ) + ) + bound_llm = MagicMock() + bound_llm.ainvoke = AsyncMock(side_effect=[tool_call_response, final_response]) + mock_llm.bind_tools = MagicMock(return_value=bound_llm) + mock_llm.ainvoke = AsyncMock(return_value=repaired_response) callback = MagicMock() agent = ShallowResearcherAgent( @@ -718,8 +825,69 @@ async def test_missing_citation_fallback_skips_ambiguous_multi_source_registry(s assert len(sources) >= 2 assert sources[0].citation_key == "mcp_time__get_current_time" assert any(source.url == "https://docs.nvidia.com/cuda/" for source in sources) - assert result.messages[-1].content == "CUDA is a parallel computing platform." - callback.emit_final_report.assert_called_once_with(result.messages[-1].content, cited_urls=[]) + assert "CUDA is a parallel computing platform [2]" in result.messages[-1].content + assert "mcp_time__get_current_time" in result.messages[-1].content + callback.emit_final_report.assert_called_once_with( + result.messages[-1].content, + cited_urls=["https://docs.nvidia.com/cuda/"], + ) + assert bound_llm.ainvoke.await_count == 2 + mock_llm.ainvoke.assert_awaited_once() + assert mock_llm.bind_tools.call_count == 2 + repair_call = mock_llm.ainvoke.await_args + assert repair_call.kwargs["config"]["tags"] == [SUPPRESS_OUTPUT_ARTIFACT_TAG] + assert isinstance(repair_call.args[0][0], SystemMessage) + + @pytest.mark.asyncio + async def test_failed_multi_source_repair_is_not_published(self, mock_llm_provider, mock_llm): + """A single unsuccessful repair fails closed without emitting an uncited report.""" + populate_from_config( + [ + { + "id": "mcp_time", + "name": "MCP Time", + "description": "Get current time and timezone information through MCP.", + "tools": ["mcp_time"], + }, + { + "id": "web_search", + "name": "Web Search", + "description": "Search the web for real-time information.", + "tools": ["web_search_with_urls"], + }, + ], + group_names={"mcp_time"}, + ) + tool_call_response = AIMessage( + content="", + tool_calls=[ + {"name": "mcp_time__get_current_time", "args": {"timezone": "Asia/Tokyo"}, "id": "1"}, + {"name": "web_search_with_urls", "args": {"query": "CUDA"}, "id": "2"}, + ], + ) + source_only_draft = AIMessage( + content=( + "CUDA is a parallel computing platform.\n\n" + "**References:**\n- [1] CUDA Toolkit Documentation - https://docs.nvidia.com/cuda/" + ) + ) + mock_llm.ainvoke = AsyncMock(side_effect=[tool_call_response, source_only_draft, source_only_draft]) + callback = MagicMock() + agent = ShallowResearcherAgent( + llm_provider=mock_llm_provider, + tools=[mcp_time__get_current_time, web_search_with_urls], + callbacks=[callback], + ) + + with pytest.raises(CitationIntegrityError, match="citation_integrity_lost"): + await agent.run( + ShallowResearchAgentState(messages=[HumanMessage(content="What is CUDA? Also note the time.")]) + ) + + assert mock_llm.ainvoke.await_count == 3 + for invocation in mock_llm.ainvoke.await_args_list: + assert invocation.kwargs["config"]["tags"] == [SUPPRESS_OUTPUT_ARTIFACT_TAG] + callback.emit_final_report.assert_not_called() @pytest.mark.asyncio async def test_registered_exact_data_source_tool_without_urls_is_captured(self, mock_llm_provider, mock_llm): @@ -830,7 +998,7 @@ async def test_final_verification_report_is_published_with_its_cited_urls(self, """The final report body and authoritative URL set come from the same verification pass.""" source_url = "https://docs.nvidia.com/cuda/" draft_report = f"CUDA is a parallel computing platform [1].\n\n## Sources\n[1] CUDA Docs: {source_url}" - verified_report = f"CUDA is a parallel computing platform.\n\n## Sources\n[1] CUDA Docs: {source_url}" + verified_report = draft_report tool_call_response = AIMessage( content="", tool_calls=[{"name": "web_search_with_urls", "args": {"query": "CUDA"}, "id": "1"}], @@ -844,12 +1012,12 @@ async def test_final_verification_report_is_published_with_its_cited_urls(self, ) first_verification = MagicMock( verified_report=draft_report, - valid_citations=[{"url": source_url}], + valid_citations=[{"number": 1, "url": source_url}], removed_citations=[], ) final_verification = MagicMock( verified_report=verified_report, - valid_citations=[{"url": source_url}], + valid_citations=[{"number": 1, "url": source_url}], removed_citations=[], ) @@ -862,6 +1030,37 @@ async def test_final_verification_report_is_published_with_its_cited_urls(self, assert result.messages[-1].content == verified_report callback.emit_final_report.assert_called_once_with(verified_report, cited_urls=[source_url]) + @pytest.mark.asyncio + async def test_finalization_cannot_publish_a_report_after_losing_inline_citations( + self, mock_llm_provider, mock_llm + ): + """A later sanitization regression must fail closed instead of publishing source-only prose.""" + source_url = "https://docs.nvidia.com/cuda/" + cited_report = f"CUDA is a parallel computing platform [1].\n\n## Sources\n[1] CUDA Docs: {source_url}" + source_only_report = f"CUDA is a parallel computing platform.\n\n## Sources\n[1] CUDA Docs: {source_url}" + tool_call_response = AIMessage( + content="", + tool_calls=[{"name": "web_search_with_urls", "args": {"query": "CUDA"}, "id": "1"}], + ) + mock_llm.ainvoke = AsyncMock(side_effect=[tool_call_response, AIMessage(content=cited_report)]) + callback = MagicMock() + agent = ShallowResearcherAgent( + llm_provider=mock_llm_provider, + tools=[web_search_with_urls], + callbacks=[callback], + ) + + with ( + patch( + "aiq_agent.agents.shallow_researcher.agent.sanitize_report", + return_value=MagicMock(sanitized_report=source_only_report), + ), + pytest.raises(CitationIntegrityError, match="citation_integrity_lost"), + ): + await agent.run(ShallowResearchAgentState(messages=[HumanMessage(content="What is CUDA?")])) + + callback.emit_final_report.assert_not_called() + @pytest.mark.asyncio async def test_invalid_citation_removed_end_to_end(self, mock_llm_provider, mock_llm): """Citations not backed by registry sources are removed from output.""" @@ -1132,3 +1331,87 @@ def test_no_leftover_header_passes_through(self): result = _append_minimal_citation(report, self._tool_source()) assert result == "Body sentence [1].\n\n**References:**\n- [1] mcp_time__get_current_time" + + def test_replaces_source_only_section_and_adds_inline_marker(self): + report = "Body sentence.\n\n## Sources\n[1] mcp_time__get_current_time" + + result = _append_minimal_citation(report, self._tool_source()) + + assert result == "Body sentence [1].\n\n**References:**\n- [1] mcp_time__get_current_time" + + @pytest.mark.parametrize("heading", ["###### References", "Sources:", "**References:**"]) + def test_replaces_canonical_heading_variants(self, heading): + report = f"Body sentence.\n\n{heading}\n[1] mcp_time__get_current_time" + + result = _append_minimal_citation(report, self._tool_source()) + + assert result.count("**References:**") == 1 + assert result == "Body sentence [1].\n\n**References:**\n- [1] mcp_time__get_current_time" + + @pytest.mark.parametrize( + "report", + [ + "## Sources of renewable energy\nSolar is renewable.", + "## References to previous work\nPrior work remains relevant.", + "Sources: market revenue grew last year.", + "Sources:\nMarket revenue grew last year.", + ], + ) + def test_preserves_source_like_answer_content(self, report): + result = _append_minimal_citation(report, self._tool_source()) + + assert report.removesuffix(".") in result + assert result.endswith("**References:**\n- [1] mcp_time__get_current_time") + + def test_replaces_reference_definitions_without_discarding_trailing_answer(self): + report = "Opening answer.\n\n## References\n- [1] Old source - https://example.com/old\n\nClosing answer." + + result = _append_minimal_citation(report, self._tool_source()) + + assert "Opening answer." in result + assert "Closing answer [1]." in result + assert "Old source" not in result + + def test_preserves_marker_first_answer_immediately_after_reference_block(self): + report = "## Sources\n[1] mcp_time__get_current_time\n[1] CUDA is a parallel computing platform." + + result = _append_minimal_citation(report, self._tool_source()) + + assert "CUDA is a parallel computing platform" in result + assert result.count("mcp_time__get_current_time") == 1 + + +class TestCitationIntegrity: + """Tests for the final inline-plus-source publication invariant.""" + + @pytest.mark.parametrize("heading", ["###### References", "Sources:", "**References:**"]) + def test_accepts_inline_marker_outside_reference_definitions(self, heading): + report = f"{heading}\n- [1] Source - https://example.com/source\n\nAnswer [1]." + + assert _has_citation_integrity(report, [{"number": 1, "url": "https://example.com/source"}]) + + def test_does_not_treat_reference_definition_as_inline_citation(self): + report = "Answer without a marker.\n\nSources:\n- [1] Source - https://example.com/source" + + assert not _has_citation_integrity(report, [{"number": 1, "url": "https://example.com/source"}]) + + def test_preserves_answer_line_that_begins_with_an_inline_marker(self): + report = "[1] CUDA is a parallel computing platform.\n\nSources:\n- [1] CUDA - https://example.com/cuda" + + assert _has_citation_integrity(report, [{"number": 1, "url": "https://example.com/cuda"}]) + + def test_does_not_count_definitions_across_multiple_source_sections(self): + report = ( + "Answer without a marker.\n\n" + "## Sources\n" + "- [1] First - https://example.com/first\n\n" + "## References\n" + "- [1] Second - https://example.com/second" + ) + + assert not _has_citation_integrity(report, [{"number": 1, "url": "https://example.com/first"}]) + + def test_preserves_marker_first_answer_after_reference_definitions(self): + report = "## Sources\n[1] mcp_time__get_current_time\n[1] CUDA is a parallel computing platform." + + assert _has_citation_integrity(report, [{"number": 1, "citation_key": "mcp_time__get_current_time"}]) diff --git a/tests/aiq_agent/test_default_model_profiles.py b/tests/aiq_agent/test_default_model_profiles.py index d13eb8186..45ab1452c 100644 --- a/tests/aiq_agent/test_default_model_profiles.py +++ b/tests/aiq_agent/test_default_model_profiles.py @@ -2,11 +2,16 @@ # SPDX-License-Identifier: Apache-2.0 from pathlib import Path +from unittest.mock import MagicMock import pytest import yaml +from aiq_agent.agents.shallow_researcher.agent import ShallowResearcherAgent +from aiq_agent.common import LLMProvider + REPO_ROOT = Path(__file__).resolve().parents[2] +SHARED_SHALLOW_PROMPT = REPO_ROOT / "src/aiq_agent/agents/shallow_researcher/prompts/researcher.j2" ULTRA_MODEL = "nvidia/nemotron-3-ultra-550b-a55b" LIGHTNING_MODEL = "nvidia/nemotron-3.5-lightning-30b-a3b" @@ -19,6 +24,10 @@ ) CONFIG_PATHS = tuple(sorted(path for pattern in CONFIG_GLOBS for path in REPO_ROOT.glob(pattern))) FRESHQA_CONFIG_PATHS = tuple(sorted(REPO_ROOT.glob("frontends/benchmarks/freshqa/configs/*.yml"))) +SHALLOW_PROFILE_PATHS = ( + REPO_ROOT / "configs/config_web_default_guardrails.yml", + REPO_ROOT / "configs/config_frontier_models.yml", +) DEPRECATED_REFERENCES = ( "/".join(("nvidia", "nemotron-3-super-120b-a12b")), @@ -143,6 +152,18 @@ def test_freshqa_research_tools_are_registered_data_sources(config_path: Path): assert set(function.get("tools", [])) <= source_tools +@pytest.mark.parametrize("config_path", SHALLOW_PROFILE_PATHS, ids=lambda path: path.name) +def test_shallow_profiles_use_the_shared_citation_prompt(config_path: Path): + """Default Lightning and frontier Luna must share the hardened prompt and runtime path.""" + config = _load_config(config_path) + shallow = config["functions"]["shallow_research_agent"] + agent = ShallowResearcherAgent(llm_provider=MagicMock(spec=LLMProvider), tools=[]) + + assert shallow["_type"] == "shallow_research_agent" + assert "system_prompt" not in shallow + assert agent.system_prompt == SHARED_SHALLOW_PROMPT.read_text(encoding="utf-8") + + def test_deprecated_model_and_endpoint_references_are_absent(): violations: list[str] = []