fix(agent): separate citations from final answers

packages/dbgpt-app/src/dbgpt_app/openapi/api_v1/agentic_data_api.py

@@ -1,3 +1,4 @@
+import asyncio
 import io
 import json
 import logging
@@ -30,6 +31,7 @@
 from dbgpt_serve.datasource.manages import ConnectorManager
 from dbgpt_serve.utils.auth import UserRequest, get_user_from_headers
 
+from .react_final import AgentFinalAnswer, FinalAnswerAssembler
 from .subagent.dispatcher import DISPATCH_PROMPT_SECTION, make_dispatch_tool
 from .subagent.history import (
     build_subagent_history_snapshot,
@@ -970,9 +972,108 @@ def _sse_event(payload: Dict[str, Any]) -> str:
     return f"data: {json.dumps(payload, ensure_ascii=False)}\n\n"
 
 
+def _sse_event_type(event: Any) -> Optional[str]:
+    """Read an SSE event type without trusting arbitrary streamed text."""
+    if not isinstance(event, str):
+        return None
+    first_line = event.splitlines()[0] if event else ""
+    if not first_line.startswith("data:"):
+        return None
+    try:
+        payload = json.loads(first_line.removeprefix("data:").strip())
+    except (json.JSONDecodeError, TypeError, ValueError):
+        return None
+    event_type = payload.get("type") if isinstance(payload, dict) else None
+    return event_type if isinstance(event_type, str) else None
+
+
+def _react_terminal_events(
+    storage_conv: Any,
+    history_payload: str,
+    final_answer: AgentFinalAnswer,
+) -> Tuple[str, str]:
+    """Persist one ReAct round without risking its terminal SSE events."""
+    try:
+        storage_conv.add_view_message(history_payload)
+        storage_conv.end_current_round()
+        storage_conv.save_to_storage()
+    except Exception:
+        logger.exception("Failed to persist ReAct agent history")
+
+    return (
+        _sse_event(final_answer.to_sse_payload()),
+        _sse_event({"type": "done"}),
+    )
+
+
+async def _cancel_and_await_agent_task(task: "asyncio.Task[Any]") -> None:
+    """Cancel a running agent task and always consume its terminal result."""
+    was_done = task.done()
+    if not was_done:
+        task.cancel()
+    try:
+        await task
+    except asyncio.CancelledError:
+        pass
+    except Exception:
+        if not was_done:
+            logger.exception("ReAct agent task failed during stream cleanup")
+
+
 async def _react_agent_stream(
     dialogue: ConversationVo,
     tool_mode: str = "full",
+) -> AsyncGenerator[str, None]:
+    """Stream ReAct events while owning the lifetime of its background task."""
+    agent_task_holder: List["asyncio.Task[Any]"] = []
+    final_emitted = False
+    done_emitted = False
+    try:
+        async for event in _react_agent_stream_impl(
+            dialogue,
+            tool_mode=tool_mode,
+            agent_task_holder=agent_task_holder,
+        ):
+            event_type = _sse_event_type(event)
+            final_emitted = final_emitted or event_type == "final"
+            done_emitted = done_emitted or event_type == "done"
+            yield event
+    except Exception:
+        logger.exception("ReAct agent stream failed before normal completion")
+        if not final_emitted and not done_emitted:
+            yield _sse_event(
+                AgentFinalAnswer(
+                    content="抱歉，回答生成过程中发生错误，请重试。"
+                ).to_sse_payload()
+            )
+        if not done_emitted:
+            yield _sse_event({"type": "done"})
+    finally:
+        if agent_task_holder:
+            await _cancel_and_await_agent_task(agent_task_holder[0])
+
+
+class _AgentStreamingResponse(StreamingResponse):
+    """Streaming response that explicitly closes its owned body iterator."""
+
+    async def __call__(self, scope: Any, receive: Any, send: Any) -> None:
+        try:
+            await super().__call__(scope, receive, send)
+        finally:
+            close = getattr(self.body_iterator, "aclose", None)
+            if callable(close):
+                try:
+                    await close()
+                except asyncio.CancelledError:
+                    raise
+                except Exception:
+                    logger.exception("Failed to close ReAct agent stream iterator")
+
+
+async def _react_agent_stream_impl(
+    dialogue: ConversationVo,
+    tool_mode: str = "full",
+    agent_task_holder: Optional[List["asyncio.Task[Any]"]] = None,
 ) -> AsyncGenerator[str, None]:
     """Core ReAct agent streaming logic.
 
@@ -982,8 +1083,6 @@ async def _react_agent_stream(
                    "knowledge" — only knowledge base tools + todowrite + terminate,
                    optimized for pure knowledge-chat scenarios.
     """
-    import asyncio
-
     from dbgpt.agent import AgentContext, AgentMemory, AgentMessage
     from dbgpt.agent.claude_skill import get_registry, load_skills_from_dir
     from dbgpt.agent.core.memory.gpts import (
@@ -2694,6 +2793,8 @@ async def run_agent():
         )
 
     agent_task = asyncio.create_task(run_agent())
+    if agent_task_holder is not None:
+        agent_task_holder.append(agent_task)
     round_step_map: Dict[int, str] = {}
     pending_thoughts: Dict[
         int, List[str]
@@ -2704,8 +2805,7 @@ async def run_agent():
     history_steps: List[Dict[str, Any]] = []
     current_history_step: Optional[Dict[str, Any]] = None
     subagent_history: Dict[str, Dict[str, Any]] = {}
-    # Collect chunks from knowledge tool observations for citation
-    _cited_chunks: List[Dict[str, Any]] = []
+    final_answer_assembler = FinalAnswerAssembler()
 
     # Emit pre-loaded skill as an SSE step before agent starts processing
     if pre_matched_skill:
@@ -3078,6 +3178,7 @@ async def run_agent():
             observation_text = action_output.get("observations") or action_output.get(
                 "content"
             )
+            status = "done" if action_output.get("is_exe_success", True) else "failed"
             if observation_text:
                 raw_chunks = emit_tool_chunks(react_step_id, observation_text)
                 if raw_chunks:
@@ -3086,7 +3187,7 @@ async def run_agent():
                 else:
                     for chunk in chunk_text(str(observation_text), max_len=600):
                         yield step_chunk(react_step_id, "text", chunk)
-                # --- History & citation: collect from observation ---
+                # --- History: collect the tool output for replay ---
                 if current_history_step is not None:
                     parsed_obs = None
                     if isinstance(observation_text, str):
@@ -3106,22 +3207,24 @@ async def run_agent():
                                         "content": content,
                                     }
                                 )
