deepdata/my_ops/markitdown_parser/process.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
MarkItDown 多格式文档解析算子
基于 Microsoft MarkItDown 库，将办公文档统一转换为 Markdown 文本。
支持格式：docx, doc, xlsx, xls, pptx, ppt, txt, csv, xml, html, htm

⚠️ 图片处理优先级 (智能回退机制)：
1. 【优先上传】开启 upload_images 时，上传图片到外部 API，返回可访问的网络链接
2. 【Base64 内联】开启 embed_images_base64 时，转为 Base64 嵌入 Markdown
3. 【本地存储】上述均关闭时，保存到本地 images 目录
4. 【保留原样】下载/上传失败，或同源图片保持原链接
"""
import os
import sys
import time
import base64
import re
import hashlib
import mimetypes
import requests
import subprocess
import importlib
from io import BytesIO
from pathlib import Path
from typing import Dict, Any, List
from urllib.parse import urlparse
from PIL import Image

from loguru import logger

try:
    from deepdata.core.base_op import Mapper
except ImportError:
    print("⚠️ 本地调试模式：未找到 'deepdata' 库，正在启用模拟 (Mock) Mapper...")
    class Mapper:
        def __init__(self, *args, **kwargs):
            self.text_key = kwargs.get("text_key", "text")
            self.data_key = kwargs.get("data_key", "data")
            self.filename_key = kwargs.get("fileName_key", "fileName")
            self.filetype_key = kwargs.get("fileType_key", "fileType")
            self.filepath_key = kwargs.get("filePath_key", "filePath")
            self.export_path_key = kwargs.get("export_path_key", "export_path")

SUPPORTED_FORMATS = {
    ".docx", ".doc",
    ".xlsx", ".xls",
    ".pptx", ".ppt",
    ".txt", ".csv", ".xml",
    #".html", ".htm",
}

IMG_LINK_REGEX = re.compile(r'!\[([^\]]*)\]\((.*?)\)')

def ensure_import(module_name: str, package_name: str = None):
    """
    安全导入模块，若不存在则尝试自动安装。
    Args:
        module_name: Python import 名 (如 'docx', 'pptx')
        package_name: Pip 安装名 (如 'python-docx'), 默认与 module_name 相同
    """
    pkg = package_name or module_name
    try:
        return importlib.import_module(module_name)
    except ImportError:
        logger.warning(f"⚠️ 模块 '{module_name}' 未找到，正在尝试自动安装 '{pkg}'...")
        try:
            subprocess.check_call([sys.executable, "-m", "pip", "install", pkg])
            logger.info(f"✅ 成功安装 '{pkg}'")
            return importlib.import_module(module_name)
        except Exception as e:
            logger.error(f"❌ 自动安装 '{pkg}' 失败: {e}")
            raise ImportError(f"缺少依赖包 '{pkg}' 且无法自动安装")

def extract_docx_images_base64(docx_file: str, max_width: int = 1200) -> List[str]:
    """从 docx 提取图片列表 (按顺序)"""
    docx = ensure_import("docx", "python-docx")
    Document = docx.Document
    
    b64_list = []
    doc = Document(docx_file)
    for rel in doc.part.rels.values():
        if "image" in rel.target_ref:
            try:
                img_bytes = rel.target_part.blob
                img = Image.open(BytesIO(img_bytes))
                _resize_and_append(img, b64_list, max_width)
            except: continue
    return b64_list

def extract_pptx_images_base64(pptx_file: str, max_width: int = 1200) -> List[str]:
    """从 pptx 提取图片列表 (按幻灯片顺序)"""
    pptx = ensure_import("pptx", "python-pptx")
    Presentation = pptx.Presentation

    prs = Presentation(pptx_file)
    b64_list = []
    for slide in prs.slides:
        for shape in slide.shapes:
            if hasattr(shape, "image") and shape.image is not None:
                try:
                    img_bytes = shape.image.blob
                    img = Image.open(BytesIO(img_bytes))
                    _resize_and_append(img, b64_list, max_width)
                except: continue
    return b64_list

def _resize_and_append(img: Image.Image, target_list: list, max_width: int):
    if img.width > max_width:
        ratio = max_width / img.width
        new_size = (max_width, int(img.height * ratio))
        img = img.resize(new_size, Image.Resampling.LANCZOS)
    buf = BytesIO()
    img.save(buf, format="PNG", optimize=True)
    b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
    target_list.append(f"data:image/png;base64,{b64}")

def download_image(url: str) -> bytes:
    headers = {"User-Agent": "Mozilla/5.0 Chrome/120.0"}
    resp = requests.get(url, headers=headers, timeout=30)
    resp.raise_for_status()
    return resp.content

class MarkItDownParser(Mapper):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        raw_formats = kwargs.get("doc_formats", [])
        if isinstance(raw_formats, str):
            self.doc_formats = [f.strip().lower().lstrip(".") for f in raw_formats.split(",") if f.strip()]
        elif isinstance(raw_formats, list):
            self.doc_formats = [str(f).lower().lstrip(".") for f in raw_formats]
        else:
            self.doc_formats = [f.lstrip(".") for f in SUPPORTED_FORMATS]

        self.embed_images_base64 = kwargs.get("embed_images_base64", "true") == "true"
        self.download_remote_images = kwargs.get("download_remote_images", "true") == "true"
        self.max_image_width = int(kwargs.get("max_image_width", 1200))
        
        # 🆕 上传到外部图床配置
        self.upload_images = kwargs.get("upload_images", "false") == "true"
        self.upload_api_url = kwargs.get("upload_api_url", "").strip()
        self.preview_base_url = kwargs.get("preview_base_url", "").rstrip("/")
        
        print(f"self.upload_images={self.upload_images}")
            
        self._converter = None

    def _get_converter(self):
        if self._converter is None:
            markitdown = ensure_import("markitdown")
            MarkItDown = markitdown.MarkItDown
            self._converter = MarkItDown()
        return self._converter

    def _is_target_format(self, filename: str) -> bool:
        ext = Path(filename).suffix.lower()
        return ext.lstrip(".") in self.doc_formats if ext else False

    def _load_office_image_cache(self, filepath: str, file_ext: str) -> List[str]:
        """根据文件类型加载图片缓存列表"""
        logger.info("📷 正在预加载 Office 内部图片...")
        if file_ext == ".docx":
            return extract_docx_images_base64(filepath, self.max_image_width)
        elif file_ext in (".pptx", ".ppt"):
            return extract_pptx_images_base64(filepath, self.max_image_width)
        return []

    def _upload_image_to_api(self, img_bytes: bytes, filename: str) -> str:
        """上传图片到外部 API 并返回完整预览 URL"""
        if not self.upload_images or not self.upload_api_url:
            return ""
            
        ct = mimetypes.guess_type(filename)[0] or "image/jpeg"
        files = {"file": (filename, BytesIO(img_bytes), ct)}
        
        try:
            resp = requests.post(self.upload_api_url, files=files, timeout=30)
            resp.raise_for_status()
            
            try:
                data = resp.json()
                raw_path = data.get("path") or data.get("url") or data.get("data")
                path = str(raw_path) if raw_path else ""
            except:
                path = resp.text.strip()
            
            # 🛡️ 强力清洗路径，防止双斜杠 //
            clean_path = path.strip('\"\' ').lstrip('/')
            clean_base = self.preview_base_url.strip(' /')
            
            if not clean_base or not clean_path:
                raise ValueError("Base URL 或路径解析为空")
            
            return f"{clean_base}/{clean_path}"
        except Exception as e:
            logger.warning(f"⚠️ 图片上传失败 [{filename}]: {e}")
            return ""

    def _is_same_origin(self, url1: str, url2: str) -> bool:
        """判断两个 URL 是否同源"""
        if not url1 or not url2:
            return False
        try:
            parsed1 = urlparse(url1)
            parsed2 = urlparse(url2)
            return (parsed1.scheme == parsed2.scheme and 
                    parsed1.hostname == parsed2.hostname and 
                    parsed1.port == parsed2.port)
        except Exception:
            return False

    def _process_all_images(self, md_content: str, export_path: str, source_dir: str, image_cache: List[str]) -> str:
        """统一处理 Markdown 中的所有图片链接"""
        image_exts = {'.png', '.jpg', '.jpeg', '.gif', '.bmp', '.webp', '.tiff'}

        def replace_func(match):
            alt = match.group(1)
            link = match.group(2).strip()

            # 1. 处理占位符 (常见于 DOCX，从缓存取)
            if link.startswith("data:image/png;base64") and "..." in link:
                if image_cache:
                    b64_data = image_cache.pop(0)
                    return self._handle_image_bytes(b64_data, alt, link, export_path)
                return match.group(0)

            # 2. 处理已经是真实 Base64 的情况
            if link.startswith("data:image/") and "..." not in link:
                if self.upload_images and self.upload_api_url:
                    try:
                        if "," in link:
                            _, b64_data = link.split(",", 1)
                        else:
                            b64_data = link
                        img_bytes = base64.b64decode(b64_data)
                        upload_url = self._upload_image_to_api(img_bytes, "embedded_image.png")
                        if upload_url:
                            return f"![{alt}]({upload_url})"
                        else:
                            logger.warning(f"⚠️ Base64 图片上传失败，保留原 Base64: {link[:50]}...")
                    except Exception as e:
                        logger.error(f"⚠️ Base64 图片处理异常: {e}")
                return match.group(0) 
            
            # 3. 处理远程链接
            if link.startswith("http"):
                if not self.download_remote_images:
                    if self.upload_images and self.preview_base_url and self._is_same_origin(link, self.preview_base_url):
                        return match.group(0)
                    return match.group(0)
                    
                try:
                    img_bytes = download_image(link)
                    if self.upload_images and self.upload_api_url:
                        upload_url = self._upload_image_to_api(img_bytes, Path(urlparse(link).path).name or "remote.jpg")
                        if upload_url:
                            return f"![{alt}]({upload_url})"
                        else:
                            logger.warning(f"⚠️ 远程图片上传失败 [{link}]，尝试下载或保留")
                    return self._handle_image_bytes(img_bytes, alt, link, export_path)
                except Exception as e:
                    logger.warning(f"⚠️ 下载失败 [{link}]: {e}")
                    return match.group(0)
            
            # 4. 处理本地/相对路径
            possible_paths = [
                os.path.join(source_dir, link),
                os.path.join(export_path, link),
                os.path.join(source_dir, "media", link)
            ]
            
            found_bytes = None
            for p in possible_paths:
                if os.path.exists(p):
                    with open(p, "rb") as f: found_bytes = f.read()
                    break
            
            if found_bytes:
                return self._handle_image_bytes(found_bytes, alt, link, export_path)
            
            # 5. 【智能回退】文件不在磁盘上
            link_ext = Path(urlparse(link).path).suffix.lower()
            if link_ext in image_exts:
                if image_cache:
                    logger.debug(f"⚠️ 文件不存在 ({link})，从缓存池中获取图片替换")
                    b64_data = image_cache.pop(0)
                    return self._handle_image_bytes(b64_data, alt, link, export_path)
            
            return match.group(0)

        return IMG_LINK_REGEX.sub(replace_func, md_content)

    def _handle_image_bytes(self, img_data: Any, alt: str, original_ref: str, export_path: str) -> str:
        """处理图片二进制或 Base64，根据配置转为上传链接/Base64/本地文件"""
        
        # 0. 优先处理：解析图片数据为 bytes
        img_bytes = None
        try:
            if isinstance(img_data, str) and img_data.startswith("data:"):
                if "," in img_data:
                    _, b64_data = img_data.split(",", 1)
                else:
                    b64_data = img_data
                img_bytes = base64.b64decode(b64_data)
            elif isinstance(img_data, bytes):
                img_bytes = img_data
        except Exception as e:
            logger.warning(f"⚠️ 图片数据解析失败: {e}")
            return f"![{alt}]({original_ref})"

        if not img_bytes:
            return f"![{alt}]({original_ref})"

        # 🟢 优先级 1: 尝试上传到外部 API
        if self.upload_images and self.upload_api_url:
            # 尝试生成一个合理的文件名
            filename = "image.png"
            if isinstance(original_ref, str) and not original_ref.startswith("data:"):
                filename = Path(original_ref).name or "image.png"
            
            upload_url = self._upload_image_to_api(img_bytes, filename)
            if upload_url:
                return f"![{alt}]({upload_url})"
            logger.warning(f"⚠️ 图片上传失败 (ref: {original_ref[:30]})，回退到 Base64 或本地")

        # 🟡 优先级 2: Base64 内联
        if self.embed_images_base64:
            try:
                img = Image.open(BytesIO(img_bytes))
                if img.width > self.max_image_width:
                    ratio = self.max_image_width / img.width
                    img = img.resize((self.max_image_width, int(img.height * ratio)), Image.Resampling.LANCZOS)
                buf = BytesIO()
                img.save(buf, format="PNG")
                b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
                mime = mimetypes.guess_type(original_ref)[0] or "image/png"
                return f"![{alt}](data:{mime};base64,{b64})"
            except Exception as e:
                logger.warning(f"⚠️ Base64 转换失败: {e}")
        
        # 🔴 优先级 3: 本地存储
        else:
            try:
                ext = mimetypes.guess_extension(mimetypes.guess_type(original_ref)[0] or 'image/png') or '.jpg'
                fname = hashlib.md5(original_ref.encode()).hexdigest()[:12] + ext
                save_rel = os.path.join("images", fname)
                full_save = os.path.join(export_path, save_rel)
                os.makedirs(os.path.dirname(full_save), exist_ok=True)
                with open(full_save, "wb") as f:
                    f.write(img_bytes)
                return f"![{alt}]({save_rel})"
            except Exception as e:
                logger.warning(f"⚠️ 本地存储失败: {e}")

        # 最后的兜底
        return f"![{alt}]({original_ref})"
    
    def execute(self, sample: Dict[str, Any]) -> Dict[str, Any]:
        filename = sample.get(self.filename_key, "")
        filepath = sample.get(self.filepath_key, "")
        export_path = sample.get(self.export_path_key, "./output")
        file_ext = Path(filename).suffix.lower()
        
        if not self._is_target_format(filename):
            logger.info(f"⏭️ 跳过 {filename} (不在处理列表中)")
            return sample
        if not filepath or not os.path.exists(filepath):
            logger.warning(f"⚠️ 文件不存在: {filepath}")
            return sample

        start_time = time.time()
        try:
            logger.info(f"🚀 开始解析: {filename}")
            result = self._get_converter().convert(filepath)
            md_content = result.text_content
            if not md_content:
                return sample

            # 加载 Office 图片缓存 (如果是 Office 格式)
            image_cache = []
            if file_ext in (".docx", ".pptx", ".ppt"):
                image_cache = self._load_office_image_cache(filepath, file_ext)
                if image_cache:
                    logger.info(f"✅ 缓存中加载到 {len(image_cache)} 张图片")

            # 统一处理所有图片
            md_content = self._process_all_images(md_content, export_path, os.path.dirname(filepath), image_cache)

            sample[self.text_key] = md_content
            sample["target_type"] = "md"
            logger.info(f"✅ 解析成功: {filename} | 耗时 {time.time() - start_time:.2f}s")
        except Exception as e:
            logger.exception(f"❌ 解析失败 {filename}: {e}")
            raise
        return sample

# ==========================================================
# 🧪 本地调试与测试入口
# ==========================================================
if __name__ == "__main__":
    print("=" * 60)
    print("🧪 MarkItDownParser 本地测试")
    print("=" * 60)
    current_dir = os.path.dirname(os.path.abspath(__file__))
    test_files = []
    for ext in SUPPORTED_FORMATS:
        test_files.extend(Path(current_dir).glob(f"*{ext}"))
    if not test_files:
        print(f"❌ 找不到测试文件")
        sys.exit(1)
    test_files = sorted(set(test_files))
    print(f"📂 找到 {len(test_files)} 个测试文件:")
    for f in test_files: print(f"   📄 {f.name}")

    output_dir = os.path.join(current_dir, "test_output")
    os.makedirs(output_dir, exist_ok=True)
    mock_kwargs = {
        "doc_formats": list(SUPPORTED_FORMATS),
        "embed_images_base64": "true",
        "download_remote_images": "true",
        "max_image_width": 1200,
        "upload_images": "true",
        "upload_api_url": "http://mfs.ailab.ccoop.cc:18585/upload/form",
        "preview_base_url": "http://mfs.ailab.ccoop.cc:18585"
    }
    success, fail = 0, 0
    for f in test_files:
        print(f"\n📄 处理: {f.name}")
        sample = {"fileName": f.name, "filePath": str(f), "export_path": output_dir, "text": "", "target_type": ""}
        try:
            res = MarkItDownParser(**mock_kwargs).execute(sample)
            md = res.get("text", "")
            if md:
                out = os.path.join(output_dir, f.stem + ".md")
                with open(out, "w", encoding="utf-8") as fo: fo.write(md)
                if "base64..." in md:
                    print(f"   ⚠️ 成功但仍有占位符 (缓存耗尽)")
                elif "data:image" in md:
                    print(f"   ✅ 成功 (含 Base64 图片)")
                else:
                    print(f"   ✅ 成功 | MD长度: {len(md)}")
                success += 1
            else: fail += 1
        except Exception as e:
            print(f"   ❌ 失败: {e}")
            fail += 1
    print(f"\n📊 测试完成: {success} 成功, {fail} 失败")