nex_docus/backend/tests/test_chat_citations.py

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()