diff --git a/src/procleaning_mcp/__pycache__/server.cpython-311.pyc b/src/procleaning_mcp/__pycache__/server.cpython-311.pyc index 19409b2..dbb03d0 100644 Binary files a/src/procleaning_mcp/__pycache__/server.cpython-311.pyc and b/src/procleaning_mcp/__pycache__/server.cpython-311.pyc differ diff --git a/src/procleaning_mcp/server.py b/src/procleaning_mcp/server.py index 70e43b2..3532129 100644 --- a/src/procleaning_mcp/server.py +++ b/src/procleaning_mcp/server.py @@ -117,6 +117,12 @@ async def _save_ocr_result( return output_md +# ============================================================ +# 📦 Image extensions list +# ============================================================ + +IMAGE_EXTENSIONS = (".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".tif", ".webp") + # ============================================================ # 🛠 Tool 1: Submit PDF for OCR (returns immediately) # ============================================================ @@ -286,6 +292,175 @@ async def submit_pdf_for_ocr( return json.dumps(final_result, ensure_ascii=False, indent=2) +# ============================================================ +# 🛠 Tool 2: Submit Image for OCR (returns immediately) +# ============================================================ + +@mcp.tool() +async def submit_image_for_ocr( + input_file: str = None, + input_dir: str = None, + output_dir: str = None, + api_key: str = None, + base_url: str = "https://agent.imqimacau.com", +) -> str: + """ + Submit Image file(s) (jpg, png, bmp, tiff) for OCR processing. Returns task_id(s) immediately. + Then use get_ocr_result(task_id) to check status and retrieve results. + + ⚠️ This tool does NOT wait for OCR completion. It submits and returns. + Use get_ocr_result() to poll for completion. + + 此工具只提交圖片,立馬回傳 task_id。然後用 get_ocr_result() 查詢結果。 + + Args: + input_file: Path to a single image file + input_dir: Directory containing image files + output_dir: Directory to save OCR results (auto-created) + api_key: API key for the OCR service. Uses env var 'OCR_API_KEY' if not provided. + base_url: Base URL of the OCR API. + + Returns: + JSON string with task_id(s) and metadata. + """ + if not api_key: + api_key = os.environ.get("OCR_API_KEY", "") + + if input_file and input_dir: + return json.dumps( + {"error": "input_file and input_dir are mutually exclusive."}, + ensure_ascii=False, indent=2 + ) + + if not input_file and not input_dir: + return json.dumps( + {"error": "Either input_file or input_dir must be provided."}, + ensure_ascii=False, indent=2 + ) + + if input_file: + if not os.path.isfile(input_file): + return json.dumps( + {"error": f"Input file does not exist: {input_file}"}, + ensure_ascii=False, indent=2 + ) + if not output_dir: + output_dir = os.path.join(os.path.dirname(input_file) or ".", "ocr_output") + image_files = [input_file] + else: + if not os.path.isdir(input_dir): + return json.dumps( + {"error": f"Input directory does not exist: {input_dir}"}, + ensure_ascii=False, indent=2 + ) + if not output_dir: + output_dir = os.path.join(input_dir, "ocr_output") + image_files = [ + os.path.join(input_dir, f) + for f in os.listdir(input_dir) + if f.lower().endswith(IMAGE_EXTENSIONS) + ] + + if not image_files: + return json.dumps( + {"message": f"No image files found in {input_dir}", "files_processed": 0}, + ensure_ascii=False, indent=2 + ) + + tasks = _load_tasks() + results = [] + success_count = 0 + fail_count = 0 + + async with aiohttp.ClientSession() as session: + for image_file in image_files: + filename = os.path.basename(image_file) + base_name = os.path.splitext(filename)[0] + + try: + # Read and encode image + with open(image_file, "rb") as f: + file_base64 = base64.b64encode(f.read()).decode() + + # Submit task + async with session.post( + f"{base_url}/task/submit", + headers={ + "X-API-Key": api_key, + "Content-Type": "application/json", + }, + json={ + "file_base64": file_base64, + "file_type": "image", + "enable_ai_description": False, + "output_type": "ocr_only", + }, + timeout=aiohttp.ClientTimeout(total=30), + ) as response: + if response.status != 200: + fail_count += 1 + results.append({ + "file": filename, + "status": "failed", + "error": f"API error: {response.status}", + }) + continue + + data = await response.json() + task_id = data.get("task_id") + if not task_id: + fail_count += 1 + results.append({ + "file": filename, + "status": "failed", + "error": "No task_id returned", + }) + continue + + # Register task for later polling + tasks[task_id] = { + "filename": filename, + "base_name": base_name, + "output_dir": output_dir, + "base_url": base_url, + "api_key": api_key, + "status": "pending", + "submitted_at": datetime.now().isoformat(), + } + _save_tasks(tasks) + + success_count += 1 + results.append({ + "file": filename, + "status": "submitted", + "task_id": task_id, + }) + + except Exception as e: + fail_count += 1 + results.append({ + "file": filename, + "status": "failed", + "error": str(e), + }) + + # Delay between files + if len(image_files) > 1: + await asyncio.sleep(2) + + final_result = { + "summary": { + "total": len(image_files), + "submitted": success_count, + "failed": fail_count, + "note": "Use get_ocr_result(task_id) to poll for results", + }, + "files": results, + } + + return json.dumps(final_result, ensure_ascii=False, indent=2) + + # ============================================================ # 🛠 Tool 2: Get OCR Result (polls and returns) # ============================================================