diff --git a/README.md b/README.md index 0fad899..02da7d5 100644 --- a/README.md +++ b/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. diff --git a/src/procleaning_mcp/server.py b/src/procleaning_mcp/server.py index 2954849..5482c7e 100644 --- a/src/procleaning_mcp/server.py +++ b/src/procleaning_mcp/server.py @@ -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.""" - 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']}" - ) - 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) +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: - proc.wait(timeout=5) - except psutil.TimeoutExpired: - return ( - f"Warning: PID {pid} ({proc_name}) did not terminate within 5s. " - f"Consider using SIGKILL." + response = requests.get( + f"{api_url}/task/result/{task_id}", + headers={"X-API-Key": api_key}, + timeout=30, ) - 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}" - - -@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) - except Exception: + await asyncio.sleep(2) continue - if not zombies: - return "No zombie processes found." + if response.status_code != 200: + await asyncio.sleep(2) + continue - 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) - else: - return f"Found {len(zombies)} zombie process(es): {', '.join(str(z['pid']) for z in zombies)}" + try: + data = response.json() + except Exception: + await asyncio.sleep(2) + continue + + status = data.get("status") + if status == "completed": + return data.get("result", {}) + elif status == "failed": + return {"error": data.get("error", "Unknown error")} + else: + 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) - try: - result = subprocess.run( - ["sudo", "systemctl", "restart", service_name], - capture_output=True, text=True, timeout=30, + 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, ) - 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." - except Exception as e: - return f"Error restarting service: {e}" + os.makedirs(output_dir, exist_ok=True) -def main(): - """Run the MCP server.""" - mcp.run() + # 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, + ) -if __name__ == "__main__": - main() + 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: + # 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 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: + fail_count += 1 + results.append({"file": filename, "status": "failed", "error": str(e)}) + + # Delay between files + time.sleep(2) + + final_result = { + "summary": { + "total": len(pdf_files), + "success": success_count, + "fail": fail_count, + "output_dir": output_dir, + }, + "files": results, + } + + return json.dumps(final_result, ensure_ascii=False, indent=2)