document-management/backend/app.py

"""Info Extractor Web 后端.

提供 Flask API 服务供 Vue 前端调用：
- 文档上传与文本提取（PDF / Word）
- 关键信息提取（流式返回思考过程 + 标注结果）
- 关键信息定义管理
"""

from __future__ import annotations

import json
import logging
import os
import queue
import re
import tempfile
import threading
import time
import uuid
from pathlib import Path
from typing import Any, Iterator

from flask import Flask, Response, jsonify, request
from flask_cors import CORS

from dms import init_dms

# 文档解析依赖
try:
    import fitz  # PyMuPDF
except ImportError:
    fitz = None

try:
    from docx import Document
except ImportError:
    Document = None

from src.extractor import InfoExtractor

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
)
logger = logging.getLogger(__name__)

app = Flask(__name__)
CORS(app)

# DMS业务使用独立的 /api/v1 Blueprint。初始化仅注册DMS基础设施，
# 不改变既有 /api/* AI接口的响应、认证或错误处理行为。
init_dms(app)

# 从环境变量读取后端配置
# 例如：INFO_EXTRACTOR_API_KEY=your-api-key INFO_EXTRACTOR_MODEL_NAME=glm-5.2 INFO_EXTRACTOR_BASE_URL=https://... python app.py
API_KEY = os.environ.get("INFO_EXTRACTOR_API_KEY", "")
MODEL_NAME = os.environ.get("INFO_EXTRACTOR_MODEL_NAME", "glm-5.2")
BASE_URL = os.environ.get("INFO_EXTRACTOR_BASE_URL", "https://open.bigmodel.cn/api/coding/paas/v4")

logger.info("Backend config: MODEL_NAME=%s, BASE_URL=%s, API_KEY_PREFIX=%s", MODEL_NAME, BASE_URL, API_KEY[:6] if API_KEY else "")

# 上传文件临时存储目录
UPLOAD_DIR = Path(tempfile.gettempdir()) / "info_extractor_uploads"
UPLOAD_DIR.mkdir(parents=True, exist_ok=True)

# 关键信息定义列表持久化文件
DATA_DIR = Path(__file__).parent / "data"
DATA_DIR.mkdir(parents=True, exist_ok=True)
DEFINITIONS_FILE = DATA_DIR / "definitions.json"

DEFAULT_DEFINITION: dict[str, Any] = {
    "metadata": {
        "node_label": "PDJH",
        "node_name": "派对计划",
        "node_type": "concept",
    },
    "subnodes": {
        "JBXX": {
            "metadata": {
                "node_label": "JBXX",
                "node_name": "基本信息",
                "node_type": "concept",
            },
            "subnodes": {
                "PDMC": {
                    "metadata": {
                        "node_label": "PDMC",
                        "node_name": "派对名称",
                        "node_type": "property",
                        "value_type": "string",
                        "is_list": False,
                    }
                },
                "PDKSSJ": {
                    "metadata": {
                        "node_label": "PDKSSJ",
                        "node_name": "派对开始时间",
                        "node_type": "property",
                        "value_type": "date_time",
                        "is_list": False,
                    }
                },
                "PDZRS": {
                    "metadata": {
                        "node_label": "PDZRS",
                        "node_name": "派对总人数",
                        "node_type": "property",
                        "value_type": "integer",
                        "is_list": False,
                    }
                },
                "PDZYS": {
                    "metadata": {
                        "node_label": "PDZYS",
                        "node_name": "派对总预算",
                        "node_type": "property",
                        "value_type": "float",
                        "is_list": False,
                    }
                },
                "SFBA": {
                    "metadata": {
                        "node_label": "SFBA",
                        "node_name": "是否备案",
                        "node_type": "property",
                        "value_type": "boolean",
                        "is_list": False,
                    }
                },
            },
        },
        "RYXX": {
            "metadata": {
                "node_label": "RYXX",
                "node_name": "人员信息",
                "node_type": "concept",
            },
            "subnodes": {
                "ZCH": {
                    "metadata": {
                        "node_label": "ZCH",
                        "node_name": "总策划",
                        "node_type": "property",
                        "value_type": "class",
                        "is_list": False,
                    },
                    "subnodes": {
                        "XM": {
                            "metadata": {
                                "node_label": "XM",
                                "node_name": "姓名",
                                "node_type": "property",
                                "value_type": "string",
                                "is_list": False,
                            }
                        },
                        "XB": {
                            "metadata": {
                                "node_label": "XB",
                                "node_name": "性别",
                                "node_type": "property",
                                "value_type": "string",
                                "is_list": False,
                            }
                        },
                    },
                },
                "RYQD": {
                    "metadata": {
                        "node_label": "RYQD",
                        "node_name": "人员清单",
                        "node_type": "property",
                        "value_type": "class",
                        "is_list": True,
                    },
                    "subnodes": {
                        "XM": {
                            "metadata": {
                                "node_label": "XM",
                                "node_name": "姓名",
                                "node_type": "property",
                                "value_type": "string",
                                "is_list": False,
                            }
                        },
                        "XB": {
                            "metadata": {
                                "node_label": "XB",
                                "node_name": "性别",
                                "node_type": "property",
                                "value_type": "string",
                                "is_list": False,
                            }
                        },
                    },
                },
            },
        },
        "HDXX": {
            "metadata": {
                "node_label": "HDXX",
                "node_name": "活动信息",
                "node_type": "concept",
            },
            "subnodes": {
                "PBYXQD": {
                    "metadata": {
                        "node_label": "PBYXQD",
                        "node_name": "破冰游戏清单",
                        "node_type": "property",
                        "value_type": "string",
                        "is_list": True,
                    }
                }
            },
        },
    },
}


