feat: support single file input (input_file) alongside directory mode
This commit is contained in:
22
README.md
22
README.md
@@ -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
|
||||||
|
|||||||
Binary file not shown.
@@ -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 = {
|
||||||
|
|||||||
Reference in New Issue
Block a user