209 lines
8.8 KiB
Python
209 lines
8.8 KiB
Python
import unittest
|
||
from types import SimpleNamespace
|
||
from unittest.mock import AsyncMock, patch
|
||
|
||
from app.api.v1.chat import _canonicalize_message_citations, _compact_cited_refs
|
||
from app.services.rag_service import RAGService
|
||
|
||
|
||
class ChatCitationTest(unittest.TestCase):
|
||
def test_legacy_interrupt_marker_is_read_only_compat(self):
|
||
from app.api.v1.chat import (
|
||
_is_interrupted_message,
|
||
_legacy_interrupted_content,
|
||
_strip_legacy_interrupt_marker,
|
||
)
|
||
|
||
self.assertTrue(_legacy_interrupted_content("interrupt"))
|
||
self.assertTrue(_legacy_interrupted_content("内容\n\ninterrupt"))
|
||
self.assertFalse(_legacy_interrupted_content("正常内容"))
|
||
|
||
self.assertEqual(_strip_legacy_interrupt_marker("interrupt"), "")
|
||
self.assertEqual(_strip_legacy_interrupt_marker("内容\n\ninterrupt"), "内容")
|
||
self.assertEqual(_strip_legacy_interrupt_marker("正常内容"), "正常内容")
|
||
|
||
# 旧数据（status 缺失/pending + 旧标记）按已中断处理，但不修改数据库
|
||
legacy = SimpleNamespace(role="assistant", status="pending", content="内容\n\ninterrupt")
|
||
self.assertTrue(_is_interrupted_message(legacy))
|
||
# 新数据一律以 status 为准，不再比对内容
|
||
fresh = SimpleNamespace(role="assistant", status="completed", content="内容\n\ninterrupt")
|
||
self.assertFalse(_is_interrupted_message(fresh))
|
||
interrupted = SimpleNamespace(role="assistant", status="interrupted", content="")
|
||
self.assertTrue(_is_interrupted_message(interrupted))
|
||
|
||
def test_duplicate_file_citations_are_renumbered_once(self):
|
||
content, refs = _canonicalize_message_citations(
|
||
"泰坦属于土星的卫星[2][4][5]。",
|
||
[
|
||
{"citation_id": 2, "file_path": "内容导航.md", "excerpt": "片段一"},
|
||
{"citation_id": 4, "file_path": "内容导航.md", "excerpt": "片段二"},
|
||
{"citation_id": 5, "file_path": "内容导航.md", "excerpt": "片段三"},
|
||
],
|
||
)
|
||
|
||
self.assertEqual(content, "泰坦属于土星的卫星[1]。")
|
||
self.assertEqual(len(refs), 1)
|
||
self.assertEqual(refs[0]["citation_id"], 1)
|
||
self.assertEqual(refs[0]["excerpt"], "片段一\n\n片段二\n\n片段三")
|
||
|
||
def test_overlapping_blocks_are_not_duplicated(self):
|
||
self.assertEqual(
|
||
RAGService._merge_text_blocks("泰坦属于土星", "泰坦属于土星"),
|
||
"泰坦属于土星",
|
||
)
|
||
|
||
def test_single_used_reference_is_compacted_to_one(self):
|
||
content, refs = _compact_cited_refs(
|
||
"泰坦是土星的卫星[3]，土卫二也属于土星[3]。",
|
||
[
|
||
{"citation_id": 1, "file_path": "其他一.md"},
|
||
{"citation_id": 2, "file_path": "其他二.md"},
|
||
{"citation_id": 3, "file_path": "内容导航.md"},
|
||
],
|
||
)
|
||
|
||
self.assertEqual(content, "泰坦是土星的卫星[1]，土卫二也属于土星[1]。")
|
||
self.assertEqual(refs, [{
|
||
"citation_id": 1,
|
||
"file_path": "内容导航.md",
|
||
"anchor_text": "",
|
||
"excerpt": "",
|
||
"content": "",
|
||
"chunk_index": None,
|
||
"hit_terms": [],
|
||
}])
|
||
|
||
def test_used_references_follow_first_appearance_order(self):
|
||
content, refs = _compact_cited_refs(
|
||
"先使用第三份[3]，再使用第一份[1]。",
|
||
[
|
||
{"citation_id": 1, "file_path": "第一份.md"},
|
||
{"citation_id": 2, "file_path": "第二份.md"},
|
||
{"citation_id": 3, "file_path": "第三份.md"},
|
||
],
|
||
)
|
||
|
||
self.assertEqual(content, "先使用第三份[1]，再使用第一份[2]。")
|
||
self.assertEqual(
|
||
[(ref["citation_id"], ref["file_path"]) for ref in refs],
|
||
[(1, "第三份.md"), (2, "第一份.md")],
|
||
)
|
||
|
||
|
||
class RAGConversationContextTest(unittest.IsolatedAsyncioTestCase):
|
||
async def test_previous_assistant_answer_is_not_sent_to_model(self):
|
||
db = AsyncMock()
|
||
db.execute.return_value = SimpleNamespace(
|
||
scalar_one_or_none=lambda: SimpleNamespace(config_id=1)
|
||
)
|
||
history = [
|
||
{"role": "user", "content": "包含了哪几次阿波罗计划？"},
|
||
{"role": "assistant", "content": "上一轮完整答案不应再次发送给模型。"},
|
||
]
|
||
|
||
_, system_prompt, messages = await RAGService._prepare_generation(
|
||
db,
|
||
"项目中的文档包含哪些土星的卫星？",
|
||
1,
|
||
[],
|
||
history,
|
||
)
|
||
|
||
self.assertEqual(messages, [{
|
||
"role": "user",
|
||
"content": "项目中的文档包含哪些土星的卫星？",
|
||
}])
|
||
self.assertIn("包含了哪几次阿波罗计划？", system_prompt)
|
||
self.assertNotIn("上一轮完整答案不应再次发送给模型。", system_prompt)
|
||
self.assertIn("禁止复述、总结或继续回答先前问题", system_prompt)
|
||
|
||
|
||
class CitationQuoteExtractionTest(unittest.TestCase):
|
||
def test_split_sentences(self):
|
||
self.assertEqual(
|
||
RAGService._split_sentences("第一句。第二句！第三句？\n第四句；第五句话。"),
|
||
["第一句。", "第二句！", "第三句？", "第四句；", "第五句话。"],
|
||
)
|
||
|
||
def test_extract_citation_claims(self):
|
||
answer = (
|
||
"土星拥有众多卫星[1]。泰坦是其中最大的一颗[1]。\n"
|
||
"关键参数：直径约 120536 公里[1]。\n"
|
||
"该产品型号为 X200[1]，电池容量 5000mAh[1]。"
|
||
)
|
||
claims = RAGService._extract_citation_claims(answer)
|
||
self.assertEqual(claims[1], [
|
||
"土星拥有众多卫星",
|
||
"泰坦是其中最大的一颗",
|
||
"关键参数：直径约 120536 公里",
|
||
"该产品型号为 X200",
|
||
"电池容量 5000mAh",
|
||
])
|
||
|
||
def test_sentence_candidates_prefer_term_overlap(self):
|
||
candidates = RAGService._sentence_candidates(
|
||
"土星是太阳系第二大行星。\n土卫六又称泰坦，是土星最大的卫星。\n"
|
||
"木星拥有最多卫星。",
|
||
["泰坦是最大的卫星"],
|
||
max_candidates=3,
|
||
)
|
||
self.assertIn("土卫六又称泰坦，是土星最大的卫星。", candidates[:1])
|
||
|
||
def test_cosine_similarity(self):
|
||
self.assertAlmostEqual(RAGService._cosine_similarity([1.0, 0.0], [1.0, 0.0]), 1.0)
|
||
self.assertAlmostEqual(RAGService._cosine_similarity([1.0, 0.0], [0.0, 1.0]), 0.0)
|
||
self.assertEqual(RAGService._cosine_similarity([], [1.0]), 0.0)
|
||
|
||
|
||
class CitationQuoteAlignmentTest(unittest.IsolatedAsyncioTestCase):
|
||
async def test_align_citation_quotes_picks_similar_sentence(self):
|
||
refs = [{
|
||
"citation_id": 1,
|
||
"file_path": "产品资料.md",
|
||
"anchor_text": "型号 X200",
|
||
"excerpt": "型号 X200 支持 5G。\n电池容量为 5000mAh。\n屏幕尺寸 6.7 英寸。",
|
||
"content": "",
|
||
}]
|
||
answer = "该产品型号为 X200[1]，电池续航出色[1]。"
|
||
|
||
async def fake_embeddings(db, texts, config=None):
|
||
return [
|
||
[1.0, 0.0, 0.0] if ("电池" in text or "5000mAh" in text) else [0.0, 1.0, 0.0]
|
||
for text in texts
|
||
]
|
||
|
||
with patch("app.services.rag_service.zvec_service") as mock_zvec:
|
||
mock_zvec.generate_embeddings = fake_embeddings
|
||
result = await RAGService.align_citation_quotes(None, answer, refs)
|
||
|
||
self.assertIn("quotes", result[0])
|
||
self.assertTrue(any(
|
||
"5000mAh" in q["text"] or "电池" in q["text"]
|
||
for q in result[0]["quotes"]
|
||
))
|
||
# 按出现顺序回填：第一处对应型号，第二处对应电池
|
||
occurrences = result[0]["quote_occurrences"]
|
||
self.assertEqual(len(occurrences), 2)
|
||
self.assertEqual(occurrences[0]["claim"], "该产品型号为 X200")
|
||
self.assertEqual(occurrences[1]["claim"], "电池续航出色")
|
||
self.assertTrue(any("电池" in q["text"] for q in occurrences[1]["quotes"]))
|
||
|
||
async def test_align_citation_quotes_falls_back_when_embedding_unavailable(self):
|
||
refs = [{
|
||
"citation_id": 1,
|
||
"file_path": "a.md",
|
||
"excerpt": "某句支撑内容。",
|
||
"content": "",
|
||
}]
|
||
|
||
with patch("app.services.rag_service.zvec_service") as mock_zvec:
|
||
mock_zvec.generate_embeddings = AsyncMock(return_value=None)
|
||
result = await RAGService.align_citation_quotes(None, "这是论断[1]。", refs)
|
||
|
||
self.assertEqual(result, refs)
|
||
self.assertNotIn("quotes", result[0])
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|