def _load_definitions() -> list[dict[str, Any]]:
    """从 JSON 文件加载关键信息定义列表."""
    if not DEFINITIONS_FILE.exists():
        return []
    try:
        with open(DEFINITIONS_FILE, "r", encoding="utf-8") as f:
            data = json.load(f)
        if isinstance(data, list):
            return data
        if isinstance(data, dict):
            # 兼容旧版单条定义
            return [_wrap_definition(data)]
    except Exception as exc:
        logger.warning("加载定义文件失败: %s", exc)
    return []


def _save_definitions(definitions: list[dict[str, Any]]) -> None:
    """保存关键信息定义列表到 JSON 文件."""
    with open(DEFINITIONS_FILE, "w", encoding="utf-8") as f:
        json.dump(definitions, f, ensure_ascii=False, indent=2)


def _wrap_definition(definition: dict[str, Any], name: str = "", description: str = "") -> dict[str, Any]:
    """将原始定义包装为列表项."""
    now = time.strftime("%Y-%m-%d %H:%M:%S")
    metadata = definition.get("metadata", {})
    if not name:
        name = metadata.get("node_name", "未命名定义")
    return {
        "id": str(uuid.uuid4()),
        "name": name,
        "description": description,
        "definition": definition,
        "created_at": now,
        "updated_at": now,
    }


# 全局定义列表
_definitions: list[dict[str, Any]] = _load_definitions()
if not _definitions:
    _definitions = [_wrap_definition(DEFAULT_DEFINITION, name="派对计划", description="示例关键信息定义")]
    _save_definitions(_definitions)


def _extract_text_from_pdf(file_path: str) -> str:
    """从 PDF 中提取文本."""
    if fitz is None:
        raise RuntimeError("PyMuPDF 未安装，无法解析 PDF")

    text_parts: list[str] = []
    with fitz.open(file_path) as doc:
        for page in doc:
            text_parts.append(page.get_text())
    return "\n\n".join(text_parts)


def _extract_text_from_docx(file_path: str) -> str:
    """从 Word 文档中提取文本."""
    if Document is None:
        raise RuntimeError("python-docx 未安装，无法解析 Word 文档")

    doc = Document(file_path)
    text_parts: list[str] = []
    for para in doc.paragraphs:
        if para.text.strip():
            text_parts.append(para.text)
    return "\n\n".join(text_parts)


def _to_markdown(text: str) -> str:
    """简单地将正文文本转为 markdown 格式.

    当前只做基本分段和标题检测，未来可接入更强的文档解析模型。
    """
    lines = text.split("\n")
    result: list[str] = []
    for line in lines:
        stripped = line.strip()
        if not stripped:
            result.append("")
            continue
        # 简单启发式：短且没有标点的行可能是标题
        if len(stripped) < 30 and not re.search(r"[。，；：！？]", stripped):
            result.append(f"## {stripped}")
        else:
            result.append(stripped)
    return "\n\n".join(result)


