bid-doctor/server/src/bid_doctor/document/service.py

"""文档、解析任务（FR-M2-01、02，FR-M3-01）。"""

from __future__ import annotations

import json
from datetime import UTC, datetime
from typing import Any
from uuid import UUID

from sqlalchemy import func, select
from sqlalchemy.ext.asyncio import AsyncSession

from bid_doctor.audit import service as audit
from bid_doctor.billing import service as billing
from bid_doctor.core import keys, queue
from bid_doctor.core.crypto import open_object, seal_object, sha256
from bid_doctor.core.db import TenantContext
from bid_doctor.core.errors import AppError
from bid_doctor.core.ids import uuid7
from bid_doctor.core.storage import project_prefix, storage
from bid_doctor.document import citations, uploads
from bid_doctor.document.citations import document_label
from bid_doctor.document.models import Citation, ParseJob, SourceDocument
from bid_doctor.document.pipeline.text_layer import Page
from bid_doctor.project import service as projects

PARSABLE = {"tender", "addendum", "clarification"}
PARSE_TASK = "parse_document"


def doc_prefix(doc: SourceDocument) -> str:
    return f"{project_prefix(doc.tenant_id, doc.project_id)}/doc/{doc.id}"


async def create(
    s: AsyncSession, ctx: TenantContext, project_id: UUID, upload_id: UUID, role: str, volume_name: str | None
) -> SourceDocument:
    p = await projects.get(s, ctx, project_id)
    projects.ensure_editable(p)
    meta, data = await uploads.take(ctx, upload_id)
    assert meta.detected_format is not None
    seq = (
        await s.execute(
            select(func.coalesce(func.max(SourceDocument.seq_no), -1) + 1).where(
                SourceDocument.project_id == project_id, SourceDocument.role == role
            )
        )
    ).scalar_one()
    if role == "tender":
        seq = 0
    doc_id = uuid7()
    doc = SourceDocument(
        id=doc_id,
        tenant_id=ctx.tenant_id,
        project_id=project_id,
        role=role,
        seq_no=seq,
        volume_name=volume_name,
        lineage_id=doc_id,
        original_filename=meta.filename,
        detected_format=meta.detected_format,
        size_bytes=meta.size_bytes,
        original_key="",
        original_sha256=sha256(data),
        page_count=meta.page_count,
        uploaded_by=ctx.user_id,
    )
    doc.original_key = f"{doc_prefix(doc)}/original"
    dk = await keys.data_key(s, ctx, project_id)
    storage().put(doc.original_key, seal_object(dk, data, doc.original_key))
    s.add(doc)
    await audit.record(
        s,
        ctx,
        "document.create",
        "source_document",
        doc.id,
        project_id=project_id,
        after={"role": role, "filename": meta.filename},
    )
    await s.flush()
    await s.refresh(doc)
    await uploads.discard(ctx, upload_id)
    return doc


async def get(s: AsyncSession, document_id: UUID) -> SourceDocument:
    doc = (await s.execute(select(SourceDocument).where(SourceDocument.id == document_id))).scalar_one_or_none()
    if doc is None:
        raise AppError("not_found")
    return doc


async def start_parse(s: AsyncSession, ctx: TenantContext, document_id: UUID, expected_charge: int | None) -> ParseJob:
    """预扣 1 次解析 → 建任务 → 入队，在同一个事务里（架构 §3.4）。"""
    doc = await get(s, document_id)
    if doc.role not in PARSABLE:
        raise AppError("document_role_not_parsable")
    projects.ensure_editable(await projects.get(s, ctx, doc.project_id))
    if expected_charge is not None and expected_charge != 1:
        raise AppError("billing_changed", charge=1)  # 用户看到的扣点和实际不一致：拒绝，不扣点（API4）
    running = (
        await s.execute(
            select(ParseJob.id).where(ParseJob.document_id == doc.id, ParseJob.status.in_(("queued", "running")))
        )
    ).scalar_one_or_none()
    if running:
        raise AppError("job_already_running", job_id=str(running))
    attempt = (
        await s.execute(
            select(func.coalesce(func.max(ParseJob.attempt_no), 0) + 1).where(ParseJob.document_id == doc.id)
        )
    ).scalar_one()
    job = ParseJob(
        id=uuid7(),
        tenant_id=ctx.tenant_id,
        project_id=doc.project_id,
        document_id=doc.id,
        attempt_no=attempt,
        created_by=ctx.user_id,
    )
    s.add(job)
    await s.flush()
    job.quota_hold_id = await billing.hold(
        s, ctx, unit="parse", ref_type="parse_job", ref_id=job.id, project_id=doc.project_id
    )
    job.queue_job_id = await queue.enqueue(
        s, PARSE_TASK, {"tenant_id": str(ctx.tenant_id), "job_id": str(job.id)}, lock=f"doc:{doc.id}"
    )
    await audit.record(s, ctx, "parse_job.start", "parse_job", job.id, project_id=doc.project_id)
    await s.flush()
    await s.refresh(job)
    return job


