66 lines
2.5 KiB
Python
66 lines
2.5 KiB
Python
import unittest
|
|
|
|
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": "",
|
|
}])
|
|
|
|
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")],
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|