第一次提交
This commit is contained in:
@@ -0,0 +1,166 @@
|
||||
"""上传图片 → OCR → 结构化提取 → 生成草稿题目。作为 BackgroundTask 运行。
|
||||
|
||||
注意：BackgroundTask 在请求返回后执行，需要自己开新的 DB session。
|
||||
"""
|
||||
import json
|
||||
import logging
|
||||
|
||||
from pydantic import ValidationError
|
||||
from sqlalchemy import select
|
||||
|
||||
from app.db import SessionLocal
|
||||
from app.models.image import Image
|
||||
from app.models.job import JobStatus, ProcessingJob
|
||||
from app.models.question import Question, QuestionType
|
||||
from app.models.user import UserSettings
|
||||
from app.schemas.question import QuestionCreate
|
||||
from app.services.llm.client import LLMClient, LLMConfigError
|
||||
from app.services.llm.prompts import EXTRACT_SYSTEM, EXTRACT_USER_TEMPLATE
|
||||
from app.services.ocr.pipeline import run_ocr
|
||||
from app.services.storage import download_image_async
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _parse_extract_json(raw: str) -> dict:
|
||||
"""从模型输出里解析 JSON（容忍代码块包裹）。"""
|
||||
text = raw.strip()
|
||||
if text.startswith("```"):
|
||||
text = text.split("```", 2)[1]
|
||||
if text.startswith("json"):
|
||||
text = text[4:]
|
||||
return json.loads(text.strip())
|
||||
|
||||
|
||||
async def run_extract_job(job_id: str, user_id: str, object_key: str, mime: str):
|
||||
async with SessionLocal() as db:
|
||||
job = await db.get(ProcessingJob, job_id)
|
||||
if job is None:
|
||||
logger.error("job %s 不存在", job_id)
|
||||
return
|
||||
try:
|
||||
# 加载用户 AI 配置
|
||||
settings_row = await db.scalar(
|
||||
select(UserSettings).where(UserSettings.user_id == user_id)
|
||||
)
|
||||
llm = LLMClient(settings_row)
|
||||
|
||||
# 1. OCR
|
||||
job.status = JobStatus.ocr_running
|
||||
await db.commit()
|
||||
image_bytes = await download_image_async(object_key)
|
||||
ocr = await run_ocr(image_bytes, mime, llm)
|
||||
|
||||
# 2. 结构化提取
|
||||
job.status = JobStatus.ai_running
|
||||
job.engine_used = ocr.engine
|
||||
await db.commit()
|
||||
extract = await llm.chat_text(
|
||||
EXTRACT_SYSTEM,
|
||||
EXTRACT_USER_TEMPLATE.format(content=ocr.markdown),
|
||||
json_mode=True,
|
||||
)
|
||||
questions = _build_questions(user_id, extract.content, ocr, job.image_id)
|
||||
|
||||
# 缓存原始 OCR 文本，便于日后追溯或重新拆分
|
||||
if job.image_id:
|
||||
image = await db.get(Image, job.image_id)
|
||||
if image is not None:
|
||||
image.ocr_markdown = ocr.markdown
|
||||
|
||||
db.add_all(questions)
|
||||
await db.flush()
|
||||
job.result_question_ids = [q.id for q in questions]
|
||||
job.status = JobStatus.done
|
||||
await db.commit()
|
||||
logger.info(
|
||||
"job %s 完成，engine=%s，拆出 %d 条笔记",
|
||||
job_id,
|
||||
ocr.engine,
|
||||
len(questions),
|
||||
)
|
||||
except LLMConfigError as e:
|
||||
await _fail(db, job, f"AI 未配置：{e}")
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.exception("job %s 失败", job_id)
|
||||
await _fail(db, job, str(e))
|
||||
|
||||
|
||||
def _build_questions(
|
||||
user_id: str, extract_raw: str, ocr, image_id: str | None
|
||||
) -> list[Question]:
|
||||
"""把提取结果拆成多条笔记。
|
||||
|
||||
一张照片常含多道题，每道题存成独立笔记。整体解析失败时不丢内容，
|
||||
降级为一条 unclassified 笔记（原文进题干），用户可自行整理。
|
||||
"""
|
||||
try:
|
||||
data = _parse_extract_json(extract_raw)
|
||||
raw_items = data.get("questions")
|
||||
if not isinstance(raw_items, list) or not raw_items:
|
||||
raise ValueError("提取结果里没有 questions 数组")
|
||||
except (json.JSONDecodeError, KeyError, ValueError) as e:
|
||||
logger.warning("提取结果解析失败，降级为一条未分类笔记：%s", e)
|
||||
return [_fallback_question(user_id, ocr, image_id)]
|
||||
|
||||
questions: list[Question] = []
|
||||
for idx, item in enumerate(raw_items):
|
||||
try:
|
||||
questions.append(_build_one(user_id, item, ocr, image_id))
|
||||
except (KeyError, ValueError, ValidationError, TypeError) as e:
|
||||
# 单条坏掉不影响其他条：这条退化为未分类，原文留在题干里
|
||||
logger.warning("第 %d 条解析失败，存为未分类：%s", idx + 1, e)
|
||||
stem = ""
|
||||
if isinstance(item, dict):
|
||||
stem = str(item.get("stem_markdown") or "")
|
||||
questions.append(
|
||||
_fallback_question(user_id, ocr, image_id, stem_override=stem or None)
|
||||
)
|
||||
|
||||
return questions or [_fallback_question(user_id, ocr, image_id)]
|
||||
|
||||
|
||||
def _build_one(user_id: str, item: dict, ocr, image_id: str | None) -> Question:
|
||||
parsed = QuestionCreate(
|
||||
# 模型没给或给了不认识的题型，就当未分类，别硬猜
|
||||
type=_coerce_type(item.get("type")),
|
||||
stem_markdown=item.get("stem_markdown") or ocr.markdown,
|
||||
options=item.get("options"),
|
||||
correct_answer=item.get("correct_answer"),
|
||||
difficulty=item.get("difficulty"),
|
||||
)
|
||||
return Question(
|
||||
user_id=user_id,
|
||||
type=parsed.type,
|
||||
stem_markdown=parsed.stem_markdown,
|
||||
options=[o.model_dump() for o in parsed.options] if parsed.options else None,
|
||||
correct_answer=parsed.correct_answer,
|
||||
difficulty=parsed.difficulty,
|
||||
source_image_id=image_id,
|
||||
ocr_engine=ocr.engine,
|
||||
)
|
||||
|
||||
|
||||
def _coerce_type(raw) -> QuestionType:
|
||||
try:
|
||||
return QuestionType(raw)
|
||||
except ValueError:
|
||||
return QuestionType.unclassified
|
||||
|
||||
|
||||
def _fallback_question(
|
||||
user_id: str, ocr, image_id: str | None, stem_override: str | None = None
|
||||
) -> Question:
|
||||
return Question(
|
||||
user_id=user_id,
|
||||
type=QuestionType.unclassified,
|
||||
stem_markdown=stem_override or ocr.markdown,
|
||||
source_image_id=image_id,
|
||||
ocr_engine=ocr.engine,
|
||||
)
|
||||
|
||||
|
||||
async def _fail(db, job: ProcessingJob, msg: str):
|
||||
job.status = JobStatus.failed
|
||||
job.error_message = msg[:1000]
|
||||
await db.commit()
|
||||
Reference in New Issue
Block a user