#!/usr/bin/env python # coding=utf-8 import os import time import requests import zipfile import json from config import ( MINERU_API_URL, MINERU_TOKEN, OSS_ACCESS_KEY_ID, OSS_ACCESS_KEY_SECRET, OSS_BUCKET_NAME, OSS_ENDPOINT, TEMP_DIR, ) import sys sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from shared.oss_upload import upload_file_to_oss def create_parse_task(url, model_version="vlm"): """ 创建解析任务 :param url: 文件URL :param model_version: 模型版本，默认为vlm :return: 任务ID """ print(f"开始创建解析任务: {url}") api_url = f"{MINERU_API_URL}/extract/task" headers = { "Content-Type": "application/json", "Authorization": f"Bearer {MINERU_TOKEN}", } data = {"url": url, "model_version": model_version} try: response = requests.post(api_url, headers=headers, json=data) response.raise_for_status() result = response.json() if result.get("code") == 0: task_id = result.get("data", {}).get("task_id") print(f"任务创建成功: {task_id}") return task_id else: print(f"任务创建失败: {result.get('msg')}") return None except Exception as e: print(f"创建任务失败: {e}") import traceback traceback.print_exc() return None def get_task_status(task_id): """ 获取任务状态 :param task_id: 任务ID :return: 任务状态 """ print(f"查询任务状态: {task_id}") api_url = f"{MINERU_API_URL}/extract/task/{task_id}" headers = { "Content-Type": "application/json", "Authorization": f"Bearer {MINERU_TOKEN}", } try: response = requests.get(api_url, headers=headers) response.raise_for_status() result = response.json() if result.get("code") == 0: status = result.get("data", {}) print(f"任务状态: {status.get('state')}") return status else: print(f"查询状态失败: {result.get('msg')}") return None except Exception as e: print(f"查询状态失败: {e}") import traceback traceback.print_exc() return None def poll_task_status(task_id, max_retries=60, interval=5): """ 轮询任务状态 :param task_id: 任务ID :param max_retries: 最大重试次数 :param interval: 轮询间隔（秒） :return: 任务完成状态 """ print(f"开始轮询任务状态: {task_id}") for i in range(max_retries): status = get_task_status(task_id) if status: state = status.get("state") if state == "done": print("任务完成!") return status elif state == "failed": print(f"任务失败: {status.get('err_msg')}") return None elif state in ["pending", "running", "converting"]: print(f"任务正在进行中 ({state})，{interval}秒后重试...") time.sleep(interval) else: print(f"未知状态: {state}") time.sleep(interval) else: print(f"获取状态失败，{interval}秒后重试...") time.sleep(interval) print("轮询超时，任务可能仍在处理中") return None def download_and_extract_result(zip_url, output_dir): """ 下载并提取解析结果 :param zip_url: 结果压缩包URL :param output_dir: 输出目录 :return: 提取的文件列表 """ print(f"开始下载解析结果: {zip_url}") # 确保输出目录存在 os.makedirs(output_dir, exist_ok=True) # 下载压缩包 zip_path = os.path.join(output_dir, "result.zip") try: response = requests.get(zip_url, stream=True) response.raise_for_status() with open(zip_path, "wb") as f: for chunk in response.iter_content(chunk_size=8192): f.write(chunk) print(f"压缩包下载成功: {zip_path}") # 提取压缩包 extracted_files = [] with zipfile.ZipFile(zip_path, "r") as zip_ref: zip_ref.extractall(output_dir) print(f"压缩包提取成功: {output_dir}") extracted_files = zip_ref.namelist() # 删除压缩包 os.remove(zip_path) print(f"删除临时压缩包: {zip_path}") return extracted_files except Exception as e: print(f"下载或提取失败: {e}") import traceback traceback.print_exc() return None def get_markdown_result(extracted_files, output_dir): """ 获取Markdown格式的解析结果 :param extracted_files: 提取的文件列表 :param output_dir: 输出目录 :return: Markdown内容 """ print("查找Markdown格式的解析结果") for file_name in extracted_files: if file_name.endswith(".md"): md_path = os.path.join(output_dir, file_name) print(f"找到Markdown文件: {md_path}") try: with open(md_path, "r", encoding="utf-8") as f: md_content = f.read() print(f"Markdown文件读取成功，长度: {len(md_content)} 字符") return md_content except Exception as e: print(f"读取Markdown文件失败: {e}") import traceback traceback.print_exc() return None print("未找到Markdown格式的解析结果") return None def parse_local_file(file_path, model_version="vlm"): """ 解析本地文件 :param file_path: 本地文件路径 :param model_version: 模型版本，默认为vlm :return: Markdown内容 """ print(f"开始解析本地文件: {file_path}") # 检查文件是否存在 if not os.path.exists(file_path): print(f"文件不存在: {file_path}") return None # 检查文件大小 file_size = os.path.getsize(file_path) if file_size > 200 * 1024 * 1024: # 200MB print(f"文件大小超出限制: {file_size} bytes (最大200MB)") return None # 生成对象名称 timestamp = int(time.time()) file_name = os.path.basename(file_path) object_name = f"mineru/{timestamp}_{file_name}" # 上传文件到OSS oss_url = upload_file_to_oss( file_path, OSS_BUCKET_NAME, object_name, OSS_ACCESS_KEY_ID, OSS_ACCESS_KEY_SECRET, OSS_ENDPOINT, ) if not oss_url: print("文件上传失败，无法继续解析") return None # 创建解析任务 task_id = create_parse_task(oss_url, model_version) if not task_id: print("任务创建失败，无法继续解析") return None # 轮询任务状态 task_status = poll_task_status(task_id) if not task_status: print("任务执行失败，无法获取解析结果") return None # 获取结果URL zip_url = task_status.get("full_zip_url") if not zip_url: print("未找到解析结果URL") return None # 生成输出目录 output_dir = os.path.join(TEMP_DIR, task_id) # 下载并提取结果 extracted_files = download_and_extract_result(zip_url, output_dir) if not extracted_files: print("下载或提取结果失败") return None # 获取Markdown结果 md_content = get_markdown_result(extracted_files, output_dir) if not md_content: print("未找到Markdown格式的解析结果") return None print("文件解析完成，成功获取Markdown格式的结果") return {"content": md_content, "output_dir": output_dir} def main(): """ 主函数 """ import sys if len(sys.argv) != 2: print("使用方法: python mineru_parser.py <本地文件路径>") print("示例: python mineru_parser.py ./example.pdf") sys.exit(1) file_path = sys.argv[1] # 解析文件 result = parse_local_file(file_path) if result: md_content = result["content"] output_dir = result["output_dir"] # 保存Markdown结果 output_file = os.path.join(TEMP_DIR, f"{os.path.basename(file_path)}.md") with open(output_file, "w", encoding="utf-8") as f: f.write(md_content) # Print JSON result for caller to parse print( "JSON_RESULT:" + json.dumps( { "status": "success", "markdown_file": os.path.abspath(output_file), "output_dir": os.path.abspath(output_dir), } ) ) else: print("文件解析失败") sys.exit(1) if __name__ == "__main__": main()