Compare commits

..
29 Commits
Author SHA1 Message Date
moosecrap 627ab3aec6 ZIT prompting guide and tweaks 2026-08-05 22:32:01 -07:00
moosecrap 74eaa2a3e2 requirements.txt 2026-08-05 22:31:53 -07:00
moosecrap 747bf5e02b Wikipedia redirect 2026-08-03 22:10:44 -07:00
moosecrap 8c18d0d40e Wikipedia text cleanup 2026-08-03 16:27:59 -07:00
moosecrap bad4a24426 Model preset tweaks, tag csv 2026-08-03 16:27:46 -07:00
moosecrap 3e523e9725 Tag search fixes, WIP on Danbooru wiki 2026-07-28 16:07:18 -07:00
moosecrap b107aff150 User agent, readme 2026-07-28 13:15:12 -07:00
moosecrap 1109a33edb Model preset tweak 2026-07-27 05:30:08 -07:00
moosecrap e67aebbd3d Wiki rename 2026-07-27 05:02:24 -07:00
moosecrap 2d30bd2155 Model switch fix 2026-07-27 04:47:54 -07:00
moosecrap f6f021f2a3 Forge parameter fix 2026-07-27 04:38:34 -07:00
moosecrap 05a2a5010c Shutdown 2026-07-27 04:02:07 -07:00
moosecrap f423764645 Proxy config 2026-07-27 02:47:00 -07:00
moosecrap c0c3a4c405 Tag search limit 2026-07-27 01:45:52 -07:00
moosecrap f0e5267828 Return type standardization 2026-07-27 00:35:20 -07:00
moosecrap baa48d7cec Readme 2026-07-26 23:02:46 -07:00
moosecrap 81fef17af2 Tag not found update 2026-07-26 22:59:25 -07:00
moosecrap aa59df60d9 Model info and res preset fix 2026-07-26 22:42:20 -07:00
moosecrap e9ee89f8a3 Contact sheet small size fixes 2026-07-26 22:17:34 -07:00
moosecrap f676a07a71 Model switch but not quite working yet 2026-07-26 22:17:15 -07:00
moosecrap 51a4708092 ZIT presets 2026-07-26 19:41:16 -07:00
moosecrap 79c975dd28 Tag search 2026-07-26 13:13:38 -07:00
moosecrap 62c06b6254 More guidance tuning 2026-07-26 03:40:07 -07:00
moosecrap ee6465d994 Wiki description fix 2026-07-26 03:17:42 -07:00
moosecrap c24a37652c Prompting guide 2026-07-26 02:28:24 -07:00
moosecrap 8beb7fe2a7 NoobAIXL settings 2026-07-25 22:18:38 -07:00
moosecrap 689cd1244c Resolutions, remove flux.1 2026-07-25 21:19:27 -07:00
moosecrap ca33869007 Cleanup 2026-07-25 21:00:51 -07:00
moosecrap f808f4d599 Stable Diffusion WebUI integration. File CORP proxy. 2026-07-25 20:22:51 -07:00
19 changed files with 141644 additions and 94 deletions
+57 -17
View File
@@ -1,10 +1,15 @@
# MooseCP Image Server # MooseCP Image Server
A Model Context Protocol (MCP) server designed to provide LLMs with efficient, token-optimized visual access to image directories. Instead of dumping full-resolution images (which waste tokens and cause context overflow), MooseCP provides a hierarchical workflow: **List $\rightarrow$ Scan $\rightarrow$ Preview $\rightarrow$ Inspect**. A Model Context Protocol (MCP) server designed to provide LLMs with efficient, token-optimized visual access to image directories and AI generation capabilities. Instead of dumping full-resolution images (which waste tokens and cause context overflow), MooseCP provides a hierarchical workflow: **List $\rightarrow$ Scan $\rightarrow$ Preview $\rightarrow$ Inspect**.
## ⚠️ AI SLOP DISCLAIMER ## ⚠️ AI SLOP DISCLAIMER
This entire project was vibe-coded by an AI. It is 100% slop code. Use it at your own risk. This entire project was vibe-coded by an AI. It is 100% slop code. Use it at your own risk.
## 🚨 SECURITY WARNING
**This server is designed to be run on `localhost` ONLY.**
It contains tools (such as `read_image` and `list_directory`) that allow the LLM to read arbitrary files from your filesystem. If you expose this server to the network or a public IP, any user or compromised AI could potentially read sensitive system files (e.g., SSH keys, `/etc/passwd`).
**NEVER run this server on a public-facing IP without implementing strict path validation.**
## Tools Overview ## Tools Overview
### 📁 `list_directory` ### 📁 `list_directory`
@@ -32,6 +37,22 @@ Extracts AI generation parameters from PNG files.
Returns the full-resolution image. Returns the full-resolution image.
* **Best for**: Final confirmation or deep visual analysis where every pixel counts. * **Best for**: Final confirmation or deep visual analysis where every pixel counts.
### 🎨 `generate_image`
Triggers an image generation on a local Stable Diffusion WebUI Forge instance.
* **Workflow**: Always call `get_model_info` first to determine the correct prompting style (e.g., tag-based vs. natural language).
* **Features**: Supports model-specific presets, resolution presets, and standard parameter overrides.
* **Output**: Returns a Base64 image for AI analysis and a proxy URL for direct embedding in the chat.
### ️ `get_model_info`
Provides the "manual" for available generation models.
* **Best for**: Learning the prompting style, recommended settings, and available resolution presets for a specific model.
* **Workflow**: Call without arguments to see the catalog; call with `model_name` for the detailed guide.
### 🏷️ `search_tags`
Searches the Danbooru tag database for recognized tags and aliases.
* **Best for**: Finding the correct booru-style tags, checking tag popularity, and resolving aliases (e.g., 'lesbian' → 'yuri').
* **Workflow**: Provide a query string to get a list of the most popular matching tags.
### 🌐 `browse_wikipedia` ### 🌐 `browse_wikipedia`
Allows the model to browse Wikipedia using its API. Allows the model to browse Wikipedia using its API.
* **Best for**: Quickly retrieving summaries, structural maps (ToC), or specific section content from Wikipedia without dumping the entire page. * **Best for**: Quickly retrieving summaries, structural maps (ToC), or specific section content from Wikipedia without dumping the entire page.
@@ -47,9 +68,11 @@ This server requires Python 3.10+ and the following packages:
* `uvicorn`: ASGI server for the SSE transport. * `uvicorn`: ASGI server for the SSE transport.
* `starlette`: Lightweight ASGI framework. * `starlette`: Lightweight ASGI framework.
* `Pillow`: Image processing and thumbnail generation. * `Pillow`: Image processing and thumbnail generation.
* `requests`: For communicating with the Stable Diffusion API.
* `httpx`: For asynchronous API requests (e.g., Wikipedia).
```bash ```bash
pip install uvicorn starlette Pillow pip install -r requirements.txt
``` ```
### Setup ### Setup
@@ -62,21 +85,38 @@ pip install uvicorn starlette Pillow
## Configuration (`config.py`) ## Configuration (`config.py`)
You can tune the server's behavior in `config.py`: Tuning the server's behavior is done via `config.py`.
### Server & Logging ### 🌐 Server & Logging
- **`HOST` / `PORT`**: The network address and port the server binds to. | Parameter | Description | Default |
- **`LOG_LEVEL`**: Logging verbosity (`DEBUG`, `INFO`, `WARNING`, `ERROR`). | :--- | :--- | :--- |
- **`LOG_FILE`**: Path to the server log file. | `HOST` | The network address the server binds to. | `"127.0.0.1"` |
- **`USER_AGENT`**: The User-Agent string used for API requests (e.g., Wikipedia). Use a browser-like string to avoid 403 Forbidden errors. | `PORT` | The port the server listens on. | `8000` |
| `LOG_LEVEL` | Logging verbosity (`DEBUG`, `INFO`, `WARNING`, `ERROR`). | `"WARNING"` |
| `LOG_FILE` | Absolute path to the server log file. | `ROOT_DIR / "debug.log"` |
| `USER_AGENT` | User-Agent string for API requests (e.g., Wikipedia). | Browser-like string |
### Model & Token Tuning ### 🎨 Stable Diffusion Integration
- **`PATCH_SIZE`**: Set this to `(clip.vision.patch_size * n_merge)` for your specific model to ensure token-perfect resizing. | Parameter | Description | Default |
- **`PREVIEW_TOKEN_BUDGET`**: Controls how many tokens the `preview_image` tool aims for (default: 70). | :--- | :--- | :--- |
- **`CONTACT_SHEET_COLS` / `ROWS`**: Adjust the grid size to fit within your model's maximum context window. | `SD_URL` | Base URL of the SD WebUI/Forge instance. | `"http://127.0.0.1:7860"` |
- **`CONTACT_SHEET_THUMB_SIZE`**: The pixel size of thumbnails in the contact sheet. | `MODEL_PRESETS_PATH` | Path to the `model_presets.toml` file. | `ROOT_DIR / "model_presets.toml"` |
| `RES_PRESETS_PATH` | Path to the `resolution_presets.toml` file. | `ROOT_DIR / "resolution_presets.toml"` |
| `TAG_DATABASE_PATH` | Path to the Danbooru `tags.csv` file. | (Path to extension folder) |
| `TAG_SEARCH_LIMIT` | Number of results returned by `search_tags` (direct or similar). | `20` |
### Image & Font Settings ### 🧠 Model & Token Tuning (Optimized for Gemma 4)
- **`IMAGE_QUALITY`**: Adjust JPEG compression (1-100). | Parameter | Description | Default |
- **`SYSTEM_FONT_NAMES`**: A list of font names Pillow should attempt to find in the system path. | :--- | :--- | :--- |
- **`FALLBACK_FONT_PATHS`**: Absolute paths to `.ttf` files used if no system fonts are found. | `PATCH_SIZE` | Model's vision patch size in pixels. | `48` |
| `PREVIEW_TOKEN_BUDGET` | Target token count for `preview_image` thumbnails. | `70` |
| `CONTACT_SHEET_COLS` | Number of columns in the contact sheet grid. | `10` |
| `CONTACT_SHEET_ROWS` | Number of rows in the contact sheet grid. | `7` |
| `CONTACT_SHEET_THUMB_SIZE` | Pixel size of thumbnails in the contact sheet. | `192` |
### 🖼️ Image & Font Settings
| Parameter | Description | Default |
| :--- | :--- | :--- |
| `IMAGE_QUALITY` | JPEG compression quality (1-100). | `95` |
| `SYSTEM_FONT_NAMES` | List of font names for Pillow to try in system paths. | Arial, DejaVu, etc. |
| `FALLBACK_FONT_PATHS` | List of absolute paths to `.ttf` files. | Linux-specific paths |
+20 -3
View File
@@ -1,11 +1,28 @@
import logging import logging
from pathlib import Path
from contextvars import ContextVar
# --- Project Root ---
# Get the directory where config.py is located
ROOT_DIR = Path(__file__).parent.resolve()
# --- Server & Logs --- # --- Server & Logs ---
HOST = "127.0.0.1" HOST = "127.0.0.1"
PORT = 8000 PORT = 8000
LOG_LEVEL = "WARNING" # Options: "DEBUG", "INFO", "WARNING", "ERROR" request_host = ContextVar("request_host", default=f"{HOST}:{PORT}")
LOG_FILE = "debug.log" LOG_LEVEL = "DEBUG" # Options: "DEBUG", "INFO", "WARNING", "ERROR"
USER_AGENT = "Mozilla/5.0 (X11; Linux x86_64; rv:151.0) Gecko/20100101 Firefox/151.0" # Stealth User-Agent to bypass Wikipedia's bot detection LOG_FILE = str(ROOT_DIR / "debug.log")
USER_AGENT = "MooseCP/1.0 Local MCP Server (https://long-cat.net/)"
# --- Stable Diffusion Config ---
SD_URL = "http://127.0.0.1:7860"
MODEL_PRESETS_PATH = str(ROOT_DIR / "model_presets.toml")
RES_PRESETS_PATH = str(ROOT_DIR / "resolution_presets.toml")
TAG_DATABASE_PATH = str(ROOT_DIR / "danbooru.csv")
TAG_SEARCH_LIMIT = 20
ENABLE_TAG_WIKI = False
DANBOORU_LOGIN = "" # Your Danbooru username (Optional)
DANBOORU_API_KEY = "" # Your Danbooru API key (Optional)
# --- Model Specific Token Tuning (Tuned for Gemma 4) --- # --- Model Specific Token Tuning (Tuned for Gemma 4) ---
# Patch size is typically (clip.vision.patch_size * n_merge) # Patch size is typically (clip.vision.patch_size * n_merge)
+140782
View File
File diff suppressed because it is too large Load Diff
+36 -6
View File
@@ -3,9 +3,10 @@ import logging
import uvicorn import uvicorn
import signal import signal
import os import os
import sys
from starlette.applications import Starlette from starlette.applications import Starlette
from starlette.routing import Route from starlette.routing import Route
from starlette.responses import Response, StreamingResponse from starlette.responses import Response, StreamingResponse, FileResponse
from starlette.middleware.cors import CORSMiddleware from starlette.middleware.cors import CORSMiddleware
from mcp_logic import MCPServer from mcp_logic import MCPServer
@@ -56,6 +57,10 @@ async def messages_endpoint(request):
if not sid or sid not in mcp_logic.sessions: if not sid or sid not in mcp_logic.sessions:
return Response("Session not found", status_code=404) return Response("Session not found", status_code=404)
# Capture the host header to ensure image links work across the network
host_header = request.headers.get("host", f"{config.HOST}:{config.PORT}")
config.request_host.set(host_header)
try: try:
body = await request.json() body = await request.json()
except Exception as e: except Exception as e:
@@ -75,10 +80,35 @@ async def messages_endpoint(request):
await mcp_logic.send_to_session(sid, {"jsonrpc": "2.0", "id": req_id, "result": result}) await mcp_logic.send_to_session(sid, {"jsonrpc": "2.0", "id": req_id, "result": result})
return Response(status_code=202) return Response(status_code=202)
async def file_endpoint(request):
"""
Proxy endpoint to serve files from disk with CORP/CORS headers to bypass browser restrictions.
"""
path = request.path_params.get("path")
if not path:
return Response("Path not provided", status_code=400)
# Ensure the path is absolute (SDFiles are usually in /tmp/gradio/...)
# If the path doesn't start with /, we assume it's relative to root for simplicity in this context
full_path = path if path.startswith("/") else f"/{path}"
if not os.path.exists(full_path):
return Response(f"File not found: {full_path}", status_code=404)
return FileResponse(
full_path,
headers={
"Cross-Origin-Resource-Policy": "cross-origin",
"Access-Control-Allow-Origin": "*",
"Access-Control-Allow-Methods": "GET"
}
)
app = Starlette( app = Starlette(
routes=[ routes=[
Route("/sse", endpoint=sse_endpoint), Route("/sse", endpoint=sse_endpoint),
Route("/messages", endpoint=messages_endpoint, methods=["POST"]), Route("/messages", endpoint=messages_endpoint, methods=["POST"]),
Route("/file/{path:path}", endpoint=file_endpoint),
] ]
) )
@@ -91,7 +121,8 @@ app.add_middleware(
) )
if __name__ == "__main__": if __name__ == "__main__":
config = uvicorn.Config( # Use a separate variable for uvicorn config to avoid shadowing the config module
uvicorn_config = uvicorn.Config(
app=app, app=app,
host=config.HOST, host=config.HOST,
port=config.PORT, port=config.PORT,
@@ -99,7 +130,7 @@ if __name__ == "__main__":
timeout_graceful_shutdown=2 # Reduced from 5 to 2 seconds timeout_graceful_shutdown=2 # Reduced from 5 to 2 seconds
) )
server = uvicorn.Server(config) server = uvicorn.Server(uvicorn_config)
def signal_handler(sig, frame): def signal_handler(sig, frame):
logger.info("Shutdown signal received") logger.info("Shutdown signal received")
@@ -115,6 +146,5 @@ if __name__ == "__main__":
except KeyboardInterrupt: except KeyboardInterrupt:
pass pass
finally: finally:
# Final fallback to ensure the process actually dies logger.info("Server process exiting.")
# if the event loop is still hanging on a connection sys.exit(0)
os._exit(0)
+3 -13
View File
@@ -42,23 +42,13 @@ class MCPServer:
) )
try: try:
result_data = await tool.handler(args) content = await tool.handler(args)
except ToolError as e: except ToolError as e:
# Convert tool errors into a text response so the LLM can understand and potentially fix the input # Convert tool errors into a text response so the LLM can understand and potentially fix the input
result_data = {"text": str(e)} content = [{"type": "text", "text": str(e)}]
except Exception as e: except Exception as e:
# Catch-all for unexpected crashes to prevent server death # Catch-all for unexpected crashes to prevent server death
result_data = {"text": f"Unexpected internal error: {str(e)}"} content = [{"type": "text", "text": f"Unexpected internal error: {str(e)}"}]
# Handle results that are already lists of content (for multimodal/multi-part responses)
if isinstance(result_data, list):
content = result_data
elif isinstance(result_data, dict) and "text" in result_data:
content = [{"type": "text", "text": result_data["text"]}]
elif isinstance(result_data, dict) and result_data.get("type") == "image":
content = [result_data]
else:
content = [result_data]
return {"content": content} return {"content": content}
+54
View File
@@ -0,0 +1,54 @@
# Model Presets
# Format: ["Model Name"]
# Description and Guide are mandatory.
# Other keys should be the Labels found in the /info endpoint.
["NoobAI XL"]
description = "Anime-style model, very good at specific artist styles and character knowledge. Not aesthetic-tuned: needs precise, explicit prompting for best results. Can do any sort of NSFW."
guide = """
This model is not aesthetic tuned, it must be given explicit tags for everything. Omitted parts of the prompt will not default to something 'good'. This model shows high variance, you can often get a very different image by just resubmitting the same prompt, so do not hesitate to try again.
Has very excellent understanding of characters and artists down to extremely niche. Unless prompting an original character, their name is enough to decribe their appearance completely except for clothing.
Accepts a list of comma-separated booru-style tags. Use spaces, not underscores, for tags. **Use the `search_tags` tool to verify your tags**.
Prompts MUST follow this format: <1girl/1boy/1other/solo/couple/(can use multiple)>, <character(s)>, <series>, <artist>, <tags>
Every prompt MUST include every one of the above sections (the angle brackets are not part of the prompt).
Quality tags such as "masterpiece", "best quality", "very awa", are a LAST RESORT, they override the artist tags. If absolutely required they should be prepended.
It understands every artist, so pick one appropriate for the image. If desired, it also knows 'implicit' artist tags such as "official art" or "game cg" for the 'artist' immediately following the series name.
Natural language understanding is very limited, but can do things like "dark blue skirt" or natural language order of tags such as "lying, on bed".
Do not use a negative prompt unless explicitly required to exclude something, your first prompt should have a blank negative prompt.
Has the SDXL problem with hands, works best if hand posture is explicitly prompted.
No default background, so prompts should include "outdoors", "indoors", or something like "patterned background" etc.
CFG Scale from 3.0 - 5.5 but the default 4.0 is usually fine.
"""
filename = "noobaiXLNAIXL_vPred10Version.safetensors"
"Resolution Set" = "sdxl"
preset = "xl"
"CFG Scale" = 4.0
"Hires. fix" = true
"Denoising strength" = 0.5
"Upscale by" = 1.5
"Upscaler" = "Lanczos"
"Hires steps" = 16
"Hires CFG Scale" = 4.0
"Sampling Steps" = 32
"Sampling Method" = "Euler a"
"Schedule Type" = "Uniform"
"Rescale CFG" = 0.3
["Z-Image-Turbo"]
description = "General image generation model. Aesthetic tuned, gets good results first try. Can do softcore NSFW, e.g. underwear, breasts, asses. No full-frontal nudity."
guide = """
This model is aesthetic tuned, regenerating with the same prompt will yield essentially the same image. Change the prompt before resubmitting.
It's a CFG 1.0 turbo model, the negative prompt has no effect.
Understands natural language very well, uses Qwen 3 4B as the text encoder. The longer and more detailed prompt the better, take advantage of line breaks and formatting.
If any part of the image is left out of the prompt, it will default to generic AI slop which is most NOT what you want. Be very explicit about each character's details. Appearance, ethnicity, age, individual outfit components, facial expression, pose, action, where they are looking, position in the image, etc. should ALL be included in the prompt. Characters can be referenced by naming them and using this name later in the prompt. This also helps the model avoid mixing traits between them.
This also applies to the image itself. Composition, framing, lighting, image style, setting, background, etc. should all be explicitly specified in the prompt.
Has some idiosyncrasies so you may need to iterate the prompt a few times to get around some weird artifacts. Think things like "green eyes" making them glow green, or "blush" making the entire face glow. Also really wants to make shirts tucked in for some reason. Examine the generated image closely and edit the prompt if needed.
"""
preset = "zit"
filename = "z_image_turbo_bf16.safetensors"
modules = ["ae.safetensors", "qwen_3_4b_abliterated.safetensors"]
"Resolution Set" = "2k"
"CFG Scale" = 1.0
"Sampling Steps" = 9
"Sampling Method" = "Euler"
"Schedule Type" = "Beta"
+5
View File
@@ -0,0 +1,5 @@
uvicorn
starlette
requests
httpx
Pillow
+15
View File
@@ -0,0 +1,15 @@
# Resolution Presets
# Format: [preset_name]
# Values: "WidthxHeight"
[sdxl]
widescreen = "1344x768"
landscape = "1152x896"
square = "1024x1024"
portrait = "896x1152"
[2k]
widescreen = "2048x1152"
landscape = "2048x1536"
square = "2048x2048"
portrait = "1536x2048"
+44 -1
View File
@@ -22,7 +22,10 @@ from .contact_sheet import handle as contact_sheet_handler
from .list_directory import handle as list_directory_details_handler from .list_directory import handle as list_directory_details_handler
from .preview_image import handle as preview_image_handler from .preview_image import handle as preview_image_handler
from .get_text_context import handle as get_text_context_handler from .get_text_context import handle as get_text_context_handler
from .wikipedia import handle as wikipedia_handler from .browse_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
from .search_tags import handle as search_tags_handler
# Central registry of all available tools # Central registry of all available tools
TOOL_REGISTRY = [ TOOL_REGISTRY = [
@@ -124,4 +127,44 @@ TOOL_REGISTRY = [
}, },
handler=wikipedia_handler 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 from the model info"},
"cfg_scale": {"type": "number", "description": "CFG scale for prompt adherence. Usually should be left omitted to select the default."},
},
"required": ["model_name", "prompt", "resolution_preset"],
},
handler=generate_image_handler
),
Tool(
name="search_tags",
description="Searches the Danbooru tag database for tags matching a query. Returns the Danbooru wiki page on an exact match. On any query, returns a list of similar tags as well as aliases. Useful for finding the correct booru-style tags for anime models or getting more information.",
schema={
"type": "object",
"properties": {
"query": {"type": "string", "description": "The tag or partial tag to search for (e.g., 'white hair' or 'lesbian')."}
},
"required": ["query"],
},
handler=search_tags_handler
),
] ]
@@ -1,9 +1,8 @@
import urllib.request import httpx
import urllib.parse
import json import json
import re import re
import asyncio import asyncio
from typing import Any, Dict, Optional from typing import Any, Dict, List, Optional
from .utils import ToolError from .utils import ToolError
import config import config
@@ -13,22 +12,33 @@ def strip_html(text: str) -> str:
"""Removes HTML tags from a string using regex to provide clean text to the LLM.""" """Removes HTML tags from a string using regex to provide clean text to the LLM."""
return re.sub(r'<[^>]*>', '', text) return re.sub(r'<[^>]*>', '', text)
def _make_request(params: Dict[str, Any]) -> Dict[str, Any]: def clean_wikitext(text: str) -> str:
""" """
Synchronous helper to make the API request using urllib. Removes the most distracting elements of raw Wikitext:
urllib is used instead of httpx to avoid TLS/HTTP fingerprinting 1. HTML comments (<!-- ... -->)
that triggers 403 Forbidden responses from Wikipedia. 2. Citations (<ref /> and <ref>...</ref>)
""" """
query_string = urllib.parse.urlencode(params) # Remove HTML comments
url = f"{API_URL}?{query_string}" text = re.sub(r'<!--.*?-->', '', text, flags=re.DOTALL)
# Remove self-closing citations FIRST to prevent them being seen as opening tags
text = re.sub(r'<ref[^>]*/>', '', text)
# Remove paired citations
text = re.sub(r'<ref[^>]*>.*?</ref>', '', text, flags=re.DOTALL)
return text
async def _make_request(params: Dict[str, Any]) -> Dict[str, Any]:
"""
Asynchronous helper to make the API request using httpx.
A custom User-Agent is used to avoid bot detection.
"""
headers = { headers = {
"User-Agent": config.USER_AGENT "User-Agent": config.USER_AGENT
} }
req = urllib.request.Request(url, headers=headers) async with httpx.AsyncClient(headers=headers, timeout=15.0) as client:
with urllib.request.urlopen(req) as response: response = await client.get(API_URL, params=params)
return json.loads(response.read().decode('utf-8')) response.raise_for_status()
return response.json()
async def _search(title: str, limit: int = 5, fallback: bool = False) -> str: async def _search(title: str, limit: int = 5, fallback: bool = False) -> str:
""" """
@@ -42,15 +52,15 @@ async def _search(title: str, limit: int = 5, fallback: bool = False) -> str:
"format": "json", "format": "json",
"srlimit": limit "srlimit": limit
} }
# Run synchronous urllib call in a thread to avoid blocking the event loop # Directly await the async request
data = await asyncio.to_thread(_make_request, params) data = await _make_request(params)
search_results = data.get("query", {}).get("search", []) search_results = data.get("query", {}).get("search", [])
if not search_results: if not search_results:
return f"No Wikipedia results found for '{title}'." return f"No Wikipedia results found for '{title}'."
if fallback: if fallback:
header = f"Article not found. Please select one of the following similar articles to proceed with '{title}':" header = f"Article not found. You MUST select one of the following articles and submit another query:"
else: else:
header = f"Search results for '{title}':" header = f"Search results for '{title}':"
@@ -62,14 +72,15 @@ async def _search(title: str, limit: int = 5, fallback: bool = False) -> str:
lines.append(f"{i}. [{title_res}] - Snippet: {snippet}") lines.append(f"{i}. [{title_res}] - Snippet: {snippet}")
if fallback: if fallback:
lines.append("\n[Tip: Use the title in brackets [ ] to explore the selected page.]") lines.append("\nUse the title in brackets [ ] to explore the selected page.")
return "\n".join(lines) return "\n".join(lines)
async def _fetch_content(title: str, section_index: int) -> Optional[str]: async def _fetch_content(title: str, section_index: int) -> tuple[Optional[str], Optional[str]]:
""" """
Retrieves raw Wikitext for a specific section. Retrieves raw Wikitext for a specific section.
section_index=0 returns the lead section. section_index=0 returns the lead section.
Returns a tuple of (content, final_title).
""" """
params = { params = {
"action": "query", "action": "query",
@@ -80,27 +91,30 @@ async def _fetch_content(title: str, section_index: int) -> Optional[str]:
"redirects": 1, "redirects": 1,
"format": "json" "format": "json"
} }
data = await asyncio.to_thread(_make_request, params) data = await _make_request(params)
pages = data.get("query", {}).get("pages", {}) pages = data.get("query", {}).get("pages", {})
if not pages: if not pages:
return None return None, None
page_id = next(iter(pages)) page_id = next(iter(pages))
page = pages[page_id] page = pages[page_id]
final_title = page.get("title")
if "missing" in page: if "missing" in page:
return None return None, None
revisions = page.get("revisions", []) revisions = page.get("revisions", [])
if not revisions: if not revisions:
return None return None, final_title
return revisions[0].get("*") content = revisions[0].get("*")
return (clean_wikitext(content) if content else None), final_title
async def _fetch_toc(title: str) -> Optional[str]: async def _fetch_toc(title: str) -> tuple[Optional[str], Optional[str]]:
""" """
Retrieves the Table of Contents data and formats it hierarchically. Retrieves the Table of Contents data and formats it hierarchically.
Returns a tuple of (toc, final_title).
""" """
params = { params = {
"action": "parse", "action": "parse",
@@ -109,18 +123,19 @@ async def _fetch_toc(title: str) -> Optional[str]:
"format": "json", "format": "json",
"redirects": 1 "redirects": 1
} }
data = await asyncio.to_thread(_make_request, params) data = await _make_request(params)
parse_data = data.get("parse") parse_data = data.get("parse")
if not parse_data: if not parse_data:
return None return None, None
final_title = parse_data.get("title")
toc_data = parse_data.get("tocdata", {}) toc_data = parse_data.get("tocdata", {})
sections = toc_data.get("sections", []) sections = toc_data.get("sections", [])
if not sections: if not sections:
return "No table of contents found for this page." return "No table of contents found for this page.", final_title
lines = [f"Table of Contents for \"{parse_data.get('title', title)}\":"] lines = [f"Table of Contents for \"{final_title}\":"]
for s in sections: for s in sections:
level = s.get("tocLevel", 1) level = s.get("tocLevel", 1)
index = s.get("index") index = s.get("index")
@@ -128,9 +143,9 @@ async def _fetch_toc(title: str) -> Optional[str]:
indent = " " * (level - 1) indent = " " * (level - 1)
lines.append(f"{indent}[{index}] {line}") lines.append(f"{indent}[{index}] {line}")
return "\n".join(lines) return "\n".join(lines), final_title
async def handle(args: Dict[str, Any]) -> Dict[str, Any]: async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
""" """
Main handler for the browse_wikipedia tool. Main handler for the browse_wikipedia tool.
Supports modes: summary, toc, section, and search. Supports modes: summary, toc, section, and search.
@@ -143,18 +158,24 @@ async def handle(args: Dict[str, Any]) -> Dict[str, Any]:
section_index = args.get("section_index") section_index = args.get("section_index")
search_limit = args.get("search_limit", 5) search_limit = args.get("search_limit", 5)
final_title = None
if mode == "search": if mode == "search":
result = await _search(title, search_limit, fallback=False) result = await _search(title, search_limit, fallback=False)
elif mode == "summary": elif mode == "summary":
# Lead section + ToC is the default 'summary' to guide the AI's next steps # Lead section + ToC is the default 'summary' to guide the AI's next steps
content = await _fetch_content(title, 0) content, title_from_content = await _fetch_content(title, 0)
final_title = title_from_content
if content is None: if content is None:
result = await _search(title, search_limit, fallback=True) result = await _search(title, search_limit, fallback=True)
else: else:
toc = await _fetch_toc(title) toc, title_from_toc = await _fetch_toc(title)
if title_from_toc:
final_title = title_from_toc
result = f"{content}\n\n---\n\n{toc}" result = f"{content}\n\n---\n\n{toc}"
elif mode == "toc": elif mode == "toc":
toc = await _fetch_toc(title) toc, title_from_toc = await _fetch_toc(title)
final_title = title_from_toc
if toc is None: if toc is None:
result = await _search(title, search_limit, fallback=True) result = await _search(title, search_limit, fallback=True)
else: else:
@@ -162,7 +183,8 @@ async def handle(args: Dict[str, Any]) -> Dict[str, Any]:
elif mode == "section": elif mode == "section":
if section_index is None: if section_index is None:
raise ToolError("Missing required parameter 'section_index' for mode='section'") raise ToolError("Missing required parameter 'section_index' for mode='section'")
content = await _fetch_content(title, int(section_index)) content, title_from_content = await _fetch_content(title, int(section_index))
final_title = title_from_content
if content is None: if content is None:
result = await _search(title, search_limit, fallback=True) result = await _search(title, search_limit, fallback=True)
else: else:
@@ -170,4 +192,7 @@ async def handle(args: Dict[str, Any]) -> Dict[str, Any]:
else: else:
raise ToolError(f"Invalid mode '{mode}'. Supported modes: summary, toc, section, search") raise ToolError(f"Invalid mode '{mode}'. Supported modes: summary, toc, section, search")
return {"text": result} if final_title and final_title != title:
result = f"Redirected to \"{final_title}\"\n\n{result}"
return [{"type": "text", "text": result}]
+24 -11
View File
@@ -2,7 +2,9 @@ import base64
import logging import logging
import io import io
import datetime import datetime
import math
from pathlib import Path from pathlib import Path
from typing import Any, Dict, List from typing import Any, Dict, List
from PIL import Image, ImageDraw, ImageFont, ImageOps from PIL import Image, ImageDraw, ImageFont, ImageOps
import config import config
@@ -14,8 +16,7 @@ logger = logging.getLogger("MooseCP")
async def handle(args: Dict[str, Any]): async def handle(args: Dict[str, Any]):
""" """
Generates a contact sheet of images. Generates a contact sheet of images.
Grid: 10 columns x 7 rows (70 images total). Sized for optimal token usage according to config values.
Sized for optimal token usage (1120 tokens) on llama.cpp.
""" """
dir_path_str = args.get("path") dir_path_str = args.get("path")
page = int(args.get("page", 1)) page = int(args.get("page", 1))
@@ -44,10 +45,10 @@ async def handle(args: Dict[str, Any]):
} }
sort_label = sort_labels.get(sort_by, "Name (alphabetical)") sort_label = sort_labels.get(sort_by, "Name (alphabetical)")
# Fixed grid dimensions for token optimization # Grid dimensions for token optimization (see config.py)
COLS = config.CONTACT_SHEET_COLS MAX_COLS = config.CONTACT_SHEET_COLS
ROWS = config.CONTACT_SHEET_ROWS MAX_ROWS = config.CONTACT_SHEET_ROWS
PAGE_SIZE = COLS * ROWS PAGE_SIZE = MAX_COLS * MAX_ROWS
total_files = len(file_info_list) total_files = len(file_info_list)
@@ -56,9 +57,21 @@ async def handle(args: Dict[str, Any]):
if not paged_files: if not paged_files:
return {"text": f"No images found on page {page}."} return {"text": f"No images found on page {page}."}
thumb_size = 192 # 192 / 48 = 4 patches per side # Dynamic grid sizing to avoid blank space
canvas_w = COLS * thumb_size n_images = len(paged_files)
canvas_h = ROWS * thumb_size thumb_size = config.CONTACT_SHEET_THUMB_SIZE
if n_images <= config.CONTACT_SHEET_ROWS ** 2:
# Smallest square-ish grid that fits all
cols = math.ceil(math.sqrt(n_images))
rows = math.ceil(n_images / cols)
else:
# Lock rows, expand columns
rows = config.CONTACT_SHEET_ROWS
cols = math.ceil(n_images / config.CONTACT_SHEET_ROWS)
canvas_w = cols * thumb_size
canvas_h = rows * thumb_size
canvas = Image.new('RGB', (canvas_w, canvas_h), (30, 30, 30)) canvas = Image.new('RGB', (canvas_w, canvas_h), (30, 30, 30))
draw = ImageDraw.Draw(canvas) draw = ImageDraw.Draw(canvas)
@@ -87,8 +100,8 @@ async def handle(args: Dict[str, Any]):
start_idx = (page - 1) * PAGE_SIZE start_idx = (page - 1) * PAGE_SIZE
for i, info in enumerate(paged_files): for i, info in enumerate(paged_files):
full_path = info["path"] full_path = info["path"]
row = i // COLS row = i // cols
col = i % COLS col = i % cols
x = col * thumb_size x = col * thumb_size
y = row * thumb_size y = row * thumb_size
+245
View File
@@ -0,0 +1,245 @@
import requests
import base64
import asyncio
import config
from typing import Any, Dict, List
from tools.utils import ToolError, load_toml
async def ensure_model_state(model_name: str, model_preset: Dict[str, Any]):
"""
Checks the current server state and switches model/modules if they differ from the preset.
"""
try:
config_resp = await asyncio.to_thread(requests.get, f"{config.SD_URL}/config", timeout=10)
config_resp.raise_for_status()
cfg_data = config_resp.json()
components = cfg_data.get("components", [])
# Extract current state from components
current_state = {}
for comp in components:
elem_id = comp.get("props", {}).get("elem_id")
if elem_id:
current_state[elem_id] = comp.get("props", {}).get("value")
# 1. Check and change Checkpoint
active_ckpt = current_state.get("setting_sd_model_checkpoint")
# Use 'filename' from TOML if available, otherwise fall back to the model_name key
target_ckpt = model_preset.get("filename", model_name)
if active_ckpt != target_ckpt:
preset = model_preset.get("preset", "xl")
await asyncio.to_thread(
requests.post,
f"{config.SD_URL}/api/checkpoint_change",
json={"data": [target_ckpt, preset]},
timeout=30
)
# 2. Check and change VAE / Text Encoders
active_modules = current_state.get("setting_sd_modules", [])
target_modules = model_preset.get("modules", [])
if active_modules != target_modules:
preset = model_preset.get("preset", "xl")
await asyncio.to_thread(
requests.post,
f"{config.SD_URL}/api/modules_change",
json={"data": [target_modules, preset]},
timeout=30
)
except Exception as e:
# We log this but don't necessarily raise a ToolError unless the generation itself fails,
# as the server might still be able to generate if it's just a state-check failure.
print(f"Warning: Failed to sync model state: {e}")
async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
"""
Generates an image using the specified model and parameters.
"""
# 1. REQUIRED ARGUMENTS
prompt = args.get("prompt")
if not prompt:
raise ToolError("The 'prompt' argument is required.")
model_name = args.get("model_name")
if not model_name:
raise ToolError("The 'model_name' argument is required. Use get_model_info to see available models.")
res_preset_name = args.get("resolution_preset")
if not res_preset_name:
raise ToolError("The 'resolution_preset' argument is required. Use get_model_info to see available resolutions for the chosen model.")
# 2. FETCH CURRENT SERVER DEFAULTS & MAP LABELS
try:
info_resp = await asyncio.to_thread(requests.get, f"{config.SD_URL}/info", timeout=10)
info_resp.raise_for_status()
info_data = info_resp.json()
params_info = info_data["named_endpoints"]["/txt2img"]["parameters"]
except Exception as e:
raise ToolError(f"Failed to connect to Stable Diffusion server: {str(e)}")
# Build the label-to-index map and a sparse payload
label_map = {}
label_counts = {}
# Find the maximum param index to determine payload size
max_param_idx = 0
for p in params_info:
name = p.get("parameter_name", "")
if name.startswith("param_"):
try:
idx = int(name.replace("param_", ""))
max_param_idx = max(max_param_idx, idx)
except ValueError:
pass
# Initialize payload with Nones (size is max_idx + 1)
payload = [None] * (max_param_idx + 1)
for idx_in_list, p in enumerate(params_info):
name = p.get("parameter_name", "")
label = p.get("label")
default = p.get("parameter_default")
# Determine the absolute index in the payload
if name == "id_task":
abs_idx = 0
elif name.startswith("param_"):
try:
abs_idx = int(name.replace("param_", ""))
except ValueError:
continue
else:
# Fallback for unexpected names, though unlikely
continue
# Set the default value at the absolute index
payload[abs_idx] = default
# Handle label mapping for AI overrides
if not label or label.startswith("parameter_"):
label = name
if label in label_counts:
label_counts[label] += 1
mapped_label = f"{label} [{label_counts[label]}]"
else:
label_counts[label] = 0
mapped_label = label
label_map[mapped_label] = abs_idx
# 3. LOAD CONFIGS
models_cfg = load_toml(config.MODEL_PRESETS_PATH)
res_cfg = load_toml(config.RES_PRESETS_PATH)
if model_name not in models_cfg:
raise ToolError(f"Model '{model_name}' not found in presets. Please call get_model_info to see available models and their correct names.")
model_preset = models_cfg[model_name]
# Ensure server state (Checkpoint, VAE, etc.) matches the preset before generating
await ensure_model_state(model_name, model_preset)
# 4. MERGE PIPELINE
for key, value in model_preset.items():
if key in ["description", "guide", "Resolution Set"]:
continue
if key in label_map:
payload[label_map[key]] = value
elif key.startswith("param_"):
try:
idx = int(key.replace("param_", ""))
if 0 <= idx < len(payload):
payload[idx] = value
except ValueError:
pass
res_preset_name = args.get("resolution_preset")
if res_preset_name:
res_set_name = model_preset.get("Resolution Set")
if res_set_name and res_set_name in res_cfg:
res_set = res_cfg[res_set_name]
if res_preset_name in res_set:
res_val_str = res_set[res_preset_name]
try:
w, h = map(int, res_val_str.split('x'))
if "Width" in label_map:
payload[label_map["Width"]] = w
if "Height" in label_map:
payload[label_map["Height"]] = h
except (ValueError, AttributeError):
raise ToolError(f"Invalid resolution format for preset '{res_preset_name}': {res_val_str}. Expected 'WidthxHeight'.")
else:
raise ToolError(f"Resolution preset '{res_preset_name}' not found for this model. You MUST call get_model_info with this model_name to see the available resolution presets.")
else:
raise ToolError(f"No resolution set configured for model '{model_name}'.")
overrides = {
"prompt": "Prompt",
"negative_prompt": "Negative Prompt",
"cfg_scale": "CFG Scale",
}
for arg_key, label in overrides.items():
if arg_key in args and label in label_map:
payload[label_map[label]] = args[arg_key]
if "Prompt" in label_map:
payload[label_map["Prompt"]] = prompt
if "Seed" in label_map:
payload[label_map["Seed"]] = -1
# The "Magic Number" splice is no longer needed as gaps are
# automatically filled by the sparse-to-dense mapping.
# 6. EXECUTE GENERATION
try:
gen_payload = {"data": payload}
gen_resp = await asyncio.to_thread(requests.post, f"{config.SD_URL}/api/txt2img", json=gen_payload, timeout=300)
gen_resp.raise_for_status()
res_data = gen_resp.json()
gallery_data = res_data["data"][0]["value"]
if not gallery_data:
raise ToolError("Server returned successfully but no image was generated.")
sd_img_url = gallery_data[0]["image"]["url"]
# Extract raw path from SD URL (e.g. http://.../file=/tmp/gradio/abc.png -> /tmp/gradio/abc.png)
if "/file=" in sd_img_url:
raw_path = sd_img_url.split("/file=")[1]
else:
# Fallback if the URL format changes
raise ToolError(f"Could not extract file path from SD URL: {sd_img_url}")
# Construct proxy URL through our MCP server
# Use the captured request host to ensure links work across the network
host = config.request_host.get()
proxy_url = f"http://{host}/file{raw_path}"
# Download the image and convert to base64 for the AI's vision
img_resp = await asyncio.to_thread(requests.get, sd_img_url, timeout=30)
img_resp.raise_for_status()
b64_data = base64.b64encode(img_resp.content).decode('utf-8')
return [
{
"type": "text",
"text": f"Image successfully generated using model '{model_name}'.\n\nLocal path: {raw_path}\nTo show the image to the user, you can include this markdown link in your response: ![Generated Image]({proxy_url})"
},
{
"type": "image",
"data": b64_data,
"mimeType": "image/png"
}
]
except requests.exceptions.HTTPError as e:
raise ToolError(f"Server error during generation: {str(e)}")
except Exception as e:
raise ToolError(f"Unexpected error during generation: {str(e)}")
+79
View File
@@ -0,0 +1,79 @@
import requests
from typing import Any, Dict, List
import config
from tools.utils import ToolError, load_toml
async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
"""
Returns information about available models or detailed guides for a specific model.
"""
model_name = args.get("model_name")
models = load_toml(config.MODEL_PRESETS_PATH)
res_presets = load_toml(config.RES_PRESETS_PATH)
if not model_name:
# Return a catalog of all models
catalog = []
for name, data in models.items():
catalog.append({
"name": name,
"description": data.get("description", "No description provided.")
})
return [
{
"type": "text",
"text": "Available models:\n\n" + "\n".join([f"- {m['name']}: {m['description']}" for m in catalog]) +
"\n\nBefore generating an image, you MUST call this tool again with your selected model as the 'model_name' argument."
}
]
if model_name not in models:
# Return the full catalog if the specific model isn't found
catalog = []
for name, data in models.items():
catalog.append({
"name": name,
"description": data.get("description", "No description provided.")
})
return [
{
"type": "text",
"text": f"Model '{model_name}' not found.\n\nAvailable models:\n\n" +
"\n".join([f"- {m['name']}: {m['description']}" for m in catalog]) +
"\n\nBefore generating an image, you MUST call this tool again with selected model as the 'model_name' argument."
}
]
model_data = models[model_name]
res_set_name = model_data.get("Resolution Set")
# Get available resolution presets for this model
res_options = []
if res_set_name and res_set_name in res_presets:
res_options = list(res_presets[res_set_name].keys())
else:
res_options = ["Default (1024x1024)"]
guide = model_data.get("guide", "No detailed guide available.")
description = model_data.get("description", "")
# Filter out internal config keys to show only the human-friendly presets
presets_to_show = {k: v for k, v in model_data.items() if k not in ["description", "guide", "Resolution Set"]}
preset_text = "\n".join([f"- {k}: {v}" for k, v in presets_to_show.items()])
res_text = ", ".join(res_options)
return [
{
"type": "text",
"text": (
f"Model: {model_name}\n"
f"Description: {description}\n\n"
f"--- Prompting Guide ---\n{guide}\n\n"
f"--- Active Presets ---\n{preset_text}\n\n"
f"--- Available Resolution Presets ---\n{res_text}\n\n"
f"Use 'resolution_preset' in generate_image to choose one of these."
)
}
]
+3 -3
View File
@@ -35,7 +35,7 @@ async def handle(args: Dict[str, Any]):
# Case 1: No pattern and no lines provided -> Dump whole file # Case 1: No pattern and no lines provided -> Dump whole file
if not pattern and not lines_str: if not pattern and not lines_str:
output = [f"{i+1}: {line}" for i, line in enumerate(all_lines)] output = [f"{i+1}: {line}" for i, line in enumerate(all_lines)]
return {"text": "".join(output)} return [{"type": "text", "text": "".join(output)}]
# Determine target line numbers (1-based) # Determine target line numbers (1-based)
target_lines: Set[int] = set() target_lines: Set[int] = set()
@@ -61,7 +61,7 @@ async def handle(args: Dict[str, Any]):
raise ToolError(f"Error parsing line indices: {e}") raise ToolError(f"Error parsing line indices: {e}")
if not target_lines: if not target_lines:
return {"text": "No matching lines found."} return [{"type": "text", "text": "No matching lines found."}]
# Expand targets with context # Expand targets with context
lines_to_show: Set[int] = set() lines_to_show: Set[int] = set()
@@ -83,4 +83,4 @@ async def handle(args: Dict[str, Any]):
output.append(f"{ln}: {all_lines[ln-1]}") output.append(f"{ln}: {all_lines[ln-1]}")
last_line = ln last_line = ln
return {"text": "".join(output)} return [{"type": "text", "text": "".join(output)}]
+2 -2
View File
@@ -43,7 +43,7 @@ async def handle(args: Dict[str, Any]):
paged_items = all_items[start_idx:end_idx] paged_items = all_items[start_idx:end_idx]
if not paged_items: if not paged_items:
return {"text": f"No items found on page {page}."} return [{"type": "text", "text": f"No items found on page {page}."}]
output = [] output = []
for item in paged_items: for item in paged_items:
@@ -61,4 +61,4 @@ async def handle(args: Dict[str, Any]):
header = f"Listing of {path_str}\nPage {effective_page} of {total_pages} ({total_items} total items). Showing {start_idx+1}-{min(end_idx, total_items)}." header = f"Listing of {path_str}\nPage {effective_page} of {total_pages} ({total_items} total items). Showing {start_idx+1}-{min(end_idx, total_items)}."
return {"text": f"{header}\n\n" + "\n".join(output)} return [{"type": "text", "text": f"{header}\n\n" + "\n".join(output)}]
+1 -1
View File
@@ -25,6 +25,6 @@ async def handle(args: Dict[str, Any]):
try: try:
data = base64.b64encode(path.read_bytes()).decode("utf-8") data = base64.b64encode(path.read_bytes()).decode("utf-8")
return {"type": "image", "data": data, "mimeType": mime_type} return [{"type": "image", "data": data, "mimeType": mime_type}]
except Exception as e: except Exception as e:
raise ToolError(f"Error reading image: {str(e)}") raise ToolError(f"Error reading image: {str(e)}")
+2 -2
View File
@@ -26,11 +26,11 @@ async def handle(args: Dict[str, Any]):
# Pillow parses PNG text chunks into the .info dictionary # Pillow parses PNG text chunks into the .info dictionary
metadata = img.info metadata = img.info
if not metadata: if not metadata:
return {"text": "No metadata found in this PNG."} return [{"type": "text", "text": "No metadata found in this PNG."}]
# Format as a simple key: value list # Format as a simple key: value list
output = "\n".join([f"{k}: {v}" for k, v in metadata.items()]) output = "\n".join([f"{k}: {v}" for k, v in metadata.items()])
return {"text": f"PNG Metadata:\n{output}"} return [{"type": "text", "text": f"PNG Metadata:\n{output}"}]
except ToolError: except ToolError:
raise raise
except Exception as e: except Exception as e:
+180
View File
@@ -0,0 +1,180 @@
import csv
import os
import config
import difflib
import httpx
import re
from typing import List, Dict, Any, Optional
from tools.utils import ToolError, format_count, get_type_suffix
# Global cache to avoid reading the CSV from disk on every request
_TAG_CACHE: Optional[List[Dict[str, Any]]] = None
def _load_tags() -> List[Dict[str, Any]]:
"""Reads the Danbooru tag CSV and caches it in memory."""
tag_file_path = config.TAG_DATABASE_PATH
tags = []
if not os.path.exists(tag_file_path):
# We raise ToolError here because it's a fatal configuration issue
raise ToolError(f"Tag database not found at {tag_file_path}. Please ensure the tag-autocomplete extension is installed.")
try:
with open(tag_file_path, mode='r', encoding='utf-8') as f:
reader = csv.reader(f)
for row in reader:
if not row or len(row) < 3:
continue
name = row[0]
tag_type = row[1]
count = int(row[2]) if row[2].isdigit() else 0
aliases = row[3].lower().split(',') if len(row) > 3 else []
tags.append({
"name": name,
"name_lower": name.lower(),
"type": tag_type,
"count": count,
"aliases": aliases
})
except Exception as e:
raise ToolError(f"Error reading tag database: {str(e)}")
return tags
return tags
async def _fetch_wiki_info(tag_name: str) -> Optional[str]:
"""Fetches the wiki description for a tag from Danbooru."""
if not config.ENABLE_TAG_WIKI:
return None
url = f"https://danbooru.donmai.us/wiki_pages/{tag_name}.json"
try:
# Use Basic Auth if credentials are provided
auth = None
if config.DANBOORU_LOGIN and config.DANBOORU_API_KEY:
auth = (config.DANBOORU_LOGIN, config.DANBOORU_API_KEY)
async with httpx.AsyncClient(timeout=5.0) as client:
# Using a browser-like UA to avoid potential blocks
headers = {"User-Agent": config.USER_AGENT}
resp = await client.get(url, headers=headers, auth=auth)
if resp.status_code == 403:
return " (Wiki access blocked by Danbooru/Cloudflare)"
if resp.status_code == 200:
data = resp.json()
wiki_body = data.get("wiki_page", {}).get("body")
if wiki_body:
# Strip HTML tags for the LLM
return re.sub(r'<[^>]*>', '', wiki_body).strip()
except Exception as e:
return f" (Error fetching wiki: {str(e)})"
return None
async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
"""
Searches the Danbooru tag database for tags matching a query.
Returns a unified list prioritized by substring matches then similarity.
"""
global _TAG_CACHE
# Normalize query: treat spaces and underscores as identical
query = args.get("query", "").lower().replace(" ", "_")
if not query:
raise ToolError("The 'query' argument is required.")
# Lazy-load the tags into memory
if _TAG_CACHE is None:
_TAG_CACHE = _load_tags()
# 1. Find Substring/Alias Matches (High Priority)
substring_matches = []
for tag in _TAG_CACHE:
is_direct = query in tag["name_lower"]
is_alias = any(query in alias.strip() for alias in tag["aliases"])
if is_direct or is_alias:
# Find which alias actually matched for reporting
matched_alias = ""
if not is_direct:
matched_alias = next((a.strip() for a in tag["aliases"] if query in a.strip()), query)
substring_matches.append({
"tag": tag,
"matched_via": "name" if is_direct else "alias",
"alias_match": matched_alias
})
# Sort substring matches by count descending
substring_matches.sort(key=lambda x: x["tag"]["count"], reverse=True)
# 2. Find Similarity Matches (Low Priority)
all_names_lower = [t["name_lower"] for t in _TAG_CACHE]
similar_names_lower = difflib.get_close_matches(query, all_names_lower, n=config.TAG_SEARCH_LIMIT, cutoff=0.5)
similar_matches = []
for s_lower in similar_names_lower:
tag = next((t for t in _TAG_CACHE if t["name_lower"] == s_lower), None)
if tag:
similar_matches.append(tag)
# Build Unified List
final_results = []
# Add substring matches first
for m in substring_matches:
final_results.append(m)
if len(final_results) >= config.TAG_SEARCH_LIMIT:
break
# Fill remaining slots with similar matches
if len(final_results) < config.TAG_SEARCH_LIMIT:
mentioned_names = {m["tag"]["name_lower"] for m in final_results}
for tag in similar_matches:
if tag["name_lower"] not in mentioned_names:
final_results.append({"tag": tag, "matched_via": "similarity", "alias_match": ""})
if len(final_results) >= config.TAG_SEARCH_LIMIT:
break
if not final_results:
return [{"type": "text", "text": f"No tags found matching '{query}'."}]
# Format output lines
output_lines = []
# Special Case: Exact Match Wiki Header
# Check if the very first result is an exact match
first_res = final_results[0]
if first_res["tag"]["name_lower"] == query:
exact_tag = first_res["tag"]
count_fmt = format_count(str(exact_tag["count"]))
type_sfx = get_type_suffix(exact_tag["type"])
output_lines.append(f"Exact Match: {exact_tag['name']} ({count_fmt}){type_sfx}")
wiki_info = await _fetch_wiki_info(exact_tag["name"])
if wiki_info:
output_lines.append(f"Wiki: {wiki_info}\n")
else:
output_lines.append("") # spacer
# List the tags
for res in final_results:
tag = res["tag"]
count_fmt = format_count(str(tag["count"]))
type_sfx = get_type_suffix(tag["type"])
# If this is the exact match we already listed in the header, skip it
if first_res["tag"]["name_lower"] == query and tag["name_lower"] == query:
continue
if res["matched_via"] == "alias":
line = f"- {res['alias_match']}{tag['name']} ({count_fmt}){type_sfx}"
else:
line = f"- {tag['name']} ({count_fmt}){type_sfx}"
output_lines.append(line)
return [{"type": "text", "text": "\n".join(output_lines)}]
+32
View File
@@ -1,5 +1,6 @@
import time import time
import datetime import datetime
import tomllib
from pathlib import Path from pathlib import Path
from typing import List, Dict, Any, Optional from typing import List, Dict, Any, Optional
@@ -8,6 +9,14 @@ class ToolError(Exception):
"""Custom exception for tool-related errors to be caught by the MCP server.""" """Custom exception for tool-related errors to be caught by the MCP server."""
pass pass
def load_toml(path: str) -> Dict[str, Any]:
"""Loads a TOML file into a dictionary."""
try:
with open(path, "rb") as f:
return tomllib.load(f)
except Exception as e:
raise ToolError(f"Failed to load config file {path}: {str(e)}")
def format_relative_time(timestamp: float) -> str: def format_relative_time(timestamp: float) -> str:
"""Converts a timestamp to a human-readable relative format.""" """Converts a timestamp to a human-readable relative format."""
now = time.time() now = time.time()
@@ -108,3 +117,26 @@ def parse_indices(indices_str: str) -> List[int]:
except ValueError: except ValueError:
raise ToolError(f"Invalid index format: {part}") raise ToolError(f"Invalid index format: {part}")
return indices return indices
def format_count(count_str: str) -> str:
"""Formats large numbers into a human-readable string (e.g., 1.2M, 218k)."""
try:
count = int(count_str)
if count >= 1_000_000:
return f"{count / 1_000_000:.1f}M".replace(".0", "")
if count >= 1_000:
return f"{count / 1_000:.0f}k"
return str(count)
except (ValueError, TypeError):
return "0"
def get_type_suffix(type_val: str) -> str:
"""Returns the human-readable suffix for the tag type."""
mapping = {
"0": "", # General
"1": " [Artist]",
"3": " [Copyright]",
"4": " [Character]",
"5": " [Meta]"
}
return mapping.get(str(type_val), f" [Type {type_val}]")