deer-flow/backend/tests/test_batch_rag_evidence.py
Guo-Yixin 3182ecb1e4
feat(subagents): preserve RAG evidence in durable batch exports (#6413)
* feat(subagents): preserve RAG evidence in durable batch exports

* docs(subagents): clarify evidence omissions and migration head

* docs(persistence): update the migration chain head

* test(persistence): track the new batch evidence migration head

* test(persistence): verify batch evidence migration on PostgreSQL

* docs(subagents): retain evidence migration compatibility guidance

* fix(subagents): propagate cited evidence forwarding omissions

* fix(tests): adapt batch evidence PostgreSQL CI connection options

* fix(tests): preserve TLS policy in batch evidence migration fixture
2026-10-07 18:37:31 +08:00

450 lines
22 KiB
Python

"""RAG producer -> durable worker -> reopened SQL -> owner JSONL evidence."""
import asyncio
import importlib.util
import json
import os
import sys
import uuid
from copy import deepcopy
from datetime import UTC, datetime, timedelta
from enum import Enum
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
import pytest_asyncio
from langchain_core.messages import ToolMessage
from support.postgres import asyncpg_test_url
from deerflow.community.ragflow.formatting import format_retrieval_sources
from deerflow.config.database_config import DatabaseConfig
from deerflow.config.subagent_batches_config import SubagentBatchesConfig
from deerflow.config.subagent_runtime_config import SubagentRuntimeConfig
from deerflow.persistence.engine import close_engine, get_session_factory, init_engine_from_config
from deerflow.persistence.subagent_batches import SubagentBatchRepository
from deerflow.subagents import batch_service
from deerflow.subagents.batch_runtime import BatchSubmitRequest
from deerflow.subagents.step_events import capture_step_message
class Status(Enum):
RUNNING = "running"
COMPLETED = "completed"
FAILED = "failed"
@property
def is_terminal(self):
return self is not Status.RUNNING
def retrieval(*, text="Original evidence", name="knowledge_search"):
content, artifact = format_retrieval_sources(
{"chunks": [{"id": "chunk", "dataset_id": "kb", "document_id": "doc", "document_keyword": "Manual.pdf", "content": text}]},
dataset_names_by_id={"kb": "Engineering"},
max_chars_per_chunk=20_000,
max_total_chars=30_000,
)
captured = []
capture_step_message(ToolMessage(content=content, artifact=artifact, name=name, tool_call_id="search", id="message"), captured, set())
return content, captured
@pytest_asyncio.fixture(params=["sqlite", "postgres"])
async def env(monkeypatch, tmp_path, request):
schema = None
if request.param == "postgres":
uri = os.environ.get("TEST_POSTGRES_URI")
if not uri:
pytest.skip("requires TEST_POSTGRES_URI (real Postgres batch evidence chain)")
schema = f"batch_evidence_{uuid.uuid4().hex}"
db = DatabaseConfig(backend="postgres", postgres_url=asyncpg_test_url(uri), postgres_schema=schema)
else:
db = DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path / "db"))
service = None
try:
await init_engine_from_config(db)
if schema:
from sqlalchemy import text
async with get_session_factory()() as session:
assert (await session.execute(text("SELECT current_schema()"))).scalar_one() == schema
repo = SubagentBatchRepository(get_session_factory())
state = SimpleNamespace(repo=repo, db=db, result=SimpleNamespace(status=Status.COMPLETED, result="plain report", ai_messages=[], error=None, stop_reason=None, token_usage_records=[]))
class Executor:
def __init__(self, **kwargs):
state.owner = kwargs["user_id"]
state.scope = kwargs["knowledge_scope"]
def execute_async(self, prompt, task_id=None):
return task_id
monkeypatch.setattr(batch_service, "SubagentExecutor", Executor)
monkeypatch.setattr(batch_service, "SubagentStatus", Status)
monkeypatch.setattr(batch_service, "get_background_task_result", lambda _: state.result)
monkeypatch.setattr(batch_service, "cleanup_background_task", lambda _: None)
monkeypatch.setattr(batch_service, "resolve_subagent_model_name", lambda *args, **kwargs: "offline-model")
monkeypatch.setattr("deerflow.tools.get_available_tools", lambda **kwargs: [])
state.service = batch_service.SubagentBatchService(repository=repo, config=SubagentBatchesConfig(), runtime_config=SubagentRuntimeConfig(), app_config=SimpleNamespace())
service = state.service
yield state
finally:
try:
if service is not None:
await service.stop()
finally:
try:
await close_engine()
finally:
if schema:
import sqlalchemy as sa
from sqlalchemy.ext.asyncio import create_async_engine
engine = create_async_engine(db.app_sqlalchemy_url)
try:
async with engine.begin() as conn:
await conn.execute(sa.text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE'))
finally:
await engine.dispose()
async def execute(env):
batch = await env.service.submit(
BatchSubmitRequest(
user_id="user-1",
thread_id="thread-1",
run_id="run-1",
tool_call_id="call-1",
submission_key="submission",
title="Research",
subagent_type="general-purpose",
items=[{"key": "one", "prompt": "Research the selected knowledge base"}],
max_live_items=None,
max_running_items=None,
execution_spec={"subagent_config": {"name": "general-purpose", "description": "Researcher", "system_prompt": "Research"}, "knowledge_scope": {"version": 1, "mode": "selected", "dataset_ids": ["kb"]}},
)
)
await env.service.run_once(now=datetime.now(UTC))
await asyncio.gather(*list(env.service._executions.values()))
return batch
@pytest.mark.asyncio
@pytest.mark.parametrize("tool_name", ["knowledge_search", "task"])
async def test_cited_evidence_survives_worker_database_reopen_and_export(env, monkeypatch, tool_name):
from app.gateway.routers import subagent_batches as router
env.result.result, env.result.ai_messages = retrieval(name=tool_name)
source = env.result.ai_messages[0]["artifact"]["knowledge_sources"]["sources"][0]
expected_source = deepcopy(source)
batch = await execute(env)
assert env.owner == "user-1" and env.scope["dataset_ids"] == ["kb"]
# Actual engine teardown/reopen; neither the worker nor provider can serve evidence.
await env.service.stop()
await close_engine()
await init_engine_from_config(env.db)
repo = SubagentBatchRepository(get_session_factory())
source["text"] = "Provider content changed after retrieval"
env.result.ai_messages = []
assert "result_artifact" not in (await repo.list_items(batch["id"], user_id="user-1"))[0]
assert await repo.list_items(batch["id"], user_id="other", include_result=True) is None
monkeypatch.setattr(router, "get_current_user", AsyncMock(return_value="user-1"))
request = SimpleNamespace(app=SimpleNamespace(state=SimpleNamespace(subagent_batch_repo=repo, subagent_batches_available=False)))
response = await router.export_batch_results.__wrapped__(thread_id="thread-1", batch_id=batch["id"], request=request)
rows = [json.loads(line) async for line in response.body_iterator]
row = rows[0]
assert row["result"] == env.result.result
assert row["result_artifact"]["knowledge_sources"]["sources"] == [expected_source]
assert "Original evidence" in row["result_artifact"]["knowledge_sources"]["sources"][0]["text"]
assert "execution_spec" not in row and "prompt" not in row
@pytest.mark.asyncio
async def test_real_subagent_graph_captures_artifact_for_durable_export(env, monkeypatch, tmp_path):
from langchain_core.language_models.fake_chat_models import GenericFakeChatModel
from langchain_core.messages import AIMessage
from langchain_core.tools import StructuredTool
from deerflow.config.app_config import AppConfig
from deerflow.subagents.config import SubagentConfig
# conftest stubs executor to break a circular import. Load the production
# executor, keeping its actual graph, middleware, ToolNode and step capture.
path = Path(__file__).parents[1] / "packages/harness/deerflow/subagents/executor.py"
spec = importlib.util.spec_from_file_location("_batch_evidence_executor", path)
module = importlib.util.module_from_spec(spec)
monkeypatch.setitem(sys.modules, spec.name, module)
spec.loader.exec_module(module)
monkeypatch.setenv("DEER_FLOW_HOME", str(tmp_path / "graph"))
content, messages = retrieval()
artifact = messages[0]["artifact"]
class Model(GenericFakeChatModel):
def bind_tools(self, tools, **kwargs):
return self
model = Model(messages=iter([AIMessage(content="Searching", tool_calls=[{"name": "knowledge_search", "args": {}, "id": "search"}])]))
monkeypatch.setattr(module, "create_chat_model", lambda **kwargs: model)
tool = StructuredTool.from_function(lambda: (content, artifact), name="knowledge_search", description="Offline captured evidence", response_format="content_and_artifact", return_direct=True)
config = AppConfig.model_validate({"models": [{"name": "offline", "use": "langchain_openai:ChatOpenAI", "model": "offline"}], "sandbox": {"use": "deerflow.sandbox.local:LocalSandboxProvider"}, "summarization": {"enabled": False}})
executor = module.SubagentExecutor(config=SubagentConfig(name="researcher", description="Research", skills=[]), tools=[tool], parent_model="offline", app_config=config, thread_id="thread-1", user_id="user-1")
try:
env.result = await executor._aexecute("Return evidence")
assert env.result.status is module.SubagentStatus.COMPLETED, env.result.error
assert env.result.result == content
captured = [m for m in env.result.ai_messages if m["type"] == "tool"]
assert captured[0]["artifact"] == artifact
monkeypatch.setattr(batch_service, "SubagentStatus", module.SubagentStatus)
batch = await execute(env)
row = (await env.repo.list_items(batch["id"], user_id="user-1", include_result=True))[0]
assert row["status"] == "succeeded"
assert row["result_artifact"]["knowledge_sources"]["sources"] == artifact["knowledge_sources"]["sources"]
finally:
module._shutdown_isolated_subagent_loop()
@pytest.mark.parametrize("change", ["uncited", "unsupported-tool", "future-version", "invalid-id", "invalid-provider", "invalid-text", "invalid-locator", "overlong-name"])
def test_unrelated_or_invalid_evidence_is_not_persisted(change):
from deerflow.community.ragflow.sources import durable_source_artifact
content, messages = retrieval()
source = messages[0]["artifact"]["knowledge_sources"]["sources"][0]
if change == "uncited":
content = "A report with no citations"
elif change == "unsupported-tool":
messages[0]["name"] = "arbitrary_tool"
elif change == "future-version":
messages[0]["artifact"]["knowledge_sources"]["version"] = 2
elif change == "invalid-id":
content = content.replace(source["id"], "invalid")
source["id"] = "invalid"
elif change == "invalid-provider":
source["provider"] = "other"
elif change == "invalid-text":
source["text"] = None
elif change == "invalid-locator":
source["document_id"] = ""
else:
source["document_name"] = "x" * 513
result = durable_source_artifact(messages, content, max_chars=1000)
assert result is None or result["knowledge_sources"]["sources"] == []
@pytest.mark.parametrize("max_chars", [1, 1000, 1800, 10_000])
def test_evidence_budget_keeps_whole_records_and_reports_omissions(max_chars):
from deerflow.community.ragflow.sources import durable_source_artifact
content, messages = retrieval(text='Quoted "evidence" and 中文 ' * 200)
original = deepcopy(messages)
small_content, small_messages = retrieval(text="A small source fits")
# Both citations occur in the report; a large early source cannot starve a small later one.
result = durable_source_artifact(messages + small_messages, content + small_content, max_chars=max_chars)
assert messages == original
if result is None:
assert max_chars == 1
return
assert len(json.dumps(result, ensure_ascii=False)) <= max_chars
payload = result["knowledge_sources"]
assert len(payload["sources"]) + payload["omitted_count"] == 2
originals = [m["artifact"]["knowledge_sources"]["sources"][0] for m in messages + small_messages]
assert all(source in originals for source in payload["sources"])
if max_chars in (1000, 1800):
assert payload["sources"] == [originals[1]]
assert payload["omitted_count"] == 1
def test_forwarding_limit_counts_duplicate_omissions_once():
from deerflow.community.ragflow.sources import cited_source_artifact, durable_source_artifact
entries = [retrieval(text=f"Evidence {index}") for index in range(102)]
content = "\n".join(entry[0] for entry in entries)
messages = [message for entry in entries for message in entry[1]]
result = durable_source_artifact(messages + messages, content, max_chars=100_000)
assert len(result["knowledge_sources"]["sources"]) == 100
assert result["knowledge_sources"]["omitted_count"] == 2
# Ordinary task forwarding retains its existing artifact shape.
assert "omitted_count" not in cited_source_artifact(messages, content)["knowledge_sources"]
def forwarded_task(messages, content):
from deerflow.tools.builtins.task_tool import _task_result_command
command = _task_result_command(tool_call_id="child", status="completed", result=content, source_messages=messages)
captured = []
capture_step_message(command.update["messages"][0], captured, set())
return captured
@pytest.mark.asyncio
async def test_nested_task_omissions_survive_worker_reopen_and_export(env, monkeypatch):
from app.gateway.routers import subagent_batches as router
entries = [retrieval(text=f"Evidence {index}") for index in range(102)]
report = "\n".join(entry[0] for entry in entries)
messages = [message for entry in entries for message in entry[1]]
forwarded = forwarded_task(messages, report)
# Repeated child messages and a second delegation must not double-count.
env.result.result = report
env.result.ai_messages = forwarded_task(forwarded + forwarded, report)
snapshot = deepcopy(env.result.ai_messages)
batch = await execute(env)
await env.service.stop()
await close_engine()
await init_engine_from_config(env.db)
repo = SubagentBatchRepository(get_session_factory())
monkeypatch.setattr(router, "get_current_user", AsyncMock(return_value="user-1"))
request = SimpleNamespace(app=SimpleNamespace(state=SimpleNamespace(subagent_batch_repo=repo)))
response = await router.export_batch_results.__wrapped__(thread_id="thread-1", batch_id=batch["id"], request=request)
row = json.loads([line async for line in response.body_iterator][0])
payload = row["result_artifact"]["knowledge_sources"]
assert len(payload["sources"]) == 100
assert payload["omitted_count"] == 2
assert "omitted_source_ids" not in payload
assert env.result.ai_messages == snapshot
@pytest.mark.parametrize("selection", ["retained", "omitted", "unknown", "partial"])
def test_forwarded_omissions_follow_only_current_complete_references(selection):
from deerflow.community.ragflow.sources import durable_source_artifact
entries = [retrieval(text=f"Evidence {index}") for index in range(101)]
report = "\n".join(entry[0] for entry in entries)
messages = [message for entry in entries for message in entry[1]]
forwarded = forwarded_task(messages, report)
selected = {"retained": entries[0][0], "omitted": entries[-1][0], "unknown": "[citation:1](#knowledge-" + "f" * 32 + "-1)", "partial": entries[-1][0].split(")")[0]}[selection]
result = durable_source_artifact(forwarded, selected, max_chars=1000)
if selection in {"unknown", "partial"}:
assert result is None
else:
assert result["knowledge_sources"]["omitted_count"] == (selection == "omitted")
assert len(result["knowledge_sources"]["sources"]) == (selection == "retained")
@pytest.mark.parametrize("reverse", [False, True])
def test_recaptured_source_supersedes_forwarded_omission(reverse):
from deerflow.community.ragflow.sources import durable_source_artifact
content, messages = retrieval()
source_id = messages[0]["artifact"]["knowledge_sources"]["sources"][0]["id"]
omitted = {"type": "tool", "name": "task", "artifact": {"knowledge_sources": {"version": 1, "sources": [], "omitted_source_ids": [source_id, source_id, None, "invalid"], "omitted_count": 999}}}
inputs = [omitted, *messages] if not reverse else [*messages, omitted]
result = durable_source_artifact(inputs, content, max_chars=1000)
assert result["knowledge_sources"]["sources"] == messages[0]["artifact"]["knowledge_sources"]["sources"]
assert result["knowledge_sources"]["omitted_count"] == 0
@pytest.mark.asyncio
async def test_report_truncation_selects_only_complete_stored_citations(env):
content, messages = retrieval()
env.service._config = SubagentBatchesConfig(max_result_chars=1000, result_preview_max_chars=64)
env.result.result = "x" * 1100 + content
env.result.ai_messages = messages
batch = await execute(env)
row = (await env.repo.list_items(batch["id"], user_id="user-1", include_result=True))[0]
assert row["result"] == "x" * 1000 and row["result_truncated"] is True
assert row["result_artifact"] is None
@pytest.mark.asyncio
async def test_plain_results_keep_nullable_evidence(env):
batch = await execute(env)
row = (await env.repo.list_items(batch["id"], user_id="user-1", include_result=True))[0]
assert row["result"] == "plain report" and row["result_artifact"] is None
@pytest.mark.asyncio
@pytest.mark.parametrize("action", ["cancel", "cancel-running", "stale-lease", "failure", "retry"])
async def test_fenced_or_failed_attempts_cannot_publish_evidence(env, action):
content, messages = retrieval()
from deerflow.community.ragflow.sources import durable_source_artifact
artifact = durable_source_artifact(messages, content, max_chars=1000)
batch = await env.repo.create_batch(
batch_id="batch",
user_id="user-1",
thread_id="thread-1",
run_id=None,
tool_call_id=None,
submission_key="key",
title="R",
subagent_type="general-purpose",
items=[{"key": "item", "prompt": "Research"}],
max_live_items=1,
max_running_items=1,
max_attempts=2 if action == "stale-lease" else 1,
execution_spec={},
)
now = datetime.now(UTC)
item = (await env.repo.claim_items(now=now, lease_owner="worker", lease_seconds=60, limit=1))[0]
if action == "cancel-running":
assert await env.repo.mark_item_running(item["id"], lease_owner="worker", now=now)
if action in {"cancel", "cancel-running"}:
await env.repo.cancel_batch(batch["id"], user_id="user-1")
if action == "stale-lease":
reclaimed = await env.repo.claim_items(now=now + timedelta(seconds=61), lease_owner="replacement", lease_seconds=60, limit=1)
assert len(reclaimed) == 1 and reclaimed[0]["attempt"] == 2
accepted = await env.repo.finalize_item(
item["id"],
lease_owner="worker",
succeeded=action not in {"failure", "retry"},
result=content,
result_preview=content[:64],
result_truncated=False,
result_artifact=artifact,
error="failed",
stop_reason=None,
token_usage=None,
model_name=None,
completed_at=now,
)
assert accepted is (action not in {"cancel", "cancel-running", "stale-lease"})
if action == "retry":
# A legacy/external failure row may have old evidence; explicit retry clears it.
from deerflow.persistence.subagent_batches.model import SubagentBatchItemRow
async with get_session_factory()() as session:
stored = await session.get(SubagentBatchItemRow, item["id"])
stored.result_artifact = artifact
await session.commit()
assert await env.repo.retry_item(batch["id"], item["id"], user_id="other") is None
await env.repo.retry_item(batch["id"], item["id"], user_id="user-1")
row = (await env.repo.list_items(batch["id"], user_id="user-1", include_result=True))[0]
assert row["result_artifact"] is None
@pytest.mark.asyncio
@pytest.mark.parametrize("user,permission,thread,status", [("user-1", True, "thread-1", 200), ("user-1", False, "thread-1", 403), ("other", True, "thread-1", 404), ("user-1", True, "other-thread", 404), (None, True, "thread-1", 401)])
async def test_http_export_checks_permission_thread_and_batch_owner(env, monkeypatch, user, permission, thread, status):
import httpx
from fastapi import FastAPI
from app.gateway import deps
from app.gateway.authz import AuthContext
from app.gateway.routers import subagent_batches as router
env.result.result, env.result.ai_messages = retrieval()
batch = await execute(env)
app = FastAPI()
app.include_router(router.router)
app.state.subagent_batch_repo = env.repo
@app.middleware("http")
async def auth_boundary(request, call_next):
request.state.auth = AuthContext(user=SimpleNamespace(id=user) if user else None, permissions=["threads:read"] if permission else [])
return await call_next(request)
# Authentication and the thread store are controlled; the actual permission decorator,
# route, SQL batch ownership, streaming and serialization remain production code.
monkeypatch.setattr(deps, "get_thread_store", lambda request: SimpleNamespace(check_access=AsyncMock(return_value=thread == "thread-1")))
monkeypatch.setattr(router, "get_current_user", AsyncMock(return_value=user))
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
response = await client.get(f"/api/threads/{thread}/subagent-batches/{batch['id']}/results.jsonl")
assert response.status_code == status
if status == 200:
row = json.loads(response.text)
assert row["result_artifact"]["knowledge_sources"]["sources"][0]["text"] == "Original evidence"
else:
assert "Original evidence" not in response.text