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

@@ -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": {