ComfyUI-Lora-Manager/py/services/base_model_service.py

from abc import ABC, abstractmethod
import asyncio
from typing import Dict, List, Optional, Type, TYPE_CHECKING
import logging
import os

from ..utils.models import BaseModelMetadata
from ..utils.metadata_manager import MetadataManager
from .model_query import FilterCriteria, ModelCacheRepository, ModelFilterSet, SearchStrategy, SettingsProvider
from .settings_manager import get_settings_manager

logger = logging.getLogger(__name__)

if TYPE_CHECKING:
    from .model_update_service import ModelUpdateService

class BaseModelService(ABC):
    """Base service class for all model types"""
    
    def __init__(
        self,
        model_type: str,
        scanner,
        metadata_class: Type[BaseModelMetadata],
        *,
        cache_repository: Optional[ModelCacheRepository] = None,
        filter_set: Optional[ModelFilterSet] = None,
        search_strategy: Optional[SearchStrategy] = None,
        settings_provider: Optional[SettingsProvider] = None,
        update_service: Optional["ModelUpdateService"] = None,
    ):
        """Initialize the service.

        Args:
            model_type: Type of model (lora, checkpoint, etc.).
            scanner: Model scanner instance.
            metadata_class: Metadata class for this model type.
            cache_repository: Custom repository for cache access (primarily for tests).
            filter_set: Filter component controlling folder/tag/favorites logic.
            search_strategy: Search component for fuzzy/text matching.
            settings_provider: Settings object; defaults to the global settings manager.
            update_service: Service used to determine whether models have remote updates available.
        """
        self.model_type = model_type
        self.scanner = scanner
        self.metadata_class = metadata_class
        self.settings = settings_provider or get_settings_manager()
        self.cache_repository = cache_repository or ModelCacheRepository(scanner)
        self.filter_set = filter_set or ModelFilterSet(self.settings)
        self.search_strategy = search_strategy or SearchStrategy()
        self.update_service = update_service
    
    async def get_paginated_data(
        self,
        page: int,
        page_size: int,
        sort_by: str = 'name',
        folder: str = None,
        search: str = None,
        fuzzy_search: bool = False,
        base_models: list = None,
        tags: list = None,
        search_options: dict = None,
        hash_filters: dict = None,
        favorites_only: bool = False,
        has_update: bool = False,
        **kwargs,
    ) -> Dict:
        """Get paginated and filtered model data"""
        sort_params = self.cache_repository.parse_sort(sort_by)
        sorted_data = await self.cache_repository.fetch_sorted(sort_params)

        if hash_filters:
            filtered_data = await self._apply_hash_filters(sorted_data, hash_filters)
        else:
            filtered_data = await self._apply_common_filters(
                sorted_data,
                folder=folder,
                base_models=base_models,
                tags=tags,
                favorites_only=favorites_only,
                search_options=search_options,
            )

            if search:
                filtered_data = await self._apply_search_filters(
                    filtered_data,
                    search,
                    fuzzy_search,
                    search_options,
                )

            filtered_data = await self._apply_specific_filters(filtered_data, **kwargs)

            if has_update:
                filtered_data = await self._apply_update_filter(filtered_data)

        paginated = self._paginate(filtered_data, page, page_size)
        paginated['items'] = await self._annotate_update_flags(
            paginated['items'],
            assume_true=has_update,
        )
        return paginated

    
    async def _apply_hash_filters(self, data: List[Dict], hash_filters: Dict) -> List[Dict]:
        """Apply hash-based filtering"""
        single_hash = hash_filters.get('single_hash')
        multiple_hashes = hash_filters.get('multiple_hashes')
        
        if single_hash:
            # Filter by single hash
            single_hash = single_hash.lower()
            return [
                item for item in data
                if item.get('sha256', '').lower() == single_hash
            ]
        elif multiple_hashes:
            # Filter by multiple hashes
            hash_set = set(hash.lower() for hash in multiple_hashes)
            return [
                item for item in data
                if item.get('sha256', '').lower() in hash_set
            ]
        
        return data
    
    async def _apply_common_filters(
        self,
        data: List[Dict],
        folder: str = None,
        base_models: list = None,
        tags: list = None,
        favorites_only: bool = False,
        search_options: dict = None,
    ) -> List[Dict]:
        """Apply common filters that work across all model types"""
        normalized_options = self.search_strategy.normalize_options(search_options)
        criteria = FilterCriteria(
            folder=folder,
            base_models=base_models,
            tags=tags,
            favorites_only=favorites_only,
            search_options=normalized_options,
        )
        return self.filter_set.apply(data, criteria)
    
    async def _apply_search_filters(
        self,
        data: List[Dict],
        search: str,
        fuzzy_search: bool,
        search_options: dict,
    ) -> List[Dict]:
        """Apply search filtering"""
        normalized_options = self.search_strategy.normalize_options(search_options)
        return self.search_strategy.apply(data, search, normalized_options, fuzzy_search)
    
    async def _apply_specific_filters(self, data: List[Dict], **kwargs) -> List[Dict]:
        """Apply model-specific filters - to be overridden by subclasses if needed"""
        return data

    async def _apply_update_filter(self, data: List[Dict]) -> List[Dict]:
        """Filter models to those with remote updates available when requested."""
        if not data:
            return []
        if self.update_service is None:
            logger.warning(
                "Requested has_update filter for %s models but update service is unavailable",
                self.model_type,
            )
            return []

        id_to_items: Dict[int, List[Dict]] = {}
        ordered_ids: List[int] = []
        for item in data:
            model_id = self._extract_model_id(item)
            if model_id is None:
                continue
            if model_id not in id_to_items:
                id_to_items[model_id] = []
                ordered_ids.append(model_id)
            id_to_items[model_id].append(item)

        if not ordered_ids:
            return []

        resolved: Optional[Dict[int, bool]] = None
        bulk_method = getattr(self.update_service, "has_updates_bulk", None)
        if callable(bulk_method):
            try:
                resolved = await bulk_method(self.model_type, ordered_ids)
            except Exception as exc:
                logger.error(
                    "Failed to resolve update status in bulk for %s models (%s): %s",
                    self.model_type,
                    ordered_ids,
                    exc,
                    exc_info=True,
                )
                resolved = None

        if resolved is None:
            tasks = [
                self.update_service.has_update(self.model_type, model_id)
                for model_id in ordered_ids
            ]
            results = await asyncio.gather(*tasks, return_exceptions=True)
            resolved = {}
            for model_id, result in zip(ordered_ids, results):
                if isinstance(result, Exception):
                    logger.error(
                        "Failed to resolve update status for model %s (%s): %s",
                        model_id,
                        self.model_type,
                        result,
                    )
                    continue
                resolved[model_id] = bool(result)

        filtered: List[Dict] = []
        for item in data:
            model_id = self._extract_model_id(item)
            if model_id is not None and resolved.get(model_id, False):
                filtered.append(item)
        return filtered

    async def _annotate_update_flags(
        self,
        items: List[Dict],
        *,
        assume_true: bool = False,
    ) -> List[Dict]:
        """Attach an update_available flag to each response item.

        Items without a civitai model id default to False. When the caller already
        filtered for updates we can skip the lookup and mark everything True.
        """
        if not items:
            return []

        annotated = [dict(item) for item in items]

        if assume_true:
            for item in annotated:
                item['update_available'] = True
            return annotated

        if self.update_service is None:
            for item in annotated:
                item['update_available'] = False
            return annotated

        id_to_items: Dict[int, List[Dict]] = {}
        ordered_ids: List[int] = []
        for item in annotated:
            model_id = self._extract_model_id(item)
            if model_id is None:
                item['update_available'] = False
                continue
            if model_id not in id_to_items:
                id_to_items[model_id] = []
                ordered_ids.append(model_id)
            id_to_items[model_id].append(item)

        if not ordered_ids:
            return annotated

        resolved: Optional[Dict[int, bool]] = None
        bulk_method = getattr(self.update_service, "has_updates_bulk", None)
        if callable(bulk_method):
            try:
                resolved = await bulk_method(self.model_type, ordered_ids)
            except Exception as exc:
                logger.error(
                    "Failed to resolve update status in bulk for %s models (%s): %s",
                    self.model_type,
                    ordered_ids,
                    exc,
                    exc_info=True,
                )
                resolved = None

        if resolved is None:
            tasks = [
                self.update_service.has_update(self.model_type, model_id)
                for model_id in ordered_ids
            ]
            results = await asyncio.gather(*tasks, return_exceptions=True)
            resolved = {}
            for model_id, result in zip(ordered_ids, results):
                if isinstance(result, Exception):
                    logger.error(
                        "Failed to resolve update status for model %s (%s): %s",
                        model_id,
                        self.model_type,
                        result,
                    )
                    continue
                resolved[model_id] = bool(result)

        for model_id, items_for_id in id_to_items.items():
            flag = bool(resolved.get(model_id, False))
            for item in items_for_id:
                item['update_available'] = flag

        return annotated

    @staticmethod
    def _extract_model_id(item: Dict) -> Optional[int]:
        civitai = item.get('civitai') if isinstance(item, dict) else None
        if not isinstance(civitai, dict):
            return None
        try:
            value = civitai.get('modelId')
            if value is None:
                return None
            return int(value)
        except (TypeError, ValueError):
            return None
    
    def _paginate(self, data: List[Dict], page: int, page_size: int) -> Dict:
        """Apply pagination to filtered data"""
        total_items = len(data)
        start_idx = (page - 1) * page_size
        end_idx = min(start_idx + page_size, total_items)
        
        return {
            'items': data[start_idx:end_idx],
            'total': total_items,
            'page': page,
            'page_size': page_size,
            'total_pages': (total_items + page_size - 1) // page_size
        }
    
    @abstractmethod
    async def format_response(self, model_data: Dict) -> Dict:
        """Format model data for API response - must be implemented by subclasses"""
        pass
    
    # Common service methods that delegate to scanner
    async def get_top_tags(self, limit: int = 20) -> List[Dict]:
        """Get top tags sorted by frequency"""
        return await self.scanner.get_top_tags(limit)
    
    async def get_base_models(self, limit: int = 20) -> List[Dict]:
        """Get base models sorted by frequency"""
        return await self.scanner.get_base_models(limit)
    
    def has_hash(self, sha256: str) -> bool:
        """Check if a model with given hash exists"""
        return self.scanner.has_hash(sha256)
    
    def get_path_by_hash(self, sha256: str) -> Optional[str]:
        """Get file path for a model by its hash"""
        return self.scanner.get_path_by_hash(sha256)
    
    def get_hash_by_path(self, file_path: str) -> Optional[str]:
        """Get hash for a model by its file path"""
        return self.scanner.get_hash_by_path(file_path)
    
    async def scan_models(self, force_refresh: bool = False, rebuild_cache: bool = False):
        """Trigger model scanning"""
        return await self.scanner.get_cached_data(force_refresh=force_refresh, rebuild_cache=rebuild_cache)
    
    async def get_model_info_by_name(self, name: str):
        """Get model information by name"""
        return await self.scanner.get_model_info_by_name(name)
    
    def get_model_roots(self) -> List[str]:
        """Get model root directories"""
        return self.scanner.get_model_roots()
    
    def filter_civitai_data(self, data: Dict, minimal: bool = False) -> Dict:
        """Filter relevant fields from CivitAI data"""
        if not data:
            return {}

        fields = ["id", "modelId", "name", "trainedWords"] if minimal else [
            "id", "modelId", "name", "createdAt", "updatedAt",
            "publishedAt", "trainedWords", "baseModel", "description",
            "model", "images", "customImages", "creator"
        ]
        return {k: data[k] for k in fields if k in data}
    
    async def get_folder_tree(self, model_root: str) -> Dict:
        """Get hierarchical folder tree for a specific model root"""
        cache = await self.scanner.get_cached_data()
        
        # Build tree structure from folders
        tree = {}
        
        for folder in cache.folders:
            # Check if this folder belongs to the specified model root
            folder_belongs_to_root = False
            for root in self.scanner.get_model_roots():
                if root == model_root:
                    folder_belongs_to_root = True
                    break
            
            if not folder_belongs_to_root:
                continue
            
            # Split folder path into components
            parts = folder.split('/') if folder else []
            current_level = tree
            
            for part in parts:
                if part not in current_level:
                    current_level[part] = {}
                current_level = current_level[part]
        
        return tree
    
    async def get_unified_folder_tree(self) -> Dict:
        """Get unified folder tree across all model roots"""
        cache = await self.scanner.get_cached_data()
        
        # Build unified tree structure by analyzing all relative paths
        unified_tree = {}
        
        # Get all model roots for path normalization
        model_roots = self.scanner.get_model_roots()
        
        for folder in cache.folders:
            if not folder:  # Skip empty folders
                continue
            
            # Find which root this folder belongs to by checking the actual file paths
            # This is a simplified approach - we'll use the folder as-is since it should already be relative
            relative_path = folder
            
            # Split folder path into components
            parts = relative_path.split('/')
            current_level = unified_tree
            
            for part in parts:
                if part not in current_level:
                    current_level[part] = {}
                current_level = current_level[part]
        
        return unified_tree

    async def get_model_notes(self, model_name: str) -> Optional[str]:
        """Get notes for a specific model file"""
        cache = await self.scanner.get_cached_data()
        
        for model in cache.raw_data:
            if model['file_name'] == model_name:
                return model.get('notes', '')
        
        return None
    
    async def get_model_preview_url(self, model_name: str) -> Optional[str]:
        """Get the static preview URL for a model file"""
        cache = await self.scanner.get_cached_data()
        
        for model in cache.raw_data:
            if model['file_name'] == model_name:
                preview_url = model.get('preview_url')
                if preview_url:
                    from ..config import config
                    return config.get_preview_static_url(preview_url)
        
        return '/loras_static/images/no-preview.png'
    
    async def get_model_civitai_url(self, model_name: str) -> Dict[str, Optional[str]]:
        """Get the Civitai URL for a model file"""
        cache = await self.scanner.get_cached_data()
        
        for model in cache.raw_data:
            if model['file_name'] == model_name:
                civitai_data = model.get('civitai', {})
                model_id = civitai_data.get('modelId')
                version_id = civitai_data.get('id')
                
                if model_id:
                    civitai_url = f"https://civitai.com/models/{model_id}"
                    if version_id:
                        civitai_url += f"?modelVersionId={version_id}"
                    
                    return {
                        'civitai_url': civitai_url,
                        'model_id': str(model_id),
                        'version_id': str(version_id) if version_id else None
                    }
        
        return {'civitai_url': None, 'model_id': None, 'version_id': None}

    async def get_model_metadata(self, file_path: str) -> Optional[Dict]:
        """Load full metadata for a single model.

        Listing/search endpoints return lightweight cache entries; this method performs
        a lazy read of the on-disk metadata snapshot when callers need full detail.
        """
        metadata, should_skip = await MetadataManager.load_metadata(file_path, self.metadata_class)
        if should_skip or metadata is None:
            return None
        return self.filter_civitai_data(metadata.to_dict().get("civitai", {}))


    async def get_model_description(self, file_path: str) -> Optional[str]:
        """Return the stored modelDescription field for a model."""
        metadata, should_skip = await MetadataManager.load_metadata(file_path, self.metadata_class)
        if should_skip or metadata is None:
            return None
        return metadata.modelDescription or ''


    async def search_relative_paths(self, search_term: str, limit: int = 15) -> List[str]:
        """Search model relative file paths for autocomplete functionality"""
        cache = await self.scanner.get_cached_data()
        
        matching_paths = []
        search_lower = search_term.lower()
        
        # Get model roots for path calculation
        model_roots = self.scanner.get_model_roots()
        
        for model in cache.raw_data:
            file_path = model.get('file_path', '')
            if not file_path:
                continue
                
            # Calculate relative path from model root
            relative_path = None
            for root in model_roots:
                # Normalize paths for comparison
                normalized_root = os.path.normpath(root)
                normalized_file = os.path.normpath(file_path)
                
                if normalized_file.startswith(normalized_root):
                    # Remove root and leading separator to get relative path
                    relative_path = normalized_file[len(normalized_root):].lstrip(os.sep)
                    break
            
            if relative_path and search_lower in relative_path.lower():
                matching_paths.append(relative_path)
                
                if len(matching_paths) >= limit * 2:  # Get more for better sorting
                    break
        
        # Sort by relevance (exact matches first, then by length)
        matching_paths.sort(key=lambda x: (
            not x.lower().startswith(search_lower),  # Exact prefix matches first
            len(x),  # Then by length (shorter first)
            x.lower()  # Then alphabetically
        ))
        
        return matching_paths[:limit]