"""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 '' 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"]