281 lines
11 KiB
Python
281 lines
11 KiB
Python
import tempfile
|
|
import unittest
|
|
import os
|
|
import time
|
|
from pathlib import Path
|
|
from unittest.mock import AsyncMock
|
|
|
|
from rag.embedding import EMBEDDING_DIMENSION
|
|
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] * EMBEDDING_DIMENSION])
|
|
service = KnowledgeUploadService(
|
|
storage_dir=tmp,
|
|
milvus_client=milvus,
|
|
embedder=embedder,
|
|
publisher=publisher,
|
|
)
|
|
|
|
result = await service.confirm(
|
|
filename="faq.md",
|
|
content="Q: 什么是基金？\nA: 一种集合投资工具。".encode(),
|
|
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.assertEqual(len(list(Path(tmp).glob("*.md"))), 0)
|
|
self.assertEqual(len(list(Path(tmp).glob("*.json"))), 0)
|
|
|
|
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] * EMBEDDING_DIMENSION for _ in texts])
|
|
service = KnowledgeUploadService(
|
|
storage_dir=tmp,
|
|
milvus_client=milvus,
|
|
embedder=embedder,
|
|
publisher=AsyncMock(),
|
|
)
|
|
|
|
await service.confirm(
|
|
filename="doc.md",
|
|
content="一二三四五六七八九十十一十二".encode(),
|
|
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] * EMBEDDING_DIMENSION]),
|
|
publisher=publisher,
|
|
)
|
|
|
|
with self.assertRaises(RuntimeError):
|
|
await service.confirm(
|
|
filename="policy.md",
|
|
content=b"policy text",
|
|
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(),
|
|
)
|
|
|
|
with self.assertRaises(UploadValidationError):
|
|
await service.confirm(
|
|
filename="faq.md",
|
|
content=b"Q: Q\nA: A",
|
|
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,
|
|
)
|
|
|
|
with self.assertRaises(UploadValidationError):
|
|
await service.confirm(
|
|
filename="faq.md",
|
|
content=b"Q: Q\nA: A",
|
|
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_unsupported_extension(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
service = KnowledgeUploadService(
|
|
storage_dir=tmp,
|
|
milvus_client=AsyncMock(),
|
|
embedder=AsyncMock(),
|
|
publisher=AsyncMock(),
|
|
)
|
|
|
|
with self.assertRaises(UploadValidationError):
|
|
await service.confirm(
|
|
filename="script.exe",
|
|
content=b"bad",
|
|
title="Bad",
|
|
doc_id="bad-doc",
|
|
collection_name="fin_faq",
|
|
strategy="default",
|
|
)
|
|
|
|
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() |