158 lines
9.4 KiB
Python
158 lines
9.4 KiB
Python
from typing import Any, Callable, Dict
|
|
|
|
# Define the Tool class here so it's available to all handlers and the registry
|
|
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
|
|
}
|
|
|
|
# Import handlers from separate files
|
|
from .read_image import handle as read_image_handler
|
|
from .read_metadata import handle as read_png_metadata_handler
|
|
from .contact_sheet import handle as contact_sheet_handler
|
|
from .list_directory import handle as list_directory_details_handler
|
|
from .preview_image import handle as preview_image_handler
|
|
from .get_text_context import handle as get_text_context_handler
|
|
from .wikipedia import handle as wikipedia_handler
|
|
from .get_model_info import handle as get_model_info_handler
|
|
from .generate_image import handle as generate_image_handler
|
|
|
|
# Central registry of all available tools
|
|
TOOL_REGISTRY = [
|
|
Tool(
|
|
name="read_image",
|
|
description="Reads an image at full resolution for maximum detail. Use this for final confirmation of a selected image or when deep visual analysis of a specific file is required.",
|
|
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 prompts for AI generated images.",
|
|
schema={
|
|
"type": "object",
|
|
"properties": {"path": {"type": "string", "description": "Path to the PNG file"}},
|
|
"required": ["path"],
|
|
},
|
|
handler=read_png_metadata_handler
|
|
),
|
|
Tool(
|
|
name="contact_sheet",
|
|
description="Generates a coarse, high-density overview of images in a directory. This tool is paginated; a single page may not show all images. To perform a comprehensive scan or find specific images, you MUST iterate through multiple pages. If ANYTHING requires selecting individual images, then use indices as inputs to preview_image. Rely on this tool ONLY for a general characterization of images in a directory, NOT anything about specific images.",
|
|
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). Increment this value to navigate through the full list of images in the folder.", "default": 1},
|
|
"sort_by": {"type": "string", "description": "Sort order: 'mtime' (newest first, default), 'name' (alphabetical), or 'size' (largest first)", "default": "mtime"},
|
|
},
|
|
"required": ["path"],
|
|
},
|
|
handler=contact_sheet_handler
|
|
),
|
|
Tool(
|
|
name="list_directory",
|
|
description="Lists the contents of a directory in a simplified text format. Use this for checking file existence, sizes, or dates or for quick filesystem navigation. DO NOT use this tool for any sort of visual or image selection.",
|
|
schema={
|
|
"type": "object",
|
|
"properties": {
|
|
"path": {"type": "string", "description": "Path to the directory to list"},
|
|
"page": {"type": "integer", "description": "The page number to retrieve (1-indexed)", "default": 1},
|
|
"page_size": {"type": "integer", "description": "Number of items 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_directory_details_handler
|
|
),
|
|
Tool(
|
|
name="preview_image",
|
|
description="Provides low-resolution previews and information such as filenames, size, and resolution for specific images. It uses the minimum viable token budget, so details WILL be missed. Only use this tool for quick down-selection of images from contact_sheet or where accuracy does not matter. For final confirmation or deep analysis of a single selected image, you MUST use read_image after this tool.",
|
|
schema={
|
|
"type": "object",
|
|
"properties": {
|
|
"path": {"type": "string", "description": "Path to an image file or a directory containing images"},
|
|
"indices": {"type": "string", "description": "1-based indices or ranges (e.g., '1, 3-5, 10') to preview from the directory. Required if path is a directory."},
|
|
"sort_by": {"type": "string", "description": "Sort order to resolve indices: 'mtime' (default), 'name', or 'size'."},
|
|
},
|
|
"required": ["path"],
|
|
},
|
|
handler=preview_image_handler
|
|
),
|
|
Tool(
|
|
name="get_text_context",
|
|
description="Extracts specific sections of a file with absolute line numbers and surrounding context. This tool MUST be called immediately before edit_file to verify the exact line numbers and content of the block being replaced, preventing misalignment errors.",
|
|
schema={
|
|
"type": "object",
|
|
"properties": {
|
|
"path": {"type": "string", "description": "Path to the file"},
|
|
"pattern": {"type": "string", "description": "Regex pattern to search for"},
|
|
"lines": {"type": "string", "description": "1-based indices or ranges (e.g., '10-20, 45')"},
|
|
"context": {"type": "integer", "description": "Number of lines of context to show above and below", "default": 1},
|
|
},
|
|
"required": ["path"],
|
|
},
|
|
handler=get_text_context_handler
|
|
),
|
|
Tool(
|
|
name="browse_wikipedia",
|
|
description="Browse Wikipedia pages to retrieve information. Defaults to the page summary. Can also retrieve the table of contents or a specific section. If a page is not found, it automatically returns search results.",
|
|
schema={
|
|
"type": "object",
|
|
"properties": {
|
|
"title": {"type": "string", "description": "The title of the Wikipedia page to browse or the search query."},
|
|
"mode": {
|
|
"type": "string",
|
|
"description": "The retrieval mode. 'summary' (default) for the lead section, 'toc' for the table of contents, 'section' for a specific section, or 'search' for explicit search results.",
|
|
"enum": ["summary", "toc", "section", "search"],
|
|
"default": "summary"
|
|
},
|
|
"section_index": {"type": "integer", "description": "The linear index of the section to retrieve. Required if mode='section'. This index can be found by calling the tool in 'toc' mode."},
|
|
"search_limit": {"type": "integer", "description": "The maximum number of search results to return. Default is 5.", "default": 5},
|
|
},
|
|
"required": ["title"],
|
|
},
|
|
handler=wikipedia_handler
|
|
),
|
|
Tool(
|
|
name="get_model_info",
|
|
description="Returns a list of available txt2img models and their short descriptions. If a specific model name is provided, it returns the comprehensive prompting guide, available resolution presets, and active configuration tips for that model. CRITICAL: You must call this for any model you are unfamiliar with, as prompting styles vary wildly (e.g., tag-based vs. natural language) and using the wrong style will result in poor image quality.",
|
|
schema={
|
|
"type": "object",
|
|
"properties": {
|
|
"model_name": {"type": "string", "description": "The name of the model to get detailed info for. Leave empty to list all available models."}
|
|
},
|
|
"required": [],
|
|
},
|
|
handler=get_model_info_handler
|
|
),
|
|
Tool(
|
|
name="generate_image",
|
|
description="Generates an image using the specified model and parameters. To ensure high quality, verify the model's prompting requirements via get_model_info before calling this tool. You will usually not get the correct result the first time. If the generated image needs work, change the prompt and call this tool again. If it's good, use the provided markdown to display it to the user.",
|
|
schema={
|
|
"type": "object",
|
|
"properties": {
|
|
"model_name": {"type": "string", "description": "The name of the model to use. Required."},
|
|
"prompt": {"type": "string", "description": "The prompt for the image. Required."},
|
|
"negative_prompt": {"type": "string", "description": "The negative prompt to exclude unwanted elements."},
|
|
"resolution_preset": {"type": "string", "description": "A named resolution preset (e.g., 'square', 'portrait'). Available options depend on the model."},
|
|
"cfg_scale": {"type": "number", "description": "CFG scale for prompt adherence. Usually should be left omitted to select the default."},
|
|
},
|
|
"required": ["model_name", "prompt"],
|
|
},
|
|
handler=generate_image_handler
|
|
),
|
|
]
|