263 lines
9.5 KiB
Python
263 lines
9.5 KiB
Python
|
"""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"
|