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