chore: update gitignore; feat: 新增customer_agent业务模块与api路由

2026-09-11 11:02:15 +08:00
parent 5915319f9f
commit 2004b8fcf4
59 changed files with 3527 additions and 18 deletions
@@ -0,0 +1,84 @@
+import unittest
+from unittest.mock import AsyncMock
+
+from rag.ingestion import ingest_document_atomic
+
+
+class AtomicIngestionTests(unittest.IsolatedAsyncioTestCase):
+    async def test_validation_completes_before_embedding(self):
+        milvus = AsyncMock()
+        embedder = AsyncMock()
+
+        with self.assertRaises(ValueError):
+            await ingest_document_atomic(
+                "Q: only a question", "doc-1", "FAQ", "fin_faq", "qa_pair",
+                milvus_client=milvus, embedder=embedder,
+            )
+
+        embedder.assert_not_awaited()
+        milvus.insert.assert_not_awaited()
+
+    async def test_embedding_failure_cleans_up_document_rows(self):
+        milvus = AsyncMock()
+        embedder = AsyncMock(side_effect=RuntimeError("embedding down"))
+
+        with self.assertRaises(RuntimeError):
+            await ingest_document_atomic(
+                "plain text", "doc-2", "Policy", "fin_policy", "default",
+                milvus_client=milvus, embedder=embedder,
+            )
+
+        milvus.delete.assert_awaited_once_with(
+            collection_name="fin_policy", filter='doc_id == "doc-2"'
+        )
+
+    async def test_milvus_failure_cleans_up_partial_document_rows(self):
+        milvus = AsyncMock()
+        milvus.insert.side_effect = RuntimeError("insert down")
+        embedder = AsyncMock(return_value=[[0.0] * 768])
+
+        with self.assertRaises(RuntimeError):
+            await ingest_document_atomic(
+                "plain text", "doc-3", "Policy", "fin_policy", "default",
+                milvus_client=milvus, embedder=embedder,
+            )
+
+        milvus.delete.assert_awaited_once_with(
+            collection_name="fin_policy", filter='doc_id == "doc-3"'
+        )
+
+    async def test_success_inserts_complete_document_rows(self):
+        milvus = AsyncMock()
+        embedder = AsyncMock(return_value=[[0.0] * 768])
+
+        result = await ingest_document_atomic(
+            "plain text", "doc-4", "Policy", "fin_policy", "default",
+            milvus_client=milvus, embedder=embedder,
+        )
+
+        self.assertEqual(result["doc_id"], "doc-4")
+        milvus.delete.assert_not_awaited()
+        milvus.insert.assert_awaited_once()
+
+    async def test_cleans_text_before_embedding_and_milvus_insert(self):
+        milvus = AsyncMock()
+        embedder = AsyncMock(return_value=[[0.0] * 768, [0.0] * 768])
+
+        await ingest_document_atomic(
+            "\ufeff基金名称：示例基金\r\n\r\n\r\n风险等级：R3\t ",
+            "doc-clean",
+            "示例基金",
+            "fin_fund_doc",
+            "default",
+            milvus_client=milvus,
+            embedder=embedder,
+        )
+
+        self.assertEqual(
+            embedder.await_args.args[0],
+            ["基金名称：示例基金", "风险等级：R3"],
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()