def _clean_markdown_for_extraction(markdown: str) -> str:
    """将 markdown 还原为纯文本供 LLM 提取."""
    text = re.sub(r"^#+\s*", "", markdown, flags=re.MULTILINE)
    return text.strip()


@app.route("/api/health", methods=["GET"])
def health():
    return jsonify({"status": "ok"})


@app.route("/api/upload", methods=["POST"])
def upload_file():
    """上传 PDF 或 Word 文档，返回文档 ID 和提取的文本内容."""
    if "file" not in request.files:
        return jsonify({"error": "未提供文件"}), 400

    file = request.files["file"]
    if file.filename == "":
        return jsonify({"error": "文件名为空"}), 400

    ext = Path(file.filename).suffix.lower()
    if ext not in (".pdf", ".docx", ".doc"):
        return jsonify({"error": "仅支持 PDF 和 Word 文档"}), 400

    doc_id = str(uuid.uuid4())
    file_path = UPLOAD_DIR / f"{doc_id}{ext}"
    file.save(str(file_path))

    try:
        if ext == ".pdf":
            raw_text = _extract_text_from_pdf(str(file_path))
        else:
            raw_text = _extract_text_from_docx(str(file_path))
    except Exception as exc:
        logger.exception("文档解析失败")
        return jsonify({"error": f"文档解析失败: {exc}"}), 500

    markdown = _to_markdown(raw_text)

    return jsonify({
        "doc_id": doc_id,
        "filename": file.filename,
        "markdown": markdown,
        "raw_text": raw_text,
    })


@app.route("/api/definition", methods=["GET", "POST"])
def definition():
    """获取或保存关键信息定义列表."""
    global _definitions
    if request.method == "GET":
        return jsonify({"definitions": _definitions})

    data = request.get_json()
    if not isinstance(data, dict) or "definitions" not in data:
        return jsonify({"error": "非法的关键信息定义列表格式"}), 400
    definitions = data["definitions"]
    if not isinstance(definitions, list):
        return jsonify({"error": "definitions 必须是数组"}), 400
    now = time.strftime("%Y-%m-%d %H:%M:%S")
    for item in definitions:
        if not isinstance(item, dict) or "definition" not in item:
            return jsonify({"error": "列表项必须包含 definition 字段"}), 400
        item.setdefault("id", str(uuid.uuid4()))
        item.setdefault("name", "未命名定义")
        item.setdefault("description", "")
        item["updated_at"] = now
    _definitions = definitions
    _save_definitions(_definitions)
    return jsonify({"success": True})


@app.route("/api/definition/import", methods=["POST"])
def import_definition():
    """导入 JSON 格式定义，自动提取名称和描述."""
    data = request.get_json()
    if not isinstance(data, dict):
        return jsonify({"error": "请求体必须是 JSON 对象"}), 400

    definition = data.get("definition")
    if not isinstance(definition, dict) or "metadata" not in definition:
        return jsonify({"error": "非法的关键信息定义格式"}), 400

    name = data.get("name", "").strip()
    description = data.get("description", "").strip()
    item = _wrap_definition(definition, name=name, description=description)
    _definitions.append(item)
    _save_definitions(_definitions)
    return jsonify({"success": True, "definition": item})


