feat: 新增多格式文件上传入库与认证体系
- 新增 JWT 认证模块,支持登录/注册/用户管理 - 新增文件上传接口,支持 .txt/.md/.html/.pdf/.docx 等格式解析入库 - 新增检索结果 AI 总结功能 - 新增文本去重缓存机制 - 新增全局认证夹具简化测试 - 新增配置项与环境变量支持 - 完善文档与测试覆盖
This commit is contained in:
@@ -9,6 +9,7 @@ Redis 不可用或写入失败仅记录 warning,不影响任务执行。
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import uuid
|
||||
from datetime import UTC, datetime
|
||||
from enum import StrEnum
|
||||
@@ -25,6 +26,8 @@ logger = structlog.get_logger()
|
||||
|
||||
# Redis 任务状态 key 前缀
|
||||
REDIS_KEY_PREFIX = "ingest_task:"
|
||||
# Redis 文本去重 key 前缀(value 为已入库文档的 IngestionResult JSON)
|
||||
DEDUP_KEY_PREFIX = "dedup:sha256:"
|
||||
|
||||
|
||||
class IngestTaskStatus(StrEnum):
|
||||
@@ -62,9 +65,39 @@ class IngestTaskManager:
|
||||
self._mirror_tasks: set[asyncio.Task[None]] = set()
|
||||
|
||||
async def submit(self, doc: DocumentInput) -> str:
|
||||
"""登记入库任务并后台执行,立即返回 task_id"""
|
||||
"""登记入库任务并后台执行,立即返回 task_id
|
||||
|
||||
文本去重:基于 doc.text 的 sha256 在 Redis 中查重;命中则直接复用旧
|
||||
IngestionResult(仅置 deduplicated=True),不重跑流水线;未命中走原
|
||||
异步入库流程,完成后写入去重记录供后续命中复用。Redis 不可用时跳过
|
||||
去重,按原流程执行,不影响主流程。
|
||||
"""
|
||||
task_id = uuid.uuid4().hex
|
||||
now = _utc_now_iso()
|
||||
text_hash = hashlib.sha256(doc.text.encode("utf-8")).hexdigest()
|
||||
|
||||
# 1. 去重命中:直接置 done,复用旧结果,不调 _run
|
||||
dedup_record = await self._lookup_dedup(text_hash)
|
||||
if dedup_record is not None:
|
||||
result_dict = dict(dedup_record)
|
||||
result_dict["deduplicated"] = True
|
||||
self._tasks[task_id] = {
|
||||
"task_id": task_id,
|
||||
"status": IngestTaskStatus.DONE,
|
||||
"created_at": now,
|
||||
"updated_at": now,
|
||||
"result": result_dict,
|
||||
"error": None,
|
||||
}
|
||||
self._schedule_mirror(task_id)
|
||||
logger.info(
|
||||
"入库任务命中去重,复用既有文档",
|
||||
task_id=task_id,
|
||||
document_id=result_dict.get("document_id"),
|
||||
)
|
||||
return task_id
|
||||
|
||||
# 2. 未命中:登记 pending 并后台跑流水线
|
||||
self._tasks[task_id] = {
|
||||
"task_id": task_id,
|
||||
"status": IngestTaskStatus.PENDING,
|
||||
@@ -74,7 +107,7 @@ class IngestTaskManager:
|
||||
"error": None,
|
||||
}
|
||||
self._schedule_mirror(task_id)
|
||||
background = asyncio.create_task(self._run(task_id, doc))
|
||||
background = asyncio.create_task(self._run(task_id, doc, text_hash))
|
||||
self._background_tasks.add(background)
|
||||
background.add_done_callback(self._background_tasks.discard)
|
||||
logger.info("入库任务已登记", task_id=task_id, title=doc.title)
|
||||
@@ -110,8 +143,8 @@ class IngestTaskManager:
|
||||
raise TimeoutError(f"入库任务 {task_id} 在 {timeout}s 内未进入终态")
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
async def _run(self, task_id: str, doc: DocumentInput) -> None:
|
||||
"""后台执行入库:并发限流 + 阶段状态推进 + 结果/错误落账"""
|
||||
async def _run(self, task_id: str, doc: DocumentInput, text_hash: str) -> None:
|
||||
"""后台执行入库:并发限流 + 阶段状态推进 + 结果/错误落账 + 去重记录写入"""
|
||||
async with self._semaphore:
|
||||
try:
|
||||
result = await self._ingester.ingest(doc, progress_cb=lambda stage: self._on_progress(task_id, stage))
|
||||
@@ -128,14 +161,39 @@ class IngestTaskManager:
|
||||
except Exception as exc:
|
||||
self._finish_failed(task_id, {"stage": "unknown", "message": str(exc), "partial_summary": None})
|
||||
else:
|
||||
result_dict = result.model_dump(mode="json")
|
||||
self._tasks[task_id].update(
|
||||
status=IngestTaskStatus.DONE,
|
||||
updated_at=_utc_now_iso(),
|
||||
result=result.model_dump(mode="json"),
|
||||
result=result_dict,
|
||||
)
|
||||
self._schedule_mirror(task_id)
|
||||
await self._record_dedup(text_hash, result_dict)
|
||||
logger.info("入库任务完成", task_id=task_id)
|
||||
|
||||
async def _lookup_dedup(self, text_hash: str) -> dict[str, Any] | None:
|
||||
"""查询文本去重记录;Redis 不可用或异常时降级为未命中"""
|
||||
if self._redis is None:
|
||||
return None
|
||||
try:
|
||||
return await self._redis.get_json(f"{DEDUP_KEY_PREFIX}{text_hash}")
|
||||
except Exception:
|
||||
logger.warning("去重记录查询失败,降级为未命中", text_hash=text_hash, exc_info=True)
|
||||
return None
|
||||
|
||||
async def _record_dedup(self, text_hash: str, result_dict: dict[str, Any]) -> None:
|
||||
"""写入文本去重记录(含完整 IngestionResult),供后续命中复用;失败仅告警"""
|
||||
if self._redis is None:
|
||||
return
|
||||
try:
|
||||
await self._redis.set_json(
|
||||
f"{DEDUP_KEY_PREFIX}{text_hash}",
|
||||
result_dict,
|
||||
ttl=self._settings.ingest_task_ttl_done,
|
||||
)
|
||||
except Exception:
|
||||
logger.warning("去重记录写入失败", text_hash=text_hash, exc_info=True)
|
||||
|
||||
def _finish_failed(self, task_id: str, error: dict[str, Any]) -> None:
|
||||
"""将任务置为 failed 并记录错误信息"""
|
||||
self._tasks[task_id].update(status=IngestTaskStatus.FAILED, updated_at=_utc_now_iso(), error=error)
|
||||
|
||||
Reference in New Issue
Block a user