import unittest
from rag import retrieve


class RetrievalTests(unittest.TestCase):
    def test_returns_the_most_relevant_document(self):
        documents = ["LoRA 使用低秩矩阵微调", "RAG 先检索文档再生成答案"]
        result = retrieve("RAG 如何检索文档", documents)
        self.assertEqual(result["index"], 1)

    def test_refuses_when_no_document_reaches_threshold(self):
        result = retrieve("量子物理", ["苹果和香蕉"], threshold=0.2)
        self.assertIsNone(result)


if __name__ == "__main__":
    unittest.main()