@app.route("/api/extract", methods=["POST"])
def extract():
    """流式提取关键信息.

    前端通过 SSE 接收：
    - type=thinking: 模型思考过程/流式输出片段
    - type=annotated: 最终标注后的文本
    - type=done: 完成
    - type=error: 错误
    """
    data = request.get_json()
    article = data.get("article", "")
    mode = data.get("mode", "two_phase")
    model_name = data.get("model_name", MODEL_NAME)
    base_url = data.get("base_url", BASE_URL)
    logger.info("Extract request: model_name=%s, base_url=%s", model_name, base_url)

    if not article:
        return jsonify({"error": "文章内容不能为空"}), 400

    definition_data = data.get("definition")
    definition_id = data.get("definition_id")
    if definition_data is None:
        if definition_id:
            for item in _definitions:
                if item.get("id") == definition_id:
                    definition_data = item.get("definition")
                    break
            if definition_data is None:
                return jsonify({"error": "未找到指定的关键信息定义"}), 400
        else:
            if not _definitions:
                return jsonify({"error": "未配置关键信息定义"}), 400
            definition_data = _definitions[0].get("definition")

    def generate() -> Iterator[str]:
        if not API_KEY:
            yield f"data: {json.dumps({'type': 'error', 'message': '后端未配置 API Key'}, ensure_ascii=False)}\n\n"
            return

        thinking_queue: queue.Queue[Any] = queue.Queue()
        extract_result: dict[str, Any] = {}
        extract_error: list[Exception] = []

        def run_extraction() -> None:
            try:
                extractor = InfoExtractor(
                    model_name=model_name,
                    api_key=API_KEY,
                    base_url=base_url or None,
                    temperature=0,
                    max_retries=1,
                )
                extractor.stream_callback = thinking_queue.put

                def on_step(label: str) -> None:
                    thinking_queue.put(("__step__", label))

                extractor.step_callback = on_step
                result = extractor.extract(
                    article=article,
                    definition=definition_data,
                    mode=mode,
                    verbose=True,
                    annotate=True,
                )
                extract_result["data"] = result
            except Exception as exc:
                logger.exception("提取失败")
                extract_error.append(exc)

        extraction_thread = threading.Thread(target=run_extraction)
        extraction_thread.start()

        # 实时从队列读取思考过程并推送
        while extraction_thread.is_alive() or not thinking_queue.empty():
            try:
                chunk = thinking_queue.get(timeout=0.1)
            except queue.Empty:
                continue
            if isinstance(chunk, tuple) and chunk and chunk[0] == "__step__":
                yield f"data: {json.dumps({'type': 'step', 'content': chunk[1]}, ensure_ascii=False)}\n\n"
            else:
                yield f"data: {json.dumps({'type': 'thinking', 'content': chunk}, ensure_ascii=False)}\n\n"

        extraction_thread.join()

        if extract_error:
            yield f"data: {json.dumps({'type': 'error', 'message': str(extract_error[0])}, ensure_ascii=False)}\n\n"
            return

        result = extract_result.get("data", {})
        annotated_text = result.get("annotated_text", "") if isinstance(result, dict) else ""
        yield f"data: {json.dumps({'type': 'annotated', 'content': annotated_text}, ensure_ascii=False)}\n\n"
        yield f"data: {json.dumps({'type': 'done'}, ensure_ascii=False)}\n\n"

    response = Response(generate(), mimetype="text/event-stream")
    response.headers["Access-Control-Allow-Origin"] = "*"
    return response


# 标注结果持久化目录
ANNOTATION_DIR = DATA_DIR / "annotations"
ANNOTATION_DIR.mkdir(parents=True, exist_ok=True)


def _annotation_path(doc_key: str) -> Path:
    """根据前端传入的 doc_key 生成本地存储路径.

    doc_key 由前端生成（如 doc_id + definition_id），仅允许字母数字与下划线连字符，
    避免路径穿越风险。
    """
    safe = re.sub(r"[^A-Za-z0-9_\-.]", "_", doc_key)
    if not safe:
        safe = "default"
    return ANNOTATION_DIR / f"{safe}.json"


@app.route("/api/annotation", methods=["GET", "POST"])
def annotation():
    """保存或读取标注后的文章内容.

    GET  ?doc_key=xxx  ->  {"content": "..."}
    POST {doc_key, content}  ->  {"success": true}
    """
    if request.method == "GET":
        doc_key = request.args.get("doc_key", "")
        path = _annotation_path(doc_key)
        if not path.exists():
            return jsonify({"content": ""})
        try:
            with open(path, "r", encoding="utf-8") as f:
                data = json.load(f)
            return jsonify({"content": data.get("content", "")})
        except Exception as exc:
            logger.warning("读取标注失败: %s", exc)
            return jsonify({"content": ""})

    data = request.get_json() or {}
    doc_key = data.get("doc_key", "")
    content = data.get("content", "")
    if not doc_key:
        return jsonify({"error": "缺少 doc_key"}), 400
    path = _annotation_path(doc_key)
    try:
        with open(path, "w", encoding="utf-8") as f:
            json.dump({"content": content, "updated_at": time.strftime("%Y-%m-%d %H:%M:%S")}, f, ensure_ascii=False, indent=2)
    except Exception as exc:
        logger.exception("保存标注失败")
        return jsonify({"error": f"保存失败: {exc}"}), 500
    return jsonify({"success": True})


if __name__ == "__main__":
    app.run(
        host=os.environ.get("DMS_HOST", "0.0.0.0"),
        port=int(os.environ.get("DMS_PORT", "8755")),
        debug=False,
    )