185 lines
7.5 KiB
Python
185 lines
7.5 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_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()
|