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 @@
+"""Knowledge-base upload and ingestion services."""
@@ -0,0 +1,262 @@
+"""Frontend document upload, preview, and confirmed Milvus ingestion."""
+from __future__ import annotations
+
+import logging
+import json
+import uuid
+from inspect import isawaitable
+from pathlib import Path
+
+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.milvus_collections import KNOWLEDGE_COLLECTIONS
+from rag.milvus_delete import delete_document_vectors
+
+
+logger = logging.getLogger("service.knowledge_base.upload")
+
+
+class UploadValidationError(ValueError):
+    pass
+
+
+class KnowledgeUploadService:
+    def __init__(
+        self,
+        *,
+        storage_dir: str | Path,
+        milvus_client,
+        embedder,
+        publisher,
+        max_upload_bytes: int = 10 * 1024 * 1024,
+        document_exists=None,
+        upload_ttl_seconds: int = 1800,
+    ):
+        self.storage_dir = Path(storage_dir)
+        self.storage_dir.mkdir(parents=True, exist_ok=True)
+        self.milvus_client = milvus_client
+        self.embedder = embedder
+        self.publisher = publisher
+        self.max_upload_bytes = max_upload_bytes
+        self.document_exists = document_exists
+        self.upload_ttl_seconds = upload_ttl_seconds
+
+    async def preview(
+        self,
+        filename: str,
+        content: bytes,
+        *,
+        strategy: str,
+        chunk_size: int | None = None,
+        chunk_overlap: int | None = None,
+    ) -> dict:
+        self.cleanup_expired_uploads()
+        suffix = self._validate_upload(filename, content)
+        upload_id = uuid.uuid4().hex
+        stored_filename = f"{upload_id}{suffix}"
+        path = self.storage_dir / stored_filename
+        path.write_bytes(content)
+        try:
+            result = preview_document(
+                path,
+                strategy=strategy,
+                chunk_size=chunk_size,
+                chunk_overlap=chunk_overlap,
+            )
+        except Exception:
+            path.unlink(missing_ok=True)
+            raise
+        self._manifest_path(upload_id).write_text(
+            json.dumps(
+                {
+                    "strategy": strategy,
+                    "chunk_size": result["chunk_size"],
+                    "chunk_overlap": result["chunk_overlap"],
+                },
+                ensure_ascii=False,
+            ),
+            encoding="utf-8",
+        )
+        return {
+            "upload_id": upload_id,
+            "stored_filename": stored_filename,
+            "filename": Path(filename).name,
+            **result,
+        }
+
+    def cleanup_expired_uploads(self, *, now: float | None = None) -> int:
+        """Remove expired upload inputs and their preview manifests."""
+        if self.upload_ttl_seconds <= 0:
+            return 0
+        if now is None:
+            import time
+
+            now = time.time()
+        cutoff = now - self.upload_ttl_seconds
+        removed = 0
+        for path in self.storage_dir.iterdir():
+            if not path.is_file() or path.suffix.lower() not in SUPPORTED_EXTENSIONS:
+                continue
+            if path.stat().st_mtime > cutoff:
+                continue
+            upload_id = path.stem
+            path.unlink(missing_ok=True)
+            self._manifest_path(upload_id).unlink(missing_ok=True)
+            removed += 1
+        return removed
+
+    async def delete_document(self, doc_id: str) -> dict:
+        if not doc_id or not doc_id.strip():
+            raise UploadValidationError("doc_id不能为空")
+        await delete_document_vectors(doc_id, milvus_client=self.milvus_client)
+        await publish_knowledge_update(
+            self.publisher,
+            doc_id=doc_id,
+            collection_name="",
+            chunk_count=0,
+            action="deleted",
+        )
+        return {"doc_id": doc_id, "deleted": True}
+
+    async def list_documents(self) -> list[dict]:
+        documents = {}
+        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,
+                    },
+                )
+                document["chunk_count"] += 1
+        return list(documents.values())
+
+    async def get_document(self, doc_id: str) -> dict | None:
+        if not doc_id or not doc_id.strip():
+            raise UploadValidationError("doc_id不能为空")
+        escaped = doc_id.replace("\\", "\\\\").replace('"', '\\"')
+        for collection_name in KNOWLEDGE_COLLECTIONS:
+            rows = await self.milvus_client.query(
+                collection_name=collection_name,
+                filter=f'doc_id == "{escaped}"',
+                output_fields=["doc_id", "title", "section_title", "strategy", "text"],
+            )
+            if rows:
+                first = rows[0]
+                return {
+                    "doc_id": first.get("doc_id", doc_id),
+                    "title": first.get("title", ""),
+                    "collection_name": collection_name,
+                    "strategy": first.get("strategy", ""),
+                    "chunk_count": len(rows),
+                    "chunks": [
+                        {
+                            "section_title": row.get("section_title") or None,
+                            "text": row.get("text", ""),
+                        }
+                        for row in rows
+                    ],
+                }
+        return None
+
+    async def confirm(
+        self,
+        *,
+        upload_id: str,
+        title: str,
+        doc_id: str,
+        collection_name: str,
+        strategy: str,
+        chunk_size: int | None = None,
+        chunk_overlap: int | None = None,
+    ) -> dict:
+        self.cleanup_expired_uploads()
+        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,
+                ),
+            )
+            try:
+                await publish_knowledge_update(
+                    self.publisher,
+                    doc_id=doc_id,
+                    collection_name=collection_name,
+                    chunk_count=result["chunk_count"],
+                )
+            except Exception:
+                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)
+                except Exception:
+                    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)
+
+    def _validate_upload(self, filename: str, content: bytes) -> str:
+        suffix = Path(filename).suffix.lower()
+        if suffix not in SUPPORTED_EXTENSIONS:
+            raise UploadValidationError(f"Unsupported document type: {suffix or '<none>'}")
+        if not content:
+            raise UploadValidationError("上传文件不能为空")
+        if len(content) > self.max_upload_bytes:
+            raise UploadValidationError("上传文件超过大小限制")
+        return suffix
+
+    def _find_upload(self, upload_id: str) -> Path:
+        if not upload_id or Path(upload_id).name != upload_id:
+            raise UploadValidationError("upload_id无效")
+        paths = [
+            path for path in self.storage_dir.glob(f"{upload_id}.*")
+            if path.suffix.lower() in SUPPORTED_EXTENSIONS
+        ]
+        if len(paths) != 1:
+            raise UploadValidationError("上传文件不存在或已过期")
+        return paths[0]
+
+    def _manifest_path(self, upload_id: str) -> Path:
+        return self.storage_dir / f"{upload_id}.json"
				`@@ -0,0 +1 @@`
				`"""Knowledge-base upload and ingestion services."""`