Extract batchOCR.py into MCP tool: batch_ocr_pdf
This commit is contained in:
Tony
49
README.md
49
README.md
@@ -1,44 +1,43 @@
|
||||
# procleaning-mcp
|
||||
|
||||
Python MCP Server for process monitoring and cleaning.
|
||||
Python MCP Server for PDF OCR batch processing.
|
||||
|
||||
## Features
|
||||
|
||||
- **List Processes** - Query running processes with flexible filtering
|
||||
- **Kill Process** - Terminate processes by PID with SIGTERM or SIGKILL
|
||||
- **Kill by Name** - Batch kill processes matching a name pattern
|
||||
- **System Resource Usage** - Get CPU, memory, disk, and load average overview
|
||||
- **Find Zombies** - Detect zombie processes
|
||||
- **Restart Service** - Restart systemd services
|
||||
- **Batch OCR** - Process multiple PDF files and save results as Markdown
|
||||
|
||||
## Installation
|
||||
## Tool
|
||||
|
||||
### `batch_ocr_pdf`
|
||||
|
||||
Batch process all PDF files in a directory using the OCR API.
|
||||
|
||||
**Parameters:**
|
||||
- `input_dir`: Directory containing PDF files
|
||||
- `output_dir`: Directory to save Markdown results
|
||||
- `api_key`: OCR API key (optional, reads from `OCR_API_KEY` env var)
|
||||
- `base_url`: OCR API base URL (default: `https://agent.imqimacau.com`)
|
||||
|
||||
**Returns:** JSON with processing results (success/fail counts per file)
|
||||
|
||||
## Setup
|
||||
|
||||
```bash
|
||||
cd procleaning-mcp
|
||||
uv sync
|
||||
uv run python -m procleaning_mcp.server
|
||||
```
|
||||
|
||||
## Usage (as MCP Server)
|
||||
## Usage in Hermes
|
||||
|
||||
```bash
|
||||
uv run procleaning-mcp
|
||||
```
|
||||
|
||||
Or use with Hermes Agent by adding to `~/.hermes/config.yaml`:
|
||||
Add to `~/.hermes/config.yaml`:
|
||||
|
||||
```yaml
|
||||
mcp_servers:
|
||||
procleaning:
|
||||
command: "uv"
|
||||
args: ["run", "--directory", "/home/tony/procleaning-mcp", "procleaning-mcp"]
|
||||
args: ["run", "--project", "/home/tony/procleaning-mcp", "python", "-m", "procleaning_mcp.server"]
|
||||
timeout: 300
|
||||
```
|
||||
|
||||
## Tools
|
||||
|
||||
| Tool | Description |
|
||||
|------|-------------|
|
||||
| `list_processes` | List processes with filters (name, user, CPU%, MEM%, RSS) |
|
||||
| `kill_process` | Kill a specific process by PID |
|
||||
| `kill_by_name` | Kill all processes matching a name pattern |
|
||||
| `system_resource_usage` | System-wide resource usage overview |
|
||||
| `find_zombies` | Find zombie processes |
|
||||
| `restart_service` | Restart a systemd service |
|
||||
Then use `mcp_procleaning_batch_ocr_pdf` in any Hermes conversation.
|
||||
|
||||
@@ -1,257 +1,224 @@
|
||||
"""Process Cleaning MCP Server - Tools for monitoring and cleaning system processes."""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import signal
|
||||
import subprocess
|
||||
from dataclasses import asdict
|
||||
from datetime import datetime, timezone
|
||||
import base64
|
||||
import time
|
||||
import requests
|
||||
from datetime import datetime
|
||||
|
||||
import psutil
|
||||
from mcp.server import FastMCP
|
||||
|
||||
mcp = FastMCP(
|
||||
name="procleaning-mcp",
|
||||
instructions="Process monitoring and cleaning utilities. List, kill, and monitor system processes.",
|
||||
instructions="Process monitoring, cleaning, and PDF OCR utilities.",
|
||||
)
|
||||
|
||||
# ============================================================
|
||||
# 🔧 PDF OCR Batch Processing Tool
|
||||
# ============================================================
|
||||
|
||||
def _proc_to_info(proc: psutil.Process) -> dict | None:
|
||||
"""Convert a psutil.Process to a dict, handling exceptions."""
|
||||
async def _poll_task_result(api_url: str, task_id: str, api_key: str) -> dict:
|
||||
"""Poll the OCR task status until completion."""
|
||||
while True:
|
||||
try:
|
||||
with proc.oneshot():
|
||||
return {
|
||||
"pid": proc.pid,
|
||||
"name": proc.name(),
|
||||
"username": proc.username(),
|
||||
"cpu_percent": round(proc.cpu_percent(), 2),
|
||||
"memory_percent": round(proc.memory_percent(), 2),
|
||||
"memory_rss_mb": round(proc.memory_info().rss / 1024 / 1024, 2),
|
||||
"status": proc.status(),
|
||||
"create_time": datetime.fromtimestamp(
|
||||
proc.create_time(), tz=timezone.utc
|
||||
).isoformat(),
|
||||
"cmd_line": " ".join(proc.cmdline()) or "(unknown)",
|
||||
}
|
||||
except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
|
||||
return None
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
def list_processes(
|
||||
name_filter: str | None = None,
|
||||
username: str | None = None,
|
||||
min_cpu: float | None = None,
|
||||
min_memory: float | None = None,
|
||||
min_rss_mb: float | None = None,
|
||||
max_processes: int = 50,
|
||||
) -> str:
|
||||
"""List system processes with optional filtering by name, user, and resource usage thresholds."""
|
||||
results = []
|
||||
for proc in psutil.process_iter():
|
||||
info = _proc_to_info(proc)
|
||||
if info is None:
|
||||
continue
|
||||
|
||||
if name_filter and name_filter.lower() not in info["name"].lower():
|
||||
continue
|
||||
if username and info["username"] != username:
|
||||
continue
|
||||
if min_cpu is not None and info["cpu_percent"] < min_cpu:
|
||||
continue
|
||||
if min_memory is not None and info["memory_percent"] < min_memory:
|
||||
continue
|
||||
if min_rss_mb is not None and info["memory_rss_mb"] < min_rss_mb:
|
||||
continue
|
||||
|
||||
results.append(info)
|
||||
if len(results) >= max_processes:
|
||||
break
|
||||
|
||||
if not results:
|
||||
return "No processes found matching the given criteria."
|
||||
|
||||
lines = [f"Found {len(results)} process(es):\n"]
|
||||
for p in results:
|
||||
lines.append(
|
||||
f" PID: {p['pid']:>6} | Name: {p['name']:<25} | User: {p['username']:<12} "
|
||||
f"| CPU: {p['cpu_percent']:>5}% | MEM: {p['memory_percent']:>4}% | "
|
||||
f"RSS: {p['memory_rss_mb']:>7}MB | Status: {p['status']}"
|
||||
response = requests.get(
|
||||
f"{api_url}/task/result/{task_id}",
|
||||
headers={"X-API-Key": api_key},
|
||||
timeout=30,
|
||||
)
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
def kill_process(pid: int, signal: str = "SIGTERM", confirm: bool = True) -> str:
|
||||
"""Terminate a running process by PID. Default signal is SIGTERM (graceful); use SIGKILL for forceful termination."""
|
||||
if not confirm:
|
||||
return "Kill cancelled. Set confirm=true to proceed."
|
||||
|
||||
sig_num = getattr(signal, signal, signal.SIGTERM)
|
||||
try:
|
||||
proc = psutil.Process(pid)
|
||||
proc_name = proc.name()
|
||||
proc.kill(sig_num)
|
||||
try:
|
||||
proc.wait(timeout=5)
|
||||
except psutil.TimeoutExpired:
|
||||
return (
|
||||
f"Warning: PID {pid} ({proc_name}) did not terminate within 5s. "
|
||||
f"Consider using SIGKILL."
|
||||
)
|
||||
return f"Successfully killed process PID {pid} ({proc_name}) with {signal}."
|
||||
except psutil.NoSuchProcess:
|
||||
return f"Error: PID {pid} does not exist."
|
||||
except psutil.AccessDenied:
|
||||
return f"Error: permission denied killing PID {pid}."
|
||||
except Exception as e:
|
||||
return f"Error killing PID {pid}: {e}"
|
||||
await asyncio.sleep(2)
|
||||
continue
|
||||
|
||||
if response.status_code != 200:
|
||||
await asyncio.sleep(2)
|
||||
continue
|
||||
|
||||
@mcp.tool()
|
||||
def kill_by_name(
|
||||
name_pattern: str,
|
||||
signal: str = "SIGTERM",
|
||||
max_kill: int = 10,
|
||||
confirm: bool = True,
|
||||
) -> str:
|
||||
"""Terminate all processes matching a given name pattern. Default signal is SIGTERM (graceful)."""
|
||||
if not confirm:
|
||||
return "Kill cancelled. Set confirm=true to proceed."
|
||||
|
||||
sig_num = getattr(signal, signal, signal.SIGTERM)
|
||||
|
||||
targets = []
|
||||
for proc in psutil.process_iter():
|
||||
info = _proc_to_info(proc)
|
||||
if info and name_pattern.lower() in info["name"].lower():
|
||||
targets.append(info)
|
||||
|
||||
if not targets:
|
||||
return f"No processes found matching '{name_pattern}'."
|
||||
|
||||
targets = targets[:max_kill]
|
||||
killed = 0
|
||||
results = []
|
||||
for t in targets:
|
||||
try:
|
||||
p = psutil.Process(t["pid"])
|
||||
p.kill(sig_num)
|
||||
killed += 1
|
||||
results.append(f" Killed PID {t['pid']} ({t['name']})")
|
||||
except Exception as e:
|
||||
results.append(f" Failed to kill PID {t['pid']}: {e}")
|
||||
|
||||
return (
|
||||
f"Killed {killed}/{len(targets)} process(es) matching '{name_pattern}'.\n\n"
|
||||
+ "\n".join(results)
|
||||
)
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
def system_resource_usage(top_n: int = 10, sort_by: str = "cpu") -> str:
|
||||
"""Get current system resource usage overview including CPU, memory, disk, and top resource consumers."""
|
||||
cpu_percent = psutil.cpu_percent(interval=1)
|
||||
mem = psutil.virtual_memory()
|
||||
disk = psutil.disk_usage("/")
|
||||
load_avg = os.getloadavg()
|
||||
cpu_count = psutil.cpu_count(logical=False) or 1
|
||||
threads = psutil.cpu_count(logical=True) or 1
|
||||
|
||||
key_map = {
|
||||
"cpu": lambda p: p["cpu_percent"],
|
||||
"memory": lambda p: p["memory_percent"],
|
||||
"rss": lambda p: p["memory_rss_mb"],
|
||||
}
|
||||
sort_key = key_map.get(sort_by, key_map["cpu"])
|
||||
|
||||
procs = []
|
||||
for proc in psutil.process_iter():
|
||||
info = _proc_to_info(proc)
|
||||
if info:
|
||||
procs.append(info)
|
||||
|
||||
procs.sort(key=sort_key, reverse=True)
|
||||
top_procs = procs[:top_n]
|
||||
|
||||
summary = (
|
||||
f"System Resource Usage\n"
|
||||
f"=====================\n\n"
|
||||
f"CPU Usage: {cpu_percent}% ({cpu_count} cores, {threads} threads)\n"
|
||||
f"Memory Usage: {mem.percent}% ({mem.used / 1024**3:.1f}GB / {mem.total / 1024**3:.1f}GB)\n"
|
||||
f"Disk Usage (/): {disk.percent}% ({disk.used / 1024**3:.1f}GB / {disk.total / 1024**3:.1f}GB)\n"
|
||||
f"Load Average: {load_avg[0]:.2f}, {load_avg[1]:.2f}, {load_avg[2]:.2f}\n\n"
|
||||
f"Top {top_n} Processes (sorted by {sort_by}):"
|
||||
)
|
||||
|
||||
for i, p in enumerate(top_procs, 1):
|
||||
summary += (
|
||||
f"\n {i:2d}. PID: {p['pid']:>6} | Name: {p['name']:<25} "
|
||||
f"| CPU: {p['cpu_percent']:>5}% | MEM: {p['memory_percent']:>4}% "
|
||||
f"| RSS: {p['memory_rss_mb']:>7}MB"
|
||||
)
|
||||
|
||||
return summary
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
def find_zombies(include_details: bool = True) -> str:
|
||||
"""Find and optionally clean zombie processes. Zombies are processes in 'zombie' status."""
|
||||
zombies = []
|
||||
for proc in psutil.process_iter():
|
||||
try:
|
||||
info = _proc_to_info(proc)
|
||||
if info and info["status"] == "zombie":
|
||||
zombies.append(info)
|
||||
data = response.json()
|
||||
except Exception:
|
||||
await asyncio.sleep(2)
|
||||
continue
|
||||
|
||||
if not zombies:
|
||||
return "No zombie processes found."
|
||||
|
||||
if include_details:
|
||||
lines = [f"Found {len(zombies)} zombie process(es):"]
|
||||
for z in zombies:
|
||||
lines.append(
|
||||
f" PID: {z['pid']}, Name: {z['name']}, "
|
||||
f"CMD: {z['cmd_line'][:80]}"
|
||||
)
|
||||
return "\n".join(lines)
|
||||
status = data.get("status")
|
||||
if status == "completed":
|
||||
return data.get("result", {})
|
||||
elif status == "failed":
|
||||
return {"error": data.get("error", "Unknown error")}
|
||||
else:
|
||||
return f"Found {len(zombies)} zombie process(es): {', '.join(str(z['pid']) for z in zombies)}"
|
||||
await asyncio.sleep(2)
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
def restart_service(service_name: str, confirm: bool = True) -> str:
|
||||
"""Restart a systemd service by name. Requires sudo privileges for systemctl commands."""
|
||||
if not confirm:
|
||||
return "Restart cancelled. Set confirm=true to proceed."
|
||||
def batch_ocr_pdf(
|
||||
input_dir: str,
|
||||
output_dir: str,
|
||||
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.
|
||||
|
||||
if not service_name:
|
||||
return "Error: service_name is required."
|
||||
Args:
|
||||
input_dir: Directory containing PDF files to process
|
||||
output_dir: Directory to save OCR results (Markdown files)
|
||||
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)
|
||||
"""
|
||||
# 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):
|
||||
return json.dumps(
|
||||
{"error": f"Input directory does not exist: {input_dir}"},
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
)
|
||||
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
# Find all PDF files
|
||||
pdf_files = [
|
||||
os.path.join(input_dir, f)
|
||||
for f in os.listdir(input_dir)
|
||||
if f.lower().endswith(".pdf")
|
||||
]
|
||||
|
||||
if not pdf_files:
|
||||
return json.dumps(
|
||||
{"message": f"No PDF files found in {input_dir}", "files_processed": 0},
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
)
|
||||
|
||||
results = []
|
||||
success_count = 0
|
||||
fail_count = 0
|
||||
|
||||
for pdf_file in pdf_files:
|
||||
filename = os.path.basename(pdf_file)
|
||||
base_name = os.path.splitext(filename)[0]
|
||||
output_md = os.path.join(output_dir, f"{base_name}.md")
|
||||
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["sudo", "systemctl", "restart", service_name],
|
||||
capture_output=True, text=True, timeout=30,
|
||||
# Read and encode PDF
|
||||
with open(pdf_file, "rb") as f:
|
||||
file_base64 = base64.b64encode(f.read()).decode()
|
||||
|
||||
# Submit task
|
||||
response = requests.post(
|
||||
f"{base_url}/task/submit",
|
||||
headers={
|
||||
"X-API-Key": api_key,
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
json={
|
||||
"file_base64": file_base64,
|
||||
"file_type": "pdf",
|
||||
"enable_ai_description": False,
|
||||
"output_type": "ocr_only",
|
||||
},
|
||||
timeout=30,
|
||||
)
|
||||
if result.returncode == 0:
|
||||
return f"Successfully restarted service '{service_name}'."
|
||||
else:
|
||||
return f"Failed to restart '{service_name}':\n{result.stderr.strip()}"
|
||||
except FileNotFoundError:
|
||||
return "Error: systemctl not found. Is this a systemd-based system?"
|
||||
except subprocess.TimeoutExpired:
|
||||
return "Error: restart operation timed out."
|
||||
except PermissionError:
|
||||
return "Error: permission denied. Ensure sudo access is configured."
|
||||
|
||||
if response.status_code != 200:
|
||||
fail_count += 1
|
||||
results.append(
|
||||
{
|
||||
"file": filename,
|
||||
"status": "failed",
|
||||
"error": f"API error: {response.status_code}",
|
||||
}
|
||||
)
|
||||
continue
|
||||
|
||||
task_id = response.json().get("task_id")
|
||||
if not task_id:
|
||||
fail_count += 1
|
||||
results.append(
|
||||
{
|
||||
"file": filename,
|
||||
"status": "failed",
|
||||
"error": "No task_id returned",
|
||||
}
|
||||
)
|
||||
continue
|
||||
|
||||
# Poll for result (runs in background thread to avoid blocking)
|
||||
loop = asyncio.new_event_loop()
|
||||
full_result = loop.run_until_complete(
|
||||
_poll_task_result(base_url, task_id, api_key)
|
||||
)
|
||||
loop.close()
|
||||
|
||||
if "error" in full_result:
|
||||
fail_count += 1
|
||||
results.append(
|
||||
{
|
||||
"file": filename,
|
||||
"status": "failed",
|
||||
"error": full_result["error"],
|
||||
}
|
||||
)
|
||||
continue
|
||||
|
||||
# Save Markdown result
|
||||
ocr_text = full_result.get("ocr_text", "")
|
||||
spatial_structure = full_result.get("spatial_structure", None)
|
||||
ai_description = full_result.get("ai_description", "")
|
||||
|
||||
with open(output_md, "w", encoding="utf-8") as f:
|
||||
f.write(f"# OCR 识别结果\n\n")
|
||||
f.write(
|
||||
f"**生成时间**: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n"
|
||||
)
|
||||
f.write(f"**Task ID**: `{task_id}`\n\n")
|
||||
f.write(f"**源文件**: `{filename}`\n\n")
|
||||
f.write("---\n\n")
|
||||
|
||||
if ocr_text:
|
||||
f.write(f"## 📝 OCR 识别结果\n\n")
|
||||
f.write(ocr_text)
|
||||
|
||||
if spatial_structure:
|
||||
f.write(f"\n\n## 🗺️ 空间结构信息\n\n")
|
||||
if "block_count" in spatial_structure:
|
||||
f.write(
|
||||
f"- **图片尺寸**: {spatial_structure.get('width', 'N/A')} x {spatial_structure.get('height', 'N/A')}\n"
|
||||
)
|
||||
f.write(
|
||||
f"- **识别块数量**: {spatial_structure.get('block_count', 0)}\n"
|
||||
)
|
||||
|
||||
if ai_description:
|
||||
f.write(f"\n\n## 🤖 AI 分析结果\n\n")
|
||||
f.write(ai_description)
|
||||
|
||||
success_count += 1
|
||||
results.append({"file": filename, "status": "success", "output": output_md})
|
||||
|
||||
except Exception as e:
|
||||
return f"Error restarting service: {e}"
|
||||
fail_count += 1
|
||||
results.append({"file": filename, "status": "failed", "error": str(e)})
|
||||
|
||||
# Delay between files
|
||||
time.sleep(2)
|
||||
|
||||
def main():
|
||||
"""Run the MCP server."""
|
||||
mcp.run()
|
||||
final_result = {
|
||||
"summary": {
|
||||
"total": len(pdf_files),
|
||||
"success": success_count,
|
||||
"fail": fail_count,
|
||||
"output_dir": output_dir,
|
||||
},
|
||||
"files": results,
|
||||
}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
return json.dumps(final_result, ensure_ascii=False, indent=2)
|
||||
|
||||
Reference in New Issue
Block a user