Azornes-Comfyui-Model-Resolver/core/services/search_dependencies.py

"""Typed dependencies shared by search orchestration and providers."""

from dataclasses import dataclass
from typing import Any, Protocol

from ..contracts import ModelCatalogEntry, ModelMatch, SearchResult
from ..routes.context import RouteContext


class LoggerProtocol(Protocol):
    def debug(self, message: str, **kwargs: Any) -> Any: ...

    def info(self, message: str, **kwargs: Any) -> Any: ...

    def warning(self, message: str, **kwargs: Any) -> Any: ...

    def exception(self, message: str, **kwargs: Any) -> Any: ...


class SearchTrackerProtocol(Protocol):
    def is_cancelled(self, progress_id: str) -> bool: ...

    def update(self, *args: Any, **kwargs: Any) -> Any: ...


class ModelResultBuilderProtocol(Protocol):
    def __call__(self, source: str, **fields: Any) -> SearchResult: ...


class ClearCacheProtocol(Protocol):
    def __call__(self) -> Any: ...


class ReloadProtocol(Protocol):
    def __call__(self) -> Any: ...


class ExtractSha256Protocol(Protocol):
    def __call__(self, metadata: dict[str, Any]) -> str | None: ...


class FormatSizeProtocol(Protocol):
    def __call__(self, value: Any, include_space: bool = True) -> str: ...


class SearchProviderProtocol(Protocol):
    def __call__(
        self,
        filename: str,
        **options: Any,
    ) -> SearchResult | list[SearchResult] | None: ...


class PopularModelUrlProtocol(Protocol):
    def __call__(self, filename: str) -> ModelCatalogEntry | None: ...


class ModelListSearchProtocol(Protocol):
    def __call__(self, filename: str) -> SearchResult | None: ...


class LocalHashSearchProtocol(Protocol):
    def __call__(
        self,
        sha256: str,
        *,
        category: str | None,
        max_matches: int,
        force_rescan: bool,
    ) -> list[ModelMatch]: ...


class ResolveUrnProtocol(Protocol):
    def __call__(
        self,
        model_id: int | str,
        version_id: int | str,
    ) -> dict[str, Any] | None: ...


class ResolveArchiveVersionProtocol(Protocol):
    def __call__(
        self,
        model_id: int | str,
        version_id: int | str,
        *,
        query: str,
    ) -> SearchResult | None: ...


class GetDownloadUrlProtocol(Protocol):
    def __call__(self, version_id: int | str) -> str | None: ...


@dataclass(frozen=True, slots=True)
class SearchDependencies:
    """Validated services and provider functions used by search code."""

    search_tracker: SearchTrackerProtocol
    search_result_timestamps: dict[str, str]
    logger: LoggerProtocol
    civarchive_search_error: type[BaseException]
    asyncio: Any
    build_model_result: ModelResultBuilderProtocol
    clear_civarchive_search_cache: ClearCacheProtocol
    clear_civitai_search_cache: ClearCacheProtocol
    clear_huggingface_search_cache: ClearCacheProtocol
    clear_lora_manager_archive_search_cache: ClearCacheProtocol
    extract_sha256_from_metadata: ExtractSha256Protocol
    format_size_bytes: FormatSizeProtocol
    get_civitai_download_url: GetDownloadUrlProtocol
    get_popular_model_url: PopularModelUrlProtocol
    reload_model_list: ReloadProtocol
    reload_popular_databases: ReloadProtocol
    resolve_civarchive_model_version: ResolveArchiveVersionProtocol
    resolve_urn: ResolveUrnProtocol
    search_civarchive_for_file: SearchProviderProtocol
    search_civitai: SearchProviderProtocol
    search_civitai_for_file: SearchProviderProtocol
    search_huggingface_for_file: SearchProviderProtocol
    search_local_matches_by_hash: LocalHashSearchProtocol
    search_lora_manager_archive_for_file: SearchProviderProtocol
    search_model_list: ModelListSearchProtocol
    web: Any

    @staticmethod
    def _extension_value(extension: Any, name: str) -> Any:
        try:
            return getattr(extension, name)
        except AttributeError as exc:
            raise KeyError(f"Missing route dependency: self.{name}") from exc

    @classmethod
    def from_context(cls, context: RouteContext) -> "SearchDependencies":
        """Build dependencies and fail fast on an incomplete route context."""
        extension = context.require("self")
        return cls(
            search_tracker=cls._extension_value(extension, "search_tracker"),
            search_result_timestamps=cls._extension_value(
                extension,
                "search_result_timestamps",
            ),
            logger=cls._extension_value(extension, "logger"),
            asyncio=context.require("asyncio"),
            civarchive_search_error=context.require("CivArchiveSearchError"),
            build_model_result=context.require("build_model_result"),
            clear_civarchive_search_cache=context.require(
                "clear_civarchive_search_cache"
            ),
            clear_civitai_search_cache=context.require(
                "clear_civitai_search_cache"
            ),
            clear_huggingface_search_cache=context.require(
                "clear_huggingface_search_cache"
            ),
            clear_lora_manager_archive_search_cache=context.require(
                "clear_lora_manager_archive_search_cache"
            ),
            extract_sha256_from_metadata=context.require(
                "extract_sha256_from_metadata"
            ),
            format_size_bytes=context.require("format_size_bytes"),
            get_civitai_download_url=context.require("get_civitai_download_url"),
            get_popular_model_url=context.require("get_popular_model_url"),
            reload_model_list=context.require("reload_model_list"),
            reload_popular_databases=context.require("reload_popular_databases"),
            resolve_civarchive_model_version=context.require(
                "resolve_civarchive_model_version"
            ),
            resolve_urn=context.require("resolve_urn"),
            search_civarchive_for_file=context.require(
                "search_civarchive_for_file"
            ),
            search_civitai=context.require("search_civitai"),
            search_civitai_for_file=context.require("search_civitai_for_file"),
            search_huggingface_for_file=context.require(
                "search_huggingface_for_file"
            ),
            search_local_matches_by_hash=context.require(
                "search_local_matches_by_hash"
            ),
            search_lora_manager_archive_for_file=context.require(
                "search_lora_manager_archive_for_file"
            ),
            search_model_list=context.require("search_model_list"),
            web=context.require("web"),
        )