orbitkit 0.8.58__tar.gz → 0.8.60__tar.gz

{orbitkit-0.8.58/orbitkit.egg-info → orbitkit-0.8.60}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: orbitkit
-Version: 0.8.58
+Version: 0.8.60
 Summary: This project is only for Orbit Tech internal use.
 Home-page: https://github.com/clown-0726/orbitkit
 Author: Lilu Cao

orbitkit-0.8.60/orbitkit/VERSION

@@ -0,0 +1 @@
+0.8.60

{orbitkit-0.8.58 → orbitkit-0.8.60}/orbitkit/airflow_handler/file_flow_entry_process.py

@@ -151,7 +151,7 @@ class FilingOfficialProcessor:
 
         return file_flow_info
 
-    def send_task(self, file_flow_info, tags, is_important, priority, spider_name_source_type):
+    def send_task(self, file_flow_info, tags, is_important, priority, spider_name_source_type, queue_type):
         for step_str, records in file_flow_info.items():
             steps = step_str.split('@__@')
             start_stage = steps[0]
@@ -172,12 +172,14 @@ class FilingOfficialProcessor:
                                                                             tag=x_spider_name,
                                                                             priority=priority,
                                                                             source_type=spider_name_source_type[
-                                                                                x_spider_name])
+                                                                                x_spider_name],
+                                                                            queue_type=queue_type)
             else:
                 status, ids, message = self.file_handler.entry_point(records=records, start_stage=start_stage,
                                                                      target_stage=target_stage, tags=tags,tag=x_spider_name,
                                                                      priority=priority,
-                                                                     source_type=spider_name_source_type[x_spider_name])
+                                                                     source_type=spider_name_source_type[x_spider_name],
+                                                                     queue_type=queue_type)
             self.all_stat_count['file_flow'] += len(records)
             logger.info(f"{len(records)}--{start_stage}-{target_stage}-{x_spider_name} status: {status}, message: {message}")
 
@@ -213,7 +215,7 @@ class FilingOfficialProcessor:
 
     async def process_task_entry(self, source: str,
                            query: dict, tags: list[str], priority: str,
-                           is_important: bool = False, custom_step: Optional[list[str]] = None, important_level = None, db_name: str = None):
+                           is_important: bool = False, custom_step: Optional[list[str]] = None, important_level = None, db_name: str = None, queue_type: str = 'Default'):
 
         if not important_level or not isinstance(important_level, int):
             important_level = 0
@@ -251,7 +253,7 @@ class FilingOfficialProcessor:
                 self.all_stat_count['doc_error'] += len(doc_error_list)
                 self.all_stat_count['step_error'] += len(except_id_list)
                 self.all_stat_count['xbrl'] += len(xbrl_data)
-                self.send_task(file_flow_info, tags, is_important, priority, spider_name_source_type)
+                self.send_task(file_flow_info, tags, is_important, priority, spider_name_source_type, queue_type)
                 self.send_xbrl_data_to_mongo(xbrl_data)
                 self.update_doc_status_to_convert(collection, doc_error_list)
                 process_data.clear()
@@ -264,7 +266,7 @@ class FilingOfficialProcessor:
             self.all_stat_count['doc_error'] += len(doc_error_list)
             self.all_stat_count['step_error'] += len(except_id_list)
             self.all_stat_count['xbrl'] += len(xbrl_data)
-            self.send_task(file_flow_info, tags, is_important, priority, spider_name_source_type)
+            self.send_task(file_flow_info, tags, is_important, priority, spider_name_source_type, queue_type)
             self.send_xbrl_data_to_mongo(xbrl_data)
             self.update_doc_status_to_convert(collection, doc_error_list)
             process_data.clear()

{orbitkit-0.8.58 → orbitkit-0.8.60}/orbitkit/airflow_handler/file_handler_v2.py

@@ -141,7 +141,8 @@ class FileFlowHandleV2:
             'updated_at': now,
             'tags': params['tags'],
             'tag': params['tag'],
-            'important_level': record.get('important_level', 0)
+            'important_level': record.get('important_level', 0),
+            'queue': params['queue']
         }
 
         step = {
@@ -225,6 +226,7 @@ class FileFlowHandleV2:
             tag: str = None,
             priority: str = '1',
             source_type: Optional[str] = None,
+            queue_type: str = 'Default'
     ) -> Tuple[bool, Any, str]:
         """
             普通任务接口
@@ -252,7 +254,7 @@ class FileFlowHandleV2:
         """
         return self._file_flow_entry_point_internal(
             records, start_stage, target_stage, tags, priority, source_type,tag,
-            urgent=False, clean_exist_data=False
+            urgent=False, clean_exist_data=False, queue_type=queue_type
         )
 
     def entry_point_urgent(
@@ -264,6 +266,7 @@ class FileFlowHandleV2:
             tag: str = None,
             priority: str = '1',
             source_type: Optional[str] = None,
+            queue_type: str = 'Default'
     ) -> Tuple[bool, Any, str]:
         """
             加急任务接口
@@ -291,7 +294,7 @@ class FileFlowHandleV2:
         """
         return self._file_flow_entry_point_internal(
             records, start_stage, target_stage, tags, priority, source_type, tag,
-            urgent=True, clean_exist_data=True
+            urgent=True, clean_exist_data=True, queue_type=queue_type
         )
 
     def _file_flow_entry_point_internal(
@@ -304,7 +307,8 @@ class FileFlowHandleV2:
             source_type: Optional[str],
             tag: str,
             urgent: bool,
-            clean_exist_data: bool = False
+            clean_exist_data: bool = False,
+            queue_type: str = 'Default'
     ) -> Tuple[bool, Any, str]:
         """核心处理逻辑"""
         params = {
@@ -315,7 +319,8 @@ class FileFlowHandleV2:
             'current_stage': start_stage,
             'source_type': source_type,
             'tag': tag,
-            'urgent': urgent
+            'urgent': urgent,
+            'queue': queue_type
         }
 
         is_valid, msg = self._validate_params(params)

