chore: update gitignore; feat: 新增customer_agent业务模块与api路由
This commit is contained in:
@@ -0,0 +1,290 @@
|
||||
import tempfile
|
||||
import unittest
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from rag.events import KNOWLEDGE_UPDATE_EVENT
|
||||
from service.knowledge_base.upload import KnowledgeUploadService, UploadValidationError
|
||||
|
||||
|
||||
class KnowledgeUploadTests(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_preview_saves_upload_and_returns_cleaning_and_chunk_preview(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
service = KnowledgeUploadService(
|
||||
storage_dir=tmp,
|
||||
milvus_client=AsyncMock(),
|
||||
embedder=AsyncMock(),
|
||||
publisher=AsyncMock(),
|
||||
)
|
||||
|
||||
result = await service.preview(
|
||||
filename="notice.md",
|
||||
content="\ufeff# Notice\r\n\r\n正文".encode(),
|
||||
strategy="chapter_semantic",
|
||||
)
|
||||
|
||||
self.assertTrue(result["upload_id"])
|
||||
self.assertTrue(result["chunks"])
|
||||
self.assertTrue(Path(tmp, result["stored_filename"]).is_file())
|
||||
self.assertIn("cleaning_warnings", result)
|
||||
|
||||
async def test_confirm_ingests_then_publishes_update_and_removes_temp_file(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
publisher = AsyncMock()
|
||||
milvus = AsyncMock()
|
||||
embedder = AsyncMock(return_value=[[0.0] * 768])
|
||||
service = KnowledgeUploadService(
|
||||
storage_dir=tmp,
|
||||
milvus_client=milvus,
|
||||
embedder=embedder,
|
||||
publisher=publisher,
|
||||
)
|
||||
preview = await service.preview(
|
||||
filename="faq.md",
|
||||
content="Q: 什么是基金？\nA: 一种集合投资工具。".encode(),
|
||||
strategy="qa_pair",
|
||||
)
|
||||
|
||||
result = await service.confirm(
|
||||
upload_id=preview["upload_id"],
|
||||
title="FAQ",
|
||||
doc_id="doc-upload-1",
|
||||
collection_name="fin_faq",
|
||||
strategy="qa_pair",
|
||||
)
|
||||
|
||||
self.assertEqual(result["doc_id"], "doc-upload-1")
|
||||
publisher.assert_awaited_once()
|
||||
self.assertEqual(publisher.await_args.args[0], KNOWLEDGE_UPDATE_EVENT)
|
||||
self.assertFalse(Path(tmp, preview["stored_filename"]).exists())
|
||||
|
||||
async def test_confirm_passes_custom_chunk_config_to_ingestion(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
milvus = AsyncMock()
|
||||
embedder = AsyncMock(side_effect=lambda texts: [[0.0] * 768 for _ in texts])
|
||||
service = KnowledgeUploadService(
|
||||
storage_dir=tmp,
|
||||
milvus_client=milvus,
|
||||
embedder=embedder,
|
||||
publisher=AsyncMock(),
|
||||
)
|
||||
preview = await service.preview(
|
||||
"doc.md", "一二三四五六七八九十十一十二".encode(), strategy="default"
|
||||
)
|
||||
|
||||
await service.confirm(
|
||||
upload_id=preview["upload_id"],
|
||||
title="Doc",
|
||||
doc_id="doc-config",
|
||||
collection_name="fin_fund_doc",
|
||||
strategy="default",
|
||||
chunk_size=10,
|
||||
chunk_overlap=2,
|
||||
)
|
||||
|
||||
self.assertEqual(len(embedder.await_args.args[0]), 2)
|
||||
|
||||
async def test_confirm_removes_vectors_when_update_event_publish_fails(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
milvus = AsyncMock()
|
||||
publisher = AsyncMock(side_effect=RuntimeError("redis unavailable"))
|
||||
service = KnowledgeUploadService(
|
||||
storage_dir=tmp,
|
||||
milvus_client=milvus,
|
||||
embedder=AsyncMock(return_value=[[0.0] * 768]),
|
||||
publisher=publisher,
|
||||
)
|
||||
preview = await service.preview(
|
||||
"policy.md", b"policy text", strategy="default"
|
||||
)
|
||||
|
||||
with self.assertRaises(RuntimeError):
|
||||
await service.confirm(
|
||||
upload_id=preview["upload_id"],
|
||||
title="Policy",
|
||||
doc_id="doc-event-failure",
|
||||
collection_name="fin_policy",
|
||||
strategy="default",
|
||||
)
|
||||
|
||||
self.assertEqual(milvus.delete.await_count, 3)
|
||||
|
||||
async def test_rejects_unsupported_extension_and_oversized_file(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
service = KnowledgeUploadService(
|
||||
storage_dir=tmp,
|
||||
max_upload_bytes=4,
|
||||
milvus_client=AsyncMock(),
|
||||
embedder=AsyncMock(),
|
||||
publisher=AsyncMock(),
|
||||
)
|
||||
|
||||
with self.assertRaises(UploadValidationError):
|
||||
await service.preview("file.exe", b"ok", strategy="default")
|
||||
with self.assertRaises(UploadValidationError):
|
||||
await service.preview("file.md", b"12345", strategy="default")
|
||||
|
||||
async def test_confirm_rejects_unknown_collection_or_strategy_mismatch(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
service = KnowledgeUploadService(
|
||||
storage_dir=tmp,
|
||||
milvus_client=AsyncMock(),
|
||||
embedder=AsyncMock(),
|
||||
publisher=AsyncMock(),
|
||||
)
|
||||
preview = await service.preview(
|
||||
"faq.md", b"Q: Q\nA: A", strategy="qa_pair"
|
||||
)
|
||||
|
||||
with self.assertRaises(UploadValidationError):
|
||||
await service.confirm(
|
||||
upload_id=preview["upload_id"], title="FAQ", doc_id="d1",
|
||||
collection_name="evil_collection", strategy="qa_pair",
|
||||
)
|
||||
|
||||
async def test_confirm_rejects_duplicate_doc_id_before_embedding(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
embedder = AsyncMock()
|
||||
service = KnowledgeUploadService(
|
||||
storage_dir=tmp,
|
||||
milvus_client=AsyncMock(),
|
||||
embedder=embedder,
|
||||
publisher=AsyncMock(),
|
||||
document_exists=lambda doc_id: True,
|
||||
)
|
||||
preview = await service.preview(
|
||||
"faq.md", b"Q: Q\nA: A", strategy="qa_pair"
|
||||
)
|
||||
|
||||
with self.assertRaises(UploadValidationError):
|
||||
await service.confirm(
|
||||
upload_id=preview["upload_id"], title="FAQ", doc_id="d1",
|
||||
collection_name="fin_faq", strategy="qa_pair",
|
||||
)
|
||||
embedder.assert_not_awaited()
|
||||
|
||||
async def test_preview_storage_cleanup_removes_expired_upload_and_manifest(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
service = KnowledgeUploadService(
|
||||
storage_dir=tmp,
|
||||
upload_ttl_seconds=60,
|
||||
milvus_client=AsyncMock(),
|
||||
embedder=AsyncMock(),
|
||||
publisher=AsyncMock(),
|
||||
)
|
||||
preview = await service.preview("faq.md", b"Q: Q\nA: A", strategy="qa_pair")
|
||||
path = Path(tmp, preview["stored_filename"])
|
||||
manifest = Path(tmp, f"{preview['upload_id']}.json")
|
||||
old = time.time() - 120
|
||||
os.utime(path, (old, old))
|
||||
os.utime(manifest, (old, old))
|
||||
|
||||
removed = service.cleanup_expired_uploads(now=time.time())
|
||||
|
||||
self.assertEqual(removed, 1)
|
||||
self.assertFalse(path.exists())
|
||||
self.assertFalse(manifest.exists())
|
||||
|
||||
async def test_confirm_rejects_expired_upload(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
service = KnowledgeUploadService(
|
||||
storage_dir=tmp,
|
||||
upload_ttl_seconds=60,
|
||||
milvus_client=AsyncMock(),
|
||||
embedder=AsyncMock(),
|
||||
publisher=AsyncMock(),
|
||||
)
|
||||
preview = await service.preview("faq.md", b"Q: Q\nA: A", strategy="qa_pair")
|
||||
path = Path(tmp, preview["stored_filename"])
|
||||
old = time.time() - 120
|
||||
os.utime(path, (old, old))
|
||||
os.utime(Path(tmp, f"{preview['upload_id']}.json"), (old, old))
|
||||
|
||||
with self.assertRaises(UploadValidationError):
|
||||
await service.confirm(
|
||||
upload_id=preview["upload_id"],
|
||||
title="FAQ",
|
||||
doc_id="expired-doc",
|
||||
collection_name="fin_faq",
|
||||
strategy="qa_pair",
|
||||
)
|
||||
|
||||
async def test_delete_document_removes_vectors_and_publishes_update(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
milvus = AsyncMock()
|
||||
publisher = AsyncMock()
|
||||
service = KnowledgeUploadService(
|
||||
storage_dir=tmp,
|
||||
milvus_client=milvus,
|
||||
embedder=AsyncMock(),
|
||||
publisher=publisher,
|
||||
)
|
||||
|
||||
result = await service.delete_document("doc-delete-1")
|
||||
|
||||
self.assertEqual(result, {"doc_id": "doc-delete-1", "deleted": True})
|
||||
self.assertEqual(milvus.delete.await_count, 3)
|
||||
publisher.assert_awaited_once()
|
||||
self.assertEqual(publisher.await_args.args[0], KNOWLEDGE_UPDATE_EVENT)
|
||||
self.assertEqual(publisher.await_args.args[1]["doc_id"], "doc-delete-1")
|
||||
self.assertEqual(publisher.await_args.args[1]["chunk_count"], 0)
|
||||
self.assertEqual(publisher.await_args.args[1]["action"], "deleted")
|
||||
|
||||
async def test_list_documents_groups_chunks_across_collections(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
milvus = AsyncMock()
|
||||
|
||||
async def query(**kwargs):
|
||||
if kwargs["collection_name"] == "fin_faq":
|
||||
return [
|
||||
{
|
||||
"doc_id": "doc-1",
|
||||
"title": "FAQ",
|
||||
"section_title": "基金基础",
|
||||
"strategy": "qa_pair",
|
||||
},
|
||||
{
|
||||
"doc_id": "doc-1",
|
||||
"title": "FAQ",
|
||||
"section_title": "购买流程",
|
||||
"strategy": "qa_pair",
|
||||
},
|
||||
]
|
||||
return []
|
||||
|
||||
milvus.query.side_effect = query
|
||||
service = KnowledgeUploadService(
|
||||
storage_dir=tmp,
|
||||
milvus_client=milvus,
|
||||
embedder=AsyncMock(),
|
||||
publisher=AsyncMock(),
|
||||
)
|
||||
|
||||
result = await service.list_documents()
|
||||
|
||||
self.assertEqual(result, [{
|
||||
"doc_id": "doc-1",
|
||||
"title": "FAQ",
|
||||
"collection_name": "fin_faq",
|
||||
"strategy": "qa_pair",
|
||||
"chunk_count": 2,
|
||||
}])
|
||||
|
||||
async def test_get_document_returns_not_found_when_doc_id_is_missing(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
milvus = AsyncMock()
|
||||
milvus.query.return_value = []
|
||||
service = KnowledgeUploadService(
|
||||
storage_dir=tmp,
|
||||
milvus_client=milvus,
|
||||
embedder=AsyncMock(),
|
||||
publisher=AsyncMock(),
|
||||
)
|
||||
|
||||
self.assertIsNone(await service.get_document("missing-doc"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user