mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-10 13:45:44 +08:00
fix(dialog): restore decorated answer in async_ask final SSE event (#13917)
## What's the problem Both `async_chat()` and `async_ask()` call `decorate_answer()` to build the final SSE payload — it inserts citation markers (`##N$$`) into the answer text and prunes `doc_aggs` to only the cited documents. Immediately after, both functions overwrite `final["answer"]` with `""`: ```python # async_chat(), line ~774 (issue #13828) final = decorate_answer(thought + full_answer) final["final"] = True final["audio_binary"] = None final["answer"] = "" # discards decorated text yield final # async_ask(), line ~1444 (same bug, different path) final = decorate_answer(full_answer) final["final"] = True final["answer"] = "" # discards decorated text yield final ``` The client receives filtered references (built for a citation-decorated answer it never sees) while displaying the raw, undecorated streaming text. Citations can never match. ## Root cause `final["answer"] = ""` was left over from an earlier design where clients were meant to reconstruct the full answer purely from delta events. Once `decorate_answer()` started placing citation markers, this blank-out broke the contract: the final event is where the decorated answer should land. ## Fix Remove the two blank-override lines — one in `async_chat()`, one in `async_ask()`: ```diff - final["answer"] = "" ``` `decorate_answer()` already sets `final["answer"]` to the correct decorated string; there is nothing to override. ## Relation to #13828 Issue #13828 and PR #13835 identify the bug in `async_chat()`. This PR absorbs that fix and also corrects the identical pattern in `async_ask()` (used by the `/retrieval` route in `chat_api.py`), which PR #13835 does not touch. ## Regression test Added `test/unit_test/api/db/services/test_dialog_service_final_answer.py` with three tests: | Test | Purpose | |------|---------| | `test_buggy_pattern_drops_answer` | Documents the old behaviour: blank-override empties the final answer | | `test_fixed_pattern_preserves_decorated_answer` | Core invariant: final event carries the decorated text from `decorate_answer()` | | `test_final_event_reference_matches_decorated_result` | Citation markers in the answer must match the pruned `doc_aggs` in the same event | Local run result: ``` test_dialog_service_final_answer.py::test_buggy_pattern_drops_answer PASSED test_dialog_service_final_answer.py::test_fixed_pattern_preserves_decorated_answer PASSED test_dialog_service_final_answer.py::test_final_event_reference_matches_decorated_result PASSED 3 passed in 0.04s ``` `ruff check` passes with no issues on all changed files. --------- Co-authored-by: edenfunf <edenfunf@gmail.com> Co-authored-by: Yingfeng <yingfeng.zhang@gmail.com>
This commit is contained in:
Eden
@@ -795,7 +795,6 @@ async def async_chat(dialog, messages, stream=True, **kwargs):
|
||||
final = decorate_answer(thought + full_answer)
|
||||
final["final"] = True
|
||||
final["audio_binary"] = None
|
||||
final["answer"] = ""
|
||||
yield final
|
||||
else:
|
||||
if llm_type == "chat":
|
||||
@@ -1465,7 +1464,6 @@ async def async_ask(question, kb_ids, tenant_id, chat_llm_name=None, search_conf
|
||||
full_answer = last_state.full_text if last_state else ""
|
||||
final = decorate_answer(full_answer)
|
||||
final["final"] = True
|
||||
final["answer"] = ""
|
||||
yield final
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,358 @@
|
||||
#
|
||||
# Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
#
|
||||
"""
|
||||
Regression tests for the bug where async_ask() and async_chat() blanked out
|
||||
final["answer"] in the last SSE event, discarding the decorated answer text
|
||||
that contains citation markers.
|
||||
|
||||
Both functions call decorate_answer() which inserts citation markers and prunes
|
||||
doc_aggs to cited documents, then overwrite final["answer"] = "" — discarding
|
||||
the decorated text before the client receives it.
|
||||
|
||||
The fix removes those two blank-override lines. Tests here drive the actual
|
||||
production functions (with heavy dependencies stubbed) to ensure regression
|
||||
protection is real: the suite would fail if the lines were re-introduced.
|
||||
|
||||
Related: PR #13835 (async_chat), this PR (async_ask + async_chat).
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import sys
|
||||
import types
|
||||
import warnings
|
||||
from copy import deepcopy
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
warnings.filterwarnings(
|
||||
"ignore",
|
||||
message="pkg_resources is deprecated as an API.*",
|
||||
category=UserWarning,
|
||||
)
|
||||
|
||||
|
||||
def _install_cv2_stub_if_unavailable():
|
||||
try:
|
||||
import cv2 # noqa: F401
|
||||
return
|
||||
except Exception:
|
||||
pass
|
||||
stub = types.ModuleType("cv2")
|
||||
stub.INTER_LINEAR = 1
|
||||
stub.INTER_CUBIC = 2
|
||||
stub.BORDER_CONSTANT = 0
|
||||
stub.BORDER_REPLICATE = 1
|
||||
stub.COLOR_BGR2RGB = 0
|
||||
stub.COLOR_BGR2GRAY = 1
|
||||
stub.COLOR_GRAY2BGR = 2
|
||||
stub.IMREAD_IGNORE_ORIENTATION = 128
|
||||
stub.IMREAD_COLOR = 1
|
||||
stub.RETR_LIST = 1
|
||||
stub.CHAIN_APPROX_SIMPLE = 2
|
||||
|
||||
def _module_getattr(name):
|
||||
if name.isupper():
|
||||
return 0
|
||||
raise RuntimeError(f"cv2.{name} is unavailable in this test environment")
|
||||
|
||||
stub.__getattr__ = _module_getattr
|
||||
sys.modules["cv2"] = stub
|
||||
|
||||
|
||||
_install_cv2_stub_if_unavailable()
|
||||
|
||||
from api.db.services import dialog_service # noqa: E402
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Shared stubs
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_KBINFOS = {
|
||||
"chunks": [
|
||||
{
|
||||
"doc_id": "doc-1",
|
||||
"content_ltks": "ragflow is a rag engine",
|
||||
"content_with_weight": "RAGFlow is a RAG engine.",
|
||||
"vector": [0.1, 0.2, 0.3],
|
||||
"docnm_kwd": "intro.pdf",
|
||||
},
|
||||
],
|
||||
"doc_aggs": [{"doc_id": "doc-1", "doc_name": "intro.pdf", "count": 1}],
|
||||
"total": 1,
|
||||
}
|
||||
|
||||
_KB = SimpleNamespace(
|
||||
id="kb-1",
|
||||
embd_id="text-embedding-ada-002@OpenAI",
|
||||
tenant_embd_id="text-embedding-ada-002@OpenAI",
|
||||
tenant_id="tenant-1",
|
||||
chunk_num=1,
|
||||
name="Test KB",
|
||||
parser_id="general",
|
||||
)
|
||||
|
||||
_LLM_CONFIG = {
|
||||
"llm_name": "gpt-4o",
|
||||
"llm_factory": "OpenAI",
|
||||
"model_type": "chat",
|
||||
"max_tokens": 8192,
|
||||
}
|
||||
|
||||
|
||||
class _StreamingChatModel:
|
||||
"""Yields a single-chunk full answer, no citations."""
|
||||
|
||||
def __init__(self, answer: str):
|
||||
self.answer = answer
|
||||
self.max_length = 8192
|
||||
|
||||
async def async_chat_streamly_delta(self, system_prompt, messages, gen_conf, **_kwargs):
|
||||
yield self.answer
|
||||
|
||||
async def async_chat(self, system_prompt, messages, gen_conf, **_kwargs):
|
||||
return self.answer
|
||||
|
||||
|
||||
class _StubRetriever:
|
||||
async def retrieval(self, *_args, **_kwargs):
|
||||
return deepcopy(_KBINFOS)
|
||||
|
||||
def retrieval_by_children(self, chunks, tenant_ids):
|
||||
return chunks
|
||||
|
||||
def insert_citations(self, answer, content_ltks, vectors, embd_mdl, **_kwargs):
|
||||
# Return the answer unchanged; no citation markers inserted.
|
||||
return answer, set()
|
||||
|
||||
|
||||
def _collect(async_gen):
|
||||
async def _run():
|
||||
return [ev async for ev in async_gen]
|
||||
return asyncio.run(_run())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests for async_ask (production code path)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.p2
|
||||
def test_async_ask_final_event_carries_decorated_answer(monkeypatch):
|
||||
"""
|
||||
Drive the real dialog_service.async_ask() and verify that the final SSE
|
||||
event (final=True) exposes the answer produced by decorate_answer(), not
|
||||
an empty string.
|
||||
|
||||
Regression guard: if `final["answer"] = ""` is re-introduced at line ~1444,
|
||||
this test fails.
|
||||
"""
|
||||
llm_answer = "RAGFlow is a RAG engine built for document understanding."
|
||||
chat_mdl = _StreamingChatModel(llm_answer)
|
||||
retriever = _StubRetriever()
|
||||
|
||||
monkeypatch.setattr(
|
||||
dialog_service.KnowledgebaseService, "get_by_ids", lambda _ids: [_KB]
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
dialog_service, "get_model_config_by_type_and_name",
|
||||
lambda _tid, _type, _name: _LLM_CONFIG,
|
||||
)
|
||||
monkeypatch.setattr(dialog_service, "LLMBundle", lambda _tid, _cfg: chat_mdl)
|
||||
monkeypatch.setattr(dialog_service.settings, "retriever", retriever, raising=False)
|
||||
monkeypatch.setattr(dialog_service.settings, "kg_retriever", retriever, raising=False)
|
||||
monkeypatch.setattr(
|
||||
dialog_service.DocMetadataService, "get_flatted_meta_by_kbs", lambda _ids: {}
|
||||
)
|
||||
monkeypatch.setattr(dialog_service, "label_question", lambda _q, _kbs: "")
|
||||
# kb_prompt calls DocumentService.get_by_ids which needs a live DB; stub it out.
|
||||
monkeypatch.setattr(
|
||||
dialog_service, "kb_prompt",
|
||||
lambda _kbinfos, _max_tokens, **_kw: ["RAGFlow is a RAG engine."],
|
||||
)
|
||||
|
||||
events = _collect(
|
||||
dialog_service.async_ask(
|
||||
question="What is RAGFlow?",
|
||||
kb_ids=["kb-1"],
|
||||
tenant_id="tenant-1",
|
||||
)
|
||||
)
|
||||
|
||||
assert events, "async_ask must yield at least one event"
|
||||
|
||||
final_events = [e for e in events if e.get("final") is True]
|
||||
assert len(final_events) == 1, (
|
||||
f"Expected exactly one final event, got {len(final_events)}: {final_events}"
|
||||
)
|
||||
final = final_events[0]
|
||||
|
||||
assert final["answer"] != "", (
|
||||
"Final event answer must not be blank — decorate_answer() result was discarded.\n"
|
||||
"This is the regression: final['answer'] = '' was removed from async_ask()."
|
||||
)
|
||||
assert llm_answer in final["answer"], (
|
||||
f"LLM answer text expected in final event, got: {final['answer']!r}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.p2
|
||||
def test_async_ask_delta_events_carry_incremental_text_only(monkeypatch):
|
||||
"""
|
||||
Intermediate delta events must have empty reference dicts.
|
||||
Only the final event should carry the populated reference from decorate_answer().
|
||||
"""
|
||||
chat_mdl = _StreamingChatModel("Incremental text for delta test.")
|
||||
retriever = _StubRetriever()
|
||||
|
||||
monkeypatch.setattr(
|
||||
dialog_service.KnowledgebaseService, "get_by_ids", lambda _ids: [_KB]
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
dialog_service, "get_model_config_by_type_and_name",
|
||||
lambda _tid, _type, _name: _LLM_CONFIG,
|
||||
)
|
||||
monkeypatch.setattr(dialog_service, "LLMBundle", lambda _tid, _cfg: chat_mdl)
|
||||
monkeypatch.setattr(dialog_service.settings, "retriever", retriever, raising=False)
|
||||
monkeypatch.setattr(dialog_service.settings, "kg_retriever", retriever, raising=False)
|
||||
monkeypatch.setattr(
|
||||
dialog_service.DocMetadataService, "get_flatted_meta_by_kbs", lambda _ids: {}
|
||||
)
|
||||
monkeypatch.setattr(dialog_service, "label_question", lambda _q, _kbs: "")
|
||||
monkeypatch.setattr(
|
||||
dialog_service, "kb_prompt",
|
||||
lambda _kbinfos, _max_tokens, **_kw: ["RAGFlow is a RAG engine."],
|
||||
)
|
||||
|
||||
events = _collect(
|
||||
dialog_service.async_ask(
|
||||
question="Describe RAGFlow briefly.",
|
||||
kb_ids=["kb-1"],
|
||||
tenant_id="tenant-1",
|
||||
)
|
||||
)
|
||||
|
||||
delta_events = [e for e in events if not e.get("final")]
|
||||
final_events = [e for e in events if e.get("final") is True]
|
||||
|
||||
assert len(final_events) == 1, f"Expected exactly one final event, got {len(final_events)}"
|
||||
for ev in delta_events:
|
||||
assert ev["reference"] == {}, f"Delta event must have empty reference, got: {ev['reference']}"
|
||||
|
||||
assert "chunks" in final_events[0]["reference"], (
|
||||
"Final event reference must contain chunk data from decorate_answer()"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests for async_chat (production code path)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _make_dialog(chat_mdl_stub):
|
||||
"""Build a minimal dialog SimpleNamespace for async_chat()."""
|
||||
return SimpleNamespace(
|
||||
id="dialog-1",
|
||||
kb_ids=["kb-1"],
|
||||
tenant_id="tenant-1",
|
||||
tenant_llm_id=None,
|
||||
llm_id="gpt-4o",
|
||||
llm_setting={"temperature": 0.1},
|
||||
prompt_type="simple",
|
||||
prompt_config={
|
||||
"system": "You are helpful. {knowledge}",
|
||||
"parameters": [{"key": "knowledge", "optional": False}],
|
||||
"quote": True,
|
||||
"empty_response": "",
|
||||
"reasoning": False,
|
||||
"refine_multiturn": False,
|
||||
"cross_languages": False,
|
||||
"keyword": False,
|
||||
"toc_enhance": False,
|
||||
"tavily_api_key": "",
|
||||
"use_kg": False,
|
||||
"tts": False,
|
||||
},
|
||||
meta_data_filter={},
|
||||
similarity_threshold=0.2,
|
||||
vector_similarity_weight=0.3,
|
||||
top_n=6,
|
||||
top_k=1024,
|
||||
rerank_id="",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.p2
|
||||
def test_async_chat_final_event_carries_decorated_answer(monkeypatch):
|
||||
"""
|
||||
Drive the real dialog_service.async_chat() streaming path and verify that
|
||||
the final SSE event (final=True) exposes the answer from decorate_answer(),
|
||||
not an empty string.
|
||||
|
||||
Regression guard: if `final["answer"] = ""` is re-introduced at line ~774,
|
||||
this test fails.
|
||||
"""
|
||||
llm_answer = "RAGFlow handles document parsing with deep understanding."
|
||||
chat_mdl = _StreamingChatModel(llm_answer)
|
||||
retriever = _StubRetriever()
|
||||
|
||||
# Stub out the heavy service/model calls
|
||||
monkeypatch.setattr(
|
||||
dialog_service.TenantLLMService, "llm_id2llm_type", lambda _llm_id: "chat"
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
dialog_service.TenantLLMService, "get_model_config",
|
||||
lambda _tid, _type, _llm_id: _LLM_CONFIG,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
dialog_service.TenantLangfuseService, "filter_by_tenant",
|
||||
lambda tenant_id: None,
|
||||
)
|
||||
# get_models returns (kbs, embd_mdl, rerank_mdl, chat_mdl, tts_mdl)
|
||||
monkeypatch.setattr(
|
||||
dialog_service, "get_models",
|
||||
lambda _dialog: ([_KB], chat_mdl, None, chat_mdl, None),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
dialog_service.KnowledgebaseService, "get_field_map", lambda _kb_ids: {}
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
dialog_service.KnowledgebaseService, "get_by_ids", lambda _ids: [_KB]
|
||||
)
|
||||
monkeypatch.setattr(dialog_service.settings, "retriever", retriever, raising=False)
|
||||
monkeypatch.setattr(dialog_service, "label_question", lambda _q, _kbs: "")
|
||||
monkeypatch.setattr(
|
||||
dialog_service, "kb_prompt",
|
||||
lambda _kbinfos, _max_tokens, **_kw: ["RAGFlow is a RAG engine."],
|
||||
)
|
||||
|
||||
dialog = _make_dialog(chat_mdl)
|
||||
messages = [{"role": "user", "content": "What is RAGFlow?"}]
|
||||
|
||||
events = _collect(dialog_service.async_chat(dialog, messages, stream=True, quote=True))
|
||||
|
||||
final_events = [e for e in events if e.get("final") is True]
|
||||
assert len(final_events) == 1, (
|
||||
f"Expected exactly one final event, got {len(final_events)}: {final_events}"
|
||||
)
|
||||
final = final_events[0]
|
||||
|
||||
assert final["answer"] != "", (
|
||||
"Final event answer must not be blank — decorate_answer() result was discarded.\n"
|
||||
"This is the regression: final['answer'] = '' was removed from async_chat()."
|
||||
)
|
||||
assert llm_answer in final["answer"], (
|
||||
f"LLM answer text expected in final event, got: {final['answer']!r}"
|
||||
)
|
||||
Reference in New Issue
Block a user