feat: add submit_image_for_ocr tool for image OCR support

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)
 # ============================================================