import os import base64 import mimetypes import logging import asyncio import io import datetime from typing import Any, Callable, Dict, List, Union from PIL import Image, ImageDraw, ImageFont logger = logging.getLogger("MattCP") class Tool: def __init__(self, name: str, description: str, schema: Dict[str, Any], handler: Callable): self.name = name self.description = description self.schema = schema self.handler = handler def to_mcp_dict(self): return { "name": self.name, "description": self.description, "inputSchema": self.schema } async def read_image_handler(args: Dict[str, Any]): path = args.get("path") if not path: return {"text": "Error: Missing path argument"} mime_type, _ = mimetypes.guess_type(path) if not mime_type or not mime_type.startswith("image/"): return {"text": f"Error: File {path} is not a supported image format."} try: with open(path, "rb") as f: data = base64.b64encode(f.read()).decode("utf-8") return {"type": "image", "data": data, "mimeType": mime_type} except Exception as e: return {"text": f"Error reading image: {str(e)}"} async def read_png_metadata_handler(args: Dict[str, Any]): path = args.get("path") if not path: return {"text": "Error: Missing path argument"} try: with Image.open(path) as img: if img.format != 'PNG': return {"text": "Error: File is not a PNG image."} metadata = img.info if not metadata: return {"text": "No metadata found in this PNG."} output = "\n".join([f"{k}: {v}" for k, v in metadata.items()]) return {"text": f"PNG Metadata:\n{output}"} except Exception as e: return {"text": f"Error reading PNG metadata: {str(e)}"} async def list_thumbnails_handler(args: Dict[str, Any]): dir_path = args.get("path") page = int(args.get("page", 1)) page_size = int(args.get("page_size", 64)) sort_by = args.get("sort_by", "mtime") if not dir_path: return {"text": "Error: Missing path argument"} if not os.path.isdir(dir_path): return {"text": f"Error: {dir_path} is not a directory."} exts = ('.png', '.jpg', '.jpeg', '.webp', '.bmp') file_info_list = [] for f in os.listdir(dir_path): if f.lower().endswith(exts): full_path = os.path.join(dir_path, f) try: stats = os.stat(full_path) file_info_list.append({ "name": f, "path": full_path, "mtime": stats.st_mtime, "size": stats.st_size }) except Exception as e: logger.error(f"Could not stat {f}: {e}") if not file_info_list: return {"text": "No supported images found in the directory."} if sort_by == "mtime": file_info_list.sort(key=lambda x: x["mtime"], reverse=True) elif sort_by == "size": file_info_list.sort(key=lambda x: x["size"], reverse=True) else: file_info_list.sort(key=lambda x: x["name"].lower()) total_files = len(file_info_list) start_idx = (page - 1) * page_size end_idx = start_idx + page_size paged_files = file_info_list[start_idx:end_idx] if not paged_files: return {"text": f"No images found on page {page}."} thumb_size = 160 padding = 15 label_height = 35 num_files = len(paged_files) cols = int(num_files**0.5) if num_files > 0 else 1 if cols == 0: cols = 1 rows = (num_files + cols - 1) // cols canvas_w = cols * (thumb_size + padding) + padding canvas_h = rows * (thumb_size + label_height + padding) + padding canvas = Image.new('RGB', (canvas_w, canvas_h), (30, 30, 30)) draw = ImageDraw.Draw(canvas) font = None font_paths = [ "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", "/usr/share/fonts/truetype/liberation/LiberationSans-Bold.ttf", "/usr/share/fonts/truetype/freefont/FreeSansBold.ttf", ] for path in font_paths: if os.path.exists(path): try: font = ImageFont.truetype(path, 20) break except Exception: continue if font is None: font = ImageFont.load_default() file_details = [] for i, info in enumerate(paged_files): filename = info["name"] full_path = info["path"] row = i // cols col = i % cols x = padding + col * (thumb_size + padding) y = padding + row * (thumb_size + label_height + padding) try: with Image.open(full_path) as img: width, height = img.size img.thumbnail((thumb_size, thumb_size)) off_x = (thumb_size - img.width) // 2 off_y = (thumb_size - img.height) // 2 canvas.paste(img, (x + off_x, y + off_y)) mod_time = datetime.datetime.fromtimestamp(info["mtime"]).strftime('%Y-%m-%d %H:%M') size_kb = info["size"] / 1024 file_details.append(f"{i+1}. {filename} | {width}x{height} | {size_kb:.1f}KB | {mod_time}") text = str(i + 1) text_w = len(text) * 10 draw.rectangle([x, y + thumb_size + 2, x + text_w + 4, y + thumb_size + 24], fill=(60, 60, 60)) draw.text((x + 2, y + thumb_size + 2), text, fill=(255, 255, 255), font=font) except Exception as e: logger.error(f"Failed to process {filename}: {e}") draw.text((x, y + thumb_size // 2), "Error", fill=(255, 0, 0), font=font) file_details.append(f"{i+1}. {filename} | ERROR") buf = io.BytesIO() canvas.save(buf, format='PNG') img_data = base64.b64encode(buf.getvalue()).decode("utf-8") total_pages = (total_files + page_size - 1) // page_size header_text = f"Directory: {dir_path}\nPage {page} of {total_pages} ({total_files} total images). Showing {start_idx+1}-{min(end_idx, total_files)}." details_text = "\n".join(file_details) return [ {"type": "text", "text": f"{header_text}\n\nFile List:\n{details_text}"}, {"type": "image", "data": img_data, "mimeType": "image/png"} ] async def list_directory_details_handler(args: Dict[str, Any]): path = args.get("path") if not path: return {"text": "Error: Missing path argument"} if not os.path.isdir(path): return {"text": f"Error: {path} is not a directory."} try: entries = os.listdir(path) except Exception as e: return {"text": f"Error reading directory: {e}"} dirs = [] files = [] for entry in entries: full_path = os.path.join(path, entry) try: stats = os.stat(full_path) mtime = datetime.datetime.fromtimestamp(stats.st_mtime).strftime('%Y-%m-%d %H:%M') if os.path.isdir(full_path): # Count items in dir (non-recursive) try: count = len(os.listdir(full_path)) except: count = 0 dirs.append({ "name": entry, "info": f"{count} items", "mtime": mtime }) else: size_kb = stats.st_size / 1024 files.append({ "name": entry, "info": f"{size_kb:.1f}KB", "mtime": mtime }) except Exception as e: logger.error(f"Could not stat {entry}: {e}") # Sort alphabetically dirs.sort(key=lambda x: x["name"].lower()) files.sort(key=lambda x: x["name"].lower()) output = [] for d in dirs: output.append(f"[DIR] {d['name']} | {d['info']} | {d['mtime']}") for f in files: output.append(f" {f['name']} | {f['info']} | {f['mtime']}") if not output: return {"text": f"Directory {path} is empty."} return {"text": f"Listing of {path}:\n\n" + "\n".join(output)} TOOL_REGISTRY = [ Tool( name="read_image", description="Reads an image from the disk and returns it as an image content object.", schema={ "type": "object", "properties": {"path": {"type": "string", "description": "Path to the image file"}}, "required": ["path"], }, handler=read_image_handler ), Tool( name="read_png_metadata", description="Reads the metadata (tEXt, zTXt, iTXt) from a PNG image to fetch info such as AI prompts.", schema={ "type": "object", "properties": {"path": {"type": "string", "description": "Path to the PNG file"}}, "required": ["path"], }, handler=read_png_metadata_handler ), Tool( name="list_thumbnails", description="Generates a thumbnail grid of images in a directory. Supports pagination and sorting.", schema={ "type": "object", "properties": { "path": {"type": "string", "description": "Path to the directory containing images"}, "page": {"type": "integer", "description": "The page number to retrieve (1-indexed)", "default": 1}, "page_size": {"type": "integer", "description": "Number of images per page", "default": 64}, "sort_by": {"type": "string", "description": "Sort order: 'mtime' (newest first, default), 'name' (alphabetical), or 'size' (largest first)", "default": "mtime"}, }, "required": ["path"], }, handler=list_thumbnails_handler ), Tool( name="list_directory", description="Lists the contents of a directory in a Nemo-like format. Directories first, then files, with sizes/counts and mtime.", schema={ "type": "object", "properties": { "path": {"type": "string", "description": "Path to the directory to list"}, }, "required": ["path"], }, handler=list_directory_details_handler ), ]