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

Cleanup and added settings in for SageUtils. You can enable and disable Ollama and LM Studio from settings.

Browse Source
This commit is contained in:
arcum42
2025-07-24 16:54:33 -07:00
parent 1a1a46456d
commit efb86cbfd5
9 changed files with 387 additions and 852 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CUSTOM_ROUTES.md
-250
View File
@@ -1,250 +0,0 @@
# SageUtils Custom API Routes
SageUtils adds custom HTTP routes to ComfyUI that expose the SageCache data. These routes allow you to access model metadata, file hashes, and cache statistics programmatically.
## Available Endpoints
All endpoints are available under the `/sage_cache/` prefix and can also be accessed with the `/api/sage_cache/` prefix.
### 1. Get Cache Info
**GET** `/sage_cache/info`
Returns the complete contents of SageCache.info as JSON. This contains model metadata, Civitai information, and cache details for all models.
**Example Response:**
```json
{
"abc123hash": {
"hash": "abc123hash",
"civitai": "true",
"modelId": "12345",
"name": "Model Name",
"baseModel": "SDXL",
"lastUsed": "2025-06-22T10:30:00"
}
}
```
### 2. Get Cache Hash Mapping
**GET** `/sage_cache/hash`
Returns the complete contents of SageCache.hash as JSON. This contains the mapping from file paths to their SHA256 hashes.
**Example Response:**
```json
{
"/path/to/model1.safetensors": "abc123hash",
"/path/to/model2.ckpt": "def456hash"
}
```
### 3. Get Cache Statistics
**GET** `/sage_cache/stats`
Returns statistics about the SageCache including file counts, Civitai status, and cache file locations.
**Example Response:**
```json
{
"total_files": 25,
"total_info_entries": 23,
"cache_files": {
"hash_path": "/path/to/sage_cache_hash.json",
"info_path": "/path/to/sage_cache_info.json",
"main_path": "/path/to/sage_cache.json",
"ollama_models_path": "/path/to/sage_cache_ollama.json"
},
"civitai_stats": {
"found_on_civitai": 18,
"not_found_on_civitai": 5
}
}
```
### 4. Get File Information by Hash
**GET** `/sage_cache/file/{file_hash}`
Returns detailed information for a specific file hash, including all file paths that use this hash.
**Example Response:**
```json
{
"hash": "abc123hash",
"info": {
"hash": "abc123hash",
"civitai": "true",
"modelId": "12345",
"name": "Model Name",
"baseModel": "SDXL",
"lastUsed": "2025-06-22T10:30:00"
},
"file_paths": [
"/path/to/model1.safetensors"
]
}
```
### 5. Get File Information by Path
**GET** `/sage_cache/path?file_path={file_path}`
Returns detailed information for a file by its path. Looks up the hash from the file path in cache.hash, then retrieves the info from cache.info.
**Parameters:**
- `file_path` (query parameter, required): The full file path to look up
**Example Response:**
```json
{
"file_path": "/path/to/model1.safetensors",
"hash": "abc123hash",
"info": {
"hash": "abc123hash",
"civitai": "true",
"modelId": "12345",
"name": "Model Name",
"baseModel": "SDXL",
"lastUsed": "2025-06-22T10:30:00"
},
"all_paths_with_same_hash": [
"/path/to/model1.safetensors"
]
}
```
## Usage Examples
### Using cURL
```bash
# Get all cache info
curl http://localhost:8188/sage_cache/info
# Get hash mappings
curl http://localhost:8188/sage_cache/hash
# Get cache statistics
curl http://localhost:8188/sage_cache/stats
# Get specific file info by hash
curl http://localhost:8188/sage_cache/file/abc123hash456def
# Get specific file info by path
curl "http://localhost:8188/sage_cache/path?file_path=/path/to/model1.safetensors"
```
### Using JavaScript/Fetch
```javascript
// Get cache statistics
async function getCacheStats() {
const response = await fetch('/api/sage_cache/stats');
const stats = await response.json();
console.log('Cache stats:', stats);
}
// Get info for a specific hash
async function getFileInfo(hash) {
const response = await fetch(`/api/sage_cache/file/${hash}`);
const info = await response.json();
console.log('File info:', info);
}
// Get info for a specific file path
async function getFileInfoByPath(filePath) {
const response = await fetch(`/api/sage_cache/path?file_path=${encodeURIComponent(filePath)}`);
const info = await response.json();
console.log('File info by path:', info);
}
```
### Using Python
```python
import requests
# Get all cache info
response = requests.get('http://localhost:8188/sage_cache/info')
cache_info = response.json()
# Get cache statistics
response = requests.get('http://localhost:8188/sage_cache/stats')
stats = response.json()
print(f"Total files: {stats['total_files']}")
print(f"Found on Civitai: {stats['civitai_stats']['found_on_civitai']}")
# Get file info by hash
file_hash = 'abc123hash'
response = requests.get(f'http://localhost:8188/sage_cache/file/{file_hash}')
file_info = response.json()
print(f"File info for hash {file_hash}: {file_info}")
# Get file info by path
file_path = '/path/to/model1.safetensors'
response = requests.get('http://localhost:8188/sage_cache/path', params={'file_path': file_path})
file_info = response.json()
print(f"File info for path {file_path}: {file_info}")
```
## Error Handling
All endpoints return appropriate HTTP status codes:
- **200**: Success
- **400**: Bad Request (e.g., missing file hash)
- **404**: Not Found (e.g., hash doesn't exist)
- **500**: Internal Server Error
Error responses include an `error` field with a descriptive message:
```json
{
"error": "No information found for hash: invalid_hash"
}
```
The implementation includes graceful fallback if PromptServer is not available during development or ComfyUI startup issues.
```
## Implementation Details
### Route Registration
The custom routes are registered using `PromptServer.instance.routes` to add new endpoints to ComfyUI's existing web server. Routes are automatically registered when the custom node imports during ComfyUI startup.
### Cache Access
Each request calls `cache.load()` to ensure fresh data is returned, respecting any changes made to the cache files while ComfyUI is running.
### File Structure
```text
comfyui_sageutils/
├── server_routes.py # Custom route implementations
├── CUSTOM_ROUTES.md # This documentation
├── test_routes.py # Testing script
├── demo_routes.py # Demonstration script
├── __init__.py # Modified to import routes
└── utils/
└── model_cache.py # Contains SageCache class
```
### Technical Notes
- The implementation is read-only (no routes modify cache data)
- Routes respect ComfyUI's security model
- Compatible with existing ComfyUI routing infrastructure
- No additional dependencies required beyond what SageUtils already uses
- Routes only register if PromptServer is available (fails silently otherwise)
- All routes are also available with the `/api/` prefix for consistency with ComfyUI standards
### Testing and Validation
When ComfyUI is running properly, you should see the message "SageUtils custom routes loaded successfully!" in the console, indicating the routes are available.
Use the provided test script to validate functionality:
```bash
cd /path/to/comfyui_sageutils
python test_routes.py
```
The cache is automatically loaded when endpoints are accessed, and all endpoints are read-only and do not modify the cache. The endpoints respect the same security model as other ComfyUI routes.
FILE_MANAGEMENT.md
-118
View File
@@ -1,118 +0,0 @@
# SageUtils File and Path Management
This document explains the reorganized file and path management system in SageUtils.
## Overview
The file and path management has been centralized into three main modules:
1. **`path_manager.py`** - Centralized path management
2. **`config_manager.py`** - Configuration file management (refactored)
3. **`model_cache.py`** - Model cache management (refactored)
4. **`sage_utils.py`** - Convenience module with commonly used functions
## Key Improvements
### 1. Centralized Path Management
- All paths are now managed in `SagePathManager`
- Consistent directory creation and path resolution
- Single source of truth for all file locations
### 2. Unified File Operations
- `SageFileManager` provides atomic JSON operations
- Consistent error handling and backup strategies
- Reusable file I/O methods
### 3. Consistent Initialization
- Config files are automatically copied from assets to user directory if missing
- Cache files are created in user directory as needed
- All necessary directories are created automatically
## File Locations
### Default Asset Files (read-only templates)
- `/assets/config.json` - Default settings
- `/assets/sage_styles.json` - Default styles
- `/assets/llm_prompts.json` - Default LLM prompts
### User Files (modifiable copies)
- `/user/default/SageUtils/config.json` - User settings
- `/user/default/SageUtils/sage_styles.json` - User styles
- `/user/default/SageUtils/llm_prompts.json` - User LLM prompts
### User Override Files (optional)
- `/user/default/SageUtils/config_user.json` - User setting overrides
- `/user/default/SageUtils/sage_styles_user.json` - User style overrides
- `/user/default/SageUtils/llm_prompts_user.json` - User prompt overrides
### Cache Files
- `/user/default/SageUtils/sage_cache.json` - Legacy cache (converted automatically)
- `/user/default/SageUtils/sage_cache_hash.json` - File path to hash mapping
- `/user/default/SageUtils/sage_cache_info.json` - Hash to metadata mapping
- `/user/default/SageUtils/sage_cache_ollama.json` - Ollama model cache
### Backup Files
- `/user/default/SageUtils/backup/` - Automatic backups of all files
### Wildcard Files
- `/user/default/SageUtils/wildcards/` - User wildcard text files
### Notes Files
- `/user/default/SageUtils/notes/` - User notes and documentation files
## Usage Examples
### Using the convenience module:
```python
from .utils.sage_utils import get_user_path, load_json, save_json
# Get paths
user_config = get_user_path("my_config.json")
# Load/save JSON
data = load_json(user_config, "my config")
save_json(user_config, {"key": "value"}, "my config")
```
### Using path manager directly:
```python
from .utils.path_manager import path_manager, file_manager
# Get specific paths
backup_path = path_manager.get_backup_file_path("backup-2024.json")
asset_path = path_manager.get_asset_file_path("defaults.json")
# File operations
file_manager.ensure_user_config_file("new_config")
config = file_manager.load_config_with_overrides("config")
```
### Using config manager:
```python
from .utils.config_manager import ConfigManager
# Create and use a config manager
my_config = ConfigManager("my_config")
data = my_config.load() # Loads with user overrides
my_config.save(data, user_override=True) # Save to override file
```
## Migration Notes
### For Existing Code
- Replace direct file path references with path_manager calls
- Use file_manager for JSON operations instead of manual file handling
- Config files will be automatically migrated on first load
### Breaking Changes
- Old global path variables (`sage_users_path`, etc.) are now in `sage_utils.py`
- Manual directory creation is no longer needed
- File operations should use the centralized managers
## Benefits
1. **Consistency** - All files are handled the same way
2. **Reliability** - Atomic operations and error handling
3. **Maintainability** - Single place to modify path/file logic
4. **Backup Safety** - Automatic backups with deduplication
5. **User Experience** - Automatic setup and migration
demo_routes.py
-88
View File
@@ -1,88 +0,0 @@
#!/usr/bin/env python3
"""
Simple demonstration of the SageUtils custom routes implementation.
This shows the structure and functionality without requiring a full ComfyUI environment.
"""
def demonstrate_routes():
"""Demonstrate the custom routes implementation for SageUtils"""
print("=== SageUtils Custom Routes Demo ===\n")
print("✓ Successfully implemented custom routes for SageUtils in ComfyUI!")
print("✓ Routes are registered via server_routes.py module")
print("✓ Routes are imported in __init__.py during custom node loading")
print("✓ Implementation follows ComfyUI documentation patterns\n")
print("📁 Files created:")
print(" - server_routes.py: Custom route implementations")
print(" - CUSTOM_ROUTES.md: Documentation for the new routes")
print(" - test_routes.py: Test script for validation")
print(" - Modified __init__.py: Added route import\n")
print("🌐 Available Routes (when ComfyUI is running):")
routes = [
("GET", "/sage_cache/info", "Returns complete SageCache.info as JSON"),
("GET", "/sage_cache/hash", "Returns complete SageCache.hash mapping as JSON"),
("GET", "/sage_cache/stats", "Returns cache statistics and counts"),
("GET", "/sage_cache/file/{hash}", "Returns info for specific file hash"),
("GET", "/sage_cache/path?file_path={path}", "Returns info for file by path")
]
for method, path, description in routes:
print(f" {method:4} {path:25} - {description}")
print(f"\n🔗 All routes also available with /api prefix")
print(f" Example: http://localhost:8188/api/sage_cache/stats")
print("\n📊 Sample Response Structure:")
# Sample stats response
sample_stats = {
"total_files": 25,
"total_info_entries": 23,
"cache_files": {
"hash_path": "/user/default/SageUtils/sage_cache_hash.json",
"info_path": "/user/default/SageUtils/sage_cache_info.json"
},
"civitai_stats": {
"found_on_civitai": 18,
"not_found_on_civitai": 5
}
}
import json
print(" /sage_cache/stats response:")
print(" " + json.dumps(sample_stats, indent=4).replace('\n', '\n '))
print("\n🛡️ Error Handling:")
print(" - Graceful import handling (warns if PromptServer unavailable)")
print(" - Try/catch blocks around cache operations")
print(" - Proper HTTP status codes (400, 404, 500)")
print(" - Descriptive error messages in JSON responses")
print("\n🔧 Implementation Details:")
print(" - Uses PromptServer.instance.routes for registration")
print(" - Imports happen at custom node load time")
print(" - Cache is loaded on each request to ensure freshness")
print(" - Compatible with ComfyUI's existing route structure")
print("\n📖 Usage Examples:")
print(" # Get cache statistics")
print(" curl http://localhost:8188/sage_cache/stats")
print("")
print(" # Get all model hashes")
print(" curl http://localhost:8188/sage_cache/hash")
print("")
print(" # Get info for specific model")
print(" curl http://localhost:8188/sage_cache/file/abc123def456")
print("")
print(" # Get info for model by file path")
print(' curl "http://localhost:8188/sage_cache/path?file_path=/path/to/model.safetensors"')
print("\n✅ Implementation Complete!")
print(" The custom routes are ready to use once ComfyUI starts successfully.")
print(" They will appear alongside ComfyUI's built-in routes like /prompt, /queue, etc.")
if __name__ == "__main__":
demonstrate_routes()
js/sage.js
+3
View File
@@ -14,6 +14,9 @@ import { setupViewTextOrAnythingNode } from "./nodes/viewAnything.js";
import { setupViewNotesNode } from "./nodes/viewNotes.js";
import { setupModelInfoDisplayNode } from "./nodes/modelInfoDisplay.js";
// Import settings UI
import "./settingsUI.js";
app.registerExtension({
name: "arcum42.sage.utils",
async setup() {
js/settingsUI.js
+181
View File
@@ -0,0 +1,181 @@
// SageUtils Settings Integration
// Integrates SageUtils settings with ComfyUI's native settings system
import { app } from "../../../../scripts/app.js";
import { api } from "../../../../scripts/api.js";
// All of the setting types are based on PrimeVue components.
// https://primevue.org/
// Props described in the PrimeVue documentation can be defined
// for ComfyUI settings by adding them in an attrs field.
// Helper functions for API communication
async function loadSageSettings() {
try {
const response = await api.fetchApi('/sage_utils/settings');
if (response.ok) {
const data = await response.json();
if (data.success) {
return data.settings;
} else {
console.error('Server returned error:', data.error);
if (data.details) {
console.error('Error details:', data.details);
}
}
} else {
console.error('HTTP error:', response.status, response.statusText);
}
throw new Error('Failed to load settings from server');
} catch (error) {
console.error('Error loading SageUtils settings:', error);
// Return fallback settings structure
return {
enable_ollama: { current_value: true },
enable_lmstudio: { current_value: true },
ollama_use_custom_url: { current_value: false },
ollama_custom_url: { current_value: "" },
lmstudio_use_custom_url: { current_value: false },
lmstudio_custom_url: { current_value: "" }
};
}
}
async function saveSageSetting(key, value) {
try {
const response = await api.fetchApi('/sage_utils/settings', {
method: 'POST',
headers: {
'Content-Type': 'application/json',
},
body: JSON.stringify({ [key]: value })
});
const data = await response.json();
if (data.success) {
console.log(`Successfully saved setting ${key}: ${value}`);
if (data.errors && data.errors.length > 0) {
console.warn('Warnings during setting save:', data.errors);
}
} else {
console.error(`Failed to save setting ${key}:`, data.error);
}
return data.success;
} catch (error) {
console.error(`Error saving SageUtils setting ${key}:`, error);
return false;
}
}
// Register SageUtils settings with ComfyUI's native settings system
app.registerExtension({
name: "arcum42.sage.utils.settings",
async setup() {
console.log("Setting up SageUtils settings integration...");
// Load current settings from server
const serverSettings = await loadSageSettings();
if (!serverSettings) {
console.warn('Could not load SageUtils settings from server - settings will use defaults');
return;
}
console.log('Loaded SageUtils settings from server:', serverSettings);
// Set current values for all settings from server
// Map backend keys to the new frontend setting IDs
const keyToIdMap = {
'enable_lmstudio': 'SageUtils.LLM Providers.enable_lmstudio',
'enable_ollama': 'SageUtils.LLM Providers.enable_ollama',
'ollama_custom_url': 'SageUtils.Local Custom Ollama URL.ollama_custom_url',
'ollama_use_custom_url': 'SageUtils.Local Custom Ollama URL.ollama_use_custom_url',
'lmstudio_custom_url': 'SageUtils.Local Custom LM Studio URL.lmstudio_custom_url',
'lmstudio_use_custom_url': 'SageUtils.Local Custom LM Studio URL.lmstudio_use_custom_url'
};
for (const [key, settingInfo] of Object.entries(serverSettings)) {
const settingId = keyToIdMap[key];
if (settingId && settingInfo.current_value !== undefined) {
try {
await app.extensionManager.setting.set(settingId, settingInfo.current_value);
console.log(`Set initial value for ${settingId}:`, settingInfo.current_value);
} catch (error) {
console.warn(`Could not set initial value for ${settingId}:`, error);
}
}
}
console.log(`SageUtils settings integration completed`);
},
settings: [
{
id: "SageUtils.LLM Providers.enable_lmstudio",
name: "Enable LM Studio Integration",
type: "boolean",
defaultValue: true,
tooltip: "Enable LM Studio LLM integration",
onChange: async (newVal, oldVal) => {
console.log(`LM Studio integration changed from ${oldVal} to ${newVal}`);
await saveSageSetting('enable_lmstudio', newVal);
}
},
{
id: "SageUtils.LLM Providers.enable_ollama",
name: "Enable Ollama Integration",
type: "boolean",
defaultValue: true,
tooltip: "Enable Ollama LLM integration",
onChange: async (newVal, oldVal) => {
console.log(`Ollama integration changed from ${oldVal} to ${newVal}`);
await saveSageSetting('enable_ollama', newVal);
}
},
{
id: "SageUtils.Local Custom Ollama URL.ollama_custom_url",
name: "Address",
type: "text",
defaultValue: "",
tooltip: "Custom URL for Ollama service (e.g., 'http://localhost:11434')",
onChange: async (newVal, oldVal) => {
console.log(`Ollama custom URL changed from '${oldVal}' to '${newVal}'`);
await saveSageSetting('ollama_custom_url', newVal);
}
},
{
id: "SageUtils.Local Custom Ollama URL.ollama_use_custom_url",
name: "Enable Custom URL",
type: "boolean",
defaultValue: false,
tooltip: "Use a custom URL for Ollama instead of the default",
onChange: async (newVal, oldVal) => {
console.log(`Ollama custom URL setting changed from ${oldVal} to ${newVal}`);
await saveSageSetting('ollama_use_custom_url', newVal);
}
},
{
id: "SageUtils.Local Custom LM Studio URL.lmstudio_custom_url",
name: "Address",
type: "text",
defaultValue: "",
tooltip: "Custom URL for LM Studio service (e.g., 'http://localhost:1234')",
onChange: async (newVal, oldVal) => {
console.log(`LM Studio custom URL changed from '${oldVal}' to '${newVal}'`);
await saveSageSetting('lmstudio_custom_url', newVal);
}
},
{
id: "SageUtils.Local Custom LM Studio URL.lmstudio_use_custom_url",
name: "Enable Custom URL",
type: "boolean",
defaultValue: false,
tooltip: "Use a custom URL for LM Studio instead of the default",
onChange: async (newVal, oldVal) => {
console.log(`LM Studio custom URL setting changed from ${oldVal} to ${newVal}`);
await saveSageSetting('lmstudio_use_custom_url', newVal);
}
}
]
});
server_routes.py
+112 -1
View File
@@ -1,17 +1,128 @@
"""
Custom routes for SageUtils to expose cache data via HTTP endpoints.
Custom routes for SageUtils to expose cache data and settings via HTTP endpoints.
"""
try:
from server import PromptServer
from aiohttp import web
from .utils.model_cache import cache
from .utils.settings import get_settings, SETTINGS_SCHEMA
# Check if PromptServer instance is available
if hasattr(PromptServer, 'instance') and PromptServer.instance is not None:
# Get the PromptServer instance
routes = PromptServer.instance.routes
# Settings management routes
@routes.get('/sage_utils/settings')
async def get_sage_settings(request):
"""
Returns all SageUtils settings with their current values and schema information.
"""
try:
import json
settings = get_settings()
settings_info = settings.list_all_settings()
# Double-check that the result is JSON serializable
json.dumps(settings_info) # This will raise an exception if not serializable
return web.json_response({
"success": True,
"settings": settings_info
})
except (TypeError, ValueError) as e:
return web.json_response(
{"success": False, "error": f"JSON serialization error: {str(e)}"},
status=500
)
except Exception as e:
import traceback
error_details = traceback.format_exc()
print(f"SageUtils settings error: {error_details}")
return web.json_response(
{"success": False, "error": f"Failed to retrieve settings: {str(e)}", "details": error_details},
status=500
)
@routes.post('/sage_utils/settings')
async def update_sage_settings(request):
"""
Updates SageUtils settings. Expects JSON body with setting key-value pairs.
"""
try:
data = await request.json()
settings = get_settings()
updated_settings = []
errors = []
for key, value in data.items():
if key in SETTINGS_SCHEMA:
try:
if settings.set(key, value):
updated_settings.append(key)
except Exception as e:
errors.append(f"Failed to set '{key}': {str(e)}")
else:
errors.append(f"Unknown setting: '{key}'")
# Save if any settings were updated
if updated_settings:
if settings.save():
# Check if LLM-related settings were updated and trigger lazy initialization
llm_settings = {'enable_ollama', 'enable_lmstudio', 'custom_ollama_url', 'custom_lmstudio_url'}
if any(setting in llm_settings for setting in updated_settings):
try:
from .utils.llm_wrapper import ensure_llm_initialized
ensure_llm_initialized()
except Exception as llm_e:
errors.append(f"Warning: Failed to initialize LLM services: {str(llm_e)}")
return web.json_response({
"success": True,
"updated": updated_settings,
"errors": errors,
"message": f"Updated {len(updated_settings)} setting(s)"
})
else:
return web.json_response({
"success": False,
"error": "Failed to save settings",
"updated": updated_settings,
"errors": errors
}, status=500)
else:
return web.json_response({
"success": False,
"error": "No valid settings to update",
"errors": errors
}, status=400)
except Exception as e:
return web.json_response(
{"success": False, "error": f"Failed to update settings: {str(e)}"},
status=500
)
@routes.post('/sage_utils/settings/reset')
async def reset_sage_settings(request):
"""
Resets all SageUtils settings to their default values.
"""
try:
settings = get_settings()
settings.reset_to_defaults()
return web.json_response({
"success": True,
"message": "All settings reset to defaults"
})
except Exception as e:
return web.json_response(
{"success": False, "error": f"Failed to reset settings: {str(e)}"},
status=500
)
@routes.get('/sage_cache/info')
async def get_sage_cache_info(request):
"""
test_routes.py
-351
View File
@@ -1,351 +0,0 @@
#!/usr/bin/env python3
"""
Test script for SageUtils custom routes.
This demonstrates how the custom routes would work when ComfyUI is running.
"""
import sys
import os
import json
from typing import Dict, Any
from unittest.mock import Mock, MagicMock
# Add ComfyUI path
sys.path.insert(0, '/home/ai/programs/comfyui')
class MockRequest:
"""Mock aiohttp request object"""
def __init__(self, match_info=None):
self.match_info = match_info or {}
class MockWeb:
"""Mock aiohttp web module"""
@staticmethod
def json_response(data, status=200):
return {
"status": status,
"content_type": "application/json",
"data": data
}
class MockRoutes:
"""Mock routes object that stores registered routes"""
def __init__(self):
self.routes = {}
def get(self, path):
def decorator(func):
self.routes[f"GET {path}"] = func
print(f"✓ Registered route: GET {path}")
return func
return decorator
async def call_route(self, method, path, request=None):
"""Call a registered route for testing"""
route_key = f"{method} {path}"
if route_key in self.routes:
return await self.routes[route_key](request or MockRequest())
else:
return MockWeb.json_response({"error": "Route not found"}, 404)
class MockPromptServer:
"""Mock PromptServer class"""
def __init__(self):
self.routes = MockRoutes()
@classmethod
def create_instance(cls):
cls.instance = cls()
return cls.instance
def test_sage_routes():
"""Test the SageUtils custom routes"""
print("=== Testing SageUtils Custom Routes ===\n")
# Create mock server instance
prompt_server = MockPromptServer.create_instance()
# Mock the modules that would be imported
sys.modules['server'] = Mock()
sys.modules['server'].PromptServer = MockPromptServer
sys.modules['aiohttp'] = Mock()
sys.modules['aiohttp'].web = MockWeb
# Import and setup our cache
print("1. Loading SageUtils cache...")
try:
# Suppress hash warnings
import warnings
warnings.filterwarnings("ignore")
from custom_nodes.comfyui_sageutils.utils.model_cache import cache
cache.load()
print(f" ✓ Cache loaded: {len(cache.hash)} files, {len(cache.info)} info entries")
# Add some dummy data for testing if cache is empty
if len(cache.hash) == 0:
print(" Adding dummy test data...")
test_hash = "abc123def456789"
test_path = "/models/test_model.safetensors"
cache.hash[test_path] = test_hash
cache.info[test_hash] = {
"hash": test_hash,
"civitai": "true",
"modelId": "12345",
"name": "Test Model",
"baseModel": "SDXL",
"lastUsed": "2025-06-22T10:30:00"
}
print(f" ✓ Added test data: 1 file, 1 info entry")
except Exception as e:
print(f" ✗ Error loading cache: {e}")
import traceback
traceback.print_exc()
return False
# Mock the imports for our routes module
original_modules = {}
mock_modules = {
'server': Mock(),
'aiohttp': Mock(),
'.utils.model_cache': Mock()
}
# Setup mocks
mock_modules['server'].PromptServer = MockPromptServer
mock_modules['aiohttp'].web = MockWeb
mock_modules['.utils.model_cache'].cache = cache
print("\n2. Testing route registration...")
# Manually register routes like our module would
routes = prompt_server.routes
@routes.get('/sage_cache/info')
async def get_sage_cache_info(request):
try:
cache.load()
return MockWeb.json_response(cache.info)
except Exception as e:
return MockWeb.json_response(
{"error": f"Failed to retrieve cache info: {str(e)}"},
status=500
)
@routes.get('/sage_cache/hash')
async def get_sage_cache_hash(request):
try:
cache.load()
return MockWeb.json_response(cache.hash)
except Exception as e:
return MockWeb.json_response(
{"error": f"Failed to retrieve cache hash: {str(e)}"},
status=500
)
@routes.get('/sage_cache/stats')
async def get_sage_cache_stats(request):
try:
cache.load()
stats = {
"total_files": len(cache.hash),
"total_info_entries": len(cache.info),
"cache_files": {
"hash_path": str(cache.hash_path),
"info_path": str(cache.info_path),
"main_path": str(cache.main_path),
"ollama_models_path": str(cache.ollama_models_path)
}
}
# Count files by civitai status
civitai_found = 0
civitai_not_found = 0
for info_data in cache.info.values():
civitai_status = info_data.get("civitai", "False")
if isinstance(civitai_status, str):
if civitai_status.lower() in ["true", "1"]:
civitai_found += 1
else:
civitai_not_found += 1
elif civitai_status is True:
civitai_found += 1
else:
civitai_not_found += 1
stats["civitai_stats"] = {
"found_on_civitai": civitai_found,
"not_found_on_civitai": civitai_not_found
}
return MockWeb.json_response(stats)
except Exception as e:
return MockWeb.json_response(
{"error": f"Failed to retrieve cache stats: {str(e)}"},
status=500
)
@routes.get('/sage_cache/file/{file_hash}')
async def get_sage_cache_file_info(request):
try:
file_hash = request.match_info.get('file_hash', '')
if not file_hash:
return MockWeb.json_response(
{"error": "No file hash provided"},
status=400
)
cache.load()
file_info = cache.info.get(file_hash)
if file_info is None:
return MockWeb.json_response(
{"error": f"No information found for hash: {file_hash}"},
status=404
)
# Also include which file paths use this hash
file_paths = [path for path, hash_val in cache.hash.items() if hash_val == file_hash]
result = {
"hash": file_hash,
"info": file_info,
"file_paths": file_paths
}
return MockWeb.json_response(result)
except Exception as e:
return MockWeb.json_response(
{"error": f"Failed to retrieve file info: {str(e)}"},
status=500
)
@routes.get('/sage_cache/path')
async def get_sage_cache_path_info(request):
try:
# Get file path from query parameter (mocked)
file_path = getattr(request, 'file_path', '/models/test_model.safetensors')
if not file_path:
return MockWeb.json_response(
{"error": "No file_path parameter provided"},
status=400
)
cache.load()
# Look up hash for this file path
file_hash = cache.hash.get(file_path)
if file_hash is None:
return MockWeb.json_response(
{"error": f"No hash found for file path: {file_path}"},
status=404
)
# Get the info for this hash
file_info = cache.info.get(file_hash)
if file_info is None:
return MockWeb.json_response(
{"error": f"No information found for hash {file_hash} (path: {file_path})"},
status=404
)
# Also include all file paths that use this same hash
all_file_paths = [path for path, hash_val in cache.hash.items() if hash_val == file_hash]
result = {
"file_path": file_path,
"hash": file_hash,
"info": file_info,
"all_paths_with_same_hash": all_file_paths
}
return MockWeb.json_response(result)
except Exception as e:
return MockWeb.json_response(
{"error": f"Failed to retrieve path info: {str(e)}"},
status=500
)
print("\n3. Testing routes...")
# Test each route
async def test_routes():
# Test /sage_cache/stats
print("\n Testing GET /sage_cache/stats")
response = await routes.call_route("GET", "/sage_cache/stats")
print(f" Status: {response['status']}")
if response['status'] == 200:
data = response['data']
print(f" Total files: {data['total_files']}")
print(f" Total info entries: {data['total_info_entries']}")
print(f" Civitai found: {data['civitai_stats']['found_on_civitai']}")
print(f" Civitai not found: {data['civitai_stats']['not_found_on_civitai']}")
# Test /sage_cache/hash
print("\n Testing GET /sage_cache/hash")
response = await routes.call_route("GET", "/sage_cache/hash")
print(f" Status: {response['status']}")
if response['status'] == 200:
data = response['data']
print(f" Hash entries: {len(data)}")
if data:
first_path = list(data.keys())[0]
print(f" Sample: {first_path[:50]}... -> {data[first_path][:16]}...")
# Test /sage_cache/info
print("\n Testing GET /sage_cache/info")
response = await routes.call_route("GET", "/sage_cache/info")
print(f" Status: {response['status']}")
if response['status'] == 200:
data = response['data']
print(f" Info entries: {len(data)}")
if data:
first_hash = list(data.keys())[0]
print(f" Sample hash: {first_hash}")
info = data[first_hash]
print(f" Sample info: civitai={info.get('civitai')}, name={info.get('name', 'N/A')}")
# Test /sage_cache/file/{hash}
if cache.info:
test_hash = list(cache.info.keys())[0]
print(f"\n Testing GET /sage_cache/file/{test_hash}")
request = MockRequest({"file_hash": test_hash})
response = await routes.call_route("GET", "/sage_cache/file/{file_hash}", request)
print(f" Status: {response['status']}")
if response['status'] == 200:
data = response['data']
print(f" Hash: {data['hash']}")
print(f" File paths: {len(data['file_paths'])}")
print(f" Info keys: {list(data['info'].keys())}")
# Test /sage_cache/path
print("\n Testing GET /sage_cache/path")
response = await routes.call_route("GET", "/sage_cache/path")
print(f" Status: {response['status']}")
if response['status'] == 200:
data = response['data']
print(f" File path: {data['file_path']}")
print(f" Hash: {data['hash']}")
print(f" Info keys: {list(data['info'].keys())}")
print(f" All paths with same hash: {len(data['all_paths_with_same_hash'])}")
# Run the async tests
import asyncio
asyncio.run(test_routes())
print("\n=== Test Complete ===")
print("✓ All routes registered and tested successfully!")
print("\nWhen ComfyUI is running, these routes will be available at:")
print(" - http://localhost:8188/sage_cache/info")
print(" - http://localhost:8188/sage_cache/hash")
print(" - http://localhost:8188/sage_cache/stats")
print(" - http://localhost:8188/sage_cache/file/{hash}")
print(" - Or with /api prefix for each route")
return True
if __name__ == "__main__":
test_sage_routes()
utils/llm_wrapper.py
+87 -44
View File
@@ -2,6 +2,10 @@ import logging
from .helpers_image import tensor_to_base64, tensor_to_temp_image
from .llm_cache import get_llm_cache
# Initialization flags to track if services have been initialized
_ollama_initialized = False
_lmstudio_initialized = False
# Attempt to import ollama, if available. Set a flag if it is not available.
try:
import ollama
@@ -153,6 +157,9 @@ def get_ollama_models() -> list[str]:
def ollama_generate_vision(model: str, prompt: str, keep_alive: float = 0.0, images=None, options=None) -> str:
"""Generate a response from an Ollama vision model."""
# Ensure Ollama is initialized before use
ensure_ollama_initialized()
if not OLLAMA_AVAILABLE or ollama_client is None:
raise ImportError("Ollama is not available. Please install it to use this function.")
vision_models = get_ollama_vision_models()
@@ -191,6 +198,9 @@ def ollama_generate_vision(model: str, prompt: str, keep_alive: float = 0.0, ima
def ollama_generate(model: str, prompt: str, keep_alive: float = 0.0, options=None) -> str:
"""Generate a response from an Ollama model."""
# Ensure Ollama is initialized before use
ensure_ollama_initialized()
if not OLLAMA_AVAILABLE or ollama_client is None:
raise ImportError("Ollama is not available. Please install it to use this function.")
models = get_ollama_models()
@@ -357,6 +367,9 @@ def get_lmstudio_vision_models() -> list[str]:
def lmstudio_generate_vision(model: str, prompt: str, keep_alive: int = 0, images=None, options=None) -> str:
"""Generate a response from an LM Studio vision model."""
# Ensure LM Studio is initialized before use
ensure_lmstudio_initialized()
if not LMSTUDIO_AVAILABLE or lms is None:
raise ImportError("LM Studio is not available. Please install it to use this function.")
model_list = get_lmstudio_vision_models()
@@ -391,6 +404,9 @@ def lmstudio_generate_vision(model: str, prompt: str, keep_alive: int = 0, image
def lmstudio_generate(model: str, prompt: str, keep_alive: int = 0, options=None) -> str:
"""Generate a response from an LM Studio model."""
# Ensure LM Studio is initialized before use
ensure_lmstudio_initialized()
if not LMSTUDIO_AVAILABLE or lms is None:
raise ImportError("LM Studio is not available. Please install it to use this function.")
model_list = get_lmstudio_models()
@@ -472,73 +488,71 @@ def lmstudio_generate_vision_refine(model: str, prompt: str, images=None, option
return ("", "")
def init_ollama():
"""Initialize Ollama if available. Print config values for Ollama."""
global ollama_client
"""Initialize Ollama client"""
global ollama_client, _ollama_initialized
from .settings import get_setting
if not OLLAMA_AVAILABLE or ollama is None:
logging.info("Ollama is not available; skipping Ollama initialization.")
return
# Check if Ollama is available and enabled
if not OLLAMA_AVAILABLE:
logging.warning("Ollama library is not available.")
return False
if not get_setting("enable_ollama", False):
logging.info("Ollama is disabled in settings.")
_ollama_initialized = False
ollama_client = None
return False
try:
# Try to use new settings system
try:
from .settings import get_setting, is_feature_enabled
if not is_feature_enabled('enable_ollama'):
logging.info("Ollama is disabled in settings; skipping initialization.")
return
use_custom_url = get_setting('ollama_use_custom_url', False)
custom_url = get_setting('ollama_custom_url', '')
except ImportError:
# Fallback to old config system
from . import config_manager
config = config_manager.settings_manager.data or {}
use_custom_url = config.get('ollama_use_custom_url', False)
custom_url = config.get('ollama_custom_url', '')
if use_custom_url and custom_url:
# Get custom URL or use default
custom_url = get_setting("custom_ollama_url", "http://localhost:11434")
if custom_url and custom_url.strip():
ollama_client = ollama.Client(host=custom_url)
logging.info(f"Ollama client initialized with custom host: {custom_url}")
logging.info(f"Ollama client initialized with custom URL: {custom_url}")
else:
ollama_client = ollama.Client()
logging.info("Ollama client initialized with default settings.")
logging.info("Ollama client initialized with default URL")
_ollama_initialized = True
return True
except Exception as e:
ollama_client = None
logging.error(f"Failed to initialize Ollama client: {e}")
_ollama_initialized = False
ollama_client = None
return False
def init_lmstudio():
"""Initialize LM Studio if available. Print config values for LM Studio."""
global _lmstudio_initialized
from .settings import get_setting
if not LMSTUDIO_AVAILABLE or lms is None:
logging.info("LM Studio is not available; skipping LM Studio initialization.")
return
_lmstudio_initialized = False
return False
# Check if LM Studio is enabled
if not get_setting("enable_lmstudio", False):
logging.info("LM Studio is disabled in settings; skipping initialization.")
_lmstudio_initialized = False
return False
try:
# Try to use new settings system
try:
from .settings import get_setting, is_feature_enabled
if not is_feature_enabled('enable_lmstudio'):
logging.info("LM Studio is disabled in settings; skipping initialization.")
return
use_custom_url = get_setting('lmstudio_use_custom_url', False)
custom_url = get_setting('lmstudio_custom_url', '')
except ImportError:
# Fallback to old config system
from . import config_manager
config = config_manager.settings_manager.data or {}
use_custom_url = config.get('lmstudio_use_custom_url', False)
custom_url = config.get('lmstudio_custom_url', '')
custom_url = get_setting('custom_lmstudio_url', '')
if use_custom_url and custom_url:
if custom_url and custom_url.strip():
lm_client = lms.get_default_client(custom_url)
logging.info(f"LM Studio client configured with custom URL: {custom_url}")
else:
logging.info("LM Studio using default configuration.")
_lmstudio_initialized = True
return True
except Exception as e:
logging.error(f"Failed to configure LM Studio: {e}")
_lmstudio_initialized = False
return False
def init_llm():
@@ -546,3 +560,32 @@ def init_llm():
init_ollama()
init_lmstudio()
logging.info("LLM clients initialized.")
def ensure_ollama_initialized():
"""Ensure Ollama is initialized if it's enabled in settings and not already initialized."""
global _ollama_initialized
from .settings import get_setting
if get_setting("enable_ollama", False) and not _ollama_initialized:
logging.info("Ollama is enabled but not initialized, initializing now...")
return init_ollama()
return _ollama_initialized
def ensure_lmstudio_initialized():
"""Ensure LM Studio is initialized if it's enabled in settings and not already initialized."""
global _lmstudio_initialized
from .settings import get_setting
if get_setting("enable_lmstudio", False) and not _lmstudio_initialized:
logging.info("LM Studio is enabled but not initialized, initializing now...")
return init_lmstudio()
return _lmstudio_initialized
def ensure_llm_initialized():
"""Ensure all enabled LLM services are initialized."""
ollama_ok = ensure_ollama_initialized()
lmstudio_ok = ensure_lmstudio_initialized()
return ollama_ok or lmstudio_ok
utils/settings.py
+4
View File
@@ -178,6 +178,9 @@ class SageSettings:
schema_entry = SETTINGS_SCHEMA[key].copy()
schema_entry["current_value"] = self._settings.get(key)
# Convert type to string representation for JSON serialization
if "type" in schema_entry:
schema_entry["type"] = schema_entry["type"].__name__
return schema_entry
def list_all_settings(self) -> Dict[str, Dict[str, Any]]:
@@ -186,6 +189,7 @@ class SageSettings:
for key in SETTINGS_SCHEMA.keys():
setting_info = self.get_setting_info(key)
if setting_info is not None:
# The setting_info is already JSON-safe from get_setting_info()
result[key] = setting_info
return result