From 6a8f0867d9a27aeac3ce157ddd8898496802d9d4 Mon Sep 17 00:00:00 2001 From: Will Miao <13051207myq@gmail.com> Date: Thu, 11 Sep 2025 22:37:46 +0800 Subject: [PATCH] refactor: migrate auto_organize_models logic to service layer with dependency injection --- py/routes/base_model_routes.py | 358 ++--------------- py/services/model_file_service.py | 447 +++++++++++++++++++++ py/services/websocket_progress_callback.py | 11 + 3 files changed, 487 insertions(+), 329 deletions(-) create mode 100644 py/services/model_file_service.py create mode 100644 py/services/websocket_progress_callback.py diff --git a/py/routes/base_model_routes.py b/py/routes/base_model_routes.py index ce0e6709..be0f2695 100644 --- a/py/routes/base_model_routes.py +++ b/py/routes/base_model_routes.py @@ -12,8 +12,8 @@ from ..utils.routes_common import ModelRouteUtils from ..services.websocket_manager import ws_manager from ..services.settings_manager import settings from ..services.server_i18n import server_i18n -from ..utils.utils import calculate_relative_path_for_model -from ..utils.constants import AUTO_ORGANIZE_BATCH_SIZE +from ..services.model_file_service import ModelFileService, ModelMoveService +from ..services.websocket_progress_callback import WebSocketProgressCallback from ..config import config logger = logging.getLogger(__name__) @@ -33,6 +33,11 @@ class BaseModelRoutes(ABC): loader=jinja2.FileSystemLoader(config.templates_path), autoescape=True ) + + # Initialize file services with dependency injection + self.model_file_service = ModelFileService(service.scanner, service.model_type) + self.model_move_service = ModelMoveService(service.scanner) + self.websocket_progress_callback = WebSocketProgressCallback() def setup_routes(self, app: web.Application, prefix: str): """Setup common routes for the model type @@ -46,6 +51,7 @@ class BaseModelRoutes(ABC): app.router.add_post(f'/api/{prefix}/delete', self.delete_model) app.router.add_post(f'/api/{prefix}/exclude', self.exclude_model) app.router.add_post(f'/api/{prefix}/fetch-civitai', self.fetch_civitai) + app.router.add_post(f'/api/{prefix}/fetch-all-civitai', self.fetch_all_civitai) app.router.add_post(f'/api/{prefix}/relink-civitai', self.relink_civitai) app.router.add_post(f'/api/{prefix}/replace-preview', self.replace_preview) app.router.add_post(f'/api/{prefix}/save-metadata', self.save_metadata) @@ -84,8 +90,6 @@ class BaseModelRoutes(ABC): app.router.add_get(f'/api/cancel-download-get', self.cancel_download_get) app.router.add_get(f'/api/download-progress/{{download_id}}', self.get_download_progress) - # CivitAI integration routes - app.router.add_post(f'/api/{prefix}/fetch-all-civitai', self.fetch_all_civitai) # app.router.add_get(f'/api/civitai/versions/{{model_id}}', self.get_civitai_versions) # Add generic page route @@ -707,33 +711,17 @@ class BaseModelRoutes(ABC): data = await request.json() file_path = data.get('file_path') target_path = data.get('target_path') + if not file_path or not target_path: return web.Response(text='File path and target path are required', status=400) - import os - source_dir = os.path.dirname(file_path) - if os.path.normpath(source_dir) == os.path.normpath(target_path): - logger.info(f"Source and target directories are the same: {source_dir}") - return web.json_response({ - 'success': True, - 'message': 'Source and target directories are the same', - 'original_file_path': file_path, - 'new_file_path': file_path - }) - - new_file_path = await self.service.scanner.move_model(file_path, target_path) - if new_file_path: - return web.json_response({ - 'success': True, - 'original_file_path': file_path, - 'new_file_path': new_file_path - }) + + result = await self.model_move_service.move_model(file_path, target_path) + + if result['success']: + return web.json_response(result) else: - return web.json_response({ - 'success': False, - 'error': 'Failed to move model', - 'original_file_path': file_path, - 'new_file_path': None - }, status=500) + return web.json_response(result, status=500) + except Exception as e: logger.error(f"Error moving model: {e}", exc_info=True) return web.Response(text=str(e), status=500) @@ -744,45 +732,13 @@ class BaseModelRoutes(ABC): data = await request.json() file_paths = data.get('file_paths', []) target_path = data.get('target_path') + if not file_paths or not target_path: return web.Response(text='File paths and target path are required', status=400) - results = [] - import os - for file_path in file_paths: - source_dir = os.path.dirname(file_path) - if os.path.normpath(source_dir) == os.path.normpath(target_path): - results.append({ - "original_file_path": file_path, - "new_file_path": file_path, - "success": True, - "message": "Source and target directories are the same" - }) - continue - - new_file_path = await self.service.scanner.move_model(file_path, target_path) - if new_file_path: - results.append({ - "original_file_path": file_path, - "new_file_path": new_file_path, - "success": True, - "message": "Success" - }) - else: - results.append({ - "original_file_path": file_path, - "new_file_path": None, - "success": False, - "message": "Failed to move model" - }) - success_count = sum(1 for r in results if r["success"]) - failure_count = len(results) - success_count - return web.json_response({ - 'success': True, - 'message': f'Moved {success_count} of {len(file_paths)} models', - 'results': results, - 'success_count': success_count, - 'failure_count': failure_count - }) + + result = await self.model_move_service.move_models_bulk(file_paths, target_path) + return web.json_response(result) + except Exception as e: logger.error(f"Error moving models in bulk: {e}", exc_info=True) return web.Response(text=str(e), status=500) @@ -816,12 +772,18 @@ class BaseModelRoutes(ABC): pass # Continue with all models if no valid JSON async with auto_organize_lock: - return await self._perform_auto_organize(file_paths) + # Use the service layer for business logic + result = await self.model_file_service.auto_organize_models( + file_paths=file_paths, + progress_callback=self.websocket_progress_callback + ) + + return web.json_response(result.to_dict()) except Exception as e: logger.error(f"Error in auto_organize_models: {e}", exc_info=True) - # Send error message via WebSocket and cleanup + # Send error message via WebSocket await ws_manager.broadcast_auto_organize_progress({ 'type': 'auto_organize_progress', 'status': 'error', @@ -833,268 +795,6 @@ class BaseModelRoutes(ABC): 'error': str(e) }, status=500) - async def _perform_auto_organize(self, file_paths=None) -> web.Response: - """Perform the actual auto-organize operation - - Args: - file_paths: Optional list of specific file paths to organize. - If None, organizes all models. - """ - try: - # Get all models from cache - cache = await self.service.scanner.get_cached_data() - all_models = cache.raw_data - - # Filter models if specific file paths are provided - if file_paths: - all_models = [model for model in all_models if model.get('file_path') in file_paths] - operation_type = 'bulk' - else: - operation_type = 'all' - - # Get model roots for this scanner - model_roots = self.service.get_model_roots() - if not model_roots: - await ws_manager.broadcast_auto_organize_progress({ - 'type': 'auto_organize_progress', - 'status': 'error', - 'error': 'No model roots configured', - 'operation_type': operation_type - }) - return web.json_response({ - 'success': False, - 'error': 'No model roots configured' - }, status=400) - - # Check if flat structure is configured for this model type - path_template = settings.get_download_path_template(self.service.model_type) - is_flat_structure = not path_template - - # Prepare results tracking - results = [] - total_models = len(all_models) - processed = 0 - success_count = 0 - failure_count = 0 - skipped_count = 0 - - # Send initial progress via WebSocket - await ws_manager.broadcast_auto_organize_progress({ - 'type': 'auto_organize_progress', - 'status': 'started', - 'total': total_models, - 'processed': 0, - 'success': 0, - 'failures': 0, - 'skipped': 0, - 'operation_type': operation_type - }) - - # Process models in batches - for i in range(0, total_models, AUTO_ORGANIZE_BATCH_SIZE): - batch = all_models[i:i + AUTO_ORGANIZE_BATCH_SIZE] - - for model in batch: - try: - file_path = model.get('file_path') - if not file_path: - if len(results) < 100: # Limit detailed results - results.append({ - "model": model.get('model_name', 'Unknown'), - "success": False, - "message": "No file path found" - }) - failure_count += 1 - processed += 1 - continue - - # Find which model root this file belongs to - current_root = None - for root in model_roots: - # Normalize paths for comparison - normalized_root = os.path.normpath(root).replace(os.sep, '/') - normalized_file = os.path.normpath(file_path).replace(os.sep, '/') - - if normalized_file.startswith(normalized_root): - current_root = root - break - - if not current_root: - if len(results) < 100: # Limit detailed results - results.append({ - "model": model.get('model_name', 'Unknown'), - "success": False, - "message": "Model file not found in any configured root directory" - }) - failure_count += 1 - processed += 1 - continue - - # Handle flat structure case - if is_flat_structure: - current_dir = os.path.dirname(file_path) - # Check if already in root directory - if os.path.normpath(current_dir) == os.path.normpath(current_root): - skipped_count += 1 - processed += 1 - continue - - # Move to root directory for flat structure - target_dir = current_root - else: - # Calculate new relative path based on settings - new_relative_path = calculate_relative_path_for_model(model, self.service.model_type) - - # If no relative path calculated (insufficient metadata), skip - if not new_relative_path: - if len(results) < 100: # Limit detailed results - results.append({ - "model": model.get('model_name', 'Unknown'), - "success": False, - "message": "Skipped - insufficient metadata for organization" - }) - skipped_count += 1 - processed += 1 - continue - - # Calculate target directory - target_dir = os.path.join(current_root, new_relative_path).replace(os.sep, '/') - - current_dir = os.path.dirname(file_path) - - # Skip if already in correct location - if current_dir.replace(os.sep, '/') == target_dir.replace(os.sep, '/'): - skipped_count += 1 - processed += 1 - continue - - # Check if target file would conflict - file_name = os.path.basename(file_path) - target_file_path = os.path.join(target_dir, file_name) - - if os.path.exists(target_file_path): - if len(results) < 100: # Limit detailed results - results.append({ - "model": model.get('model_name', 'Unknown'), - "success": False, - "message": f"Target file already exists: {target_file_path}" - }) - failure_count += 1 - processed += 1 - continue - - # Perform the move - success = await self.service.scanner.move_model(file_path, target_dir) - - if success: - success_count += 1 - else: - if len(results) < 100: # Limit detailed results - results.append({ - "model": model.get('model_name', 'Unknown'), - "success": False, - "message": "Failed to move model" - }) - failure_count += 1 - - processed += 1 - - except Exception as e: - logger.error(f"Error processing model {model.get('model_name', 'Unknown')}: {e}", exc_info=True) - if len(results) < 100: # Limit detailed results - results.append({ - "model": model.get('model_name', 'Unknown'), - "success": False, - "message": f"Error: {str(e)}" - }) - failure_count += 1 - processed += 1 - - # Send progress update after each batch - await ws_manager.broadcast_auto_organize_progress({ - 'type': 'auto_organize_progress', - 'status': 'processing', - 'total': total_models, - 'processed': processed, - 'success': success_count, - 'failures': failure_count, - 'skipped': skipped_count, - 'operation_type': operation_type - }) - - # Small delay between batches to prevent overwhelming the system - await asyncio.sleep(0.1) - - # Send completion message - await ws_manager.broadcast_auto_organize_progress({ - 'type': 'auto_organize_progress', - 'status': 'cleaning', - 'total': total_models, - 'processed': processed, - 'success': success_count, - 'failures': failure_count, - 'skipped': skipped_count, - 'message': 'Cleaning up empty directories...', - 'operation_type': operation_type - }) - - # Clean up empty directories after organizing - from ..utils.utils import remove_empty_dirs - cleanup_counts = {} - for root in model_roots: - removed = remove_empty_dirs(root) - cleanup_counts[root] = removed - - # Send cleanup completed message - await ws_manager.broadcast_auto_organize_progress({ - 'type': 'auto_organize_progress', - 'status': 'completed', - 'total': total_models, - 'processed': processed, - 'success': success_count, - 'failures': failure_count, - 'skipped': skipped_count, - 'cleanup': cleanup_counts, - 'operation_type': operation_type - }) - - # Prepare response with limited details - response_data = { - 'success': True, - 'message': f'Auto-organize {operation_type} completed: {success_count} moved, {skipped_count} skipped, {failure_count} failed out of {total_models} total', - 'summary': { - 'total': total_models, - 'success': success_count, - 'skipped': skipped_count, - 'failures': failure_count, - 'organization_type': 'flat' if is_flat_structure else 'structured', - 'cleaned_dirs': cleanup_counts, - 'operation_type': operation_type - } - } - - # Only include detailed results if under limit - if len(results) <= 100: - response_data['results'] = results - else: - response_data['results_truncated'] = True - response_data['sample_results'] = results[:50] # Show first 50 as sample - - return web.json_response(response_data) - - except Exception as e: - logger.error(f"Error in _perform_auto_organize: {e}", exc_info=True) - - # Send error message via WebSocket - await ws_manager.broadcast_auto_organize_progress({ - 'type': 'auto_organize_progress', - 'status': 'error', - 'error': str(e), - 'operation_type': operation_type if 'operation_type' in locals() else 'unknown' - }) - - raise e - async def get_auto_organize_progress(self, request: web.Request) -> web.Response: """Get current auto-organize progress for polling""" try: diff --git a/py/services/model_file_service.py b/py/services/model_file_service.py new file mode 100644 index 00000000..705423ae --- /dev/null +++ b/py/services/model_file_service.py @@ -0,0 +1,447 @@ +import asyncio +import os +import logging +from typing import List, Dict, Callable, Optional, Any +from abc import ABC, abstractmethod + +from ..utils.utils import calculate_relative_path_for_model, remove_empty_dirs +from ..utils.constants import AUTO_ORGANIZE_BATCH_SIZE +from ..services.settings_manager import settings + +logger = logging.getLogger(__name__) + + +class ProgressCallback(ABC): + """Abstract callback interface for progress reporting""" + + @abstractmethod + async def on_progress(self, progress_data: Dict[str, Any]) -> None: + """Called when progress is updated""" + pass + + +class AutoOrganizeResult: + """Result object for auto-organize operations""" + + def __init__(self): + self.total: int = 0 + self.processed: int = 0 + self.success_count: int = 0 + self.failure_count: int = 0 + self.skipped_count: int = 0 + self.operation_type: str = 'unknown' + self.cleanup_counts: Dict[str, int] = {} + self.results: List[Dict[str, Any]] = [] + self.results_truncated: bool = False + self.sample_results: List[Dict[str, Any]] = [] + self.is_flat_structure: bool = False + + def to_dict(self) -> Dict[str, Any]: + """Convert result to dictionary""" + result = { + 'success': True, + 'message': f'Auto-organize {self.operation_type} completed: {self.success_count} moved, {self.skipped_count} skipped, {self.failure_count} failed out of {self.total} total', + 'summary': { + 'total': self.total, + 'success': self.success_count, + 'skipped': self.skipped_count, + 'failures': self.failure_count, + 'organization_type': 'flat' if self.is_flat_structure else 'structured', + 'cleaned_dirs': self.cleanup_counts, + 'operation_type': self.operation_type + } + } + + if self.results_truncated: + result['results_truncated'] = True + result['sample_results'] = self.sample_results + else: + result['results'] = self.results + + return result + + +class ModelFileService: + """Service for handling model file operations and organization""" + + def __init__(self, scanner, model_type: str): + """Initialize the service + + Args: + scanner: Model scanner instance + model_type: Type of model (e.g., 'lora', 'checkpoint') + """ + self.scanner = scanner + self.model_type = model_type + + def get_model_roots(self) -> List[str]: + """Get model root directories""" + return self.scanner.get_model_roots() + + async def auto_organize_models( + self, + file_paths: Optional[List[str]] = None, + progress_callback: Optional[ProgressCallback] = None + ) -> AutoOrganizeResult: + """Auto-organize models based on current settings + + Args: + file_paths: Optional list of specific file paths to organize. + If None, organizes all models. + progress_callback: Optional callback for progress updates + + Returns: + AutoOrganizeResult object with operation results + """ + result = AutoOrganizeResult() + + try: + # Get all models from cache + cache = await self.scanner.get_cached_data() + all_models = cache.raw_data + + # Filter models if specific file paths are provided + if file_paths: + all_models = [model for model in all_models if model.get('file_path') in file_paths] + result.operation_type = 'bulk' + else: + result.operation_type = 'all' + + # Get model roots for this scanner + model_roots = self.get_model_roots() + if not model_roots: + raise ValueError('No model roots configured') + + # Check if flat structure is configured for this model type + path_template = settings.get_download_path_template(self.model_type) + result.is_flat_structure = not path_template + + # Initialize tracking + result.total = len(all_models) + + # Send initial progress + if progress_callback: + await progress_callback.on_progress({ + 'type': 'auto_organize_progress', + 'status': 'started', + 'total': result.total, + 'processed': 0, + 'success': 0, + 'failures': 0, + 'skipped': 0, + 'operation_type': result.operation_type + }) + + # Process models in batches + await self._process_models_in_batches( + all_models, + model_roots, + result, + progress_callback + ) + + # Send cleanup progress + if progress_callback: + await progress_callback.on_progress({ + 'type': 'auto_organize_progress', + 'status': 'cleaning', + 'total': result.total, + 'processed': result.processed, + 'success': result.success_count, + 'failures': result.failure_count, + 'skipped': result.skipped_count, + 'message': 'Cleaning up empty directories...', + 'operation_type': result.operation_type + }) + + # Clean up empty directories + result.cleanup_counts = await self._cleanup_empty_directories(model_roots) + + # Send completion message + if progress_callback: + await progress_callback.on_progress({ + 'type': 'auto_organize_progress', + 'status': 'completed', + 'total': result.total, + 'processed': result.processed, + 'success': result.success_count, + 'failures': result.failure_count, + 'skipped': result.skipped_count, + 'cleanup': result.cleanup_counts, + 'operation_type': result.operation_type + }) + + return result + + except Exception as e: + logger.error(f"Error in auto_organize_models: {e}", exc_info=True) + + # Send error message + if progress_callback: + await progress_callback.on_progress({ + 'type': 'auto_organize_progress', + 'status': 'error', + 'error': str(e), + 'operation_type': result.operation_type + }) + + raise e + + async def _process_models_in_batches( + self, + all_models: List[Dict[str, Any]], + model_roots: List[str], + result: AutoOrganizeResult, + progress_callback: Optional[ProgressCallback] + ) -> None: + """Process models in batches to avoid overwhelming the system""" + + for i in range(0, result.total, AUTO_ORGANIZE_BATCH_SIZE): + batch = all_models[i:i + AUTO_ORGANIZE_BATCH_SIZE] + + for model in batch: + await self._process_single_model(model, model_roots, result) + result.processed += 1 + + # Send progress update after each batch + if progress_callback: + await progress_callback.on_progress({ + 'type': 'auto_organize_progress', + 'status': 'processing', + 'total': result.total, + 'processed': result.processed, + 'success': result.success_count, + 'failures': result.failure_count, + 'skipped': result.skipped_count, + 'operation_type': result.operation_type + }) + + # Small delay between batches + await asyncio.sleep(0.1) + + async def _process_single_model( + self, + model: Dict[str, Any], + model_roots: List[str], + result: AutoOrganizeResult + ) -> None: + """Process a single model for organization""" + try: + file_path = model.get('file_path') + model_name = model.get('model_name', 'Unknown') + + if not file_path: + self._add_result(result, model_name, False, "No file path found") + result.failure_count += 1 + return + + # Find which model root this file belongs to + current_root = self._find_model_root(file_path, model_roots) + if not current_root: + self._add_result(result, model_name, False, + "Model file not found in any configured root directory") + result.failure_count += 1 + return + + # Determine target directory + target_dir = await self._calculate_target_directory( + model, current_root, result.is_flat_structure + ) + + if target_dir is None: + self._add_result(result, model_name, False, + "Skipped - insufficient metadata for organization") + result.skipped_count += 1 + return + + current_dir = os.path.dirname(file_path) + + # Skip if already in correct location + if current_dir.replace(os.sep, '/') == target_dir.replace(os.sep, '/'): + result.skipped_count += 1 + return + + # Check for conflicts + file_name = os.path.basename(file_path) + target_file_path = os.path.join(target_dir, file_name) + + if os.path.exists(target_file_path): + self._add_result(result, model_name, False, + f"Target file already exists: {target_file_path}") + result.failure_count += 1 + return + + # Perform the move + success = await self.scanner.move_model(file_path, target_dir) + + if success: + result.success_count += 1 + else: + self._add_result(result, model_name, False, "Failed to move model") + result.failure_count += 1 + + except Exception as e: + logger.error(f"Error processing model {model.get('model_name', 'Unknown')}: {e}", exc_info=True) + self._add_result(result, model.get('model_name', 'Unknown'), False, f"Error: {str(e)}") + result.failure_count += 1 + + def _find_model_root(self, file_path: str, model_roots: List[str]) -> Optional[str]: + """Find which model root the file belongs to""" + for root in model_roots: + # Normalize paths for comparison + normalized_root = os.path.normpath(root).replace(os.sep, '/') + normalized_file = os.path.normpath(file_path).replace(os.sep, '/') + + if normalized_file.startswith(normalized_root): + return root + return None + + async def _calculate_target_directory( + self, + model: Dict[str, Any], + current_root: str, + is_flat_structure: bool + ) -> Optional[str]: + """Calculate the target directory for a model""" + if is_flat_structure: + file_path = model.get('file_path') + current_dir = os.path.dirname(file_path) + + # Check if already in root directory + if os.path.normpath(current_dir) == os.path.normpath(current_root): + return None # Signal to skip + + return current_root + else: + # Calculate new relative path based on settings + new_relative_path = calculate_relative_path_for_model(model, self.model_type) + + if not new_relative_path: + return None # Signal to skip + + return os.path.join(current_root, new_relative_path).replace(os.sep, '/') + + def _add_result( + self, + result: AutoOrganizeResult, + model_name: str, + success: bool, + message: str + ) -> None: + """Add a result entry if under the limit""" + if len(result.results) < 100: # Limit detailed results + result.results.append({ + "model": model_name, + "success": success, + "message": message + }) + elif len(result.results) == 100: + # Mark as truncated and save sample + result.results_truncated = True + result.sample_results = result.results[:50] + + async def _cleanup_empty_directories(self, model_roots: List[str]) -> Dict[str, int]: + """Clean up empty directories after organizing""" + cleanup_counts = {} + for root in model_roots: + removed = remove_empty_dirs(root) + cleanup_counts[root] = removed + return cleanup_counts + + +class ModelMoveService: + """Service for handling individual model moves""" + + def __init__(self, scanner): + """Initialize the service + + Args: + scanner: Model scanner instance + """ + self.scanner = scanner + + async def move_model(self, file_path: str, target_path: str) -> Dict[str, Any]: + """Move a single model file + + Args: + file_path: Source file path + target_path: Target directory path + + Returns: + Dictionary with move result + """ + try: + source_dir = os.path.dirname(file_path) + if os.path.normpath(source_dir) == os.path.normpath(target_path): + logger.info(f"Source and target directories are the same: {source_dir}") + return { + 'success': True, + 'message': 'Source and target directories are the same', + 'original_file_path': file_path, + 'new_file_path': file_path + } + + new_file_path = await self.scanner.move_model(file_path, target_path) + if new_file_path: + return { + 'success': True, + 'original_file_path': file_path, + 'new_file_path': new_file_path + } + else: + return { + 'success': False, + 'error': 'Failed to move model', + 'original_file_path': file_path, + 'new_file_path': None + } + except Exception as e: + logger.error(f"Error moving model: {e}", exc_info=True) + return { + 'success': False, + 'error': str(e), + 'original_file_path': file_path, + 'new_file_path': None + } + + async def move_models_bulk(self, file_paths: List[str], target_path: str) -> Dict[str, Any]: + """Move multiple model files + + Args: + file_paths: List of source file paths + target_path: Target directory path + + Returns: + Dictionary with bulk move results + """ + try: + results = [] + + for file_path in file_paths: + result = await self.move_model(file_path, target_path) + results.append({ + "original_file_path": file_path, + "new_file_path": result.get('new_file_path'), + "success": result['success'], + "message": result.get('message', result.get('error', 'Unknown')) + }) + + success_count = sum(1 for r in results if r["success"]) + failure_count = len(results) - success_count + + return { + 'success': True, + 'message': f'Moved {success_count} of {len(file_paths)} models', + 'results': results, + 'success_count': success_count, + 'failure_count': failure_count + } + except Exception as e: + logger.error(f"Error moving models in bulk: {e}", exc_info=True) + return { + 'success': False, + 'error': str(e), + 'results': [], + 'success_count': 0, + 'failure_count': len(file_paths) + } \ No newline at end of file diff --git a/py/services/websocket_progress_callback.py b/py/services/websocket_progress_callback.py new file mode 100644 index 00000000..1a390f30 --- /dev/null +++ b/py/services/websocket_progress_callback.py @@ -0,0 +1,11 @@ +from typing import Dict, Any +from .model_file_service import ProgressCallback +from .websocket_manager import ws_manager + + +class WebSocketProgressCallback(ProgressCallback): + """WebSocket implementation of progress callback""" + + async def on_progress(self, progress_data: Dict[str, Any]) -> None: + """Send progress data via WebSocket""" + await ws_manager.broadcast_auto_organize_progress(progress_data) \ No newline at end of file