vent_ai_platform/app/services/document_templates.py

from __future__ import annotations

from pathlib import Path

from sqlalchemy import select
from sqlalchemy.orm import Session

from app.core.errors import APIError
from app.db.models.document import DocumentExtractionTemplate


DOCUMENT_TYPE_BY_EXTENSION = {
    ".pdf": "pdf",
    ".doc": "word",
    ".docx": "word",
    ".xls": "excel",
    ".xlsx": "excel",
    ".csv": "excel",
}


DEFAULT_EXTRACTION_TEMPLATES = [
    {
        "code": "pdf_document_v1",
        "name": "PDF 文档关键信息",
        "document_type": "pdf",
        "description": "抽取 PDF 规程、报告、措施和验收资料中的关键业务信息。",
        "prompt": (
            "你是一名矿山通防业务资料抽取专家。请根据 PDF 转出的 Markdown，"
            "按照字段定义提取信息；找不到时填 null，不得编造。"
        ),
        "extraction_schema": {
            "fields": [
                {"key": "document_title", "name": "资料名称", "type": "text", "required": True},
                {"key": "document_code", "name": "资料编号", "type": "text", "required": False},
                {"key": "issuing_department", "name": "编制或发布单位", "type": "text", "required": False},
                {"key": "effective_date", "name": "生效日期", "type": "date", "required": False},
                {"key": "applicable_scope", "name": "适用范围", "type": "text", "required": False},
                {"key": "key_requirements", "name": "关键要求", "type": "list", "required": True},
                {"key": "key_parameters", "name": "关键参数", "type": "object", "required": False},
                {"key": "risk_points", "name": "风险点", "type": "list", "required": False},
                {"key": "conclusion", "name": "结论或执行要求", "type": "text", "required": False},
            ]
        },
    },
    {
        "code": "word_document_v1",
        "name": "Word 文档关键信息",
        "document_type": "word",
        "description": "抽取 Word 措施、制度、设计和审批资料中的关键业务信息。",
        "prompt": (
            "你是一名矿山通防业务资料抽取专家。请根据 Word 转出的 Markdown，"
            "重点识别编制依据、责任部门、执行要求、审批信息和时间节点；找不到时填 null。"
        ),
        "extraction_schema": {
            "fields": [
                {"key": "document_title", "name": "资料名称", "type": "text", "required": True},
                {"key": "document_type", "name": "资料类型", "type": "text", "required": False},
                {"key": "responsible_department", "name": "责任部门", "type": "text", "required": False},
                {"key": "applicable_area", "name": "适用区域", "type": "text", "required": False},
                {"key": "effective_date", "name": "执行日期", "type": "date", "required": False},
                {"key": "requirements", "name": "执行要求", "type": "list", "required": True},
                {"key": "approval_info", "name": "审批信息", "type": "object", "required": False},
                {"key": "attachments", "name": "附件清单", "type": "list", "required": False},
                {"key": "conclusion", "name": "结论或落款", "type": "text", "required": False},
            ]
        },
    },
    {
        "code": "excel_document_v1",
        "name": "Excel 表格关键信息",
        "document_type": "excel",
        "description": "抽取 Excel 报表、台账和测风记录中的指标、范围和异常信息。",
        "prompt": (
            "你是一名矿山通防业务数据抽取专家。请根据 Excel 转出的 Markdown 表格，"
            "识别表头含义、统计周期、关键指标和异常记录；找不到时填 null，数值保持原始单位。"
        ),
        "extraction_schema": {
            "fields": [
                {"key": "workbook_title", "name": "报表名称", "type": "text", "required": True},
                {"key": "sheet_names", "name": "工作表名称", "type": "list", "required": False},
                {"key": "date_range", "name": "统计周期", "type": "text", "required": False},
                {"key": "row_count", "name": "数据行数", "type": "number", "required": False},
                {"key": "key_metrics", "name": "关键指标", "type": "object", "required": True},
                {"key": "abnormal_rows", "name": "异常记录", "type": "list", "required": False},
                {"key": "summary", "name": "数据摘要", "type": "text", "required": False},
            ]
        },
    },
]


def detect_document_type(file_ext: str | None) -> str:
    extension = Path(file_ext or "").suffix.lower()
    if not extension:
        extension = f".{str(file_ext or '').strip().lstrip('.').lower()}"
    document_type = DOCUMENT_TYPE_BY_EXTENSION.get(extension)
    if document_type is None:
        supported = "、".join(sorted(DOCUMENT_TYPE_BY_EXTENSION))
        raise APIError(
            422,
            "unsupported_document_type",
            f"暂不支持该文档类型，可用扩展名：{supported}",
        )
    return document_type


def ensure_default_templates(db: Session) -> None:
    for item in DEFAULT_EXTRACTION_TEMPLATES:
        template = db.scalar(
            select(DocumentExtractionTemplate).where(
                DocumentExtractionTemplate.code == item["code"]
            )
        )
        if template is None:
            template = DocumentExtractionTemplate(**item)
            db.add(template)
        else:
            template.name = item["name"]
            template.document_type = item["document_type"]
            template.description = item["description"]
            template.prompt = item["prompt"]
            template.extraction_schema = item["extraction_schema"]
            template.enabled = True
    db.commit()


def get_template_for_type(
    db: Session,
    document_type: str,
) -> DocumentExtractionTemplate:
    template = db.scalar(
        select(DocumentExtractionTemplate)
        .where(
            DocumentExtractionTemplate.document_type == document_type,
            DocumentExtractionTemplate.enabled.is_(True),
        )
        .order_by(DocumentExtractionTemplate.updated_at.desc())
    )
    if template is None:
        raise APIError(
            422,
            "extraction_template_not_found",
            f"未配置 {document_type} 类型的抽取模板",
        )
    return template