mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-19 11:06:18 +00:00
34bbeb1806
* feat(knowledge): add verifiable RAGFlow source citations * docs(knowledge): scope RAGFlow guidance to its own directory * fix(knowledge): preserve citations through rendering and budgets
128 lines
6.4 KiB
Python
128 lines
6.4 KiB
Python
"""Citation records and model-visible evidence must stay paired after budgeting."""
|
|
|
|
from copy import deepcopy
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
from langchain_core.messages import ToolMessage
|
|
from langgraph.types import Command
|
|
|
|
from deerflow.agents.middlewares.tool_output_budget_middleware import ToolOutputBudgetMiddleware, _patch_model_messages, _patch_result
|
|
from deerflow.community.ragflow.formatting import format_retrieval_sources
|
|
from deerflow.community.ragflow.sources import budget_source_artifact
|
|
from deerflow.config.tool_output_config import ToolOutputConfig
|
|
|
|
|
|
def message(name="knowledge_search"):
|
|
chunks = [{"id": f"chunk-{i}", "dataset_id": "kb", "document_id": "doc", "document_keyword": "Manual.pdf", "content": f"Evidence {i}: " + "x" * 4800} for i in range(8)]
|
|
content, artifact = format_retrieval_sources({"chunks": chunks}, dataset_names_by_id={"kb": "Engineering"}, max_chars_per_chunk=5000, max_total_chars=40000)
|
|
if name == "task":
|
|
links = " ".join(f"[Manual.pdf](#knowledge-{source['id']})" for source in artifact["knowledge_sources"]["sources"])
|
|
content = "Task completed. Findings: " + links + "\n" + "report " * 6000
|
|
artifact["other"] = "preserve me"
|
|
return ToolMessage(content=content, artifact=artifact, name=name, tool_call_id="call-1")
|
|
|
|
|
|
def assert_paired(original, result, limit):
|
|
assert len(result.content) <= limit
|
|
assert result.artifact["other"] == "preserve me"
|
|
sources = result.artifact.get("knowledge_sources", {}).get("sources", [])
|
|
assert sources
|
|
assert len(sources) < len(original.artifact["knowledge_sources"]["sources"])
|
|
for source in original.artifact["knowledge_sources"]["sources"]:
|
|
if source in sources:
|
|
assert f"](#knowledge-{source['id']})" in result.content
|
|
assert source["text"] in result.content
|
|
else:
|
|
assert f"#knowledge-{source['id']}" not in result.content
|
|
assert "omitted" in result.content
|
|
|
|
|
|
@pytest.mark.parametrize("name", ["knowledge_search", "task"])
|
|
@pytest.mark.parametrize("mode", ["externalize", "fallback", "override", "history", "no_storage"])
|
|
def test_complete_citations_survive_budget_paths(tmp_path, name, mode):
|
|
original = message(name)
|
|
snapshot = deepcopy(original)
|
|
limit = 12000 if mode == "externalize" else 7000
|
|
config = ToolOutputConfig(**({"fallback_max_chars": limit, "externalize_min_chars": 0} if mode in {"fallback", "history"} else {"tool_overrides": {name: limit}}))
|
|
if mode == "history":
|
|
result = _patch_model_messages([original], config)[0]
|
|
else:
|
|
command = Command(update={"messages": [original], "unrelated": 42})
|
|
patched = _patch_result(command, config, str(tmp_path) if mode not in {"fallback", "no_storage"} else None)
|
|
assert patched.update["unrelated"] == 42
|
|
result = patched.update["messages"][0]
|
|
assert_paired(original, result, limit)
|
|
assert original == snapshot
|
|
transforms = result.additional_kwargs["deerflow_tool_transforms"]
|
|
assert transforms[-1]["kind"] == "truncated"
|
|
if mode in {"externalize", "override"}:
|
|
assert transforms[-2]["kind"] == "externalized"
|
|
# A subsequent model hook must not rewrite the same evidence again.
|
|
assert _patch_model_messages([result], config) is None
|
|
|
|
|
|
@pytest.mark.parametrize("name", ["knowledge_search", "task"])
|
|
@pytest.mark.parametrize("limit", [1, 60, 159, 160, 500])
|
|
def test_tiny_budget_does_not_leave_orphaned_or_partial_citations(name, limit):
|
|
original = message(name)
|
|
config = ToolOutputConfig(externalize_min_chars=0, fallback_max_chars=limit)
|
|
result = _patch_result(original, config, None)
|
|
assert len(result.content) <= limit
|
|
assert "#knowledge-" not in result.content
|
|
assert "knowledge_sources" not in result.artifact
|
|
assert result.artifact["other"] == "preserve me"
|
|
|
|
|
|
@pytest.mark.parametrize("config", [ToolOutputConfig(enabled=False), ToolOutputConfig(exempt_tools=["knowledge_search"]), ToolOutputConfig(externalize_min_chars=0, fallback_max_chars=0)])
|
|
def test_disabled_and_exempt_budget_preserves_sources(config):
|
|
original = message()
|
|
request = SimpleNamespace(tool_call={"name": "knowledge_search", "id": "call-1"}, runtime=SimpleNamespace(state={}))
|
|
result = ToolOutputBudgetMiddleware(config).wrap_tool_call(request, lambda _: original)
|
|
assert result is original
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_tool_hook_preserves_sources():
|
|
original = message()
|
|
request = SimpleNamespace(tool_call={"name": "knowledge_search", "id": "call-1"}, runtime=SimpleNamespace(state={}))
|
|
|
|
async def handler(_):
|
|
return original
|
|
|
|
result = await ToolOutputBudgetMiddleware(ToolOutputConfig(externalize_min_chars=0, fallback_max_chars=7000)).awrap_tool_call(request, handler)
|
|
assert_paired(original, result, 7000)
|
|
|
|
|
|
@pytest.mark.parametrize("change", ["other_tool", "unknown_version", "malformed_records", "no_artifact"])
|
|
def test_unrelated_results_keep_the_generic_budget_behavior(change):
|
|
original = message()
|
|
if change == "other_tool":
|
|
original.name = "web_search"
|
|
elif change == "unknown_version":
|
|
original.artifact["knowledge_sources"]["version"] = 2
|
|
elif change == "malformed_records":
|
|
original.artifact["knowledge_sources"]["sources"] = [None, {"id": "invalid", "text": 42}]
|
|
else:
|
|
original.artifact = None
|
|
config = ToolOutputConfig(externalize_min_chars=0, fallback_max_chars=7000)
|
|
control = _patch_result(original.model_copy(update={"artifact": None}), config, None)
|
|
result = _patch_result(original, config, None)
|
|
assert result.content == control.content
|
|
assert result.artifact == original.artifact
|
|
|
|
|
|
def test_small_source_result_is_not_rewritten():
|
|
original = message()
|
|
config = ToolOutputConfig(externalize_min_chars=50000, fallback_max_chars=50000)
|
|
assert _patch_result(original, config, None) is original
|
|
|
|
|
|
def test_delegated_budget_retains_the_report_reference_and_complete_evidence():
|
|
original = message("task")
|
|
summary = "Task completed. " + "synopsis " * 500 + "\nRead the full report: /mnt/user-data/outputs/.tool-results/report.txt"
|
|
content, artifact = budget_source_artifact(original.content, original.artifact, 7000, summary=summary)
|
|
assert content.startswith("Task completed.")
|
|
assert "/mnt/user-data/outputs/.tool-results/report.txt" in content
|
|
assert_paired(original, original.model_copy(update={"content": content, "artifact": artifact}), 7000)
|