84 lines
2.3 KiB
Python
84 lines
2.3 KiB
Python
"""
|
||
文件/长任务处理进度：阶段常量与上报助手。
|
||
|
||
进度约定：
|
||
- 0–30：HTTP 上传（主要由前端 onUploadProgress 展示；入库后后端从 30 起）
|
||
- 30–90：解析 / OCR / 分块 / 向量化
|
||
- 90–100：摘要与收尾
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
from typing import Awaitable, Callable, Optional, Type
|
||
|
||
import asyncpg
|
||
|
||
# ---------- 阶段 ----------
|
||
STAGE_UPLOADING = "uploading"
|
||
STAGE_PROCESSING = "processing"
|
||
STAGE_DOWNLOADING = "downloading"
|
||
STAGE_PARSING = "parsing"
|
||
STAGE_OCR = "ocr"
|
||
STAGE_SPLITTING = "splitting"
|
||
STAGE_EMBEDDING = "embedding"
|
||
STAGE_SUMMARIZING = "summarizing"
|
||
STAGE_EXTRACTING = "extracting"
|
||
STAGE_INDEXING = "indexing"
|
||
STAGE_COMPLETED = "completed"
|
||
STAGE_FAILED = "failed"
|
||
|
||
ProgressCallback = Callable[[int, str], Awaitable[None]]
|
||
|
||
STAGE_LABELS: dict[str, str] = {
|
||
STAGE_UPLOADING: "上传中",
|
||
STAGE_PROCESSING: "处理中",
|
||
STAGE_DOWNLOADING: "下载文件",
|
||
STAGE_PARSING: "解析文档",
|
||
STAGE_OCR: "文字识别",
|
||
STAGE_SPLITTING: "文本分块",
|
||
STAGE_EMBEDDING: "向量化",
|
||
STAGE_SUMMARIZING: "生成摘要",
|
||
STAGE_EXTRACTING: "抽取实体关系",
|
||
STAGE_INDEXING: "建立索引",
|
||
STAGE_COMPLETED: "已完成",
|
||
STAGE_FAILED: "失败",
|
||
}
|
||
|
||
|
||
async def _noop_progress(_percent: int, _stage: str) -> None:
|
||
pass
|
||
|
||
|
||
class FileProgressReporter:
|
||
"""节流上报，避免同一进度重复写库。"""
|
||
|
||
def __init__(self, callback: Optional[ProgressCallback] = None):
|
||
self._callback = callback or _noop_progress
|
||
self._last_percent = -1
|
||
|
||
async def report(self, percent: int, stage: str) -> None:
|
||
percent = max(0, min(100, int(percent)))
|
||
if percent <= self._last_percent:
|
||
return
|
||
self._last_percent = percent
|
||
await self._callback(percent, stage)
|
||
|
||
|
||
def make_db_progress_callback(
|
||
pool,
|
||
file_id: int,
|
||
service_cls: Type,
|
||
) -> ProgressCallback:
|
||
"""为知识库/聊天文件创建写库进度回调。"""
|
||
|
||
async def _callback(percent: int, stage: str) -> None:
|
||
async with pool.acquire() as conn:
|
||
await service_cls.update_file_progress(conn, file_id, percent, stage)
|
||
|
||
return _callback
|
||
|
||
|
||
def stage_label(stage: Optional[str]) -> str:
|
||
if not stage:
|
||
return "处理中"
|
||
return STAGE_LABELS.get(stage, stage)
|