fix: make shallow citation enforcement configurable

docs/source/architecture/agents/shallow-researcher.md

@@ -100,6 +100,7 @@ Configured through `ShallowResearchAgentConfig` (NeMo Agent Toolkit type name: `
 | `tools` | `list[FunctionRef \| FunctionGroupRef]` | `[]` | Tools available for research (web search, document search, etc.) |
 | `max_llm_turns` | `int` | `10` | Maximum LLM interaction turns |
 | `max_tool_iterations` | `int` | `5` | Maximum tool calls before forcing synthesis |
+| `enforce_citations` | `bool` | `false` | Fail instead of returning a sanitized generated answer when citation integrity cannot be preserved |
 | `verbose` | `bool` | `false` | Enable verbose logging |
 
 **Example YAML:**
@@ -113,6 +114,7 @@ functions:
       - web_search_tool
     max_llm_turns: 10
     max_tool_iterations: 5
+    enforce_citations: false
     verbose: true
 ```
 

docs/source/customization/configuration-reference.md

@@ -428,7 +428,7 @@ functions:
 
 ### `shallow_research_agent`
 
-Fast, single-pass research agent that produces citation-backed answers in one tool-calling loop.
+Fast, single-pass research agent that attempts to produce citation-backed answers in one tool-calling loop.
 
 ```yaml
 functions:
@@ -440,6 +440,7 @@ functions:
       - knowledge_search
     max_llm_turns: 10
     max_tool_iterations: 5
+    enforce_citations: false
     verbose: true
 ```
 
@@ -449,6 +450,7 @@ functions:
 | `tools` | `list[str]` | `[]` | Search tools available to the agent. |
 | `max_llm_turns` | `int` | `10` | Maximum number of LLM turns (includes both reasoning and tool-calling steps). |
 | `max_tool_iterations` | `int` | `5` | Maximum tool-calling iterations before forcing synthesis. |
+| `enforce_citations` | `bool` | `false` | Fail the run when citation integrity cannot be preserved. When `false`, AI-Q returns the generated answer after sanitization instead of failing solely on the citation contract. |
 | `verbose` | `bool` | `false` | Enable verbose logging. |
 
 ### `deep_research_agent`

frontends/aiq_api/src/aiq_api/jobs/runner.py

@@ -1312,6 +1312,7 @@ def _create_agent_instance(
             llm_provider=llm_provider,
             tools=tools,
             max_tool_iterations=getattr(fn_config, "max_tool_iterations", 5),
+            enforce_citations=getattr(fn_config, "enforce_citations", False),
             callbacks=callbacks,
         )
     except TypeError:

src/aiq_agent/agents/shallow_researcher/agent.py

@@ -259,6 +259,7 @@ def __init__(
         max_llm_turns: int = 10,
         max_tool_iterations: int = 5,
         citation_repair_timeout: float = 60.0,
+        enforce_citations: bool = False,
         callbacks: list[Any] | None = None,
     ) -> None:
         """
@@ -274,13 +275,17 @@ def __init__(
                                 synthesis (default 5).
             citation_repair_timeout: Maximum seconds for the one-shot citation
                                      repair call (default 60).
+            enforce_citations: Whether missing or invalid citation integrity
+                               should fail the run instead of returning the
+                               generated answer (default False).
             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.enforce_citations = enforce_citations
         self.callbacks = callbacks or []
 
         # Load prompts
@@ -586,13 +591,33 @@ async def run(self, state: ShallowResearchAgentState) -> ShallowResearchAgentSta
                 enable_logging=False,
             )
             generated_answer = sanitize_report(content).sanitized_report if content is not None else None
-            raise EmptySourceRegistryError(
-                "shallow research",
-                unavailable_tools=unavailable,
-                available_count=available_count,
-                reason=classify_empty_source_registry_reason(state.data_sources, available_count, unavailable),
-                generated_answer=generated_answer,
+            if self.enforce_citations or generated_answer is None:
+                raise EmptySourceRegistryError(
+                    "shallow research",
+                    unavailable_tools=unavailable,
+                    available_count=available_count,
+                    reason=classify_empty_source_registry_reason(state.data_sources, available_count, unavailable),
+                    generated_answer=generated_answer,
+                )
+
+            logger.info(
+                "Shallow research completed without captured sources; returning generated answer because "
+                "enforce_citations is false (available_tools=%d unavailable_tools=%d)",
+                available_count,
+                len(unavailable),
             )
+            content = generated_answer
+            if last_msg is not None:
+                for cb in self.callbacks:
+                    if hasattr(cb, "emit_final_report"):
+                        cb.emit_final_report(content, cited_urls=[])
+                        break
+
+                if hasattr(last_msg, "model_copy"):
+                    validated_result["messages"][-1] = last_msg.model_copy(update={"content": content})
+                else:
+                    validated_result["messages"][-1] = type(last_msg)(content=content)
+            return ShallowResearchAgentState.model_validate(validated_result)
 
         if validated_result.get("messages"):
             if content is not None:
@@ -611,7 +636,7 @@ async def run(self, state: ShallowResearchAgentState) -> ShallowResearchAgentSta
                     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:
+                    elif not citation_integrity and self.enforce_citations:
                         logger.info(
                             "Shallow report is missing citation integrity; attempting one bounded repair "
                             "(registered_sources=%d)",
@@ -627,6 +652,12 @@ async def run(self, state: ShallowResearchAgentState) -> ShallowResearchAgentSta
                                 len(sources),
                                 len(repair_verification.valid_citations),
                             )
+                    elif not citation_integrity:
+                        logger.info(
+                            "Shallow report is missing citation integrity; returning generated answer because "
+                            "enforce_citations is false (registered_sources=%d)",
+                            len(sources),
+                        )
                 # Step 2: sanitize report (strip body URLs, shortened URLs, unsafe URLs)
                 sanitization = sanitize_report(content)
                 content = sanitization.sanitized_report
@@ -641,7 +672,8 @@ async def run(self, state: ShallowResearchAgentState) -> ShallowResearchAgentSta
                         len(registry.all_sources()),
                         len(final_verification.valid_citations),
                     )
-                    raise CitationIntegrityError()
+                    if self.enforce_citations:
+                        raise CitationIntegrityError()
                 content = _format_chat_references(final_verification.verified_report)
                 final_cited_urls = list(
                     dict.fromkeys(

src/aiq_agent/agents/shallow_researcher/register.py

@@ -59,6 +59,10 @@ class ShallowResearchAgentConfig(FunctionBaseConfig, name="shallow_research_agen
     )
     max_llm_turns: int = Field(default=10, description="Maximum number of LLM turns")
     max_tool_iterations: int = Field(default=5, description="Maximum tool-calling iterations before forcing synthesis")
+    enforce_citations: bool = Field(
+        default=False,
+        description="Fail instead of returning a generated answer when citation integrity cannot be preserved.",
+    )
     verbose: bool = Field(default=False, description="Whether to enable verbose logging")
 
 
@@ -178,6 +182,7 @@ async def _run(state: ShallowResearchAgentState) -> ShallowResearchAgentState:
                     tools=selected_tools,
                     max_llm_turns=config.max_llm_turns,
                     max_tool_iterations=config.max_tool_iterations,
+                    enforce_citations=config.enforce_citations,
                     callbacks=callbacks,
                 )
 

tests/aiq_agent/agents/shallow_researcher/test_agent.py

@@ -110,6 +110,7 @@ def test_init_with_defaults(self, mock_llm_provider, real_tool):
         assert agent.max_llm_turns == 10
         assert agent.max_tool_iterations == 5
         assert agent.citation_repair_timeout == 60.0
+        assert agent.enforce_citations is False
         assert agent.callbacks == []
         assert agent.system_prompt is not None
 
@@ -130,10 +131,12 @@ def test_init_with_custom_limits(self, mock_llm_provider, real_tool):
             tools=[real_tool],
             max_llm_turns=5,
             max_tool_iterations=3,
+            enforce_citations=True,
         )
 
         assert agent.max_llm_turns == 5
         assert agent.max_tool_iterations == 3
+        assert agent.enforce_citations is True
 
     def test_init_with_callbacks(self, mock_llm_provider, real_tool):
         """Test ShallowResearcherAgent initialization with callbacks."""
@@ -611,6 +614,7 @@ async def test_explicit_tool_not_declared_as_data_source_is_not_captured(self, m
         agent = ShallowResearcherAgent(
             llm_provider=mock_llm_provider,
             tools=[mcp_time__get_current_time],
+            enforce_citations=True,
         )
 
         state = ShallowResearchAgentState(messages=[HumanMessage(content="What time is it in Tokyo?")])
@@ -718,6 +722,7 @@ async def test_missing_uploaded_document_citation_is_repaired_before_publication
         agent = ShallowResearcherAgent(
             llm_provider=mock_llm_provider,
             tools=[knowledge_search],
+            enforce_citations=True,
             callbacks=[callback],
         )
 
@@ -816,6 +821,7 @@ async def test_missing_citations_are_repaired_once_for_multiple_sources(self, mo
         agent = ShallowResearcherAgent(
             llm_provider=mock_llm_provider,
             tools=[mcp_time__get_current_time, web_search_with_urls],
+            enforce_citations=True,
             callbacks=[callback],
         )
 
@@ -839,6 +845,58 @@ async def test_missing_citations_are_repaired_once_for_multiple_sources(self, mo
         assert repair_call.kwargs["config"]["tags"] == [SUPPRESS_OUTPUT_ARTIFACT_TAG]
         assert isinstance(repair_call.args[0][0], SystemMessage)
 
+    @pytest.mark.asyncio
+    async def test_default_multi_source_citation_integrity_failure_returns_draft(self, mock_llm_provider, mock_llm):
+        """Default citation mode returns the answer without running the strict repair pass."""
+        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])
+        callback = MagicMock()
+        agent = ShallowResearcherAgent(
+            llm_provider=mock_llm_provider,
+            tools=[mcp_time__get_current_time, web_search_with_urls],
+            callbacks=[callback],
+        )
+
+        result = await agent.run(
+            ShallowResearchAgentState(messages=[HumanMessage(content="What is CUDA? Also note the time.")])
+        )
+
+        assert "CUDA is a parallel computing platform." in result.messages[-1].content
+        assert mock_llm.ainvoke.await_count == 2
+        callback.emit_final_report.assert_called_once_with(
+            result.messages[-1].content,
+            cited_urls=["https://docs.nvidia.com/cuda/"],
+        )
+
     @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."""
@@ -877,6 +935,7 @@ async def test_failed_multi_source_repair_is_not_published(self, mock_llm_provid
         agent = ShallowResearcherAgent(
             llm_provider=mock_llm_provider,
             tools=[mcp_time__get_current_time, web_search_with_urls],
+            enforce_citations=True,
             callbacks=[callback],
         )
 
@@ -1052,6 +1111,7 @@ async def test_finalization_cannot_publish_a_report_after_losing_inline_citation
         agent = ShallowResearcherAgent(
             llm_provider=mock_llm_provider,
             tools=[web_search_with_urls],
+            enforce_citations=True,
             callbacks=[callback],
         )
 
@@ -1173,6 +1233,7 @@ async def test_run_clears_registry_in_standalone_mode(self, mock_llm_provider, m
         agent = ShallowResearcherAgent(
             llm_provider=mock_llm_provider,
             tools=[],
+            enforce_citations=True,
         )
         # Pre-populate the instance registry (simulating stale data)
         agent.source_registry.add(SourceEntry(url="https://stale.example.com"))
@@ -1182,6 +1243,20 @@ async def test_run_clears_registry_in_standalone_mode(self, mock_llm_provider, m
         with pytest.raises(EmptySourceRegistryError):
             await agent.run(state)
 
+    @pytest.mark.asyncio
+    async def test_default_empty_registry_returns_sanitized_answer(self, mock_llm_provider, mock_llm):
+        """By default, citation-source failures return the generated answer instead of raising."""
+        from aiq_agent.common.citation_verification import set_session_registry
+
+        set_session_registry(None)
+        mock_llm.ainvoke = AsyncMock(return_value=AIMessage(content="Draft answer with https://private.example/path"))
+        agent = ShallowResearcherAgent(llm_provider=mock_llm_provider, tools=[])
+        state = ShallowResearchAgentState(messages=[HumanMessage(content="Test")])
+
+        result = await agent.run(state)
+
+        assert result.messages[-1].content == "Draft answer with "
+
     @pytest.mark.parametrize(
         ("data_sources", "expected_reason"),
         [
@@ -1199,7 +1274,7 @@ async def test_empty_registry_classification_preserves_sanitized_answer(
         expected_reason,
     ):
         mock_llm.ainvoke = AsyncMock(return_value=AIMessage(content="Draft answer with https://private.example/path"))
-        agent = ShallowResearcherAgent(llm_provider=mock_llm_provider, tools=[])
+        agent = ShallowResearcherAgent(llm_provider=mock_llm_provider, tools=[], enforce_citations=True)
         state = ShallowResearchAgentState(
             messages=[HumanMessage(content="Test")],
             data_sources=data_sources,
@@ -1235,7 +1310,11 @@ async def test_enabled_source_empty_result_preserves_generated_answer(self, mock
         )
         generated_answer = "No supporting sources were found, so I cannot provide a sourced answer."
         mock_llm.ainvoke = AsyncMock(side_effect=[tool_call, AIMessage(content=generated_answer)])
-        agent = ShallowResearcherAgent(llm_provider=mock_llm_provider, tools=[empty_web_search_tool])
+        agent = ShallowResearcherAgent(
+            llm_provider=mock_llm_provider,
+            tools=[empty_web_search_tool],
+            enforce_citations=True,
+        )
         state = ShallowResearchAgentState(
             messages=[HumanMessage(content="Summarize quantum computing.")],
             data_sources=["web"],

tests/aiq_agent/jobs/test_runner.py

@@ -2923,6 +2923,50 @@ def __init__(
         assert agent.resource_limits is fn_config.resource_limits
         assert agent.resource_limits.max_source_tool_calls == 8
 
+    def test_create_agent_instance_passes_shallow_research_config(self):
+        """Async workers pass shallow citation enforcement through the constructor."""
+        from aiq_agent.agents.shallow_researcher.register import ShallowResearchAgentConfig
+        from aiq_api.jobs.runner import _create_agent_instance
+
+        class FakeShallowResearcherAgent:
+            def __init__(
+                self,
+                *,
+                llm_provider,
+                tools,
+                max_tool_iterations=5,
+                enforce_citations=False,
+                callbacks=None,
+            ):
+                self.llm_provider = llm_provider
+                self.tools = tools
+                self.max_tool_iterations = max_tool_iterations
+                self.enforce_citations = enforce_citations
+                self.callbacks = callbacks
+
+        assert ShallowResearchAgentConfig(llm="llm").enforce_citations is False
+        fn_config = ShallowResearchAgentConfig(
+            llm="llm",
+            max_tool_iterations=2,
+            enforce_citations=True,
+        )
+
+        agent = _create_agent_instance(
+            agent_cls=FakeShallowResearcherAgent,
+            llm_provider="provider",
+            llm="llm",
+            tools=["tool"],
+            fn_config=fn_config,
+            verbose=False,
+            callbacks=["callback"],
+        )
+
+        assert agent.llm_provider == "provider"
+        assert agent.tools == ["tool"]
+        assert agent.max_tool_iterations == 2
+        assert agent.enforce_citations is True
+        assert agent.callbacks == ["callback"]
+
     def test_create_agent_instance_allows_non_deep_agent_to_reuse_deep_config(self):
         """Async workers should not treat shared DeepResearchAgentConfig as a constructor contract."""
         from aiq_agent.agents.deep_researcher.register import DeepResearchAgentConfig