347 lines
12 KiB
Python
347 lines
12 KiB
Python
"""Frontend document upload, preview, and confirmed Milvus ingestion."""
|
|
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.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
|
|
|
|
|
|
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 = {}
|
|
page_size = 4096
|
|
for collection_name in KNOWLEDGE_COLLECTIONS:
|
|
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,
|
|
)
|
|
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:
|
|
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,
|
|
*,
|
|
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:
|
|
suffix = self._validate_upload(filename, content)
|
|
if collection_name not in KNOWLEDGE_COLLECTIONS:
|
|
raise UploadValidationError("知识库集合不在允许范围内")
|
|
|
|
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,
|
|
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:
|
|
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()
|
|
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" |