async def request_cancel(s: AsyncSession, job_id: UUID) -> ParseJob:
    job = await get_job(s, job_id)
    if job.status not in ("queued", "running"):
        raise AppError("job_finished")
    job.cancel_requested_at = datetime.now(UTC)
    await s.flush()
    return job


async def get_job(s: AsyncSession, job_id: UUID) -> ParseJob:
    job = (await s.execute(select(ParseJob).where(ParseJob.id == job_id))).scalar_one_or_none()
    if job is None:
        raise AppError("not_found")
    return job


ERROR_MESSAGES = {
    "file_corrupted": "这份文件无法转换，可能已加密或损坏",
    "file_encrypted": "文件已加密，请上传未加密的版本",
    "ocr_required": "这份文件大部分页面没有文字层（可能是扫描件），扫描件识别还没有上线",
    "plan_limit_exceeded": "页数超出套餐上限",
    "converter_unavailable": "文档转换服务暂时不可用",
    "model_unavailable": "模型服务暂时不可用",
    "internal_error": "系统出错了",
}


def job_to_api(job: ParseJob) -> dict[str, Any]:
    error = None
    if job.status in ("failed", "cancelled") and job.error_code:
        error = {
            "code": job.error_code,
            "message": ERROR_MESSAGES.get(job.error_code, job.error_code),
            "refunded": True,
        }
    return {
        "id": job.id,
        "document_id": job.document_id,
        "attempt_no": job.attempt_no,
        "status": job.status,
        "stage": job.stage,
        "progress": job.progress,
        "cancel_requested": job.cancel_requested_at is not None,
        "error": error,
        "result": (job.checkpoints or {}).get("result"),
        "created_at": job.created_at,
        "finished_at": job.finished_at,
    }


async def doc_to_api(s: AsyncSession, doc: SourceDocument) -> dict[str, Any]:
    latest = (
        await s.execute(
            select(ParseJob).where(ParseJob.document_id == doc.id).order_by(ParseJob.attempt_no.desc()).limit(1)
        )
    ).scalar_one_or_none()
    return {
        "id": doc.id,
        "role": doc.role,
        "label": document_label(doc.role, doc.seq_no),
        "seq_no": doc.seq_no,
        "volume_name": doc.volume_name,
        "lineage_id": doc.lineage_id,
        "is_latest_version": True,
        "original_filename": doc.original_filename,
        "detected_format": doc.detected_format,
        "size_bytes": doc.size_bytes,
        "page_count": doc.page_count,
        "normalized_sha256": doc.normalized_sha256.hex() if doc.normalized_sha256 else None,
        "status": doc.status,
        "latest_parse_job": job_to_api(latest) if latest else None,
        "uploaded_by": {"id": doc.uploaded_by, "name": ""},
        "created_at": doc.created_at,
    }


async def summaries(s: AsyncSession, project_ids: list[UUID]) -> dict[UUID, dict[str, Any]]:
    """工作台用：每个项目的招标文件数、正在跑的解析。"""
    out: dict[UUID, dict[str, Any]] = {pid: {"tenders": 0, "job": None} for pid in project_ids}
    if not project_ids:
        return out
    for pid, n in (
        await s.execute(
            select(SourceDocument.project_id, func.count())
            .where(SourceDocument.project_id.in_(project_ids), SourceDocument.role == "tender")
            .group_by(SourceDocument.project_id)
        )
    ).tuples():
        out[pid]["tenders"] = n
    for job in (
        await s.execute(
            select(ParseJob).where(ParseJob.project_id.in_(project_ids), ParseJob.status.in_(("queued", "running")))
        )
    ).scalars():
        out[job.project_id]["job"] = {"kind": "parse_job", "id": job.id, "progress": job.progress, "label": "解析中"}
    return out


async def text_layer(s: AsyncSession, ctx: TenantContext, doc: SourceDocument) -> list[dict[str, Any]]:
    """文字层（逐页的句子和坐标框），加密存在对象存储里（DM1）。"""
    if not doc.text_layer_key:
        raise AppError("document_not_ready")
    dk = await keys.data_key(s, ctx, doc.project_id)
    return json.loads(open_object(dk, storage().get(doc.text_layer_key), doc.text_layer_key))


async def human_citation(
    s: AsyncSession, ctx: TenantContext, project_id: UUID, document_id: UUID, sentence_keys: list[str]
) -> Citation:
    """人工在阅读器里选中句子建出处（api-design §4.4）：句子必须在这份文档的同一页上。"""
    doc = await get(s, document_id)
    if doc.project_id != project_id:
        raise AppError("not_found")
    pages = [Page.from_json(p) for p in await text_layer(s, ctx, doc)]
    by_key = {x.key: x for p in pages for x in p.sentences}
    chosen = [by_key[k] for k in dict.fromkeys(sentence_keys) if k in by_key]
    if len(chosen) != len(set(sentence_keys)) or len({x.page_no for x in chosen}) != 1:
        raise AppError("sentence_not_found")
    (c,) = await citations.from_sentences(s, ctx, doc, chosen, "human", [])
    return c