diff --git a/README.md b/README.md index 02da7d5..3f3a82e 100644 --- a/README.md +++ b/README.md @@ -5,19 +5,35 @@ Python MCP Server for PDF OCR batch processing. ## Features - **Batch OCR** - Process multiple PDF files and save results as Markdown +- **Single File OCR** - Process a single PDF file directly ## Tool ### `batch_ocr_pdf` -Batch process all PDF files in a directory using the OCR API. +OCR process PDF file(s) using the OCR API. Supports both single file and directory modes. **Parameters:** -- `input_dir`: Directory containing PDF files -- `output_dir`: Directory to save Markdown results +- `input_file`: Path to a single PDF file to process (mutually exclusive with `input_dir`) +- `input_dir`: Directory containing PDF files to process (mutually exclusive with `input_file`) +- `output_dir`: Directory to save Markdown results (auto-created if not provided) - `api_key`: OCR API key (optional, reads from `OCR_API_KEY` env var) - `base_url`: OCR API base URL (default: `https://agent.imqimacau.com`) +**Usage Modes:** + +1. **Single File Mode**: Provide `input_file` only + ``` + input_file: /path/to/document.pdf + ``` + +2. **Batch Mode**: Provide `input_dir` only + ``` + input_dir: /path/to/pdf_directory + ``` + +**Note:** `input_file` and `input_dir` are mutually exclusive. At least one must be provided. + **Returns:** JSON with processing results (success/fail counts per file) ## Setup diff --git a/src/procleaning_mcp/__pycache__/server.cpython-311.pyc b/src/procleaning_mcp/__pycache__/server.cpython-311.pyc index 7a192d4..83cf035 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 5482c7e..8c1c55c 100644 --- a/src/procleaning_mcp/server.py +++ b/src/procleaning_mcp/server.py @@ -53,43 +53,79 @@ async def _poll_task_result(api_url: str, task_id: str, api_key: str) -> dict: @mcp.tool() def batch_ocr_pdf( - input_dir: str, - output_dir: str, + input_file: str = None, + input_dir: str = None, + output_dir: str = None, api_key: str = None, base_url: str = "https://agent.imqimacau.com", ) -> str: """ - Batch OCR process all PDF files in the input directory. - Submits each PDF to the OCR API and saves results as Markdown files. + OCR process PDF file(s) — either a single file or an entire directory. + Submits PDF to the OCR API and saves results as Markdown files. Args: - input_dir: Directory containing PDF files to process - output_dir: Directory to save OCR results (Markdown files) + input_file: Path to a single PDF file to process (mutually exclusive with input_dir) + input_dir: Directory containing PDF files to process (mutually exclusive with input_file) + output_dir: Directory to save OCR results (Markdown files). Auto-created if not exists. api_key: API key for the OCR service (optional, uses environment variable if not provided) base_url: Base URL of the OCR API (default: https://agent.imqimacau.com) Returns: - JSON string with batch processing results (success/fail counts and per-file status) + JSON string with processing results (success/fail counts and per-file status) + + Note: + input_file and input_dir are mutually exclusive. If neither is provided, defaults to input_dir. """ # Use environment variable if api_key not provided if not api_key: api_key = os.environ.get("OCR_API_KEY", "") - if not os.path.exists(input_dir): + # Validate mutual exclusivity + if input_file and input_dir: return json.dumps( - {"error": f"Input directory does not exist: {input_dir}"}, + {"error": "input_file and input_dir are mutually exclusive. Please provide only one."}, ensure_ascii=False, indent=2, ) - os.makedirs(output_dir, exist_ok=True) + # If neither provided, default to input_dir behavior with an error message + 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, + ) - # Find all PDF files - pdf_files = [ - os.path.join(input_dir, f) - for f in os.listdir(input_dir) - if f.lower().endswith(".pdf") - ] + # Determine file list based on input mode + if input_file: + # Single file mode + if not os.path.isfile(input_file): + return json.dumps( + {"error": f"Input file does not exist: {input_file}"}, + ensure_ascii=False, + indent=2, + ) + # Auto-create output_dir based on input file location if not provided + if not output_dir: + output_dir = os.path.join(os.path.dirname(input_file) or ".", "ocr_output") + pdf_files = [input_file] + else: + # Directory mode + 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") + pdf_files = [ + os.path.join(input_dir, f) + for f in os.listdir(input_dir) + if f.lower().endswith(".pdf") + ] + + os.makedirs(output_dir, exist_ok=True) if not pdf_files: return json.dumps( @@ -208,8 +244,9 @@ def batch_ocr_pdf( fail_count += 1 results.append({"file": filename, "status": "failed", "error": str(e)}) - # Delay between files - time.sleep(2) + # Delay between files (only in batch mode) + if len(pdf_files) > 1: + time.sleep(2) final_result = { "summary": {