orbitkit-0.8.60/orbitkit/pdf_extractor/mineru_demo.py

@@ -0,0 +1,133 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+MinerU PDF 解析工具使用示例
+"""
+import os
+import tempfile
+import boto3
+from urllib.parse import urlparse
+from pdf_extractor_minerU_v1 import MinerUSingleTask
+
+
+def process_pdf(s3_path: str, api_token: str,
+                supported_versions: list = None, model_version: str = "vlm",
+                is_ocr: bool = True, enable_formula: bool = True,
+                enable_table: bool = True, timeout: int = 1800,
+                keep_all_files: bool = True) -> str:
+    """
+    处理 S3 上的 PDF 文件
+
+    :param s3_path: S3 路径 (如 s3://bucket/path/to/file.pdf)
+    :param api_token: MinerU API Token
+    :param supported_versions: 支持的版本列表，如 ["2.7.3"]，None 表示不检查
+    :param model_version: 模型版本 "pipeline" | "vlm" | "MinerU-HTML"
+    :param is_ocr: 是否启用 OCR
+    :param enable_formula: 是否启用公式识别
+    :param enable_table: 是否启用表格识别
+    :param timeout: 超时时间（秒）
+    :param keep_all_files: 是否保留所有文件
+    :return: S3 输出路径 (txt-vector/...)
+    """
+
+    # 版本检查回调
+    def check_version(version_info):
+        if supported_versions is None:
+            return
+        version = version_info.get("_version_name", "")
+        if version not in supported_versions:
+            raise Exception(
+                f"不支持的 MinerU API 版本: {version}\n"
+                f"当前支持: {supported_versions}\n"
+                f"请更新代码以支持新版本的 block 类型"
+            )
+        print(f"版本检查通过: {version}")
+
+    # 解析 S3 路径
+    parsed = urlparse(s3_path)
+    bucket = parsed.netloc
+    key = parsed.path.lstrip('/')
+    filename = os.path.basename(key)
+
+    # 初始化客户端
+    version_checker = check_version if supported_versions else None
+    client = MinerUSingleTask(api_token, version_checker=version_checker)
+    s3_client = boto3.client('s3')
+
+    print("=" * 60)
+    print("MinerU PDF 解析工具")
+    print("=" * 60)
+    print(f"输入: {s3_path}")
+
+    # 第零步：从 S3 下载 PDF 到本地
+    print("\n[步骤0] 从 S3 下载 PDF...")
+    with tempfile.TemporaryDirectory() as temp_dir:
+        local_pdf_path = os.path.join(temp_dir, filename)
+        s3_client.download_file(bucket, key, local_pdf_path)
+        print(f"下载完成: {local_pdf_path}")
+
+        # 第一步：上传文件
+        print("\n[步骤1] 上传文件到 MinerU...")
+        tasks = client.submit_task(
+            local_pdf_path,
+            model_version=model_version,
+            is_ocr=is_ocr,
+            enable_formula=enable_formula,
+            enable_table=enable_table
+        )
+
+        # 第二步：等待任务完成
+        print("\n[步骤2] 等待任务完成...")
+        results = client.wait_for_complete(tasks, timeout=timeout)
+
+        # 第三步：下载并处理结果（输出到 txt-vector/{原S3路径}）
+        print("\n[步骤3] 下载并处理结果...")
+        output_path = client.download_and_process(
+            results,
+            s3_path,  # 使用原 S3 路径生成输出目录
+            s3_client,
+            keep_all_files=keep_all_files
+        )
+
+    print("\n" + "=" * 60)
+    print(f"完成! 结果已上传到: {output_path}/")
+    print("=" * 60)
+
+    return output_path
+
+
+if __name__ == '__main__':
+    # ==================== 配置参数 ====================
+
+    # MinerU API Token（从 https://mineru.net 获取）
+    API_TOKEN = "你的API Token"
+
+    # PDF 文件的 S3 路径
+    S3_PATH = "s3://your-bucket/path/to/file.pdf"
+
+    # 支持的 API 版本（None 表示不检查版本）
+    SUPPORTED_VERSIONS = ["2.7.3"]
+
+    # 模型版本: "pipeline" | "vlm" | "MinerU-HTML"
+    MODEL_VERSION = "vlm"
+
+    # 其他选项
+    IS_OCR = True
+    ENABLE_FORMULA = True
+    ENABLE_TABLE = True
+    TIMEOUT = 1800
+    KEEP_ALL_FILES = True
+
+    # ==================== 执行 ====================
+
+    process_pdf(
+        s3_path=S3_PATH,
+        api_token=API_TOKEN,
+        supported_versions=SUPPORTED_VERSIONS,
+        model_version=MODEL_VERSION,
+        is_ocr=IS_OCR,
+        enable_formula=ENABLE_FORMULA,
+        enable_table=ENABLE_TABLE,
+        timeout=TIMEOUT,
+        keep_all_files=KEEP_ALL_FILES
+    )