AI260626/Mutual_Fund
17
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: 客服agent的切片结构调整重构

Browse Source
This commit is contained in:
ouyangyang_0626
2026-09-11 17:31:16 +08:00
parent b1b764eac9
commit e561c97a4e
20 changed files with 3309 additions and 230 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_embedding.py
+4 -4
View File
@@ -1,18 +1,18 @@
import unittest
from rag.embedding import EmbeddingError, embed_texts
from rag.embedding import EMBEDDING_DIMENSION, EmbeddingError, embed_texts
class EmbeddingTests(unittest.IsolatedAsyncioTestCase):
async def test_returns_768_dimension_vectors(self):
async def test_returns_configured_dimension_vectors(self):
class Client:
async def embed(self, texts):
return [[0.1] * 768 for _ in texts]
return [[0.1] * EMBEDDING_DIMENSION for _ in texts]
vectors = await embed_texts(["基金知识"], client=Client())
self.assertEqual(len(vectors), 1)
self.assertEqual(len(vectors[0]), 768)
self.assertEqual(len(vectors[0]), EMBEDDING_DIMENSION)
async def test_rejects_wrong_embedding_dimension(self):
class Client:
tests/test_ingestion.py
+4 -3
View File
@@ -1,6 +1,7 @@
import unittest
from unittest.mock import AsyncMock
from rag.embedding import EMBEDDING_DIMENSION
from rag.ingestion import ingest_document_atomic
@@ -35,7 +36,7 @@ class AtomicIngestionTests(unittest.IsolatedAsyncioTestCase):
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])
embedder = AsyncMock(return_value=[[0.0] * EMBEDDING_DIMENSION])
with self.assertRaises(RuntimeError):
await ingest_document_atomic(
@@ -49,7 +50,7 @@ class AtomicIngestionTests(unittest.IsolatedAsyncioTestCase):
async def test_success_inserts_complete_document_rows(self):
milvus = AsyncMock()
embedder = AsyncMock(return_value=[[0.0] * 768])
embedder = AsyncMock(return_value=[[0.0] * EMBEDDING_DIMENSION])
result = await ingest_document_atomic(
"plain text", "doc-4", "Policy", "fin_policy", "default",
@@ -62,7 +63,7 @@ class AtomicIngestionTests(unittest.IsolatedAsyncioTestCase):
async def test_cleans_text_before_embedding_and_milvus_insert(self):
milvus = AsyncMock()
embedder = AsyncMock(return_value=[[0.0] * 768, [0.0] * 768])
embedder = AsyncMock(return_value=[[0.0] * EMBEDDING_DIMENSION, [0.0] * EMBEDDING_DIMENSION])
await ingest_document_atomic(
"\ufeff基金名称：示例基金\r\n\r\n\r\n风险等级：R3\t ",
tests/test_knowledge_router.py
+5 -5
View File
@@ -9,6 +9,7 @@ from api.chat.knowledge import (
list_documents,
preview_document_upload,
)
from rag.embedding import EMBEDDING_DIMENSION
from service.knowledge_base.upload import KnowledgeUploadService
from types import SimpleNamespace
@@ -65,14 +66,13 @@ class KnowledgeRouterTests(unittest.IsolatedAsyncioTestCase):
service = KnowledgeUploadService(
storage_dir=tmp,
milvus_client=AsyncMock(),
embedder=AsyncMock(return_value=[[0.0] * 768]),
embedder=AsyncMock(return_value=[[0.0] * EMBEDDING_DIMENSION]),
publisher=AsyncMock(),
)
preview = await service.preview("faq.md", b"Q: Q\nA: A", strategy="qa_pair")
request = FakeRequest(
service,
json_data={
"upload_id": preview["upload_id"],
form_data={
"file": FakeUpload("faq.md", b"Q: Q\nA: A"),
"title": "FAQ",
"doc_id": "doc-1",
"collection_name": "fin_faq",
@@ -144,4 +144,4 @@ class KnowledgeRouterTests(unittest.IsolatedAsyncioTestCase):
if __name__ == "__main__":
unittest.main()
unittest.main()
tests/test_knowledge_upload.py
+31 -40
View File
@@ -5,6 +5,7 @@ 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
@@ -34,21 +35,17 @@ class KnowledgeUploadTests(unittest.IsolatedAsyncioTestCase):
with tempfile.TemporaryDirectory() as tmp:
publisher = AsyncMock()
milvus = AsyncMock()
embedder = AsyncMock(return_value=[[0.0] * 768])
embedder = AsyncMock(return_value=[[0.0] * EMBEDDING_DIMENSION])
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"],
filename="faq.md",
content="Q: 什么是基金？\nA: 一种集合投资工具。".encode(),
title="FAQ",
doc_id="doc-upload-1",
collection_name="fin_faq",
@@ -58,24 +55,23 @@ class KnowledgeUploadTests(unittest.IsolatedAsyncioTestCase):
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())
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] * 768 for _ in texts])
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(),
)
preview = await service.preview(
"doc.md", "一二三四五六七八九十十一十二".encode(), strategy="default"
)
await service.confirm(
upload_id=preview["upload_id"],
filename="doc.md",
content="一二三四五六七八九十十一十二".encode(),
title="Doc",
doc_id="doc-config",
collection_name="fin_fund_doc",
@@ -93,16 +89,14 @@ class KnowledgeUploadTests(unittest.IsolatedAsyncioTestCase):
service = KnowledgeUploadService(
storage_dir=tmp,
milvus_client=milvus,
embedder=AsyncMock(return_value=[[0.0] * 768]),
embedder=AsyncMock(return_value=[[0.0] * EMBEDDING_DIMENSION]),
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"],
filename="policy.md",
content=b"policy text",
title="Policy",
doc_id="doc-event-failure",
collection_name="fin_policy",
@@ -134,14 +128,15 @@ class KnowledgeUploadTests(unittest.IsolatedAsyncioTestCase):
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",
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):
@@ -154,14 +149,15 @@ class KnowledgeUploadTests(unittest.IsolatedAsyncioTestCase):
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",
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()
@@ -187,28 +183,23 @@ class KnowledgeUploadTests(unittest.IsolatedAsyncioTestCase):
self.assertFalse(path.exists())
self.assertFalse(manifest.exists())
async def test_confirm_rejects_expired_upload(self):
async def test_confirm_rejects_unsupported_extension(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",
filename="script.exe",
content=b"bad",
title="Bad",
doc_id="bad-doc",
collection_name="fin_faq",
strategy="qa_pair",
strategy="default",
)
async def test_delete_document_removes_vectors_and_publishes_update(self):
@@ -287,4 +278,4 @@ class KnowledgeUploadTests(unittest.IsolatedAsyncioTestCase):
if __name__ == "__main__":
unittest.main()
unittest.main()
tests/test_milvus_collections.py
+24 -7
View File
@@ -2,7 +2,20 @@ import unittest
from unittest.mock import AsyncMock, call
from unittest.mock import patch
from rag.milvus_collections import KNOWLEDGE_COLLECTIONS, ensure_collections
from rag.milvus_collections import (
EMBEDDING_DIMENSION,
KNOWLEDGE_COLLECTIONS,
ensure_collections,
)
def _existing_client(dim: int) -> AsyncMock:
client = AsyncMock()
client.has_collection.return_value = True
client.describe_collection.return_value = {
"fields": [{"name": "vector", "params": {"dim": dim}}]
}
return client
class MilvusCollectionTests(unittest.IsolatedAsyncioTestCase):
@@ -19,14 +32,19 @@ class MilvusCollectionTests(unittest.IsolatedAsyncioTestCase):
self.assertEqual(client.create_collection.await_count, len(KNOWLEDGE_COLLECTIONS))
async def test_reuses_existing_collections(self):
client = AsyncMock()
client.has_collection.return_value = True
client = _existing_client(EMBEDDING_DIMENSION)
await ensure_collections(client)
client.create_collection.assert_not_awaited()
async def test_collection_schema_contains_metadata_and_768_vector(self):
async def test_existing_collection_with_wrong_dimension_raises(self):
client = _existing_client(EMBEDDING_DIMENSION + 1)
with self.assertRaises(RuntimeError):
await ensure_collections(client)
async def test_collection_schema_contains_metadata_and_1024_vector(self):
client = AsyncMock()
client.has_collection.return_value = False
@@ -36,11 +54,10 @@ class MilvusCollectionTests(unittest.IsolatedAsyncioTestCase):
fields = {field["name"] for field in schema.to_dict()["fields"]}
self.assertTrue({"chunk_id", "doc_id", "title", "section_title", "text", "strategy", "vector"} <= fields)
vector = next(field for field in schema.to_dict()["fields"] if field["name"] == "vector")
self.assertEqual(vector["params"]["dim"], 768)
self.assertEqual(vector["params"]["dim"], 1024)
async def test_default_path_uses_configured_milvus_client(self):
client = AsyncMock()
client.has_collection.return_value = True
client = _existing_client(EMBEDDING_DIMENSION)
with patch("rag.milvus_collections.configured_milvus_client", return_value=client):
await ensure_collections()