219 lines
6.9 KiB
Python
219 lines
6.9 KiB
Python
"""PPT 解析后台任务。"""
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import threading
|
|
from datetime import datetime
|
|
|
|
from insurance.db.compat import db
|
|
from insurance.models.ppt_session import PptSession
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_local_locks: set[str] = set()
|
|
_local_locks_guard = threading.Lock()
|
|
_redis_locks: set[str] = set()
|
|
|
|
STALE_TASK_TIMEOUT = 600 # 10 分钟
|
|
|
|
|
|
def recover_stale_tasks():
|
|
"""启动时将长时间 parsing 的会话标记为 failed。"""
|
|
from datetime import timedelta
|
|
cutoff = datetime.now() - timedelta(seconds=STALE_TASK_TIMEOUT)
|
|
|
|
stale = PptSession.query.filter(
|
|
PptSession.status.in_(["parsing"]),
|
|
PptSession.created_at < cutoff,
|
|
).all()
|
|
for session in stale:
|
|
session.status = "error"
|
|
session.parse_error = "任务因服务重启而中断，请重新处理"
|
|
session.parse_finished_at = datetime.now()
|
|
logger.warning(f"恢复过期 PPT 解析任务: session_id={session.id}")
|
|
|
|
if stale:
|
|
db.session.commit()
|
|
logger.info(f"已恢复 {len(stale)} 个过期 PPT 解析任务")
|
|
|
|
|
|
def start_parse_task(app, session_id: str, user_id: str) -> bool:
|
|
"""启动后台解析任务，返回是否新启动。"""
|
|
if not _acquire_task_lock(session_id):
|
|
return False
|
|
|
|
thread = threading.Thread(
|
|
target=_run_parse_task,
|
|
args=(app, session_id, user_id),
|
|
daemon=True,
|
|
)
|
|
thread.start()
|
|
return True
|
|
|
|
|
|
def _run_parse_task(app, session_id: str, user_id: str):
|
|
with app.app_context():
|
|
try:
|
|
_execute_parse(session_id, user_id)
|
|
except Exception as exc:
|
|
logger.error(f"PPT 解析后台任务失败 [{session_id}]: {exc}", exc_info=True)
|
|
_mark_session_failed(session_id, str(exc))
|
|
finally:
|
|
_release_task_lock(session_id)
|
|
db.session.remove()
|
|
|
|
|
|
def _execute_parse(session_id: str, user_id: str):
|
|
from insurance.ppt.extraction import ExtractionOrchestrator
|
|
|
|
session = PptSession.query.filter_by(id=session_id, user_id=user_id).first()
|
|
if not session:
|
|
return
|
|
|
|
files = json.loads(session.files_json) if session.files_json else []
|
|
if not files:
|
|
_mark_session_failed(session_id, "没有可处理的 PDF 文件")
|
|
return
|
|
|
|
orchestrator = ExtractionOrchestrator()
|
|
extractions = []
|
|
total = len(files)
|
|
|
|
session.status = "parsing"
|
|
session.parse_progress = 0
|
|
session.parse_message = "数据结构化任务已启动"
|
|
session.parse_error = None
|
|
session.parse_started_at = datetime.now()
|
|
session.parse_finished_at = None
|
|
session.extractions_json = json.dumps([], ensure_ascii=False)
|
|
db.session.commit()
|
|
|
|
for index, file_info in enumerate(files, start=1):
|
|
filename = file_info.get("name", "")
|
|
filepath = file_info.get("path", "")
|
|
plan_type = file_info.get("type", "savings")
|
|
|
|
_update_progress(
|
|
session_id,
|
|
progress=_progress(index - 1, total),
|
|
message=f"正在处理 {filename or f'第 {index} 个文件'}",
|
|
extractions=extractions,
|
|
)
|
|
|
|
try:
|
|
result = asyncio.run(orchestrator.extract_plan(filepath, plan_type, force_reparse=True))
|
|
extractions.append(_build_extraction(file_info, filepath, result))
|
|
except Exception as exc:
|
|
logger.error(f"PDF 解析失败 [{filename}]: {exc}", exc_info=True)
|
|
extractions.append({
|
|
"pdfName": filename,
|
|
"pdfPath": filepath,
|
|
"planType": plan_type,
|
|
"status": "error",
|
|
"productName": "unknown",
|
|
"data": None,
|
|
"error": str(exc),
|
|
"yearCount": 0,
|
|
})
|
|
|
|
_update_progress(
|
|
session_id,
|
|
progress=_progress(index, total),
|
|
message=f"已完成 {index}/{total} 个文件",
|
|
extractions=extractions,
|
|
)
|
|
|
|
session = PptSession.query.filter_by(id=session_id, user_id=user_id).first()
|
|
if not session:
|
|
return
|
|
|
|
all_failed = all(e.get("status") == "error" for e in extractions)
|
|
partial_count = sum(1 for e in extractions if e.get("status") == "partial")
|
|
session.extractions_json = json.dumps(extractions, ensure_ascii=False)
|
|
session.status = "error" if all_failed else "parsed"
|
|
session.parse_progress = 100
|
|
if all_failed:
|
|
session.parse_message = "处理失败"
|
|
elif partial_count:
|
|
session.parse_message = f"处理完成，{partial_count} 个文件需补充数据"
|
|
else:
|
|
session.parse_message = "处理完成"
|
|
session.parse_error = "所有文件均处理失败" if all_failed else None
|
|
session.parse_finished_at = datetime.now()
|
|
db.session.commit()
|
|
|
|
|
|
def _build_extraction(file_info: dict, filepath: str, result) -> dict:
|
|
return {
|
|
"pdfName": file_info.get("name", ""),
|
|
"pdfPath": filepath,
|
|
"planType": result.plan_type,
|
|
"status": result.status,
|
|
"productName": result.product_name,
|
|
"data": result.data,
|
|
"error": result.error,
|
|
"yearCount": len(result.data.get("benefit_illustration", [])) if result.data else 0,
|
|
}
|
|
|
|
|
|
def _update_progress(session_id: str, progress: int, message: str, extractions: list[dict]):
|
|
session = PptSession.query.filter_by(id=session_id).first()
|
|
if not session:
|
|
return
|
|
session.status = "parsing"
|
|
session.parse_progress = progress
|
|
session.parse_message = message
|
|
session.extractions_json = json.dumps(extractions, ensure_ascii=False)
|
|
db.session.commit()
|
|
|
|
|
|
def _mark_session_failed(session_id: str, error: str):
|
|
session = PptSession.query.filter_by(id=session_id).first()
|
|
if not session:
|
|
return
|
|
session.status = "error"
|
|
session.parse_progress = 100
|
|
session.parse_message = "处理失败"
|
|
session.parse_error = error[:1000]
|
|
session.parse_finished_at = datetime.now()
|
|
db.session.commit()
|
|
|
|
|
|
def _progress(done: int, total: int) -> int:
|
|
if total <= 0:
|
|
return 0
|
|
return min(99, int(done / total * 100))
|
|
|
|
|
|
def _acquire_task_lock(session_id: str) -> bool:
|
|
redis_key = f"ppt_parse_lock:{session_id}"
|
|
try:
|
|
from insurance.db.compat import redis_client
|
|
if redis_client and redis_client.set(redis_key, "1", nx=True, ex=1800):
|
|
_redis_locks.add(session_id)
|
|
return True
|
|
if redis_client:
|
|
return False
|
|
except Exception:
|
|
pass
|
|
|
|
with _local_locks_guard:
|
|
if session_id in _local_locks:
|
|
return False
|
|
_local_locks.add(session_id)
|
|
return True
|
|
|
|
|
|
def _release_task_lock(session_id: str):
|
|
if session_id in _redis_locks:
|
|
try:
|
|
from insurance.db.compat import redis_client
|
|
if redis_client:
|
|
redis_client.delete(f"ppt_parse_lock:{session_id}")
|
|
except Exception:
|
|
pass
|
|
_redis_locks.discard(session_id)
|
|
|
|
with _local_locks_guard:
|
|
_local_locks.discard(session_id)
|