336 lines
12 KiB
Python
336 lines
12 KiB
Python
"""Tests for RAG citation batch construction."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
|
|
from deerflow.runtime.references import (
|
|
build_latest_reference_batch,
|
|
build_latest_reference_prompt,
|
|
build_reference_batches,
|
|
)
|
|
|
|
|
|
def _tool_message(call_id: str, query: str, start: int) -> dict:
|
|
return {
|
|
"type": "tool",
|
|
"name": "web_search",
|
|
"tool_call_id": call_id,
|
|
"content": json.dumps(
|
|
{
|
|
"query": query,
|
|
"results": [
|
|
{
|
|
"title": f"{query} result {start}",
|
|
"content": f"{query} content {start}",
|
|
"url": f"https://example.test/{start}",
|
|
},
|
|
{
|
|
"title": f"{query} result {start + 1}",
|
|
"content": f"{query} content {start + 1}",
|
|
"url": f"https://example.test/{start + 1}",
|
|
},
|
|
],
|
|
},
|
|
ensure_ascii=False,
|
|
),
|
|
}
|
|
|
|
|
|
def test_latest_reference_prompt_uses_one_combined_batch_for_multi_search_turn() -> None:
|
|
messages = [
|
|
{"type": "human", "id": "human-1", "content": "search two terms"},
|
|
{
|
|
"type": "ai",
|
|
"id": "ai-search-1",
|
|
"content": "",
|
|
"tool_calls": [{"id": "call-1", "name": "web_search", "args": {"query": "Trump"}}],
|
|
},
|
|
_tool_message("call-1", "Trump", 1),
|
|
{
|
|
"type": "ai",
|
|
"id": "ai-search-2",
|
|
"content": "",
|
|
"tool_calls": [{"id": "call-2", "name": "web_search", "args": {"query": "White House"}}],
|
|
},
|
|
_tool_message("call-2", "White House", 3),
|
|
]
|
|
|
|
batch = build_latest_reference_batch(messages)
|
|
assert batch is not None
|
|
assert batch["source_count"] == 4
|
|
assert [source["index"] for source in batch["sources"]] == [1, 2, 3, 4]
|
|
assert [source["title"] for source in batch["sources"]] == [
|
|
"Trump result 1",
|
|
"Trump result 2",
|
|
"White House result 3",
|
|
"White House result 4",
|
|
]
|
|
|
|
prompt, prompt_batch = build_latest_reference_prompt(messages)
|
|
assert prompt_batch == batch
|
|
assert '<reference_batch source_count="4">' in prompt
|
|
assert "[1] Trump result 1" in prompt
|
|
assert "[4] White House result 4" in prompt
|
|
assert "重新分批" in prompt
|
|
|
|
|
|
def test_reference_batches_keep_separate_user_turns_but_merge_searches_inside_turn() -> None:
|
|
messages = [
|
|
{"type": "human", "id": "human-1", "content": "first"},
|
|
{"type": "ai", "id": "ai-1", "content": "", "tool_calls": [{"id": "c1", "name": "web_search"}]},
|
|
_tool_message("c1", "first", 1),
|
|
{"type": "ai", "id": "answer-1", "content": "answer [1]"},
|
|
{"type": "human", "id": "human-2", "content": "second"},
|
|
{"type": "ai", "id": "ai-2", "content": "", "tool_calls": [{"id": "c2", "name": "web_search"}]},
|
|
_tool_message("c2", "second-a", 3),
|
|
{"type": "ai", "id": "ai-3", "content": "", "tool_calls": [{"id": "c3", "name": "web_search"}]},
|
|
_tool_message("c3", "second-b", 5),
|
|
{"type": "ai", "id": "answer-2", "content": "answer [1][4]"},
|
|
]
|
|
|
|
batches = build_reference_batches(messages)
|
|
assert [len(batch["sources"]) for batch in batches] == [2, 4]
|
|
assert [source["index"] for source in batches[1]["sources"]] == [1, 2, 3, 4]
|
|
assert batches[0]["assistant_message_id"] == "answer-1"
|
|
assert batches[1]["assistant_message_id"] == "answer-2"
|
|
|
|
|
|
def test_reference_drawer_deduplicates_across_calls_and_prefers_wiki_read_detail() -> None:
|
|
repeated_content = "唐纳德·特朗普于1946年出生，曾任美国总统并从事商业活动。"
|
|
|
|
def retrieval(
|
|
call_id: str,
|
|
tool_name: str,
|
|
results: list[dict],
|
|
) -> list[dict]:
|
|
return [
|
|
{
|
|
"type": "ai",
|
|
"id": f"ai-{call_id}",
|
|
"content": "",
|
|
"tool_calls": [
|
|
{"id": call_id, "name": tool_name, "args": {}}
|
|
],
|
|
},
|
|
{
|
|
"type": "tool",
|
|
"name": tool_name,
|
|
"tool_call_id": call_id,
|
|
"content": json.dumps(
|
|
{"ok": True, "results": results},
|
|
ensure_ascii=False,
|
|
),
|
|
},
|
|
]
|
|
|
|
messages = [
|
|
{"type": "human", "id": "human-1", "content": "介绍特朗普生平"},
|
|
*retrieval(
|
|
"raw-1",
|
|
"ordinary_knowledge_search",
|
|
[
|
|
{
|
|
"knowledge_base_id": "kb-news",
|
|
"knowledge_id": "document-trump",
|
|
"chunk_id": "provider-chunk-1",
|
|
"title": "特朗普生平.txt",
|
|
"content": repeated_content,
|
|
"score": 0.71,
|
|
},
|
|
{
|
|
"knowledge_base_id": "kb-news",
|
|
"knowledge_id": "document-trump",
|
|
"chunk_id": "provider-chunk-2",
|
|
"title": "特朗普生平.txt",
|
|
"content": "特朗普在纽约成长，这是一段与上一分片不同的原文信息。",
|
|
"score": 0.84,
|
|
},
|
|
],
|
|
),
|
|
*retrieval(
|
|
"raw-2",
|
|
"ordinary_knowledge_search",
|
|
[
|
|
{
|
|
"knowledge_base_id": "kb-news",
|
|
"knowledge_id": "document-trump",
|
|
"chunk_id": "provider-remapped-3",
|
|
"title": "特朗普生平.txt",
|
|
"content": f" {repeated_content[:18]}\n{repeated_content[18:]} ",
|
|
"score": 0.93,
|
|
}
|
|
],
|
|
),
|
|
*retrieval(
|
|
"wiki-search",
|
|
"ordinary_wiki_search",
|
|
[
|
|
{
|
|
"knowledge_base_id": "kb-news",
|
|
"page_ref": "wiki-page:v1:donald-trump",
|
|
"slug": "entity/donald-trump",
|
|
"title": "唐纳德·特朗普",
|
|
"summary": "人物索引摘要。",
|
|
}
|
|
],
|
|
),
|
|
*retrieval(
|
|
"wiki-read",
|
|
"ordinary_wiki_read_page",
|
|
[
|
|
{
|
|
"knowledge_base_id": "kb-news",
|
|
"page_ref": "wiki-page:v1:donald-trump",
|
|
"slug": "entity/donald-trump",
|
|
"title": "唐纳德·特朗普",
|
|
"summary": "人物索引摘要。",
|
|
"content_markdown": "# 唐纳德·特朗普\n\n这是 Wiki 深读后的完整人物页面正文。",
|
|
}
|
|
],
|
|
),
|
|
{"type": "ai", "id": "answer-1", "content": "完整回答"},
|
|
]
|
|
|
|
batch = build_latest_reference_batch(messages)
|
|
|
|
assert batch is not None
|
|
# This is the exact source count rendered by ChatBox's right-side drawer.
|
|
assert batch["source_count"] == 3
|
|
assert [source["index"] for source in batch["sources"]] == [1, 2, 3]
|
|
assert [source["title"] for source in batch["sources"]] == [
|
|
"特朗普生平.txt",
|
|
"特朗普生平.txt",
|
|
"唐纳德·特朗普",
|
|
]
|
|
assert batch["sources"][0]["raw"]["chunk_id"] == "provider-remapped-3"
|
|
assert batch["sources"][0]["score"] == 0.93
|
|
assert batch["sources"][1]["raw"]["chunk_id"] == "provider-chunk-2"
|
|
assert batch["sources"][2]["skillName"] == "ordinary_wiki_read_page"
|
|
assert "Wiki 深读后的完整人物页面正文" in batch["sources"][2]["content"]
|
|
|
|
|
|
def test_reference_drawer_namespaces_generic_ids_by_skill_and_result_set() -> None:
|
|
messages = [
|
|
{
|
|
"type": "human",
|
|
"id": "human-1",
|
|
"content": "综合两个技能的结果",
|
|
"additional_kwargs": {
|
|
"llmwiki_reference_sources": [
|
|
{
|
|
"id": "1",
|
|
"skillName": "skill_alpha",
|
|
"resultSetKey": "primary",
|
|
"title": "Alpha 主结果",
|
|
"content": "Alpha 技能主结果中的独立内容。",
|
|
"raw": {"id": "1"},
|
|
},
|
|
{
|
|
"id": "1",
|
|
"skillName": "skill_alpha",
|
|
"resultSetKey": "secondary",
|
|
"title": "Alpha 次结果",
|
|
"content": "Alpha 技能次结果中的另一段独立内容。",
|
|
"raw": {"id": "1"},
|
|
},
|
|
{
|
|
"id": "1",
|
|
"skillName": "skill_beta",
|
|
"resultSetKey": "primary",
|
|
"title": "Beta 主结果",
|
|
"content": "Beta 技能主结果中的独立内容。",
|
|
"raw": {"id": "1"},
|
|
},
|
|
]
|
|
},
|
|
},
|
|
{"type": "ai", "id": "answer-1", "content": "综合回答"},
|
|
]
|
|
|
|
batch = build_latest_reference_batch(messages)
|
|
|
|
assert batch is not None
|
|
assert batch["source_count"] == 3
|
|
assert [source["title"] for source in batch["sources"]] == [
|
|
"Alpha 主结果",
|
|
"Alpha 次结果",
|
|
"Beta 主结果",
|
|
]
|
|
|
|
|
|
def test_reference_drawer_replaces_wiki_lead_but_keeps_continuation_segments() -> None:
|
|
def retrieval(
|
|
call_id: str,
|
|
tool_name: str,
|
|
results: list[dict],
|
|
) -> list[dict]:
|
|
return [
|
|
{
|
|
"type": "ai",
|
|
"id": f"ai-{call_id}",
|
|
"content": "",
|
|
"tool_calls": [
|
|
{"id": call_id, "name": tool_name, "args": {}}
|
|
],
|
|
},
|
|
{
|
|
"type": "tool",
|
|
"name": tool_name,
|
|
"tool_call_id": call_id,
|
|
"content": json.dumps(
|
|
{"ok": True, "results": results},
|
|
ensure_ascii=False,
|
|
),
|
|
},
|
|
]
|
|
|
|
common = {
|
|
"knowledge_base_id": "kb-news",
|
|
"page_ref": "wiki-page:v1:donald-trump",
|
|
"slug": "entity/donald-trump",
|
|
"title": "唐纳德·特朗普",
|
|
}
|
|
messages = [
|
|
{"type": "human", "id": "human-1", "content": "完整读取人物页面"},
|
|
*retrieval(
|
|
"wiki-search",
|
|
"ordinary_wiki_search",
|
|
[{**common, "summary": "人物索引摘要。"}],
|
|
),
|
|
*retrieval(
|
|
"wiki-read-first",
|
|
"ordinary_wiki_read_page",
|
|
[
|
|
{
|
|
**common,
|
|
"content_offset": 0,
|
|
"content_markdown": "# 唐纳德·特朗普\n\n页面第一段完整正文。",
|
|
}
|
|
],
|
|
),
|
|
*retrieval(
|
|
"wiki-read-next",
|
|
"ordinary_wiki_read_page",
|
|
[
|
|
{
|
|
**common,
|
|
"content_offset": 1000,
|
|
"content_markdown": "页面续读内容，包含后续生平与经历。",
|
|
}
|
|
],
|
|
),
|
|
{"type": "ai", "id": "answer-1", "content": "完整回答"},
|
|
]
|
|
|
|
batch = build_latest_reference_batch(messages)
|
|
|
|
assert batch is not None
|
|
assert batch["source_count"] == 2
|
|
assert [source["raw"].get("content_offset", 0) for source in batch["sources"]] == [
|
|
0,
|
|
1000,
|
|
]
|
|
assert "页面第一段完整正文" in batch["sources"][0]["content"]
|
|
assert "页面续读内容" in batch["sources"][1]["content"]
|