MooseCP/tools/get_model_info.py
2026-07-25 21:00:51 -07:00

60 lines
2.3 KiB
Python

import requests
from typing import Any, Dict, List
import config
from tools.utils import ToolError, load_toml
async def handle(args: Dict[str, Any]) -> 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 {
"text": "Available models:\n\n" + "\n".join([f"- {m['name']}: {m['description']}" for m in catalog]) +
"\n\nTo get a detailed prompting guide and available resolutions for a specific model, call this tool again with the 'model_name' argument."
}
if model_name not in models:
raise ToolError(f"Model '{model_name}' not found in presets. Available models: {', '.join(models.keys())}")
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 {
"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."
)
}