infiniflow/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-10-12 02:16:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(agent): propagate deferred stream errors through Message (#18796)

Browse Source 
Fix lazy stream errors — partial output followed by failure now stops and reports error instead of persisting as success.
This commit is contained in:
Lem0nTea2002
2026-09-02 10:11:04 +08:00
committed by GitHub
parent 58decad723
commit 500fe29d62
3 changed files with 198 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
agent/canvas.py
+28 -17
View File
@@ -609,12 +609,15 @@ class Canvas(Graph):
def _node_finished(cpn_obj):
outputs = cpn_obj.output()
logged_outputs = dict(outputs)
if logged_outputs.get("_ERROR"):
logged_outputs["_ERROR"] = "<redacted>"
_logger.debug(
"[Canvas] Component '%s' (%s) finished. Outputs: %s, Error: %s",
self.get_component_name(cpn_obj._id),
self.get_component_type(cpn_obj._id),
json.dumps(outputs, ensure_ascii=False, default=str)[:500],
cpn_obj.error(),
json.dumps(logged_outputs, ensure_ascii=False, default=str)[:500],
bool(cpn_obj.error()),
)
return decorate(
"node_finished",
@@ -654,7 +657,9 @@ class Canvas(Graph):
for i in range(idx, to):
cpn = self.get_component(self.path[i])
cpn_obj = self.get_component_obj(self.path[i])
if cpn_obj.component_name.lower() == "message":
is_message = cpn_obj.component_name.lower() == "message"
streamed_message_content = None
if is_message:
if cpn_obj.get_param("auto_play"):
try:
tts_model_config = get_tenant_default_model_by_type(self._tenant_id, LLMType.TTS)
@@ -768,7 +773,6 @@ class Canvas(Graph):
await tts_queue.join()
for ev in await _drain_ready_tts():
yield ev
cpn_obj.set_output("content", _m)
finally:
for worker in tts_workers:
worker.cancel()
@@ -777,9 +781,28 @@ class Canvas(Graph):
async for ev in _stream_events():
yield ev
streamed_message_content = _m
else:
yield decorate("message", {"content": cpn_obj.output("content")})
other_branch = False
component_error = cpn_obj.error()
if component_error:
if is_message and isinstance(cpn_obj.output("content"), partial):
cpn_obj.set_output("content", None)
ex = cpn_obj.exception_handler()
if ex and ex["goto"]:
self.path.extend(ex["goto"])
other_branch = True
elif ex and ex["default_value"]:
yield decorate("message", {"content": ex["default_value"]})
yield decorate("message_end", {})
else:
self.error = component_error if "Task has been canceled" in component_error else f"Component execution failed: {cpn_obj._id}"
if is_message and not component_error:
if streamed_message_content is not None:
cpn_obj.set_output("content", streamed_message_content)
message_end = self._build_message_end(cpn_obj)
yield decorate("message_end", message_end)
@@ -790,18 +813,6 @@ class Canvas(Graph):
yield _node_finished(_cpn_obj)
partials.pop(0)
other_branch = False
if cpn_obj.error():
ex = cpn_obj.exception_handler()
if ex and ex["goto"]:
self.path.extend(ex["goto"])
other_branch = True
elif ex and ex["default_value"]:
yield decorate("message", {"content": ex["default_value"]})
yield decorate("message_end", {})
else:
self.error = cpn_obj.error()
if cpn_obj.component_name.lower() not in ("iteration", "loop"):
if isinstance(cpn_obj.output("content"), partial):
if self.error:
@@ -843,7 +854,7 @@ class Canvas(Graph):
_extend_path(cpn["downstream"])
if self.error:
logging.error(f"Runtime Error: {self.error}")
logging.error("Runtime Error: %s", self.error)
break
idx = to
agent/component/message.py
+9
View File
@@ -42,6 +42,8 @@ from common import settings
from api.db.joint_services.memory_message_service import queue_save_to_memory_task
_logger = logging.getLogger(__name__)
def _valid_message_content(content: Any) -> list[str]:
if not isinstance(content, list):
@@ -233,6 +235,13 @@ class Message(ComponentBase):
all_content += t
cnt += t
yield t
if "@" in exp:
source_component_id = exp.split("@", 1)[0]
source_component = self._canvas.get_component_obj(source_component_id)
if source_error := source_component.error():
_logger.warning("Message stream stopped after source component error: %s", source_component_id)
self.set_output("_ERROR", source_error)
return
self.set_input_value(exp, cnt)
continue
elif inspect.isawaitable(v):
test/unit_test/agent/component/test_message_stream_error.py
+161
View File
@@ -0,0 +1,161 @@
import asyncio
import logging
import time
from functools import partial
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
from agent.canvas import Canvas
from agent.component.message import Message, MessageParam
async def _collect_stream(message: Message, template: str) -> str:
return "".join([chunk async for chunk in message._stream(template)])
def _build_message(source_stream, source):
message = object.__new__(Message)
message._canvas = SimpleNamespace(
get_variable_value=lambda _: partial(source_stream),
get_component_obj=lambda _: source,
is_canceled=lambda: False,
task_id="test-task",
)
message._param = MessageParam()
message._save_to_memory = AsyncMock()
return message
@pytest.mark.p1
def test_stream_propagates_deferred_source_error(caplog):
source = SimpleNamespace(error_value=None)
source.error = lambda: source.error_value
async def source_stream():
yield "partial"
source.error_value = "provider stream failed"
message = _build_message(source_stream, source)
with caplog.at_level(logging.WARNING):
content = asyncio.run(_collect_stream(message, "{llm@content}"))
assert content == "partial"
assert message.error() == "provider stream failed"
assert message.output("content") == ""
assert message.get_input_value("llm@content") is None
message._save_to_memory.assert_not_awaited()
assert "source component error: llm" in caplog.text
assert "provider stream failed" not in caplog.text
@pytest.mark.p1
def test_stream_finalizes_successful_deferred_source():
source = SimpleNamespace(error=lambda: None)
async def source_stream():
yield "complete"
message = _build_message(source_stream, source)
content = asyncio.run(_collect_stream(message, "{llm@content}"))
assert content == "complete"
assert message.error() is None
assert message.output("content") == "complete"
assert message.get_input_value("llm@content") == "complete"
message._save_to_memory.assert_awaited_once_with("complete")
def _build_canvas_component(component_id, component_name, invoke_async):
outputs = {}
component = SimpleNamespace()
async def _invoke_async(**kwargs):
outputs["_created_time"] = time.perf_counter()
await invoke_async(component, **kwargs)
component._id = component_id
component._param = SimpleNamespace(layout_recognize=None)
component.component_name = component_name
component._invoke_async = _invoke_async
component.invoke_async = _invoke_async
component.invoke = lambda **_kwargs: None
component.reset = lambda only_output=False: None
component.get_input_elements = dict
component.get_input = dict
component.get_input_values = dict
component.output = lambda key=None: outputs.get(key) if key else dict(outputs)
component.set_output = lambda key, value: outputs.__setitem__(key, value)
component.error = lambda: outputs.get("_ERROR")
component.exception_handler = lambda: None
component.get_param = lambda _key: False
component.get_parent = lambda: None
component.thoughts = lambda: ""
return component
def _build_canvas(invoke_message):
async def invoke_begin(_component, **_kwargs):
return None
begin = _build_canvas_component("begin", "Begin", invoke_begin)
message = _build_canvas_component("message", "Message", invoke_message)
canvas = Canvas.__new__(Canvas)
canvas.__dict__.update(
_thread_pool=SimpleNamespace(_max_workers=1),
_run_token_usage={"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0, "calls": 0},
task_id="task",
path=[],
error="",
history=[],
retrieval=[],
globals={"sys.history": [], "sys.conversation_turns": 0, "sys.date": ""},
components={
"begin": {"obj": begin, "downstream": ["message"], "upstream": []},
"message": {"obj": message, "downstream": [], "upstream": ["begin"]},
},
)
canvas.is_canceled = lambda: False
canvas.get_component_name = lambda component_id: component_id
return canvas, message
@pytest.mark.p1
@pytest.mark.parametrize(
("chunk", "source_error", "completion_count"),
[
pytest.param("partial", "provider stream failed", 0, id="failed"),
pytest.param("complete", None, 1, id="successful"),
],
)
def test_canvas_message_completion_follows_stream_error(caplog, chunk, source_error, completion_count):
async def invoke_message(component, **_kwargs):
async def stream():
yield chunk
if source_error:
component.set_output("_ERROR", source_error)
component.set_output("content", partial(stream))
canvas, message = _build_canvas(invoke_message)
with caplog.at_level(logging.DEBUG):
events = asyncio.run(_collect_canvas_events(canvas))
assert [event["data"]["content"] for event in events if event["event"] == "message"] == [chunk]
assert sum(event["event"] == "message_end" for event in events) == completion_count
assert sum(event["event"] == "workflow_finished" for event in events) == completion_count
message_finished = next(event for event in events if event["event"] == "node_finished" and event["data"]["component_id"] == "message")
assert message_finished["data"]["error"] == source_error
assert message.output("content") == (None if source_error else chunk)
assert canvas.error == ("Component execution failed: message" if source_error else "")
if source_error:
assert "Runtime Error: Component execution failed: message" in caplog.text
assert source_error not in caplog.text
async def _collect_canvas_events(canvas):
return [event async for event in canvas._run_impl(query="hello", inputs={})]