""" MinerU Tianshu - Client Example 天枢客户端示例 演示如何使用 Python 客户端提交任务和查询状态 """ import asyncio import aiohttp from pathlib import Path from loguru import logger import time from typing import Dict class TianshuClient: """天枢客户端""" def __init__(self, api_url='http://localhost:8000'): self.api_url = api_url self.base_url = f"{api_url}/api/v1" async def submit_task( self, session: aiohttp.ClientSession, file_path: str, backend: str = 'pipeline', lang: str = 'ch', method: str = 'auto', formula_enable: bool = True, table_enable: bool = True, priority: int = 0 ) -> Dict: """ 提交任务 Args: session: aiohttp session file_path: 文件路径 backend: 处理后端 lang: 语言 method: 解析方法 formula_enable: 是否启用公式识别 table_enable: 是否启用表格识别 priority: 优先级 Returns: 响应字典，包含 task_id """ with open(file_path, 'rb') as f: data = aiohttp.FormData() data.add_field('file', f, filename=Path(file_path).name) data.add_field('backend', backend) data.add_field('lang', lang) data.add_field('method', method) data.add_field('formula_enable', str(formula_enable).lower()) data.add_field('table_enable', str(table_enable).lower()) data.add_field('priority', str(priority)) async with session.post(f'{self.base_url}/tasks/submit', data=data) as resp: if resp.status == 200: result = await resp.json() logger.info(f"✅ Submitted: {file_path} -> Task ID: {result['task_id']}") return result else: error = await resp.text() logger.error(f"❌ Failed to submit {file_path}: {error}") return {'success': False, 'error': error} async def get_task_status(self, session: aiohttp.ClientSession, task_id: str) -> Dict: """ 查询任务状态 Args: session: aiohttp session task_id: 任务ID Returns: 任务状态字典 """ async with session.get(f'{self.base_url}/tasks/{task_id}') as resp: if resp.status == 200: return await resp.json() else: return {'success': False, 'error': 'Task not found'} async def wait_for_task( self, session: aiohttp.ClientSession, task_id: str, timeout: int = 600, poll_interval: int = 2 ) -> Dict: """ 等待任务完成 Args: session: aiohttp session task_id: 任务ID timeout: 超时时间（秒） poll_interval: 轮询间隔（秒） Returns: 最终任务状态 """ start_time = time.time() while True: status = await self.get_task_status(session, task_id) if not status.get('success'): logger.error(f"❌ Failed to get status for task {task_id}") return status task_status = status.get('status') if task_status == 'completed': logger.info(f"✅ Task {task_id} completed!") logger.info(f" Output: {status.get('result_path')}") return status elif task_status == 'failed': logger.error(f"❌ Task {task_id} failed!") logger.error(f" Error: {status.get('error_message')}") return status elif task_status == 'cancelled': logger.warning(f"⚠️ Task {task_id} was cancelled") return status # 检查超时 if time.time() - start_time > timeout: logger.error(f"⏱️ Task {task_id} timeout after {timeout}s") return {'success': False, 'error': 'timeout'} # 等待后继续轮询 await asyncio.sleep(poll_interval) async def get_queue_stats(self, session: aiohttp.ClientSession) -> Dict: """获取队列统计""" async with session.get(f'{self.base_url}/queue/stats') as resp: return await resp.json() async def cancel_task(self, session: aiohttp.ClientSession, task_id: str) -> Dict: """取消任务""" async with session.delete(f'{self.base_url}/tasks/{task_id}') as resp: return await resp.json() async def example_single_task(): """示例1：提交单个任务并等待完成""" logger.info("=" * 60) logger.info("示例1：提交单个任务") logger.info("=" * 60) client = TianshuClient() async with aiohttp.ClientSession() as session: # 提交任务 result = await client.submit_task( session, file_path='../../demo/pdfs/demo1.pdf', backend='pipeline', lang='ch', formula_enable=True, table_enable=True ) if result.get('success'): task_id = result['task_id'] # 等待完成 logger.info(f"⏳ Waiting for task {task_id} to complete...") final_status = await client.wait_for_task(session, task_id) return final_status async def example_batch_tasks(): """示例2：批量提交多个任务并并发等待""" logger.info("=" * 60) logger.info("示例2：批量提交多个任务") logger.info("=" * 60) client = TianshuClient() # 准备任务列表 files = [ '../../demo/pdfs/demo1.pdf', '../../demo/pdfs/demo2.pdf', '../../demo/pdfs/demo3.pdf', ] async with aiohttp.ClientSession() as session: # 并发提交所有任务 logger.info(f"📤 Submitting {len(files)} tasks...") submit_tasks = [ client.submit_task(session, file) for file in files ] results = await asyncio.gather(*submit_tasks) # 提取 task_ids task_ids = [r['task_id'] for r in results if r.get('success')] logger.info(f"✅ Submitted {len(task_ids)} tasks successfully") # 并发等待所有任务完成 logger.info(f"⏳ Waiting for all tasks to complete...") wait_tasks = [ client.wait_for_task(session, task_id) for task_id in task_ids ] final_results = await asyncio.gather(*wait_tasks) # 统计结果 completed = sum(1 for r in final_results if r.get('status') == 'completed') failed = sum(1 for r in final_results if r.get('status') == 'failed') logger.info("=" * 60) logger.info(f"📊 Results: {completed} completed, {failed} failed") logger.info("=" * 60) return final_results async def example_priority_tasks(): """示例3：使用优先级队列""" logger.info("=" * 60) logger.info("示例3：优先级队列") logger.info("=" * 60) client = TianshuClient() async with aiohttp.ClientSession() as session: # 提交低优先级任务 low_priority = await client.submit_task( session, file_path='../../demo/pdfs/demo1.pdf', priority=0 ) logger.info(f"📝 Low priority task: {low_priority['task_id']}") # 提交高优先级任务 high_priority = await client.submit_task( session, file_path='../../demo/pdfs/demo2.pdf', priority=10 ) logger.info(f"🔥 High priority task: {high_priority['task_id']}") # 高优先级任务会先被处理 logger.info("⏳ 高优先级任务将优先处理...") async def example_queue_monitoring(): """示例4：监控队列状态""" logger.info("=" * 60) logger.info("示例4：监控队列状态") logger.info("=" * 60) client = TianshuClient() async with aiohttp.ClientSession() as session: # 获取队列统计 stats = await client.get_queue_stats(session) logger.info("📊 Queue Statistics:") logger.info(f" Total: {stats.get('total', 0)}") for status, count in stats.get('stats', {}).items(): logger.info(f" {status:12s}: {count}") async def main(): """主函数""" import sys if len(sys.argv) > 1: example = sys.argv[1] else: example = 'all' try: if example == 'single' or example == 'all': await example_single_task() print() if example == 'batch' or example == 'all': await example_batch_tasks() print() if example == 'priority' or example == 'all': await example_priority_tasks() print() if example == 'monitor' or example == 'all': await example_queue_monitoring() print() except Exception as e: logger.error(f"Example failed: {e}") import traceback traceback.print_exc() if __name__ == '__main__': """ 使用方法: # 运行所有示例 python client_example.py # 运行特定示例 python client_example.py single python client_example.py batch python client_example.py priority python client_example.py monitor """ asyncio.run(main())