-                                clean = _strip_html_tags(content)
-                                if len(clean) >= 10:
-                                    _cited_chunks.append({"content": clean})
                     elif isinstance(observation_text, str) and observation_text:
                         current_history_step["outputs"].append(
                             {
                                 "output_type": "text",
                                 "content": observation_text,
                             }
                         )
-                        clean = _strip_html_tags(observation_text)
-                        if len(clean) >= 10:
-                            _cited_chunks.append({"content": clean})
+
+                # Citations are opt-in: the assembler accepts only explicitly
+                # supported knowledge tools and fails closed on malformed data.
+                final_answer_assembler.observe(
+                    action,
+                    action_input_data,
+                    observation_text,
+                    succeeded=status == "done",
+                )
 
             # Mark step as done and track as last completed
-            status = "done" if action_output.get("is_exe_success", True) else "failed"
             yield step_done(react_step_id, status)
             if (
                 status == "done"
@@ -3148,20 +3251,23 @@ async def run_agent():
         error_payload = json.dumps(
             {
                 "version": 1,
+                "protocol_version": 2,
                 "type": "react-agent",
                 "final_content": err_msg,
+                "citations": [],
                 "steps": history_steps,
                 "task_plan": list(_todo_list),
                 "generated_images": react_state.get("generated_images", []),
                 "sub_agents": build_subagent_history_snapshot(subagent_history),
             },
             ensure_ascii=False,
         )
-        storage_conv.add_view_message(error_payload)
-        storage_conv.end_current_round()
-        storage_conv.save_to_storage()
-        yield _sse_event({"type": "final", "content": err_msg})
-        yield _sse_event({"type": "done"})
+        for terminal_event in _react_terminal_events(
+            storage_conv,
+            error_payload,
+            AgentFinalAnswer(content=err_msg),
+        ):
+            yield terminal_event
         return
 
     if reply.action_report and reply.action_report.terminate:
@@ -3213,133 +3319,29 @@ async def run_agent():
     else:
         final_content = reply.content or ""
 
-    # Auto-annotate citations in the final answer: match collected chunks
-    # against the answer text and insert [n] markers for traceability.
-    if _cited_chunks:
-        # Deduplicate by content while preserving order
-        seen = set()
-        unique_chunks = []
-        for c in _cited_chunks:
-            key = c.get("content", "")
-            if key not in seen:
-                seen.add(key)
-                unique_chunks.append(c)
-        _cited_chunks[:] = unique_chunks
-        final_content = _auto_annotate_citations(final_content, _cited_chunks)
-
-    # Build references XML for frontend citation panel
-    references_xml = _build_references_xml(_cited_chunks)
+    final_answer = final_answer_assembler.finalize(final_content)
 
     # Persist AI reply with structured history payload
     history_payload = json.dumps(
         {
             "version": 1,
+            "protocol_version": 2,
             "type": "react-agent",
-            "final_content": final_content,
+            "final_content": final_answer.content,
+            "citations": [citation.to_dict() for citation in final_answer.citations],
             "steps": history_steps,
             "task_plan": list(_todo_list),
             "generated_images": react_state.get("generated_images", []),
             "sub_agents": build_subagent_history_snapshot(subagent_history),
         },
         ensure_ascii=False,
     )
-    storage_conv.add_view_message(history_payload)
-    storage_conv.end_current_round()
-    storage_conv.save_to_storage()
-
-    # Emit final content with references appended so the frontend can render
-    # inline citations and the citation drawer panel.
-    final_with_refs = final_content + references_xml
-    yield _sse_event({"type": "final", "content": final_with_refs})
-    yield _sse_event({"type": "done"})
-
-
-# ---------------------------------------------------------------------------
-# Citation helpers for react-agent knowledge chat
-# ---------------------------------------------------------------------------
-
-
-def _strip_html_tags(text: str) -> str:
-    """Remove HTML/XML tags and unescape HTML entities from a string."""
-    import html
-
-    text = html.unescape(text or "")
-    text = re.sub(r"<[^>]+>", "", text)
-    text = re.sub(r"\s+", " ", text)
-    return text.strip()
-
-
-def _auto_annotate_citations(text: str, chunks: list) -> str:
-    """Insert [n] citation markers into the answer text by matching collected
-    knowledge tool observation chunks against the final answer.
-
-    Uses longest-common-substring matching (>= 10 chars) to find where each
-    chunk appears in the answer and inserts a [n] marker after it.
-    """
-    if not chunks or not text:
-        return text
-
-    # Sort by content length descending so longer (more specific) chunks
-    # are matched first.
-    indexed = sorted(
-        enumerate(chunks, start=1),
-        key=lambda x: len(x[1].get("content", "")),
-        reverse=True,
-    )
-
-    for idx, chunk_info in indexed:
-        content = (chunk_info.get("content") or "").strip()
-        if len(content) < 8:
-            continue
-        # Find longest common substring
-        match = _longest_common_substring(text, content, min_len=8, max_len=80)
-        if not match:
-            continue
-        marker = f"[{idx}]"
-        if marker not in text:
-            # Replace ALL occurrences — a single source chunk often
-            # supports multiple sentences in the answer.
-            text = text.replace(match, match + marker)
-    return text
-
-
-def _longest_common_substring(
-    text: str, chunk: str, min_len: int = 8, max_len: int = 80
-) -> str:
-    """Find the longest substring of `chunk` that appears in `text`."""
-    chunk = chunk.strip()
-    for length in range(min(max_len, len(chunk)), min_len - 1, -1):
-        for start in range(len(chunk) - length + 1):
-            sub = chunk[start : start + length]
-            if sub in text:
-                return sub
-    return ""
-
-
-def _build_references_xml(chunks: list) -> str:
-    """Build a <references> XML tag with collected chunks for the frontend
-    citation panel and inline citation hover popovers.
-
-    Uses single-quoted attribute (no HTML entity escaping) so the frontend
-    regex /<references[^>]*references='([^']*)'/ can extract and JSON.parse.
-    """
-    if not chunks:
-        return ""
-    refs = []
-    for idx, chunk_info in enumerate(chunks, start=1):
-        content = (chunk_info.get("content") or "").strip()
-        if len(content) < 10:
-            continue
-        refs.append({"index": idx, "id": idx, "content": content, "recall_score": None})
-    if not refs:
-        return ""
-    # Build with json.dumps then wrap in single quotes so the frontend can
-    # extract via regex and JSON.parse without unescaping HTML entities.
-    payload_obj = [{"name": "Knowledge Base", "chunks": refs}]
-    payload = json.dumps(payload_obj, ensure_ascii=False)
-    return (
-        '\n\n<references title="References" references=\'' + payload + "'></references>"
-    )
+    for terminal_event in _react_terminal_events(
+        storage_conv,
+        history_payload,
+        final_answer,
+    ):
+        yield terminal_event
 
 
 # ---------------------------------------------------------------------------
@@ -3564,7 +3566,7 @@ async def chat_react_agent(
         "Transfer-Encoding": "chunked",
     }
     try:
-        return StreamingResponse(
+        return _AgentStreamingResponse(
             _react_agent_stream(dialogue, tool_mode="full"),
             headers=headers,
             media_type="text/event-stream",
@@ -3605,7 +3607,7 @@ async def chat_knowledge_agent(
         "Transfer-Encoding": "chunked",
     }
     try:
-        return StreamingResponse(
+        return _AgentStreamingResponse(
             _react_agent_stream(dialogue, tool_mode="knowledge"),
             headers=headers,
             media_type="text/event-stream",