nodes/arcum42-ComfyUI_SageUtils
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

More refactoring.

Browse Source
This commit is contained in:
Arcum42
2026-07-04 12:58:32 -07:00
parent 9cc3038416
commit ba13db7f8a
4 changed files with 823 additions and 1061 deletions
Download Patch File Download Diff File Expand all files Collapse all files
routes/llm_routes.py
+255 -979
View File
File diff suppressed because it is too large Load Diff
utils/llm/providers/lmstudio/client.py
+39
View File
@@ -193,6 +193,45 @@ def get_models(enabled: bool) -> list[str]:
label='LM Studio REST models',
)
def get_models_payload(enabled: bool) -> list[dict[str, Any]]:
"""Retrieve raw model payloads from LM Studio REST."""
if is_provider_unavailable(enabled):
return []
try:
response = lmstudio_request_json_models()
return _extract_models_payload(response)
except Exception as e:
report_fetch_error(
llm_report,
'Error retrieving raw model payloads from LM Studio REST',
provider=_PROVIDER_NAME,
operation='get_models_payload',
cause=e,
fallback=[],
)
return []
def is_model_loaded(enabled: bool, model: str) -> bool:
"""Return whether an LM Studio model already has a loaded instance."""
if is_provider_unavailable(enabled):
return False
try:
for model_obj in get_models_payload(enabled):
model_name = _extract_model_name(model_obj)
if model_name != model:
continue
loaded_instances = model_obj.get('loaded_instances')
return bool(loaded_instances)
except Exception:
pass
return False
def get_vision_models(enabled: bool) -> list[str]:
"""Retrieve a list of available vision models from LM Studio REST."""
if is_provider_unavailable(enabled):
utils/llm/routes_helpers.py
+157 -82
View File
@@ -4,6 +4,7 @@ import base64
import json
import os
import tempfile
from contextvars import ContextVar
from pathlib import Path
from typing import Any, Optional
@@ -21,6 +22,7 @@ from .provider_keys import (
)
logger = get_logger('llm.routes_helpers')
_last_llm_error = ContextVar('sageutils_last_llm_error', default=None)
_BUILTIN_TOOL_PROFILES: dict[str, list[dict[str, Any]]] = {
@@ -725,16 +727,18 @@ def get_provider_availability_error(
"""Return a standardized provider availability error tuple if unavailable."""
provider = normalize_provider(provider)
if provider == LMSTUDIO_REST_KEY and not llm.LMSTUDIO_REST_AVAILABLE:
return False, 'LM Studio REST is not available', 'LLM_PROVIDER_UNAVAILABLE', 503
if provider == NATIVE_KEY:
if not llm.is_provider_available(provider, force=True):
return False, 'Native CLIP provider is not available', 'LLM_PROVIDER_UNAVAILABLE', 503
return True, None, None, None
if provider == OLLAMA_REST_KEY and not llm.OLLAMA_REST_AVAILABLE:
return False, 'Ollama REST is not available', 'LLM_PROVIDER_UNAVAILABLE', 503
if provider in {LMSTUDIO_REST_KEY, OLLAMA_REST_KEY, OPENAI_KEY}:
if not llm.is_provider_available(provider, force=True):
provider_name = provider.replace('_', ' ').title()
return False, f'{provider_name} is not available', 'LLM_PROVIDER_UNAVAILABLE', 503
return True, None, None, None
if provider == OPENAI_KEY and not llm.OPENAI_AVAILABLE:
return False, 'OpenAI provider is not available', 'LLM_PROVIDER_UNAVAILABLE', 503
return True, None, None, None
return False, f'Invalid provider: {provider}', 'LLM_INVALID_PROVIDER', 400
def validate_native_model_availability(
@@ -747,12 +751,27 @@ def validate_native_model_availability(
if provider != NATIVE_KEY:
return True, None, None, None
available_native = llm.get_native_models()
if model not in available_native:
if not llm.is_native_model_available(model):
return False, f"Native CLIP model '{model}' is not available", 'LLM_MODEL_NOT_FOUND', 404
return True, None, None, None
def get_provider_and_native_error(
provider: str,
model: str,
operation: str,
) -> tuple[bool, Optional[str], Optional[str], Optional[int]]:
"""Validate provider availability and native model selection in one check."""
is_available, error_msg, error_code, status_code = get_provider_availability_error(
provider,
operation,
)
if not is_available:
return is_available, error_msg, error_code, status_code
return validate_native_model_availability(provider, model, operation)
def format_sse_chunk(chunk_data: dict[str, Any]) -> str:
"""Format a data chunk as Server-Sent Events (SSE) format."""
return f"data: {json.dumps(chunk_data)}\n\n"
@@ -801,44 +820,142 @@ async def write_sse_error_and_close(
await response.write_eof()
async def write_provider_availability_error_if_unavailable(
def set_last_llm_error(payload: Any) -> None:
"""Store the last structured LLM error payload for route error handling."""
_last_llm_error.set(payload)
def clear_last_llm_error() -> None:
"""Clear any stale LLM error payload before route handling."""
_last_llm_error.set(None)
def pop_last_llm_error() -> Any:
"""Retrieve and clear the last structured LLM error payload."""
payload = _last_llm_error.get()
_last_llm_error.set(None)
return payload
def build_error_payload(
message: str,
*,
error_code: str = 'LLM_ERROR',
provider: str | None = None,
operation: str | None = None,
cause: str | None = None,
) -> dict[str, Any]:
payload = {
'success': False,
'error': message,
'error_code': error_code,
}
if provider:
payload['provider'] = provider
if operation:
payload['operation'] = operation
if cause:
payload['cause'] = cause
return payload
def error_response(
message: str,
status: int = 500,
*,
error_code: str = 'LLM_ERROR',
provider: str | None = None,
operation: str | None = None,
cause: str | None = None,
) -> web.Response:
return web.json_response(
build_error_payload(
message,
error_code=error_code,
provider=provider,
operation=operation,
cause=cause,
),
status=status,
)
def error_response_from_exception(
prefix: str,
exc: Exception,
status: int = 500,
default_error_code: str = 'LLM_ROUTE_ERROR',
) -> web.Response:
payload = pop_last_llm_error()
if payload:
message = payload.get('scoped_message') or payload.get('message') or f'{prefix}: {str(exc)}'
error_type = str(payload.get('error_type', 'ERROR')).upper()
return error_response(
message,
status=status,
error_code=f'LLM_{error_type}',
provider=payload.get('provider'),
operation=payload.get('operation'),
cause=payload.get('cause'),
)
return error_response(
f'{prefix}: {str(exc)}',
status=status,
error_code=default_error_code,
)
async def write_sse_error_from_exception(
response: web.StreamResponse,
exc: Exception,
operation: str,
prefix: str = 'Stream error',
default_error_code: str = 'LLM_STREAM_ERROR',
) -> None:
payload = pop_last_llm_error()
if payload:
message = str(payload.get('scoped_message') or payload.get('message') or f'{prefix}: {str(exc)}')
error_type = str(payload.get('error_type', 'STREAM_ERROR')).upper()
await write_sse_error_and_close(
response,
message,
error_code=f'LLM_{error_type}',
provider=payload.get('provider'),
operation=payload.get('operation') or operation,
cause=payload.get('cause'),
)
return
await write_sse_error_and_close(
response,
f'{prefix}: {str(exc)}',
error_code=default_error_code,
operation=operation,
)
async def write_provider_and_native_error_if_unavailable(
response: web.StreamResponse,
provider: str,
model: str,
operation: str,
) -> bool:
"""Write a provider availability SSE error and return True if unavailable."""
provider = normalize_provider(provider)
if provider == LMSTUDIO_REST_KEY and not llm.LMSTUDIO_REST_AVAILABLE:
"""Write a shared SSE error chunk for provider or native model availability failures."""
is_valid, error_msg, error_code, status_code = get_provider_and_native_error(
provider,
model,
operation,
)
if not is_valid:
await write_sse_error_and_close(
response,
'LM Studio REST is not available',
error_code='LLM_PROVIDER_UNAVAILABLE',
provider=LMSTUDIO_REST_KEY,
str(error_msg or 'Provider or model unavailable'),
error_code=error_code or 'LLM_PROVIDER_UNAVAILABLE',
provider=provider,
operation=operation,
)
return True
if provider == OLLAMA_REST_KEY and not llm.OLLAMA_REST_AVAILABLE:
await write_sse_error_and_close(
response,
'Ollama REST is not available',
error_code='LLM_PROVIDER_UNAVAILABLE',
provider=OLLAMA_REST_KEY,
operation=operation,
)
return True
if provider == OPENAI_KEY and not llm.OPENAI_AVAILABLE:
await write_sse_error_and_close(
response,
'OpenAI provider is not available',
error_code='LLM_PROVIDER_UNAVAILABLE',
provider=OPENAI_KEY,
operation=operation,
)
return True
return False
@@ -915,48 +1032,6 @@ def validate_generation_data(
def check_model_vision_capability(provider: str, model: str) -> tuple[bool, Optional[str]]:
"""
Check if a model supports vision capability before attempting vision operations.
Args:
provider: Provider name (e.g., 'ollama', 'lmstudio', 'openai')
model: Model name
Returns:
Tuple of (supports_vision, error_message)
- If supports_vision is True, error_message is None
- If supports_vision is False, error_message describes why
"""
provider = normalize_provider(provider)
try:
from . import service as llm
# Ensure LLM services are initialized for active model capability validation
llm.ensure_llm_initialized(force=True)
# Get capability map for the provider
if provider == 'lmstudio_rest':
cap_map = llm.get_lmstudio_rest_model_capabilities_map()
elif provider == 'ollama_rest':
cap_map = llm.get_ollama_rest_model_capabilities_map()
elif provider == 'openai':
cap_map = llm.get_openai_model_capabilities_map()
else:
return False, f"Unknown provider: {provider}"
# Check if model exists and supports vision
if not cap_map:
return False, f"No capability data available for {provider}"
if model not in cap_map:
available = ', '.join(sorted(cap_map.keys())[:5])
return False, f"Model '{model}' not found in {provider}. Available: {available}..."
capabilities = cap_map[model]
if not capabilities.get('vision', False):
return False, f"Model '{model}' does not support vision capability on provider {provider}"
return True, None
except Exception as e:
logger.warning(f'Error checking vision capability for {provider}:{model}: {e}')
# Degrade gracefully. If we can't check, allow the attempt.
return True, None
return llm.is_model_vision_capable(provider, model)
utils/llm/service.py
+372
View File
@@ -1,6 +1,8 @@
import time
from typing import Any, Optional
from ..logger import get_logger
from ..settings import get_setting
from .providers.settings import (
is_lmstudio_rest_enabled,
is_ollama_rest_enabled,
@@ -165,6 +167,16 @@ def _get_model_capabilities_dict(provider_key: str) -> dict[str, dict[str, objec
return {model_name: capabilities.to_dict() for model_name, capabilities in capability_map.items()}
def _status_key(provider_key: str) -> str:
return f'{provider_key}_available'
def _filter_compatible_models(model_list: list[str]) -> list[str]:
if not model_list:
return []
return [model for model in model_list if model and not model.startswith('(')]
def _generate_non_streaming(
provider_key: str,
model: str,
@@ -493,6 +505,155 @@ def _unload_provider_model(
return provider.unload_model(enabled, model)
def load_model(
provider_key: str,
model: str,
*,
keep_alive: int = 0,
) -> bool:
provider_key = normalize_provider_key(provider_key)
if provider_key == NATIVE_KEY:
native_models = get_native_models()
return model in native_models
if provider_key == LMSTUDIO_REST_KEY:
ensure_lmstudio_rest_initialized(force=True)
enabled = _provider_enabled(provider_key)
if _provider_client(provider_key).is_model_loaded(enabled, model):
return True
return _load_provider_model(provider_key, model, keep_alive=keep_alive)
if provider_key == OLLAMA_REST_KEY:
ensure_ollama_rest_initialized(force=True)
enabled = _provider_enabled(provider_key)
if _provider_client(provider_key).is_model_loaded(enabled, model):
return True
return _load_provider_model(provider_key, model, keep_alive=keep_alive)
if provider_key == OPENAI_KEY:
ensure_openai_initialized(force=True)
compatible_models = _filter_compatible_models(get_openai_models())
return model in compatible_models
raise ValueError(f'Unsupported provider for load_model: {provider_key}')
def is_model_loaded(provider_key: str, model: str) -> bool:
provider_key = normalize_provider_key(provider_key)
if provider_key == NATIVE_KEY:
return model in get_native_models()
if provider_key == LMSTUDIO_REST_KEY:
ensure_lmstudio_rest_initialized(force=True)
enabled = _provider_enabled(provider_key)
return _provider_client(provider_key).is_model_loaded(enabled, model)
if provider_key == OLLAMA_REST_KEY:
ensure_ollama_rest_initialized(force=True)
enabled = _provider_enabled(provider_key)
return _provider_client(provider_key).is_model_loaded(enabled, model)
if provider_key == OPENAI_KEY:
ensure_openai_initialized(force=True)
compatible_models = _filter_compatible_models(get_openai_models())
return model in compatible_models
raise ValueError(f'Unsupported provider for is_model_loaded: {provider_key}')
def generate_only(
provider_key: str,
model: str,
prompt: str,
*,
options=None,
system_prompt: str = '',
keep_alive=None,
):
provider_key = normalize_provider_key(provider_key)
if not is_provider_available(provider_key, force=True):
raise ValueError(f'Provider {provider_key} is not available')
if provider_key == NATIVE_KEY and not is_native_model_available(model):
raise ValueError(f"Native CLIP model '{model}' is not available")
if provider_key in {LMSTUDIO_REST_KEY, OLLAMA_REST_KEY}:
if not is_model_loaded(provider_key, model):
raise ValueError(f"Model '{model}' is not loaded for provider {provider_key}")
return generate(
provider_key,
model,
prompt,
options=options,
system_prompt=system_prompt,
keep_alive=keep_alive,
)
def is_provider_available(provider_key: str, *, force: bool = False) -> bool:
provider_key = normalize_provider_key(provider_key)
if provider_key == NATIVE_KEY:
return len(get_native_models()) > 0
if provider_key == LMSTUDIO_REST_KEY:
return ensure_lmstudio_rest_initialized(force=force) and LMSTUDIO_REST_AVAILABLE and is_lmstudio_rest_enabled()
if provider_key == OLLAMA_REST_KEY:
return ensure_ollama_rest_initialized(force=force) and OLLAMA_REST_AVAILABLE and is_ollama_rest_enabled()
if provider_key == OPENAI_KEY:
return ensure_openai_initialized(force=force) and OPENAI_AVAILABLE and is_openai_enabled()
return False
def is_native_model_available(model: str) -> bool:
return model in get_native_models()
def get_provider_model_capabilities_map(provider_key: str) -> dict[str, dict[str, object]]:
provider_key = normalize_provider_key(provider_key)
if provider_key == LMSTUDIO_REST_KEY:
return get_lmstudio_rest_model_capabilities_map()
if provider_key == OLLAMA_REST_KEY:
return get_ollama_rest_model_capabilities_map()
if provider_key == OPENAI_KEY:
return get_openai_model_capabilities_map()
return {}
def is_model_vision_capable(provider_key: str, model: str) -> tuple[bool, Optional[str]]:
provider_key = normalize_provider_key(provider_key)
if provider_key == NATIVE_KEY:
return False, 'Native provider does not support vision generation'
if provider_key not in {LMSTUDIO_REST_KEY, OLLAMA_REST_KEY, OPENAI_KEY}:
return False, f'Unknown provider: {provider_key}'
ensure_llm_initialized(force=True)
capability_map = get_provider_model_capabilities_map(provider_key)
if not capability_map:
return False, f'No capability data available for provider {provider_key}'
if model not in capability_map:
available = ', '.join(sorted(capability_map.keys())[:10])
return False, f"Model '{model}' not found for {provider_key}. Available: {available}"
capabilities = capability_map[model]
if not capabilities.get('vision', False):
return False, f"Model '{model}' does not support vision capability on provider {provider_key}"
return True, None
# ============================================================================
# MODEL DISCOVERY
# ============================================================================
@@ -798,6 +959,217 @@ def _ensure_registered_providers(force: bool = False) -> dict[str, bool]:
return results
def get_llm_status(force: bool = False) -> dict[str, dict[str, object]]:
"""Return availability and enabled state for all supported providers."""
if force:
reset_llm_initialization_state()
ensure_llm_initialized(force=force)
status = {
LMSTUDIO_REST_KEY: {
'available': LMSTUDIO_REST_AVAILABLE and is_lmstudio_rest_enabled(),
'enabled': is_lmstudio_rest_enabled(),
},
OLLAMA_REST_KEY: {
'available': OLLAMA_REST_AVAILABLE and is_ollama_rest_enabled(),
'enabled': is_ollama_rest_enabled(),
},
OPENAI_KEY: {
'available': OPENAI_AVAILABLE and is_openai_enabled(),
'enabled': is_openai_enabled(),
},
}
native_models = get_native_models()
status[NATIVE_KEY] = {
'available': len(native_models) > 0,
'enabled': True,
}
return status
def get_llm_models(force: bool = False) -> dict[str, object]:
"""Return model lists, capabilities, tool models, reasoning models, and status."""
if force:
reset_llm_initialization_state()
ensure_llm_initialized(force=force)
lmstudio_rest_enabled = is_lmstudio_rest_enabled()
ollama_rest_enabled = is_ollama_rest_enabled()
openai_enabled = is_openai_enabled()
lmstudio_rest_models = []
lmstudio_rest_tool_models = []
lmstudio_rest_reasoning_models = []
lmstudio_rest_capabilities = {}
ollama_rest_models = []
ollama_rest_tool_models = []
ollama_rest_reasoning_models = []
ollama_rest_capabilities = {}
openai_models = []
openai_tool_models = []
openai_reasoning_models = []
openai_capabilities = {}
if lmstudio_rest_enabled and LMSTUDIO_REST_AVAILABLE:
try:
models = get_lmstudio_rest_models()
lmstudio_rest_models = _filter_compatible_models(models)
lmstudio_rest_tool_models = _filter_compatible_models(get_lmstudio_rest_tool_models())
lmstudio_rest_reasoning_models = _filter_compatible_models(get_lmstudio_rest_reasoning_models())
lmstudio_rest_capabilities = get_lmstudio_rest_model_capabilities_map()
except Exception as e:
logger.warning(f'Failed to get LM Studio REST models: {e}')
if ollama_rest_enabled and OLLAMA_REST_AVAILABLE:
try:
models = get_ollama_rest_models()
ollama_rest_models = _filter_compatible_models(models)
ollama_rest_tool_models = _filter_compatible_models(get_ollama_rest_tool_models())
ollama_rest_reasoning_models = _filter_compatible_models(get_ollama_rest_reasoning_models())
ollama_rest_capabilities = get_ollama_rest_model_capabilities_map()
except Exception as e:
logger.warning(f'Failed to get Ollama REST models: {e}')
if openai_enabled and OPENAI_AVAILABLE:
try:
models = get_openai_models()
openai_models = _filter_compatible_models(models)
openai_tool_models = _filter_compatible_models(get_openai_tool_models())
openai_reasoning_models = _filter_compatible_models(get_openai_reasoning_models())
openai_capabilities = get_openai_model_capabilities_map()
except Exception as e:
logger.warning(f'Failed to get OpenAI models: {e}')
native_models = get_native_models()
return {
'models': {
LMSTUDIO_REST_KEY: lmstudio_rest_models,
OLLAMA_REST_KEY: ollama_rest_models,
OPENAI_KEY: openai_models,
NATIVE_KEY: native_models,
},
'capabilities': {
LMSTUDIO_REST_KEY: lmstudio_rest_capabilities,
OLLAMA_REST_KEY: ollama_rest_capabilities,
OPENAI_KEY: openai_capabilities,
NATIVE_KEY: {},
},
'tool_models': {
LMSTUDIO_REST_KEY: lmstudio_rest_tool_models,
OLLAMA_REST_KEY: ollama_rest_tool_models,
OPENAI_KEY: openai_tool_models,
NATIVE_KEY: [],
},
'reasoning_models': {
LMSTUDIO_REST_KEY: lmstudio_rest_reasoning_models,
OLLAMA_REST_KEY: ollama_rest_reasoning_models,
OPENAI_KEY: openai_reasoning_models,
NATIVE_KEY: [],
},
'status': {
_status_key(LMSTUDIO_REST_KEY): len(lmstudio_rest_models) > 0,
_status_key(OLLAMA_REST_KEY): len(ollama_rest_models) > 0,
_status_key(OPENAI_KEY): len(openai_models) > 0,
_status_key(NATIVE_KEY): len(native_models) > 0,
},
}
def get_llm_vision_models(force: bool = False) -> dict[str, object]:
"""Return vision model lists, capabilities, tool models, reasoning models, and status."""
if force:
reset_llm_initialization_state()
ensure_llm_initialized(force=force)
lmstudio_rest_enabled = is_lmstudio_rest_enabled()
ollama_rest_enabled = is_ollama_rest_enabled()
openai_enabled = is_openai_enabled()
lmstudio_rest_models = []
lmstudio_rest_tool_models = []
lmstudio_rest_reasoning_models = []
lmstudio_rest_capabilities = {}
ollama_rest_models = []
ollama_rest_tool_models = []
ollama_rest_reasoning_models = []
ollama_rest_capabilities = {}
openai_models = []
openai_tool_models = []
openai_reasoning_models = []
openai_capabilities = {}
if lmstudio_rest_enabled and LMSTUDIO_REST_AVAILABLE:
try:
models = get_lmstudio_rest_vision_models()
lmstudio_rest_models = _filter_compatible_models(models)
lmstudio_rest_tool_models = _filter_compatible_models(get_lmstudio_rest_tool_models())
lmstudio_rest_reasoning_models = _filter_compatible_models(get_lmstudio_rest_reasoning_models())
lmstudio_rest_capabilities = get_lmstudio_rest_model_capabilities_map()
except Exception as e:
logger.warning(f'Failed to get LM Studio REST vision models: {e}')
if ollama_rest_enabled and OLLAMA_REST_AVAILABLE:
try:
models = get_ollama_rest_vision_models()
ollama_rest_models = _filter_compatible_models(models)
ollama_rest_tool_models = _filter_compatible_models(get_ollama_rest_tool_models())
ollama_rest_reasoning_models = _filter_compatible_models(get_ollama_rest_reasoning_models())
ollama_rest_capabilities = get_ollama_rest_model_capabilities_map()
except Exception as e:
logger.warning(f'Failed to get Ollama REST vision models: {e}')
if openai_enabled and OPENAI_AVAILABLE:
try:
models = get_openai_vision_models()
openai_models = _filter_compatible_models(models)
openai_tool_models = _filter_compatible_models(get_openai_tool_models())
openai_reasoning_models = _filter_compatible_models(get_openai_reasoning_models())
openai_capabilities = get_openai_model_capabilities_map()
except Exception as e:
logger.warning(f'Failed to get OpenAI vision models: {e}')
native_models = []
return {
'models': {
LMSTUDIO_REST_KEY: lmstudio_rest_models,
OLLAMA_REST_KEY: ollama_rest_models,
OPENAI_KEY: openai_models,
NATIVE_KEY: native_models,
},
'capabilities': {
LMSTUDIO_REST_KEY: lmstudio_rest_capabilities,
OLLAMA_REST_KEY: ollama_rest_capabilities,
OPENAI_KEY: openai_capabilities,
NATIVE_KEY: {},
},
'tool_models': {
LMSTUDIO_REST_KEY: lmstudio_rest_tool_models,
OLLAMA_REST_KEY: ollama_rest_tool_models,
OPENAI_KEY: openai_tool_models,
NATIVE_KEY: [],
},
'reasoning_models': {
LMSTUDIO_REST_KEY: lmstudio_rest_reasoning_models,
OLLAMA_REST_KEY: ollama_rest_reasoning_models,
OPENAI_KEY: openai_reasoning_models,
NATIVE_KEY: [],
},
'status': {
_status_key(LMSTUDIO_REST_KEY): len(lmstudio_rest_models) > 0,
_status_key(OLLAMA_REST_KEY): len(ollama_rest_models) > 0,
_status_key(OPENAI_KEY): len(openai_models) > 0,
_status_key(NATIVE_KEY): False,
},
}
def init_llm() -> None:
"""Initialize all configured REST/OpenAI providers."""
_init_registered_providers()