96 lines
3.7 KiB
Python
96 lines
3.7 KiB
Python
import unittest
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock
|
|
|
|
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")],
|
|
)
|
|
|
|
|
|
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)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|