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

service/knowledge_base/upload.py

@@ -3,16 +3,19 @@ from __future__ import annotations
 
 import logging
 import json
+import re
+import time
 import uuid
 from inspect import isawaitable
 from pathlib import Path
+from uuid import NAMESPACE_URL, uuid5
 
 from rag.document_parser import SUPPORTED_EXTENSIONS
 from rag.events import publish_knowledge_update
-from rag.ingestion import ingest_document_atomic
-from rag.preview import preview_document
-from rag.chunk_config import resolve_chunk_config
+from rag.embedding import EMBEDDING_DIMENSION
 from rag.milvus_collections import KNOWLEDGE_COLLECTIONS
+from rag.milvus_delete import _escape_filter_value
+from rag.preview import preview_document
 from rag.milvus_delete import delete_document_vectors
 
 
@@ -123,27 +126,37 @@ class KnowledgeUploadService:
 
     async def list_documents(self) -> list[dict]:
         documents = {}
+        page_size = 4096
         for collection_name in KNOWLEDGE_COLLECTIONS:
-            rows = await self.milvus_client.query(
-                collection_name=collection_name,
-                filter="",
-                output_fields=["doc_id", "title", "strategy"],
-            )
-            for row in rows or []:
-                doc_id = row.get("doc_id", "")
-                if not doc_id:
-                    continue
-                document = documents.setdefault(
-                    doc_id,
-                    {
-                        "doc_id": doc_id,
-                        "title": row.get("title", ""),
-                        "collection_name": collection_name,
-                        "strategy": row.get("strategy", ""),
-                        "chunk_count": 0,
-                    },
+            offset = 0
+            while True:
+                rows = await self.milvus_client.query(
+                    collection_name=collection_name,
+                    filter="",
+                    output_fields=["doc_id", "title", "strategy"],
+                    limit=page_size,
+                    offset=offset,
                 )
-                document["chunk_count"] += 1
+                if not rows:
+                    break
+                for row in rows:
+                    doc_id = row.get("doc_id", "")
+                    if not doc_id:
+                        continue
+                    document = documents.setdefault(
+                        doc_id,
+                        {
+                            "doc_id": doc_id,
+                            "title": row.get("title", ""),
+                            "collection_name": collection_name,
+                            "strategy": row.get("strategy", ""),
+                            "chunk_count": 0,
+                        },
+                    )
+                    document["chunk_count"] += 1
+                if len(rows) < page_size:
+                    break
+                offset += page_size
         return list(documents.values())
 
     async def get_document(self, doc_id: str) -> dict | None:
@@ -177,47 +190,106 @@ class KnowledgeUploadService:
     async def confirm(
         self,
         *,
-        upload_id: str,
-        title: str,
-        doc_id: str,
+        filename: str,
+        content: bytes,
+        title: str | None = None,
+        doc_id: str | None = None,
         collection_name: str,
         strategy: str,
         chunk_size: int | None = None,
         chunk_overlap: int | None = None,
     ) -> dict:
-        self.cleanup_expired_uploads()
+        suffix = self._validate_upload(filename, content)
         if collection_name not in KNOWLEDGE_COLLECTIONS:
             raise UploadValidationError("知识库集合不在允许范围内")
-        path = self._find_upload(upload_id)
-        try:
-            manifest = json.loads(self._manifest_path(upload_id).read_text(encoding="utf-8"))
-            if strategy != manifest["strategy"]:
-                raise UploadValidationError("确认策略必须与预览策略一致")
-            if chunk_size is None:
-                chunk_size = manifest["chunk_size"]
-            if chunk_overlap is None:
-                chunk_overlap = manifest["chunk_overlap"]
-            if self.document_exists is not None:
-                exists = self.document_exists(doc_id)
-                if isawaitable(exists):
-                    exists = await exists
-                if exists:
-                    raise UploadValidationError("doc_id已存在，禁止重复入库")
-            from rag.document_parser import parse_document
 
-            result = await ingest_document_atomic(
-                parse_document(path),
-                doc_id,
-                title,
-                collection_name,
-                strategy,
-                milvus_client=self.milvus_client,
-                embedder=self.embedder,
-                config=resolve_chunk_config(
-                    chunk_size=chunk_size,
-                    chunk_overlap=chunk_overlap,
-                ),
+        stem = Path(filename).stem
+        title = (title or "").strip() or stem
+        if not doc_id or not doc_id.strip():
+            doc_id = self._slugify(stem) or f"doc_{int(time.time())}_{uuid.uuid4().hex[:8]}"
+        else:
+            doc_id = doc_id.strip()
+
+        if self.document_exists is not None:
+            exists = self.document_exists(doc_id)
+            if isawaitable(exists):
+                exists = await exists
+            if exists:
+                raise UploadValidationError("doc_id已存在，禁止重复入库")
+
+        upload_id = uuid.uuid4().hex
+        temp_path = self.storage_dir / f"{upload_id}{suffix}"
+        temp_path.write_bytes(content)
+
+        try:
+            preview = preview_document(
+                temp_path,
+                strategy=strategy,
+                chunk_size=chunk_size,
+                chunk_overlap=chunk_overlap,
             )
+            if not preview["chunks"]:
+                raise UploadValidationError("文档切片为空")
+
+            vectors = await self.embedder(
+                [chunk["text"] for chunk in preview["chunks"]]
+            )
+            if len(vectors) != len(preview["chunks"]) or any(
+                len(vector) != EMBEDDING_DIMENSION for vector in vectors
+            ):
+                raise UploadValidationError(
+                    f"Embedding维度必须为{EMBEDDING_DIMENSION}"
+                )
+
+            rows = [
+                {
+                    "chunk_id": str(uuid5(NAMESPACE_URL, f"{doc_id}:{index}")),
+                    "doc_id": doc_id,
+                    "title": title,
+                    "section_title": chunk.get("section_title") or "",
+                    "text": chunk["text"],
+                    "strategy": preview["actual_strategy"],
+                    "vector": vector,
+                }
+                for index, (chunk, vector) in enumerate(
+                    zip(preview["chunks"], vectors)
+                )
+            ]
+            try:
+                await self.milvus_client.insert(
+                    collection_name=collection_name, data=rows
+                )
+            except Exception:
+                logger.exception(
+                    "atomic Milvus ingestion failed for doc_id=%s", doc_id
+                )
+                try:
+                    await self.milvus_client.delete(
+                        collection_name=collection_name,
+                        filter=f'doc_id == "{_escape_filter_value(doc_id)}"',
+                    )
+                except Exception:
+                    logger.exception(
+                        "failed to clean up document rows for doc_id=%s", doc_id
+                    )
+                raise
+
+            result = {
+                "doc_id": doc_id,
+                "title": title,
+                "collection_name": collection_name,
+                "filename": Path(filename).name,
+                "strategy": preview["strategy"],
+                "actual_strategy": preview["actual_strategy"],
+                "degraded": preview["degraded"],
+                "warning": preview["warning"],
+                "cleaning_changed": preview["cleaning_changed"],
+                "cleaning_warnings": preview["cleaning_warnings"],
+                "chunk_count": len(rows),
+                "chunk_size": preview["chunk_size"],
+                "chunk_overlap": preview["chunk_overlap"],
+            }
+
             try:
                 await publish_knowledge_update(
                     self.publisher,
@@ -226,16 +298,29 @@ class KnowledgeUploadService:
                     chunk_count=result["chunk_count"],
                 )
             except Exception:
-                logger.exception("knowledge update event publish failed for doc_id=%s", doc_id)
+                logger.exception(
+                    "knowledge update event publish failed for doc_id=%s", doc_id
+                )
                 try:
-                    await delete_document_vectors(doc_id, milvus_client=self.milvus_client)
+                    await delete_document_vectors(
+                        doc_id, milvus_client=self.milvus_client
+                    )
                 except Exception:
-                    logger.exception("failed to compensate vectors for doc_id=%s", doc_id)
+                    logger.exception(
+                        "failed to compensate vectors for doc_id=%s", doc_id
+                    )
                 raise
+
             return {"event_published": True, **result}
         finally:
-            path.unlink(missing_ok=True)
-            self._manifest_path(upload_id).unlink(missing_ok=True)
+            temp_path.unlink(missing_ok=True)
+
+    @staticmethod
+    def _slugify(text: str) -> str:
+        s = text.strip().lower()
+        s = re.sub(r"[^\w\u4e00-\u9fff]+", "_", s)
+        s = re.sub(r"_+", "_", s).strip("_")
+        return s
 
     def _validate_upload(self, filename: str, content: bytes) -> str:
         suffix = Path(filename).suffix.lower()
@@ -259,4 +344,4 @@ class KnowledgeUploadService:
         return paths[0]
 
     def _manifest_path(self, upload_id: str) -> Path:
-        return self.storage_dir / f"{upload_id}.json"
+        return self.storage_dir / f"{upload_id}.json"