feat: support single file input (input_file) alongside directory mode

This commit is contained in:
Tony
2026-05-09 15:02:56 +08:00
parent 4a48ded39e
commit df11f4e79b
3 changed files with 74 additions and 21 deletions

View File

@@ -5,19 +5,35 @@ Python MCP Server for PDF OCR batch processing.
## Features ## Features
- **Batch OCR** - Process multiple PDF files and save results as Markdown - **Batch OCR** - Process multiple PDF files and save results as Markdown
- **Single File OCR** - Process a single PDF file directly
## Tool ## Tool
### `batch_ocr_pdf` ### `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:** **Parameters:**
- `input_dir`: Directory containing PDF files - `input_file`: Path to a single PDF file to process (mutually exclusive with `input_dir`)
- `output_dir`: Directory to save Markdown results - `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) - `api_key`: OCR API key (optional, reads from `OCR_API_KEY` env var)
- `base_url`: OCR API base URL (default: `https://agent.imqimacau.com`) - `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) **Returns:** JSON with processing results (success/fail counts per file)
## Setup ## Setup

View File

@@ -53,44 +53,80 @@ async def _poll_task_result(api_url: str, task_id: str, api_key: str) -> dict:
@mcp.tool() @mcp.tool()
def batch_ocr_pdf( def batch_ocr_pdf(
input_dir: str, input_file: str = None,
output_dir: str, input_dir: str = None,
output_dir: str = None,
api_key: str = None, api_key: str = None,
base_url: str = "https://agent.imqimacau.com", base_url: str = "https://agent.imqimacau.com",
) -> str: ) -> str:
""" """
Batch OCR process all PDF files in the input directory. OCR process PDF file(s) — either a single file or an entire directory.
Submits each PDF to the OCR API and saves results as Markdown files. Submits PDF to the OCR API and saves results as Markdown files.
Args: Args:
input_dir: Directory containing PDF files to process input_file: Path to a single PDF file to process (mutually exclusive with input_dir)
output_dir: Directory to save OCR results (Markdown files) 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) 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) base_url: Base URL of the OCR API (default: https://agent.imqimacau.com)
Returns: 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 # Use environment variable if api_key not provided
if not api_key: if not api_key:
api_key = os.environ.get("OCR_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": "input_file and input_dir are mutually exclusive. Please provide only one."},
ensure_ascii=False,
indent=2,
)
# 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,
)
# 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( return json.dumps(
{"error": f"Input directory does not exist: {input_dir}"}, {"error": f"Input directory does not exist: {input_dir}"},
ensure_ascii=False, ensure_ascii=False,
indent=2, indent=2,
) )
if not output_dir:
os.makedirs(output_dir, exist_ok=True) output_dir = os.path.join(input_dir, "ocr_output")
# Find all PDF files
pdf_files = [ pdf_files = [
os.path.join(input_dir, f) os.path.join(input_dir, f)
for f in os.listdir(input_dir) for f in os.listdir(input_dir)
if f.lower().endswith(".pdf") if f.lower().endswith(".pdf")
] ]
os.makedirs(output_dir, exist_ok=True)
if not pdf_files: if not pdf_files:
return json.dumps( return json.dumps(
{"message": f"No PDF files found in {input_dir}", "files_processed": 0}, {"message": f"No PDF files found in {input_dir}", "files_processed": 0},
@@ -208,7 +244,8 @@ def batch_ocr_pdf(
fail_count += 1 fail_count += 1
results.append({"file": filename, "status": "failed", "error": str(e)}) results.append({"file": filename, "status": "failed", "error": str(e)})
# Delay between files # Delay between files (only in batch mode)
if len(pdf_files) > 1:
time.sleep(2) time.sleep(2)
final_result = { final_result = {