"""Model discovery route functions for ComfyUI API."""
import httpx
from src.auth.base import ComfyAuth
from src.client.get_data import get_data
from src.client.response import ResponseGetData
from src.utils.logging import log_call
class ModelError(Exception):
"""Raised when model operations fail."""
def __init__(self, message: str, response: ResponseGetData | None = None):
super().__init__(message)
self.response = response
class NodeDefinitionError(Exception):
"""Raised when node definition operations fail."""
def __init__(self, message: str, response: ResponseGetData | None = None):
super().__init__(message)
self.response = response
@log_call(action_name="get_available_models", level_name="route")
async def get_available_models(
auth: ComfyAuth,
node_class: str = "CheckpointLoaderSimple",
*,
session: httpx.AsyncClient | None = None,
) -> ResponseGetData:
"""Get available models from ComfyUI for a specific node class.
Queries the /object_info endpoint to discover available models, checkpoints,
VAEs, LoRAs, and other resources that can be used in workflows.
Args:
auth: ComfyAuth instance for base URL and authentication
node_class: Node class name (e.g., "CheckpointLoaderSimple", "VAELoader")
session: Optional HTTPX client for connection pooling
Returns:
ResponseGetData with:
- response: Node object info including input definitions and model lists
- status: 200 on success
- is_success: True on success
Raises:
ModelError: If model discovery fails
Example:
>>> auth = NoAuth("http://127.0.0.1:8188")
>>> res = await get_available_models(
... auth=auth,
... node_class="CheckpointLoaderSimple",
... )
>>> models = res.response["CheckpointLoaderSimple"]["input"]["required"]["ckpt_name"][0]
>>> print(models)
["sd_xl_base_1.0.safetensors", "dreamshaper_8.safetensors", ...]
"""
url = f"{auth.base_url}/object_info/{node_class}"
res = await get_data(
auth=auth,
method="GET",
url=url,
session=session,
)
if not res.is_success:
raise ModelError(
f"Failed to get available models for {node_class}: HTTP {res.status}",
response=res,
)
return res
@log_call(action_name="get_node_definitions", level_name="route")
async def get_node_definitions(
auth: ComfyAuth,
node_class: str | None = None,
*,
session: httpx.AsyncClient | None = None,
) -> ResponseGetData:
"""Get node definitions from ComfyUI.
Queries the /object_info endpoint to retrieve node schemas including
input parameters, output types, categories, and other metadata.
Args:
auth: ComfyAuth instance for base URL and authentication
node_class: Optional specific node class (e.g., "KSampler").
If None, returns all available nodes.
session: Optional HTTPX client for connection pooling
Returns:
ResponseGetData with:
- response: Dict mapping node_class -> node definition
- status: 200 on success
- is_success: True on success
Raises:
NodeDefinitionError: If request fails
Example:
>>> auth = NoAuth("http://127.0.0.1:8188")
>>> # Get all nodes
>>> res = await get_node_definitions(auth=auth)
>>> print(list(res.response.keys())[:3])
["KSampler", "CLIPTextEncode", "CheckpointLoaderSimple"]
>>>
>>> # Get specific node
>>> res = await get_node_definitions(auth=auth, node_class="KSampler")
>>> print(res.response["KSampler"]["input"]["required"].keys())
dict_keys(["model", "seed", "steps", "cfg", "sampler_name", ...])
"""
if node_class:
url = f"{auth.base_url}/object_info/{node_class}"
else:
url = f"{auth.base_url}/object_info"
res = await get_data(
auth=auth,
method="GET",
url=url,
session=session,
)
if not res.is_success:
target = f" for {node_class}" if node_class else ""
raise NodeDefinitionError(
f"Failed to get node definitions{target}: HTTP {res.status}",
response=res,
)
return res
