kplam
fdb664e546
对多个文件进行代码格式化调整，将长行参数拆分为多行书写，提升代码可读性，包括： - 调整函数定义、调用的多行换行格式 - 优化列表、元组、字典的多行排版 - 新增README.md项目说明文档
239 lines
8.5 KiB
Python
239 lines
8.5 KiB
Python
"""文档 API：POST /api/v1/documents 异步入库 + 任务查询 + 文档管理（列表/详情/删除）+ 文件上传入库"""
|
||
|
||
import json
|
||
import uuid
|
||
from datetime import UTC, datetime
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
import structlog
|
||
from fastapi import APIRouter, Depends, File, Form, Query, UploadFile
|
||
from fastapi.responses import JSONResponse
|
||
|
||
from app.api.response import ApiError, ok
|
||
from app.config import Settings, settings
|
||
from app.core.auth import AuthUser, get_current_user, require_admin
|
||
from app.core.file_parser import parse_file
|
||
from app.core.ingest_tasks import IngestTaskManager
|
||
from app.core.ingestion import Ingester
|
||
from app.models.document import DocumentInput
|
||
from app.services.qdrant import QdrantService
|
||
from app.services.redis import RedisCache, get_cache
|
||
|
||
logger = structlog.get_logger()
|
||
|
||
router = APIRouter(prefix="/api/v1", tags=["document"])
|
||
|
||
# 模块级懒加载单例，避免每请求重建 Ingester 及其下游依赖
|
||
_ingester: Ingester | None = None
|
||
_qdrant: QdrantService | None = None
|
||
_task_manager: IngestTaskManager | None = None
|
||
|
||
|
||
def _get_ingester() -> Ingester:
|
||
global _ingester
|
||
if _ingester is None:
|
||
_ingester = Ingester()
|
||
return _ingester
|
||
|
||
|
||
def _get_qdrant() -> QdrantService:
|
||
global _qdrant
|
||
if _qdrant is None:
|
||
_qdrant = QdrantService()
|
||
return _qdrant
|
||
|
||
|
||
def _get_task_manager() -> IngestTaskManager:
|
||
"""入库任务管理器懒加载单例
|
||
|
||
Redis 沿用全局缓存单例（RedisCache 读写全容错，不可用时镜像写失败仅告警）；
|
||
获取缓存实例异常时传 None，退化为纯内存模式。
|
||
"""
|
||
global _task_manager
|
||
if _task_manager is None:
|
||
try:
|
||
redis: RedisCache | None = get_cache()
|
||
except Exception:
|
||
logger.warning("Redis 缓存不可用，入库任务状态仅保留在内存", exc_info=True)
|
||
redis = None
|
||
_task_manager = IngestTaskManager(
|
||
ingester=_get_ingester(), redis=redis, settings=Settings()
|
||
)
|
||
return _task_manager
|
||
|
||
|
||
def _allowed_extensions() -> set[str]:
|
||
"""解析 settings.upload_allowed_extensions 逗号分隔字符串为扩展名集合（全小写、含点号）"""
|
||
return {
|
||
ext.strip().lower()
|
||
for ext in settings.upload_allowed_extensions.split(",")
|
||
if ext.strip()
|
||
}
|
||
|
||
|
||
@router.post("/documents")
|
||
async def ingest_document(
|
||
doc: DocumentInput, user: AuthUser = Depends(get_current_user)
|
||
) -> JSONResponse:
|
||
"""文档入库入口：登记异步任务并返回 202 + task_id，入库结果经任务查询端点获取"""
|
||
if not doc.text.strip():
|
||
raise ApiError(1001, "文档内容不能为空")
|
||
task_id = await _get_task_manager().submit(doc)
|
||
return JSONResponse(
|
||
status_code=202, content=ok({"task_id": task_id, "status": "pending"})
|
||
)
|
||
|
||
|
||
@router.post("/documents/upload")
|
||
async def upload_document(
|
||
file: UploadFile = File(
|
||
..., description="上传的文件（.txt/.md/.html/.htm/.pdf/.docx）"
|
||
),
|
||
title: str = Form(default="", description="可选标题，默认取原文件名去扩展"),
|
||
source: str = Form(default="", description="可选来源标识，默认 file:{原文件名}"),
|
||
metadata: str = Form(
|
||
default="", description='可选元数据 JSON 字符串，如 \'{"author":"x"}\''
|
||
),
|
||
user: AuthUser = Depends(get_current_user),
|
||
) -> JSONResponse:
|
||
"""文件上传入库入口：校验 → 提取文本 → 落盘 → 提交异步入库流水线
|
||
|
||
与 POST /documents 共用同一 IngestTaskManager；任务状态经
|
||
GET /api/v1/documents/tasks/{task_id} 查询。
|
||
"""
|
||
original_filename = file.filename or "unnamed"
|
||
ext = Path(original_filename).suffix.lower()
|
||
|
||
# 1. 扩展名校验
|
||
allowed = _allowed_extensions()
|
||
if ext not in allowed:
|
||
raise ApiError(1001, f"不支持的文件类型: {ext or '(无扩展名)'}")
|
||
|
||
# 2. 读取字节并校验大小
|
||
content = await file.read()
|
||
max_bytes = settings.upload_max_size_mb * 1024 * 1024
|
||
if len(content) > max_bytes:
|
||
raise ApiError(1001, f"文件超过大小上限: {settings.upload_max_size_mb}MB")
|
||
|
||
# 3. 提取文本
|
||
try:
|
||
text = parse_file(original_filename, content)
|
||
except ValueError as exc:
|
||
raise ApiError(1001, str(exc)) from exc
|
||
if not text.strip():
|
||
raise ApiError(1001, "无法从文件提取文本")
|
||
|
||
# 4. 落盘（按 YYYY/MM 日期分片；失败仅 warning，不阻塞入库）
|
||
doc_id = uuid.uuid4().hex
|
||
saved_path = ""
|
||
metadata_dict: dict[str, str] = {}
|
||
try:
|
||
upload_dir = Path(settings.upload_dir).resolve()
|
||
shard_subdir = datetime.now(UTC).strftime("%Y/%m")
|
||
target_dir = upload_dir / shard_subdir
|
||
target_dir.mkdir(parents=True, exist_ok=True)
|
||
target = target_dir / f"{doc_id}_{original_filename}"
|
||
target.write_bytes(content)
|
||
saved_path = str(target)
|
||
metadata_dict.update(
|
||
{
|
||
"raw_file_path": saved_path,
|
||
"original_filename": original_filename,
|
||
"original_size_bytes": str(len(content)),
|
||
}
|
||
)
|
||
logger.info(
|
||
"上传文件已落盘", doc_id=doc_id, saved_path=saved_path, size=len(content)
|
||
)
|
||
except Exception:
|
||
logger.warning(
|
||
"上传文件落盘失败，仅做文本入库",
|
||
doc_id=doc_id,
|
||
filename=original_filename,
|
||
exc_info=True,
|
||
)
|
||
|
||
# 5. 合并用户传入的 metadata（落盘元数据优先级更高，不与用户键冲突）
|
||
if metadata.strip():
|
||
try:
|
||
user_meta = json.loads(metadata)
|
||
if isinstance(user_meta, dict):
|
||
for k, v in user_meta.items():
|
||
metadata_dict.setdefault(str(k), str(v))
|
||
except (json.JSONDecodeError, TypeError):
|
||
logger.warning("metadata 不是合法 JSON，已忽略", raw=metadata)
|
||
|
||
# 6. 默认 title / source
|
||
if not title.strip():
|
||
title = Path(original_filename).stem
|
||
if not source.strip():
|
||
source = f"file:{original_filename}"
|
||
|
||
# 7. 提交入库流水线
|
||
doc_input = DocumentInput(
|
||
text=text, title=title, source=source, metadata=metadata_dict
|
||
)
|
||
task_id = await _get_task_manager().submit(doc_input)
|
||
|
||
return JSONResponse(
|
||
status_code=202,
|
||
content=ok({"task_id": task_id, "status": "pending", "saved_path": saved_path}),
|
||
)
|
||
|
||
|
||
@router.get("/documents/tasks/{task_id}")
|
||
async def get_ingest_task(
|
||
task_id: str, user: AuthUser = Depends(get_current_user)
|
||
) -> dict[str, Any]:
|
||
"""查询入库任务状态：含 task_id/status/created_at/updated_at，done 附 result，failed 附 error"""
|
||
task = await _get_task_manager().get(task_id)
|
||
if task is None:
|
||
raise ApiError(1004, "任务不存在")
|
||
return ok(task)
|
||
|
||
|
||
@router.get("/documents")
|
||
async def list_documents(
|
||
limit: int = Query(default=20, ge=1, le=100),
|
||
offset: str | None = None,
|
||
user: AuthUser = Depends(get_current_user),
|
||
) -> dict[str, Any]:
|
||
"""分页列出文档（L1 摘要），返回 items 与下一页游标 next_offset"""
|
||
try:
|
||
items, next_offset = await _get_qdrant().scroll_l1(limit=limit, offset=offset)
|
||
except Exception as exc:
|
||
logger.error("文档列表查询失败", error=str(exc))
|
||
raise ApiError(2000, f"文档列表查询失败: {exc}") from exc
|
||
return ok({"items": items, "next_offset": next_offset})
|
||
|
||
|
||
@router.get("/documents/{doc_id}")
|
||
async def get_document(
|
||
doc_id: str, user: AuthUser = Depends(get_current_user)
|
||
) -> dict[str, Any]:
|
||
"""获取文档详情：L1 记录 + L2/L3 节点 + chunks 数量"""
|
||
try:
|
||
detail = await _get_qdrant().get_doc_detail(doc_id)
|
||
except Exception as exc:
|
||
logger.error("文档详情查询失败", doc_id=doc_id, error=str(exc))
|
||
raise ApiError(2000, f"文档详情查询失败: {exc}") from exc
|
||
if detail is None:
|
||
raise ApiError(1004, "文档不存在")
|
||
return ok(detail)
|
||
|
||
|
||
@router.delete("/documents/{doc_id}")
|
||
async def delete_document(
|
||
doc_id: str, user: AuthUser = Depends(require_admin)
|
||
) -> dict[str, Any]:
|
||
"""删除文档：四层集合中该 doc_id 的所有点；幂等，不存在也返回成功（删除数全 0）"""
|
||
try:
|
||
deleted = await _get_qdrant().delete_by_doc_id(doc_id)
|
||
except Exception as exc:
|
||
logger.error("文档删除失败", doc_id=doc_id, error=str(exc))
|
||
raise ApiError(2000, f"文档删除失败: {exc}") from exc
|
||
return ok(
|
||
{"doc_id": doc_id, "deleted": deleted, "deleted_total": sum(deleted.values())}
|
||
)
|