From ddc5333bd2f137cc262c14db55a8e9d5008eb9dd Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Fri, 21 Aug 2026 23:10:26 +0200 Subject: [PATCH] refactor: unify model data contracts --- core/contracts.py | 2309 +++++++++++++++++++++++ core/custom_nodes/base.py | 108 +- core/custom_nodes/lora_manager.py | 94 +- core/custom_nodes/registry.py | 80 +- core/download/directories.py | 5 +- core/local_hash_matches.py | 19 +- core/log_system/logger.py | 11 +- core/matcher.py | 111 +- core/metadata_audit.py | 21 +- core/metadata_builder.py | 63 +- core/metadata_model_utils.py | 40 +- core/path_templates.py | 20 +- core/path_utils.py | 35 +- core/request_utils.py | 198 ++ core/resolver.py | 422 ++--- core/routes/base_models.py | 10 +- core/routes/helpers.py | 8 +- core/scanner.py | 98 +- core/services/civitai_search_service.py | 140 +- core/services/custom_url_service.py | 127 +- core/services/directory_service.py | 19 +- core/services/download_service.py | 212 ++- core/services/hash_service.py | 178 +- core/services/loaded_models_service.py | 109 +- core/services/metadata_service.py | 64 +- core/services/model_details_service.py | 91 +- core/services/model_utils.py | 28 +- core/services/scanner_service.py | 33 +- core/services/search_dependencies.py | 41 +- core/services/search_orchestrator.py | 214 +-- core/services/search_providers.py | 325 +++- core/services/search_support_service.py | 39 +- core/services/workflow_service.py | 218 ++- core/sources/civarchive.py | 170 +- core/sources/civitai.py | 101 +- core/sources/huggingface.py | 108 +- core/sources/lora_manager_archive.py | 59 +- core/sources/model_list.py | 73 +- core/sources/popular.py | 65 +- core/type_utils.py | 134 +- core/workflow/analysis.py | 158 +- core/workflow/inventory.py | 47 +- core/workflow/references.py | 55 +- core/workflow/subgraphs.py | 161 +- core/workflow_updater.py | 143 +- tests/test_contracts.py | 1450 ++++++++++++++ tests/test_download_routes.py | 142 +- tests/test_hash_routes.py | 50 +- tests/test_huggingface.py | 96 +- tests/test_loaded_models_routes.py | 16 +- tests/test_local_hash_matches.py | 25 +- tests/test_log_system.py | 26 + tests/test_lora_manager_archive.py | 54 +- tests/test_matcher.py | 74 +- tests/test_metadata_builder.py | 4 + tests/test_metadata_model_helpers.py | 22 +- tests/test_path_security.py | 30 + tests/test_path_templates.py | 196 +- tests/test_refactoring_targets.py | 248 ++- tests/test_refactoring_unification.py | 53 +- tests/test_request_utils.py | 109 +- tests/test_resolver_robustness.py | 43 +- tests/test_scanner.py | 45 +- tests/test_scanner_routes.py | 58 +- tests/test_search_orchestration.py | 113 +- tests/test_search_service_components.py | 323 +++- tests/test_source_internals.py | 181 +- tests/test_title_matching_hashing.py | 27 +- tests/test_unified_helpers.py | 13 +- tests/test_workflow_analyzer.py | 180 +- tests/test_workflow_components.py | 91 +- tests/test_workflow_service.py | 180 +- 72 files changed, 8404 insertions(+), 2209 deletions(-) create mode 100644 core/contracts.py create mode 100644 tests/test_contracts.py create mode 100644 tests/test_log_system.py diff --git a/core/contracts.py b/core/contracts.py new file mode 100644 index 0000000..b68aee5 --- /dev/null +++ b/core/contracts.py @@ -0,0 +1,2309 @@ +"""Canonical internal data contracts used across Model Resolver features. + +The HTTP API intentionally continues to expose dictionaries. These immutable +dataclasses are used at feature boundaries so provider-specific payloads and +workflow requests are validated once, while ``to_dict``/``to_kwargs`` keep the +existing compatibility contracts intact. +""" + +import math +import re +from collections.abc import Mapping +from dataclasses import dataclass, field, replace +from types import MappingProxyType +from typing import Any, Dict, List, Optional, Sequence, Tuple + + +def _freeze_value(value: Any) -> Any: + """Recursively freeze JSON-like values stored inside a contract.""" + if isinstance(value, Mapping): + return MappingProxyType( + {key: _freeze_value(item) for key, item in value.items()} + ) + if isinstance(value, (list, tuple)): + return tuple(_freeze_value(item) for item in value) + if isinstance(value, (set, frozenset)): + return tuple(_freeze_value(item) for item in value) + return value + + +def _thaw_value(value: Any) -> Any: + """Convert frozen contract values back to JSON-compatible containers.""" + if isinstance(value, Mapping): + return {key: _thaw_value(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_thaw_value(item) for item in value] + if isinstance(value, (set, frozenset)): + return [_thaw_value(item) for item in value] + return value + + +def _as_contract_float(value: Any, name: str, default: float = 0.0) -> float: + """Normalize a canonical finite float without hiding malformed values.""" + if value is None: + return default + if isinstance(value, bool): + raise TypeError(f"{name} must be a finite number") + try: + converted = float(value) + except (TypeError, ValueError) as exc: + raise ValueError(f"{name} must be a finite number") from exc + if not math.isfinite(converted): + raise ValueError(f"{name} must be a finite number") + return converted + + +def _as_optional_bool(value: Any, name: str = "boolean") -> Optional[bool]: + """Normalize an optional boolean while rejecting unknown representations.""" + if value is None: + return None + if isinstance(value, bool): + return value + if isinstance(value, (int, float)) and value in (0, 1): + return bool(value) + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"true", "1", "yes", "on"}: + return True + if normalized in {"false", "0", "no", "off"}: + return False + raise ValueError(f"{name} must be a boolean") + + +def _as_contract_bool( + value: Any, + name: str, +) -> bool: + """Normalize a canonical boolean while rejecting unknown representations.""" + normalized = _as_optional_bool(value, name) + if normalized is None: + raise ValueError(f"{name} must be a boolean") + return normalized + + +def _as_contract_text(value: Any, name: str, default: str = "") -> str: + """Normalize a canonical text field without stringifying invalid values.""" + if value is None: + return default + if not isinstance(value, str): + raise TypeError(f"{name} must be a string") + return value + + +def _as_optional_contract_text(value: Any, name: str) -> Optional[str]: + """Normalize an optional canonical text field without string coercion.""" + if value is None: + return None + if not isinstance(value, str): + raise TypeError(f"{name} must be a string") + return value or None + + +def _as_nullable_contract_text(value: Any, name: str) -> Optional[str]: + """Validate an optional canonical text field while preserving empty strings.""" + if value is None: + return None + if not isinstance(value, str): + raise TypeError(f"{name} must be a string") + return value + + +def _as_contract_identifier(value: Any, name: str) -> int | str | None: + """Normalize an optional model/node identifier without accepting booleans.""" + if value is None: + return None + if isinstance(value, bool) or not isinstance(value, (int, str)): + raise TypeError(f"{name} must be an integer or string") + if isinstance(value, str): + return value.strip() or None + return value + + +def _as_optional_provider_int(value: Any, name: str) -> Optional[int]: + """Normalize a provider numeric identifier without accepting arbitrary values.""" + if value is None: + return None + if isinstance(value, bool): + raise TypeError(f"{name} must be an integer") + if isinstance(value, int): + return value + if isinstance(value, str): + text = value.strip() + if not text: + return None + try: + return int(text) + except ValueError as exc: + raise ValueError(f"{name} must be an integer") from exc + raise TypeError(f"{name} must be an integer") + + +def _as_optional_non_negative_int(value: Any, name: str) -> Optional[int]: + """Validate an optional non-negative integer field.""" + if value is None: + return None + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{name} must be a non-negative integer") + if value < 0: + raise ValueError(f"{name} must be a non-negative integer") + return value + + +_SIZE_UNIT_MULTIPLIERS = { + "b": 1, + "kb": 1024, + "kib": 1024, + "mb": 1024**2, + "mib": 1024**2, + "gb": 1024**3, + "gib": 1024**3, + "tb": 1024**4, + "tib": 1024**4, +} +_SIZE_LABEL_PATTERN = re.compile( + r"^([0-9]+(?:\.[0-9]+)?)\s*(b|kb|kib|mb|mib|gb|gib|tb|tib)?$", + re.IGNORECASE, +) + + +def _as_size_bytes(value: Any) -> Optional[int]: + """Normalize numeric and human-readable sizes to bytes.""" + if value is None or isinstance(value, bool): + return None + if isinstance(value, (int, float)): + if not math.isfinite(float(value)) or value < 0: + return None + return int(value) + if not isinstance(value, str): + return None + + match = _SIZE_LABEL_PATTERN.fullmatch(value.strip()) + if not match: + return None + amount = float(match.group(1)) + multiplier = _SIZE_UNIT_MULTIPLIERS.get((match.group(2) or "b").lower(), 1) + return int(amount * multiplier) + + +def _require_mapping(value: Any, name: str) -> Dict[str, Any]: + """Return a mutable mapping copy or reject malformed contract input.""" + if not isinstance(value, Mapping): + raise TypeError(f"{name} must be an object") + return dict(value) + + +def _as_contract_mapping(value: Any, name: str) -> Dict[str, Any]: + """Normalize an optional canonical mapping without discarding invalid values.""" + if value is None: + return {} + if not isinstance(value, Mapping): + raise TypeError(f"{name} must be an object") + return dict(value) + + +def _as_contract_sequence(value: Any, name: str) -> Tuple[Any, ...]: + """Normalize an optional canonical array without discarding invalid values.""" + if value is None: + return () + if not isinstance(value, (list, tuple)): + raise TypeError(f"{name} must be an array") + return tuple(value) + + +def _as_contract_size(value: Any, name: str) -> Optional[int]: + """Normalize a canonical finite byte size while rejecting malformed values.""" + if value is None: + return None + normalized = _as_size_bytes(value) + if normalized is None: + raise ValueError(f"{name} must be a non-negative size") + return normalized + + +def _require_string_mapping(value: Any, name: str) -> Dict[str, str]: + """Return a mapping with string keys and values or reject malformed input.""" + if value is None: + return {} + if not isinstance(value, Mapping): + raise TypeError(f"{name} must be an object") + if any( + not isinstance(key, str) or not isinstance(item, str) + for key, item in value.items() + ): + raise TypeError(f"{name} must contain only string keys and values") + return dict(value) + + +def _require_non_negative_int(value: Any, name: str) -> int: + """Return a non-negative integer or reject malformed contract input.""" + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{name} must be a non-negative integer") + if value < 0: + raise ValueError(f"{name} must be a non-negative integer") + return value + + +def _download_text_from_mapping( + raw: Mapping[str, Any], + field_name: str, + aliases: Sequence[str] = (), +) -> str: + """Read a download text field without bypassing invalid falsey values.""" + return _contract_text_from_mapping( + raw, + field_name, + aliases, + contract_name="DownloadSpec", + ) + + +def _contract_text_from_mapping( + raw: Mapping[str, Any], + field_name: str, + aliases: Sequence[str] = (), + *, + contract_name: str, +) -> str: + """Read a canonical text field without bypassing invalid falsey values.""" + for key in (field_name, *aliases): + if key not in raw: + continue + value = raw[key] + if value is None: + continue + if not isinstance(value, str): + raise TypeError(f"{contract_name} {key} must be a string") + if value.strip(): + return value + return "" + + +def _first_provider_value( + raw: Mapping[str, Any], + keys: Sequence[str], +) -> Any: + """Read the first meaningful provider value without truthiness fallback.""" + for key in keys: + if key not in raw: + continue + value = raw[key] + if value is None: + continue + if isinstance(value, str) and not value.strip(): + continue + return value + return None + + +def _first_provider_value_preserving_empty( + raw: Mapping[str, Any], + keys: Sequence[str], +) -> Any: + """Read the first present provider value, including explicit empty text.""" + for key in keys: + if key not in raw or raw[key] is None: + continue + return raw[key] + return None + + +def _provider_text( + raw: Mapping[str, Any], + keys: Sequence[str], + *, + field_name: str, + contract_name: str = "SearchResult", +) -> str: + """Validate a provider text alias at the external boundary.""" + value = _first_provider_value(raw, keys) + return _as_contract_text(value, f"{contract_name} {field_name}") + + +def _provider_optional_text( + raw: Mapping[str, Any], + keys: Sequence[str], + *, + field_name: str, + contract_name: str = "SearchResult", +) -> Optional[str]: + """Validate an optional provider text alias explicitly.""" + value = _first_provider_value(raw, keys) + return _as_nullable_contract_text(value, f"{contract_name} {field_name}") + + +def _is_contract_instance(value: Any, expected_type: Any) -> bool: + """Accept a contract loaded through a second package-module alias.""" + return isinstance(value, expected_type) or ( + value.__class__.__name__ == expected_type.__name__ + and callable(getattr(value, "to_dict", None)) + ) + + +_MODEL_REFERENCE_CANONICAL_KEYS = frozenset( + { + "node_id", + "node_type", + "widget_index", + "original_path", + "category", + "exists", + "subgraph_id", + "is_top_level", + } +) +_MODEL_REFERENCE_BOOLEAN_EXTRA_KEYS = frozenset( + { + "active", + "auto_download_capable", + "auto_download_candidate", + "connected", + "input_choice_matches_value", + "is_legacy_lora_manager", + "is_lora_v2", + "is_urn", + "locate_is_top_level", + "locate_via_promoted_widget", + } +) +_MODEL_REFERENCE_TEXT_EXTRA_KEYS = frozenset( + { + "custom_node_adapter", + "custom_node_original_identity", + "expected_filename", + "filename", + "full_path", + "hash_lookup_source", + "name", + "widget_name", + "workflow_sha256", + } +) +_RESOLVED_MODEL_CANONICAL_KEYS = frozenset( + { + "path", + "filename", + "relative_path", + "category", + "base_directory", + } +) +_RESOLVED_MODEL_INPUT_KEYS = _RESOLVED_MODEL_CANONICAL_KEYS | { + "resolved_path", + "name", +} +_MISSING_MODEL_RESERVED_KEYS = _MODEL_REFERENCE_CANONICAL_KEYS | { + "reference_count", + "all_node_refs", + "matches", +} +_SEARCH_RESULT_INPUT_KEYS = frozenset( + { + "source", + "model_id", + "modelId", + "model_name", + "modelName", + "version_id", + "versionId", + "modelVersionId", + "name", + "version_name", + "versionName", + "type", + "model_type", + "filename", + "file_name", + "fileName", + "url", + "page_url", + "download_url", + "downloadUrl", + "downloadURL", + "size", + "size_bytes", + "sizeKB", + "base_model", + "baseModel", + "tags", + "match_type", + "confidence", + "sha256", + "hash", + "hashes", + "trained_words", + "trainedWords", + "images", + "details_source", + "version_url", + "versionUrl", + "custom_url", + "result_mode", + "SHA256", + } +) +_SEARCH_RESULT_RESERVED_EXTRA_KEYS = _SEARCH_RESULT_INPUT_KEYS - {"hash"} + + +@dataclass(frozen=True, slots=True) +class ModelReference: + """Stable identity for a model reference extracted from a workflow.""" + + node_id: int | str | None = None + node_type: str = "" + widget_index: Optional[int] = None + original_path: str = "" + category: str = "" + exists: bool = False + subgraph_id: Optional[str] = None + is_top_level: Optional[bool] = None + extra: Mapping[str, Any] = field(default_factory=dict, repr=False, compare=False) + + def __post_init__(self) -> None: + object.__setattr__( + self, + "node_id", + _as_contract_identifier(self.node_id, "ModelReference node_id"), + ) + object.__setattr__( + self, + "widget_index", + _as_optional_non_negative_int( + self.widget_index, + "ModelReference widget_index", + ), + ) + object.__setattr__( + self, + "node_type", + _as_contract_text(self.node_type, "ModelReference node_type"), + ) + object.__setattr__( + self, + "original_path", + _as_contract_text(self.original_path, "ModelReference original_path"), + ) + object.__setattr__( + self, + "category", + _as_contract_text(self.category, "ModelReference category"), + ) + object.__setattr__( + self, + "subgraph_id", + _as_optional_contract_text( + self.subgraph_id, + "ModelReference subgraph_id", + ), + ) + object.__setattr__( + self, + "exists", + _as_contract_bool(self.exists, "ModelReference exists"), + ) + object.__setattr__( + self, + "is_top_level", + _as_optional_bool(self.is_top_level, "ModelReference is_top_level"), + ) + normalized_extra = {} + for key, item in _as_contract_mapping( + self.extra, + "ModelReference extra", + ).items(): + if key in _MODEL_REFERENCE_BOOLEAN_EXTRA_KEYS: + normalized_extra[key] = _as_contract_bool( + item, + f"ModelReference {key}", + ) + elif key in _MODEL_REFERENCE_TEXT_EXTRA_KEYS: + normalized_extra[key] = _as_nullable_contract_text( + item, + f"ModelReference {key}", + ) + else: + normalized_extra[key] = item + object.__setattr__( + self, + "extra", + _freeze_value( + { + key: item + for key, item in normalized_extra.items() + if key not in _MODEL_REFERENCE_CANONICAL_KEYS + } + ), + ) + + @classmethod + def from_mapping(cls, value: Mapping[str, Any]) -> "ModelReference": + """Build a reference from the existing workflow-reference dictionary.""" + raw = _require_mapping(value, "ModelReference") + reference = cls( + node_id=raw.get("node_id"), + node_type=raw.get("node_type"), + widget_index=raw.get("widget_index"), + original_path=raw.get("original_path"), + category=raw.get("category"), + exists=_as_contract_bool( + raw.get("exists", False), + "ModelReference exists", + ), + subgraph_id=raw.get("subgraph_id"), + is_top_level=_as_optional_bool( + raw.get("is_top_level"), + "ModelReference is_top_level", + ), + extra={ + key: item + for key, item in raw.items() + if key not in _MODEL_REFERENCE_CANONICAL_KEYS + }, + ) + if not ( + reference.node_id is not None + or reference.original_path + or reference.extra_value("full_path") + or reference.extra_value("filename") + or reference.extra_value("name") + or reference.extra_value("expected_filename") + ): + raise ValueError("ModelReference requires a node or model identity") + return reference + + def with_updates(self, **values: Any) -> "ModelReference": + """Return a typed reference with canonical or provider-specific fields updated.""" + canonical_updates = { + key: value + for key, value in values.items() + if key in _MODEL_REFERENCE_CANONICAL_KEYS + } + extra = { + key: value + for key, value in self.extra.items() + if key not in _MODEL_REFERENCE_CANONICAL_KEYS + } + extra.update( + { + key: value + for key, value in values.items() + if key not in _MODEL_REFERENCE_CANONICAL_KEYS + } + ) + return replace(self, extra=extra, **canonical_updates) + + def extra_value(self, key: str, default: Any = None) -> Any: + """Read a workflow-specific field without treating the contract as a loose dict.""" + return self.extra.get(key, default) + + def to_dict(self) -> Dict[str, Any]: + """Serialize the canonical fields while preserving provider/node extras.""" + result = _thaw_value(self.extra) + result.update( + { + "node_id": self.node_id, + "node_type": self.node_type, + "widget_index": self.widget_index, + "original_path": self.original_path, + "category": self.category, + "exists": self.exists, + } + ) + if self.subgraph_id is not None: + result["subgraph_id"] = self.subgraph_id + if self.is_top_level is not None: + result["is_top_level"] = self.is_top_level + return result + + +@dataclass(frozen=True, slots=True) +class ProviderUrlReference: + """Typed provider URL identity extracted from a supported model URL.""" + + model_id: Optional[int] = None + version_id: Optional[int] = None + sha256: Optional[str] = None + + def __post_init__(self) -> None: + object.__setattr__( + self, + "model_id", + _as_optional_provider_int( + self.model_id, + "ProviderUrlReference model_id", + ), + ) + object.__setattr__( + self, + "version_id", + _as_optional_provider_int( + self.version_id, + "ProviderUrlReference version_id", + ), + ) + sha256 = self.sha256 + if sha256 is not None: + if not isinstance(sha256, str): + raise TypeError("ProviderUrlReference sha256 must be a string") + sha256 = sha256.strip().lower() + if sha256.startswith("sha256:"): + sha256 = sha256[7:] + if not re.fullmatch(r"[0-9a-f]{64}", sha256): + raise ValueError("ProviderUrlReference sha256 must be a SHA256 hash") + object.__setattr__(self, "sha256", sha256 or None) + if self.model_id is None and self.version_id is None and self.sha256 is None: + raise ValueError("ProviderUrlReference requires a provider identity") + + @classmethod + def from_mapping(cls, value: Mapping[str, Any]) -> "ProviderUrlReference": + """Build a provider URL identity from parser output.""" + raw = _require_mapping(value, "ProviderUrlReference") + model_id = raw.get("model_id") + if model_id is None: + model_id = raw.get("modelId") + version_id = raw.get("version_id") + if version_id is None: + version_id = raw.get("versionId") + if version_id is None: + version_id = raw.get("modelVersionId") + sha256 = raw.get("sha256") + if sha256 is None: + sha256 = raw.get("SHA256") + return cls(model_id=model_id, version_id=version_id, sha256=sha256) + + def to_dict(self) -> Dict[str, Any]: + """Serialize to the legacy provider parser shape.""" + result: Dict[str, Any] = {} + if self.model_id is not None: + result["model_id"] = self.model_id + if self.version_id is not None: + result["version_id"] = self.version_id + if self.sha256 is not None: + result["sha256"] = self.sha256 + return result + + +@dataclass(frozen=True, slots=True) +class HuggingFaceFileReference: + """Typed HuggingFace repository file identity extracted from a URL.""" + + repo: str + filename: str + branch: str = "main" + + def __post_init__(self) -> None: + for field_name in ("repo", "filename", "branch"): + value = getattr(self, field_name) + if not isinstance(value, str): + raise TypeError(f"HuggingFaceFileReference {field_name} must be a string") + object.__setattr__(self, field_name, value.strip()) + if not self.repo: + raise ValueError("HuggingFaceFileReference repo is required") + if not self.filename: + raise ValueError("HuggingFaceFileReference filename is required") + if not self.branch: + object.__setattr__(self, "branch", "main") + + def to_dict(self) -> Dict[str, str]: + """Serialize to the existing parser result shape.""" + return { + "repo": self.repo, + "branch": self.branch, + "filename": self.filename, + } + +@dataclass(frozen=True, slots=True) +class ResolvedModel: + """Canonical local model record used by matching and workflow updates.""" + + path: str = "" + filename: str = "" + relative_path: str = "" + category: str = "" + base_directory: str = "" + extra: Mapping[str, Any] = field(default_factory=dict, repr=False, compare=False) + raw: Mapping[str, Any] = field(default_factory=dict, repr=False, compare=False) + + def __post_init__(self) -> None: + for field_name in ( + "path", + "filename", + "relative_path", + "category", + "base_directory", + ): + object.__setattr__( + self, + field_name, + _as_contract_text( + getattr(self, field_name), + f"ResolvedModel {field_name}", + ), + ) + object.__setattr__( + self, + "extra", + _freeze_value( + { + key: item + for key, item in _as_contract_mapping( + self.extra, + "ResolvedModel extra", + ).items() + if key not in _RESOLVED_MODEL_INPUT_KEYS + } + ), + ) + object.__setattr__( + self, + "raw", + _freeze_value(_as_contract_mapping(self.raw, "ResolvedModel raw")), + ) + + @classmethod + def from_mapping(cls, value: Mapping[str, Any]) -> "ResolvedModel": + """Normalize a scanner or frontend model mapping.""" + raw = _require_mapping(value, "ResolvedModel") + model = cls( + path=_contract_text_from_mapping( + raw, + "path", + ("resolved_path",), + contract_name="ResolvedModel", + ), + filename=_contract_text_from_mapping( + raw, + "filename", + ("name",), + contract_name="ResolvedModel", + ), + relative_path=_contract_text_from_mapping( + raw, + "relative_path", + contract_name="ResolvedModel", + ), + category=_contract_text_from_mapping( + raw, + "category", + contract_name="ResolvedModel", + ), + base_directory=_contract_text_from_mapping( + raw, + "base_directory", + contract_name="ResolvedModel", + ), + extra={ + key: item + for key, item in raw.items() + if key not in _RESOLVED_MODEL_INPUT_KEYS + }, + raw=raw, + ) + if not (model.path or model.filename or model.relative_path): + raise ValueError("ResolvedModel requires a path or filename") + return model + + def to_dict(self) -> Dict[str, Any]: + """Serialize while preserving provider/scanner-specific fields.""" + result = _thaw_value( + { + key: item + for key, item in self.raw.items() + if key not in _RESOLVED_MODEL_INPUT_KEYS + } + ) + result.update( + _thaw_value( + { + key: item + for key, item in self.extra.items() + if key not in _RESOLVED_MODEL_INPUT_KEYS + } + ) + ) + if self.path: + result["path"] = self.path + if self.filename: + result["filename"] = self.filename + if self.relative_path: + result["relative_path"] = self.relative_path + if self.category: + result["category"] = self.category + if self.base_directory: + result["base_directory"] = self.base_directory + return result + + def extra_value(self, key: str, default: Any = None) -> Any: + """Read an optional scanner/provider field explicitly.""" + return self.extra.get(key, default) + + def with_extra(self, **values: Any) -> "ResolvedModel": + """Return a copy enriched with local scan or download metadata.""" + extra = { + key: value + for key, value in self.extra.items() + if key not in _RESOLVED_MODEL_INPUT_KEYS + } + extra.update( + { + key: value + for key, value in values.items() + if key not in _RESOLVED_MODEL_INPUT_KEYS + } + ) + return replace(self, extra=extra) + + def with_updates(self, **values: Any) -> "ResolvedModel": + """Return a copy with canonical model fields replaced immutably.""" + canonical_updates = { + key: value + for key, value in values.items() + if key in _RESOLVED_MODEL_CANONICAL_KEYS + } + extra = { + key: value + for key, value in self.extra.items() + if key not in _RESOLVED_MODEL_INPUT_KEYS + } + extra.update( + { + key: value + for key, value in values.items() + if key not in _RESOLVED_MODEL_INPUT_KEYS + } + ) + return replace(self, extra=extra, **canonical_updates) + + +@dataclass(frozen=True, slots=True) +class ModelMatch: + """Canonical local match returned for a workflow model reference.""" + + model: ResolvedModel + filename: str = "" + similarity: float = 0.0 + confidence: float = 0.0 + extra: Mapping[str, Any] = field(default_factory=dict, repr=False, compare=False) + raw: Mapping[str, Any] = field(default_factory=dict, repr=False, compare=False) + + def __post_init__(self) -> None: + model = self.model + if not _is_contract_instance(model, ResolvedModel) and isinstance( + model, + Mapping, + ): + model = ResolvedModel.from_mapping(model) + if not _is_contract_instance(model, ResolvedModel): + raise TypeError("ModelMatch model must be a model object") + object.__setattr__(self, "model", model) + filename = self.filename + if filename is None or filename == "": + filename = model.filename + object.__setattr__( + self, + "filename", + _as_contract_text(filename, "ModelMatch filename"), + ) + object.__setattr__( + self, + "similarity", + _as_contract_float(self.similarity, "ModelMatch similarity"), + ) + object.__setattr__( + self, + "confidence", + _as_contract_float(self.confidence, "ModelMatch confidence"), + ) + object.__setattr__( + self, + "extra", + _freeze_value(_as_contract_mapping(self.extra, "ModelMatch extra")), + ) + object.__setattr__( + self, + "raw", + _freeze_value(_as_contract_mapping(self.raw, "ModelMatch raw")), + ) + + @classmethod + def from_mapping(cls, value: Mapping[str, Any]) -> "ModelMatch": + """Normalize a fuzzy or hash match mapping.""" + raw = _require_mapping(value, "ModelMatch") + model_value = raw.get("model") + if _is_contract_instance(model_value, ResolvedModel): + model = model_value + elif isinstance(model_value, Mapping): + model = ResolvedModel.from_mapping(model_value) + else: + raise TypeError("ModelMatch model must be a model object") + known_keys = {"model", "filename", "similarity", "confidence"} + return cls( + model=model, + filename=raw.get("filename"), + similarity=raw.get("similarity"), + confidence=raw.get("confidence"), + extra={key: item for key, item in raw.items() if key not in known_keys}, + raw=raw, + ) + + def to_dict(self) -> Dict[str, Any]: + """Serialize using the existing local-match response shape.""" + result = _thaw_value(self.raw) + result.update(_thaw_value(self.extra)) + result["model"] = self.model.to_dict() + result["filename"] = self.filename + result["similarity"] = self.similarity + result["confidence"] = self.confidence + return result + + def with_extra(self, **values: Any) -> "ModelMatch": + """Return a copy enriched with local-match metadata.""" + extra = dict(self.extra) + extra.update(values) + return replace(self, extra=extra) + + def extra_value(self, key: str, default: Any = None) -> Any: + """Read an optional local-match field explicitly.""" + return self.extra.get(key, default) + + +@dataclass(frozen=True, slots=True) +class MissingModel: + """Grouped workflow model reference with typed local matches.""" + + reference: ModelReference + reference_count: Optional[int] = None + all_node_refs: Tuple[ModelReference, ...] = () + matches: Optional[Tuple[ModelMatch, ...]] = None + extra: Mapping[str, Any] = field(default_factory=dict, repr=False, compare=False) + + def __post_init__(self) -> None: + reference = self.reference + if not _is_contract_instance(reference, ModelReference) and isinstance( + reference, + Mapping, + ): + reference = ModelReference.from_mapping(reference) + if not _is_contract_instance(reference, ModelReference): + raise TypeError("MissingModel reference must be a model reference") + object.__setattr__(self, "reference", reference) + + raw_refs = self.all_node_refs + if not isinstance(raw_refs, (list, tuple)): + raise TypeError("MissingModel all_node_refs must be an array") + if not raw_refs: + raw_refs = (reference,) + normalized_refs = [] + for item in raw_refs: + if _is_contract_instance(item, ModelReference): + normalized_refs.append(item) + elif isinstance(item, Mapping): + normalized_refs.append(ModelReference.from_mapping(item)) + else: + raise TypeError("MissingModel all_node_refs must contain model references") + object.__setattr__(self, "all_node_refs", tuple(normalized_refs)) + + if self.reference_count is None: + reference_count = len(normalized_refs) + else: + reference_count = _require_non_negative_int( + self.reference_count, + "MissingModel reference_count", + ) + if reference_count < 1: + raise ValueError( + "MissingModel reference_count must be a positive integer" + ) + if reference_count != len(normalized_refs): + raise ValueError( + "MissingModel reference_count must match all_node_refs" + ) + object.__setattr__(self, "reference_count", reference_count) + + raw_matches = self.matches + if raw_matches is None: + normalized_matches = None + else: + if not isinstance(raw_matches, (list, tuple)): + raise TypeError("MissingModel matches must be an array") + normalized_matches = [] + for item in raw_matches: + if _is_contract_instance(item, ModelMatch): + normalized_matches.append(item) + elif isinstance(item, Mapping): + normalized_matches.append(ModelMatch.from_mapping(item)) + else: + raise TypeError("MissingModel matches must contain model matches") + object.__setattr__( + self, + "matches", + None if normalized_matches is None else tuple(normalized_matches), + ) + object.__setattr__( + self, + "extra", + _freeze_value( + { + key: item + for key, item in _as_contract_mapping( + self.extra, + "MissingModel extra", + ).items() + if key not in _MISSING_MODEL_RESERVED_KEYS + } + ), + ) + + @classmethod + def from_mapping(cls, value: Mapping[str, Any]) -> "MissingModel": + """Build a grouped missing-model value from raw workflow data.""" + raw = _require_mapping(value, "MissingModel") + raw_refs_value = raw.get("all_node_refs") + if raw_refs_value is None: + raw_refs = (raw,) + elif isinstance(raw_refs_value, (list, tuple)): + raw_refs = tuple(raw_refs_value) or (raw,) + else: + raise TypeError("MissingModel all_node_refs must be an array") + + normalized_refs = [] + for item in raw_refs: + if _is_contract_instance(item, ModelReference): + normalized_refs.append(item) + elif isinstance(item, Mapping): + normalized_refs.append(ModelReference.from_mapping(item)) + else: + raise TypeError( + "MissingModel all_node_refs must contain model references" + ) + + raw_matches = raw.get("matches") + matches = None + if raw_matches is not None: + if not isinstance(raw_matches, (list, tuple)): + raise TypeError("MissingModel matches must be an array") + normalized_matches = [] + for item in raw_matches: + if _is_contract_instance(item, ModelMatch): + normalized_matches.append(item) + continue + if not isinstance(item, Mapping): + raise TypeError( + "MissingModel matches must contain model matches" + ) + normalized_matches.append(ModelMatch.from_mapping(item)) + matches = tuple(normalized_matches) + + known_keys = _MISSING_MODEL_RESERVED_KEYS + raw_reference_count = raw.get("reference_count") + if raw_reference_count is None: + reference_count = len(normalized_refs) + else: + reference_count = _require_non_negative_int( + raw_reference_count, + "MissingModel reference_count", + ) + if reference_count < 1: + raise ValueError( + "MissingModel reference_count must be a positive integer" + ) + + return cls( + reference=ModelReference.from_mapping(raw), + reference_count=reference_count, + all_node_refs=tuple(normalized_refs), + matches=matches, + extra={key: item for key, item in raw.items() if key not in known_keys}, + ) + + def with_extra(self, **values: Any) -> "MissingModel": + """Return a copy with workflow-enrichment fields merged in.""" + extra = { + key: item + for key, item in self.extra.items() + if key not in _MISSING_MODEL_RESERVED_KEYS + } + extra.update( + { + key: value + for key, value in values.items() + if key not in _MISSING_MODEL_RESERVED_KEYS + } + ) + return replace(self, extra=extra) + + def with_matches(self, matches: List[ModelMatch]) -> "MissingModel": + """Return a copy with normalized local matches attached.""" + return replace(self, matches=tuple(matches)) + + def with_references(self, references: Sequence[ModelReference]) -> "MissingModel": + """Return a copy with all workflow occurrences attached.""" + normalized = tuple(references) + if not normalized: + raise ValueError("MissingModel requires at least one reference") + return replace( + self, + reference=normalized[0], + reference_count=len(normalized), + all_node_refs=normalized, + ) + + def extra_value(self, key: str, default: Any = None) -> Any: + """Read an optional workflow-enrichment value explicitly.""" + return self.extra.get(key, default) + + def to_dict(self) -> Dict[str, Any]: + """Serialize to the workflow-analysis HTTP response shape.""" + result = self.reference.to_dict() + result["reference_count"] = self.reference_count + result["all_node_refs"] = [item.to_dict() for item in self.all_node_refs] + if self.matches is not None: + result["matches"] = [item.to_dict() for item in self.matches] + canonical = dict(result) + result = _thaw_value( + { + key: item + for key, item in self.extra.items() + if key not in _MISSING_MODEL_RESERVED_KEYS + } + ) + result.update(canonical) + return result + +@dataclass(frozen=True, slots=True) +class WorkflowAnalysisResult: + """Typed internal result serialized at the workflow HTTP boundary.""" + + missing_models: Tuple[MissingModel, ...] = () + resolved_models: Tuple[ModelReference, ...] = () + total_resolved: int = 0 + total_missing: int = 0 + total_models_analyzed: int = 0 + + def __post_init__(self) -> None: + if not isinstance(self.missing_models, (list, tuple)): + raise TypeError( + "WorkflowAnalysisResult missing_models must be an array" + ) + if not isinstance(self.resolved_models, (list, tuple)): + raise TypeError( + "WorkflowAnalysisResult resolved_models must be an array" + ) + + normalized_missing = [] + for item in self.missing_models: + if _is_contract_instance(item, MissingModel): + normalized_missing.append(item) + elif isinstance(item, Mapping): + normalized_missing.append(MissingModel.from_mapping(item)) + else: + raise TypeError("WorkflowAnalysisResult missing_models must contain missing models") + + normalized_resolved = [] + for item in self.resolved_models: + if _is_contract_instance(item, ModelReference): + normalized_resolved.append(item) + elif isinstance(item, Mapping): + normalized_resolved.append(ModelReference.from_mapping(item)) + else: + raise TypeError("WorkflowAnalysisResult resolved_models must contain model references") + + object.__setattr__(self, "missing_models", tuple(normalized_missing)) + object.__setattr__(self, "resolved_models", tuple(normalized_resolved)) + for field_name in ( + "total_resolved", + "total_missing", + "total_models_analyzed", + ): + object.__setattr__( + self, + field_name, + _require_non_negative_int( + getattr(self, field_name), + f"WorkflowAnalysisResult {field_name}", + ), + ) + + if self.total_missing != len(self.missing_models): + raise ValueError( + "WorkflowAnalysisResult total_missing must equal the number " + "of missing_models" + ) + if self.total_resolved != len(self.resolved_models): + raise ValueError( + "WorkflowAnalysisResult total_resolved must equal the number " + "of resolved_models" + ) + if self.total_models_analyzed < self.total_missing + self.total_resolved: + raise ValueError( + "WorkflowAnalysisResult total_models_analyzed must be at least " + "total_missing plus total_resolved" + ) + + def to_dict(self) -> Dict[str, Any]: + """Serialize the complete analysis response for the frontend.""" + return { + "missing_models": [item.to_dict() for item in self.missing_models], + "resolved_models": [ + {**item.to_dict(), "matches": []} + for item in self.resolved_models + ], + "total_resolved": self.total_resolved, + "total_missing": self.total_missing, + "total_models_analyzed": self.total_models_analyzed, + } + + +@dataclass(frozen=True, slots=True) +class WorkflowModelInventory: + """Typed shared workflow inventory used by resolver and HTTP services.""" + + available_models: Tuple[ResolvedModel, ...] = () + model_refs: Tuple[ModelReference, ...] = () + + def __post_init__(self) -> None: + if not isinstance(self.available_models, (list, tuple)): + raise TypeError("WorkflowModelInventory available_models must be an array") + if not isinstance(self.model_refs, (list, tuple)): + raise TypeError("WorkflowModelInventory model_refs must be an array") + object.__setattr__( + self, + "available_models", + tuple( + item + if _is_contract_instance(item, ResolvedModel) + else ResolvedModel.from_mapping(item) + for item in self.available_models + ), + ) + object.__setattr__( + self, + "model_refs", + tuple( + item + if _is_contract_instance(item, ModelReference) + else ModelReference.from_mapping(item) + for item in self.model_refs + ), + ) + + +@dataclass(frozen=True, slots=True) +class SearchResult: + """Canonical model-search result shared by all supported providers.""" + + source: str + model_id: int | str | None = None + version_id: int | str | None = None + name: str = "" + version_name: str = "" + model_type: str = "" + filename: str = "" + url: str = "" + download_url: Optional[str] = None + size: Optional[int] = None + base_model: Optional[str] = None + tags: Tuple[Any, ...] = () + match_type: str = "similar" + confidence: float = 0.0 + sha256: Optional[str] = None + hashes: Mapping[str, Any] = field(default_factory=dict) + trained_words: Tuple[Any, ...] = () + images: Tuple[Any, ...] = () + details_source: Optional[str] = None + version_url: Optional[str] = None + custom_url: bool = False + result_mode: str = "search" + extra: Mapping[str, Any] = field(default_factory=dict, repr=False, compare=False) + + def __post_init__(self) -> None: + """Defensively copy mutable payload fields at the boundary.""" + object.__setattr__( + self, + "model_id", + _as_contract_identifier(self.model_id, "SearchResult model_id"), + ) + object.__setattr__( + self, + "version_id", + _as_contract_identifier(self.version_id, "SearchResult version_id"), + ) + source = _as_contract_text(self.source, "SearchResult source").strip() + if not source: + raise ValueError("SearchResult source is required") + object.__setattr__(self, "source", source) + for field_name in ( + "name", + "version_name", + "model_type", + "filename", + "url", + "match_type", + "result_mode", + ): + object.__setattr__( + self, + field_name, + _as_contract_text( + getattr(self, field_name), + f"SearchResult {field_name}", + ), + ) + for field_name in ( + "download_url", + "base_model", + "sha256", + "details_source", + "version_url", + ): + object.__setattr__( + self, + field_name, + _as_nullable_contract_text( + getattr(self, field_name), + f"SearchResult {field_name}", + ), + ) + object.__setattr__( + self, + "tags", + _freeze_value(_as_contract_sequence(self.tags, "SearchResult tags")), + ) + object.__setattr__( + self, + "hashes", + _freeze_value(_as_contract_mapping(self.hashes, "SearchResult hashes")), + ) + object.__setattr__( + self, + "trained_words", + _freeze_value( + _as_contract_sequence( + self.trained_words, + "SearchResult trained_words", + ) + ), + ) + object.__setattr__( + self, + "images", + _freeze_value(_as_contract_sequence(self.images, "SearchResult images")), + ) + object.__setattr__( + self, + "custom_url", + _as_contract_bool(self.custom_url, "SearchResult custom_url"), + ) + object.__setattr__( + self, + "extra", + _freeze_value( + { + key: item + for key, item in _as_contract_mapping( + self.extra, + "SearchResult extra", + ).items() + if key not in _SEARCH_RESULT_RESERVED_EXTRA_KEYS + } + ), + ) + object.__setattr__( + self, + "confidence", + _as_contract_float(self.confidence, "SearchResult confidence"), + ) + object.__setattr__( + self, + "size", + _as_contract_size(self.size, "SearchResult size"), + ) + + if self.result_mode not in { + "search", + "custom_url", + "compact_custom_url", + }: + raise ValueError(f"Unsupported model result mode: {self.result_mode}") + + @classmethod + def from_mapping( + cls, + value: Mapping[str, Any], + *, + source: Optional[str] = None, + ) -> "SearchResult": + """Normalize common provider aliases into a canonical search result. + + This adapter is intentionally conservative: unknown provider fields are + retained in ``extra`` so adding metadata does not break the API. + """ + raw = _require_mapping(value, "SearchResult") + raw_hashes = raw.get("hashes") + nested_hashes = _as_contract_mapping( + raw_hashes, + "SearchResult hashes", + ) + known_keys = _SEARCH_RESULT_INPUT_KEYS + size = raw.get("size", raw.get("size_bytes")) + if size is None and raw.get("sizeKB") is not None: + try: + if isinstance(raw["sizeKB"], bool): + raise TypeError + size = int(float(raw["sizeKB"]) * 1024) + except (TypeError, ValueError): + raise ValueError( + "SearchResult size must be a non-negative size" + ) from None + + custom_url = _as_contract_bool( + raw.get("custom_url", False), + "SearchResult custom_url", + ) + result_mode = _provider_text( + raw, + ("result_mode",), + field_name="result_mode", + ) + if not result_mode and custom_url: + full_custom_url_keys = { + "version_name", + "versionName", + "type", + "model_type", + "size", + "size_bytes", + "sizeKB", + "tags", + "sha256", + "SHA256", + "hash", + "hashes", + "trained_words", + "trainedWords", + "images", + } + result_mode = ( + "custom_url" + if full_custom_url_keys.intersection(raw) + else "compact_custom_url" + ) + if not result_mode: + result_mode = "search" + + model_id = raw.get("model_id") + if model_id is None: + model_id = raw.get("modelId") + version_id = raw.get("version_id") + if version_id is None: + version_id = raw.get("versionId") + if version_id is None: + version_id = raw.get("modelVersionId") + sha256_value = _first_provider_value( + raw, + ( + "sha256", + "SHA256", + "hash", + ), + ) + if sha256_value is None: + sha256_value = _first_provider_value( + nested_hashes, + ("sha256", "SHA256"), + ) + if sha256_value is None: + sha256_value = _first_provider_value_preserving_empty( + raw, + ("sha256", "SHA256", "hash"), + ) + if sha256_value is None: + sha256_value = _first_provider_value_preserving_empty( + nested_hashes, + ("sha256", "SHA256"), + ) + + return cls( + source=( + _provider_text( + raw, + ("source",), + field_name="source", + ) + if source is None + else source + ), + model_id=model_id, + version_id=version_id, + name=_provider_text( + raw, + ("name", "model_name", "modelName"), + field_name="name", + ), + version_name=_provider_text( + raw, + ("version_name", "versionName"), + field_name="version_name", + ), + model_type=_provider_text( + raw, + ("type", "model_type"), + field_name="model_type", + ), + filename=_provider_text( + raw, + ("filename", "file_name", "fileName", "name"), + field_name="filename", + ), + url=_provider_text( + raw, + ("url", "page_url"), + field_name="url", + ), + download_url=_provider_optional_text( + raw, + ("download_url", "downloadUrl", "downloadURL"), + field_name="download_url", + ), + size=_as_contract_size(size, "SearchResult size"), + base_model=_provider_optional_text( + raw, + ("base_model", "baseModel"), + field_name="base_model", + ), + tags=_as_contract_sequence(raw.get("tags"), "SearchResult tags"), + match_type=_provider_text( + raw, + ("match_type",), + field_name="match_type", + ) or "similar", + confidence=raw.get("confidence", 0.0), + sha256=_as_nullable_contract_text( + sha256_value, + "SearchResult sha256", + ), + hashes=nested_hashes, + trained_words=_as_contract_sequence( + _first_provider_value(raw, ("trained_words", "trainedWords")), + "SearchResult trained_words", + ), + images=_as_contract_sequence(raw.get("images"), "SearchResult images"), + details_source=_provider_optional_text( + raw, + ("details_source",), + field_name="details_source", + ), + version_url=_provider_optional_text( + raw, + ("version_url", "versionUrl"), + field_name="version_url", + ), + custom_url=custom_url, + result_mode=result_mode, + extra={key: item for key, item in raw.items() if key not in known_keys}, + ) + + def to_dict(self) -> Dict[str, Any]: + """Serialize using the existing Model Resolver result contract.""" + if self.result_mode == "compact_custom_url": + result: Dict[str, Any] = { + "source": self.source, + "details_source": self.details_source or self.source, + "name": self.name, + "filename": self.filename, + "url": self.url, + "version_url": self.version_url or self.url, + "download_url": self.download_url, + "match_type": self.match_type, + "custom_url": self.custom_url, + } + if self.model_id is not None: + result["model_id"] = self.model_id + if self.version_id is not None: + result["version_id"] = self.version_id + canonical = dict(result) + result = _thaw_value( + { + key: item + for key, item in self.extra.items() + if key not in _SEARCH_RESULT_RESERVED_EXTRA_KEYS + } + ) + result.update(canonical) + return result + + result = { + "source": self.source, + "model_id": self.model_id, + "version_id": self.version_id, + "name": self.name, + "version_name": self.version_name, + "type": self.model_type, + "filename": self.filename, + "url": self.url, + "download_url": self.download_url, + "size": self.size, + "base_model": self.base_model, + "tags": _thaw_value(self.tags or ()), + "match_type": self.match_type, + "confidence": self.confidence, + "sha256": self.sha256, + "hashes": _thaw_value(self.hashes or {}), + "trained_words": _thaw_value(self.trained_words or ()), + "images": _thaw_value(self.images or ()), + } + + if self.result_mode == "custom_url": + result.pop("confidence", None) + result.update( + { + "details_source": self.details_source or self.source, + "version_url": self.version_url or self.url, + "match_type": "custom_url", + "custom_url": self.custom_url, + } + ) + else: + if self.details_source is not None: + result["details_source"] = self.details_source + if self.version_url is not None: + result["version_url"] = self.version_url + if self.custom_url: + result["custom_url"] = True + + canonical = dict(result) + result = _thaw_value( + { + key: item + for key, item in self.extra.items() + if key not in _SEARCH_RESULT_RESERVED_EXTRA_KEYS + } + ) + result.update(canonical) + return result + + def with_extra(self, **values: Any) -> "SearchResult": + """Return a copy enriched with provider-specific fields.""" + extra = { + key: item + for key, item in self.extra.items() + if key not in _SEARCH_RESULT_RESERVED_EXTRA_KEYS + } + extra.update( + { + key: value + for key, value in values.items() + if key not in _SEARCH_RESULT_RESERVED_EXTRA_KEYS + } + ) + return replace(self, extra=extra) + + def with_updates(self, **values: Any) -> "SearchResult": + """Return a copy with canonical fields replaced immutably.""" + return replace(self, **values) + + def extra_value(self, key: str, default: Any = None) -> Any: + """Read a provider-specific field without falling back to a mapping API.""" + return self.extra.get(key, default) + + +@dataclass(frozen=True, slots=True) +class ModelCatalogEntry: + """Validated entry returned by the local popular-model catalog.""" + + filename: str = "" + name: str = "" + url: str = "" + download_url: Optional[str] = None + model_type: str = "" + size: Optional[int] = None + directory: str = "" + match_type: str = "exact" + confidence: float = 0.0 + extra: Mapping[str, Any] = field(default_factory=dict, repr=False, compare=False) + + def __post_init__(self) -> None: + for field_name in ( + "filename", + "name", + "url", + "model_type", + "directory", + "match_type", + ): + object.__setattr__( + self, + field_name, + _as_contract_text( + getattr(self, field_name), + f"ModelCatalogEntry {field_name}", + ).strip(), + ) + object.__setattr__( + self, + "download_url", + _as_nullable_contract_text( + self.download_url, + "ModelCatalogEntry download_url", + ), + ) + object.__setattr__( + self, + "size", + _as_contract_size(self.size, "ModelCatalogEntry size"), + ) + object.__setattr__( + self, + "confidence", + _as_contract_float( + self.confidence, + "ModelCatalogEntry confidence", + ), + ) + object.__setattr__( + self, + "extra", + _freeze_value( + { + key: item + for key, item in _as_contract_mapping( + self.extra, + "ModelCatalogEntry extra", + ).items() + if key + not in { + "filename", + "name", + "url", + "download_url", + "type", + "model_type", + "size", + "match_type", + "confidence", + "directory", + } + } + ), + ) + + @classmethod + def from_mapping( + cls, + value: Mapping[str, Any], + *, + filename: str = "", + ) -> "ModelCatalogEntry": + """Build a catalog entry from the source database mapping.""" + raw = _require_mapping(value, "ModelCatalogEntry") + raw_filename = raw.get("filename") + if raw_filename is None: + raw_filename = filename + name = _provider_text( + raw, + ("name",), + field_name="name", + contract_name="ModelCatalogEntry", + ) + if not name: + name = _as_contract_text(raw_filename, "ModelCatalogEntry filename") + size = raw.get("size") + if size is None and raw.get("sizeKB") is not None: + size = raw.get("sizeKB") + if not isinstance(size, bool): + try: + size = float(size) * 1024 + except (TypeError, ValueError): + raise ValueError( + "ModelCatalogEntry size must be a non-negative size" + ) from None + return cls( + filename=_as_contract_text( + raw_filename, + "ModelCatalogEntry filename", + ), + name=name, + url=_provider_optional_text( + raw, + ("url", "page_url"), + field_name="url", + contract_name="ModelCatalogEntry", + ) + or "", + download_url=_provider_optional_text( + raw, + ("download_url", "downloadUrl", "downloadURL"), + field_name="download_url", + contract_name="ModelCatalogEntry", + ), + model_type=_provider_text( + raw, + ("type", "model_type"), + field_name="model_type", + contract_name="ModelCatalogEntry", + ), + size=_as_contract_size(size, "ModelCatalogEntry size"), + directory=_provider_text( + raw, + ("directory",), + field_name="directory", + contract_name="ModelCatalogEntry", + ), + match_type=_provider_text( + raw, + ("match_type",), + field_name="match_type", + contract_name="ModelCatalogEntry", + ) + or "exact", + confidence=raw.get("confidence", 0.0), + extra={ + key: item + for key, item in raw.items() + if key + not in { + "filename", + "name", + "url", + "page_url", + "download_url", + "downloadUrl", + "downloadURL", + "type", + "model_type", + "size", + "sizeKB", + "match_type", + "confidence", + "directory", + } + }, + ) + + def extra_value(self, key: str, default: Any = None) -> Any: + """Read catalog-specific metadata explicitly.""" + return self.extra.get(key, default) + + def with_extra(self, **values: Any) -> "ModelCatalogEntry": + """Return a catalog entry enriched with source-specific metadata.""" + extra = dict(self.extra) + extra.update(values) + return replace(self, extra=extra) + + def to_dict(self) -> Dict[str, Any]: + """Serialize the catalog entry for compatibility with source payloads.""" + result = _thaw_value(self.extra) + if self.filename: + result["filename"] = self.filename + if self.name: + result["name"] = self.name + if self.url: + result["url"] = self.url + if self.download_url is not None: + result["download_url"] = self.download_url + if self.model_type: + result["type"] = self.model_type + if self.size is not None: + result["size"] = self.size + if self.directory: + result["directory"] = self.directory + if self.match_type: + result["match_type"] = self.match_type + if self.confidence: + result["confidence"] = self.confidence + return result + + +@dataclass(frozen=True, slots=True) +class DownloadSpec: + """Validated internal description of a model download request.""" + + url: str + filename: str + category: str + headers: Mapping[str, str] = field(default_factory=dict, repr=False, compare=False) + subfolder: str = "" + base_directory: str = "" + metadata: Mapping[str, Any] = field(default_factory=dict, repr=False, compare=False) + expected_sha256: Optional[str] = None + + def __post_init__(self) -> None: + for field_name in ( + "url", + "filename", + "category", + "subfolder", + "base_directory", + ): + value = getattr(self, field_name) + if not isinstance(value, str): + raise TypeError(f"DownloadSpec {field_name} must be a string") + object.__setattr__(self, field_name, value.strip()) + expected_sha256 = self.expected_sha256 + if expected_sha256 is not None: + if not isinstance(expected_sha256, str): + raise TypeError( + "DownloadSpec expected_sha256 must be a string" + ) + expected_sha256 = expected_sha256.strip() or None + object.__setattr__(self, "expected_sha256", expected_sha256) + object.__setattr__( + self, + "headers", + _freeze_value(_require_string_mapping(self.headers, "DownloadSpec headers")), + ) + object.__setattr__( + self, + "metadata", + _freeze_value(_require_mapping(self.metadata, "DownloadSpec metadata")), + ) + + if not self.url: + raise ValueError("Download URL is required") + if not self.filename: + raise ValueError("Download filename is required") + if not self.category: + raise ValueError("Download category is required") + + @classmethod + def from_mapping(cls, value: Mapping[str, Any]) -> "DownloadSpec": + """Build a download specification from an HTTP/provider payload.""" + raw = _require_mapping(value, "DownloadSpec") + metadata = raw.get("metadata") + if metadata is None: + metadata = raw.get("download_metadata") + if metadata is None: + metadata = {} + expected_sha256 = raw.get("expected_sha256") + if expected_sha256 is None or ( + isinstance(expected_sha256, str) and not expected_sha256.strip() + ): + expected_sha256 = raw.get("sha256") + return cls( + url=_download_text_from_mapping(raw, "url", ("download_url",)), + filename=_download_text_from_mapping(raw, "filename", ("name",)), + category=_download_text_from_mapping(raw, "category"), + headers=raw.get("headers"), + subfolder=_download_text_from_mapping(raw, "subfolder"), + base_directory=_download_text_from_mapping(raw, "base_directory"), + metadata=metadata, + expected_sha256=expected_sha256, + ) + + def to_kwargs(self) -> Dict[str, Any]: + """Return arguments accepted by the existing background downloader.""" + metadata = _thaw_value(self.metadata) + if self.expected_sha256 and not metadata.get("sha256"): + metadata["sha256"] = self.expected_sha256 + return { + "url": self.url, + "filename": self.filename, + "category": self.category, + "headers": _thaw_value(self.headers) if self.headers else None, + "subfolder": self.subfolder, + "base_directory": self.base_directory, + "metadata": metadata, + } + + +@dataclass(frozen=True, slots=True) +class CustomNodeMetadata: + """Typed metadata passed from workflow references to custom-node adapters.""" + + adapter_id: str = "" + original_identity: str = "" + is_legacy_lora_manager: bool = False + + def __post_init__(self) -> None: + object.__setattr__( + self, + "adapter_id", + _as_contract_text( + self.adapter_id, + "CustomNodeMetadata adapter_id", + ).strip(), + ) + object.__setattr__( + self, + "original_identity", + _as_contract_text( + self.original_identity, + "CustomNodeMetadata original_identity", + ).strip(), + ) + object.__setattr__( + self, + "is_legacy_lora_manager", + _as_contract_bool( + self.is_legacy_lora_manager, + "CustomNodeMetadata is_legacy_lora_manager", + ), + ) + + @classmethod + def from_mapping(cls, value: Any) -> "CustomNodeMetadata": + """Normalize serialized adapter metadata into the typed contract.""" + if value is None: + return cls() + if _is_contract_instance(value, cls): + return value + raw = _require_mapping(value, "CustomNodeMetadata") + return cls( + adapter_id=_contract_text_from_mapping( + raw, + "custom_node_adapter", + ("adapter_id",), + contract_name="CustomNodeMetadata", + ), + original_identity=_contract_text_from_mapping( + raw, + "custom_node_original_identity", + ( + "original_lora_name", + "original_identity", + "name", + "original_path", + ), + contract_name="CustomNodeMetadata", + ), + is_legacy_lora_manager=raw.get( + "is_legacy_lora_manager", + raw.get("is_lora_v2", False), + ), + ) + + def to_dict(self) -> Dict[str, Any]: + """Serialize adapter metadata for the legacy workflow-update boundary.""" + result: Dict[str, Any] = {} + if self.adapter_id: + result["custom_node_adapter"] = self.adapter_id + if self.original_identity: + result["custom_node_original_identity"] = self.original_identity + if self.is_legacy_lora_manager: + result["is_lora_v2"] = True + return result + + +@dataclass(frozen=True, slots=True) +class Resolution: + """Validated workflow replacement passed through the typed update flow.""" + + node_id: int | str | None = None + widget_index: Optional[int] = None + resolved_path: Optional[str] = None + category: Optional[str] = None + base_directory: Optional[str] = None + resolved_model: Optional[ResolvedModel] = None + subgraph_id: Optional[str] = None + is_top_level: Optional[bool] = None + nested_key: Optional[str] = None + promoted_widget_name: Optional[str] = None + reference: Optional[ModelReference] = None + custom_node_metadata: CustomNodeMetadata = field( + default_factory=CustomNodeMetadata, + repr=False, + compare=False, + ) + + def __post_init__(self) -> None: + object.__setattr__( + self, + "node_id", + _as_contract_identifier(self.node_id, "Resolution node_id"), + ) + object.__setattr__( + self, + "widget_index", + _as_optional_non_negative_int( + self.widget_index, + "Resolution widget_index", + ), + ) + for field_name in ( + "resolved_path", + "category", + "base_directory", + "subgraph_id", + "nested_key", + "promoted_widget_name", + ): + object.__setattr__( + self, + field_name, + _as_nullable_contract_text( + getattr(self, field_name), + f"Resolution {field_name}", + ), + ) + + resolved_model = self.resolved_model + if not _is_contract_instance(resolved_model, ResolvedModel) and isinstance( + resolved_model, + Mapping, + ): + resolved_model = ResolvedModel.from_mapping(resolved_model) + if resolved_model is not None and not _is_contract_instance( + resolved_model, + ResolvedModel, + ): + raise TypeError("Resolution resolved_model must be a model object") + object.__setattr__(self, "resolved_model", resolved_model) + + reference = self.reference + if not _is_contract_instance(reference, ModelReference) and isinstance( + reference, + Mapping, + ): + reference = ModelReference.from_mapping(reference) + if reference is not None and not _is_contract_instance( + reference, + ModelReference, + ): + raise TypeError("Resolution reference must be a model reference") + object.__setattr__(self, "reference", reference) + object.__setattr__( + self, + "is_top_level", + _as_optional_bool(self.is_top_level, "Resolution is_top_level"), + ) + object.__setattr__( + self, + "custom_node_metadata", + CustomNodeMetadata.from_mapping(self.custom_node_metadata), + ) + + @classmethod + def from_mapping(cls, value: Mapping[str, Any]) -> "Resolution": + """Deserialize a workflow resolution request from JSON.""" + raw = _require_mapping(value, "Resolution") + resolved_model = raw.get("resolved_model") + if resolved_model is not None and not isinstance( + resolved_model, + Mapping, + ) and not _is_contract_instance(resolved_model, ResolvedModel): + raise ValueError("resolved_model must be an object") + if isinstance(resolved_model, Mapping): + for path_key in ("path", "resolved_path"): + path_value = resolved_model.get(path_key) + if path_value is not None and not isinstance(path_value, str): + raise ValueError( + f"resolved_model.{path_key} must be a string" + ) + _as_contract_identifier(raw.get("node_id"), "Resolution node_id") + _as_optional_non_negative_int( + raw.get("widget_index"), + "Resolution widget_index", + ) + for field_name in ( + "resolved_path", + "category", + "base_directory", + "subgraph_id", + "nested_key", + "promoted_widget_name", + ): + field_value = raw.get(field_name) + if field_value is not None and not isinstance(field_value, str): + raise TypeError(f"Resolution {field_name} must be a string") + + return cls( + node_id=raw.get("node_id"), + widget_index=raw.get("widget_index"), + resolved_path=raw.get("resolved_path"), + category=raw.get("category"), + base_directory=raw.get("base_directory"), + resolved_model=( + resolved_model + if _is_contract_instance(resolved_model, ResolvedModel) + else ResolvedModel.from_mapping(resolved_model) + if resolved_model is not None + else None + ), + subgraph_id=raw.get("subgraph_id"), + is_top_level=_as_optional_bool( + raw.get("is_top_level"), + "Resolution is_top_level", + ), + nested_key=raw.get("nested_key"), + promoted_widget_name=raw.get("promoted_widget_name"), + reference=ModelReference.from_mapping(raw), + custom_node_metadata=CustomNodeMetadata.from_mapping( + raw.get("custom_node_metadata") + ), + ) + + def with_custom_node_metadata( + self, + metadata: CustomNodeMetadata | Mapping[str, Any], + ) -> "Resolution": + """Return a copy enriched with adapter metadata for workflow updates.""" + return replace( + self, + custom_node_metadata=CustomNodeMetadata.from_mapping(metadata), + ) + + def with_base_directory(self, base_directory: Optional[str]) -> "Resolution": + """Return a copy with a derived model base directory.""" + return replace(self, base_directory=base_directory) + + def validate(self) -> None: + """Validate fields required by the workflow updater.""" + if self.node_id is None: + raise ValueError("Resolution requires node_id") + if isinstance(self.node_id, bool) or not isinstance( + self.node_id, + (int, str), + ): + raise ValueError("Resolution node_id must be an integer or string") + if isinstance(self.node_id, str) and not self.node_id.strip(): + raise ValueError("Resolution node_id cannot be empty") + + if not isinstance(self.widget_index, int) or isinstance( + self.widget_index, + bool, + ): + raise ValueError("Resolution widget_index must be a non-negative integer") + if self.widget_index < 0: + raise ValueError("Resolution widget_index must be a non-negative integer") + + if self.resolved_path is not None and not isinstance( + self.resolved_path, + str, + ): + raise ValueError("Resolution resolved_path must be a string") + if self.resolved_model is not None and not _is_contract_instance( + self.resolved_model, + ResolvedModel, + ): + raise ValueError("Resolution resolved_model must be a model object") + + resolved_model_path = self.resolved_model.path if self.resolved_model else "" + if resolved_model_path and not isinstance(resolved_model_path, str): + raise ValueError("Resolution resolved_model.path must be a string") + if not (self.resolved_path or resolved_model_path): + raise ValueError("Resolution requires resolved_path or resolved_model.path") + + for field_name in ( + "category", + "base_directory", + "subgraph_id", + "nested_key", + "promoted_widget_name", + ): + value = getattr(self, field_name) + if value is not None and not isinstance(value, str): + raise ValueError(f"Resolution {field_name} must be a string") + + if self.is_top_level is not None and not isinstance(self.is_top_level, bool): + raise ValueError("Resolution is_top_level must be a boolean") + if self.reference is not None and not _is_contract_instance( + self.reference, + ModelReference, + ): + raise ValueError("Resolution reference must be a model reference") + if not _is_contract_instance( + self.custom_node_metadata, + CustomNodeMetadata, + ): + raise ValueError( + "Resolution custom_node_metadata must be typed metadata" + ) diff --git a/core/custom_nodes/base.py b/core/custom_nodes/base.py index c39144e..e6941fe 100644 --- a/core/custom_nodes/base.py +++ b/core/custom_nodes/base.py @@ -1,7 +1,30 @@ """Shared contract for backend custom-node model adapters.""" +from collections.abc import Mapping from dataclasses import dataclass, field -from typing import Any, Callable, Dict, Optional, Tuple +from types import MappingProxyType +from typing import Any, Callable, Dict, List, Optional, Protocol, Tuple + +from ..contracts import CustomNodeMetadata, ModelReference, ResolvedModel + + +class AnalyzeReferencesCallback(Protocol): + """Typed callback used to extract references from a serialized node.""" + + def __call__( + self, + node: Dict[str, Any], + available_models: Optional[List[ResolvedModel]] = None, + *, + is_active: bool, + get_widget_name_hint: Callable[[Dict[str, Any], int], str], + ) -> Optional[List[ModelReference]]: ... + + +class HasPotentialReferenceCallback(Protocol): + """Typed callback used for cheap custom-node reference detection.""" + + def __call__(self, node: Dict[str, Any]) -> bool: ... @dataclass(frozen=True) @@ -11,9 +34,80 @@ class CustomNodeModelAdapter: adapter_id: str node_types: Tuple[str, ...] category_hint: Optional[str] = None - widget_categories: Dict[int, str] = field(default_factory=dict) - analyze_references: Optional[Callable[..., Any]] = None - has_potential_reference: Optional[Callable[..., bool]] = None - update_model_path: Optional[Callable[..., Optional[bool]]] = None - should_skip_existing: Optional[Callable[..., bool]] = None - adapt_loaded_model: Optional[Callable[..., Any]] = None + widget_categories: Mapping[int, str] = field(default_factory=dict) + analyze_references: Optional[AnalyzeReferencesCallback] = None + has_potential_reference: Optional[HasPotentialReferenceCallback] = None + update_model_path: Optional[ + Callable[ + [Dict[str, Any], int, Optional[ResolvedModel], CustomNodeMetadata], + Optional[bool], + ] + ] = None + should_skip_existing: Optional[Callable[[ModelReference], bool]] = None + adapt_loaded_model: Optional[ + Callable[[ModelReference, str, Any], Tuple[str, Any]] + ] = None + + def __post_init__(self) -> None: + if not isinstance(self.adapter_id, str): + raise TypeError("CustomNodeModelAdapter adapter_id must be a string") + adapter_id = self.adapter_id.strip() + if not adapter_id: + raise ValueError("CustomNodeModelAdapter adapter_id is required") + object.__setattr__(self, "adapter_id", adapter_id) + + if not isinstance(self.node_types, (list, tuple)): + raise TypeError("CustomNodeModelAdapter node_types must be an array") + if any(not isinstance(node_type, str) for node_type in self.node_types): + raise TypeError( + "CustomNodeModelAdapter node_types must contain only strings" + ) + normalized_node_types = tuple( + node_type.strip() for node_type in self.node_types + ) + if any(not node_type for node_type in normalized_node_types): + raise ValueError( + "CustomNodeModelAdapter node_types cannot contain empty strings" + ) + object.__setattr__(self, "node_types", normalized_node_types) + + if self.category_hint is not None: + if not isinstance(self.category_hint, str): + raise TypeError( + "CustomNodeModelAdapter category_hint must be a string" + ) + object.__setattr__(self, "category_hint", self.category_hint.strip()) + + if not isinstance(self.widget_categories, Mapping): + raise TypeError( + "CustomNodeModelAdapter widget_categories must be an object" + ) + normalized_widget_categories = {} + for widget_index, category in self.widget_categories.items(): + if isinstance(widget_index, bool) or not isinstance(widget_index, int): + raise TypeError( + "CustomNodeModelAdapter widget category keys must be integers" + ) + if not isinstance(category, str): + raise TypeError( + "CustomNodeModelAdapter widget category values must be strings" + ) + normalized_widget_categories[widget_index] = category.strip() + object.__setattr__( + self, + "widget_categories", + MappingProxyType(normalized_widget_categories), + ) + + for field_name in ( + "analyze_references", + "has_potential_reference", + "update_model_path", + "should_skip_existing", + "adapt_loaded_model", + ): + callback = getattr(self, field_name) + if callback is not None and not callable(callback): + raise TypeError( + f"CustomNodeModelAdapter {field_name} must be callable" + ) diff --git a/core/custom_nodes/lora_manager.py b/core/custom_nodes/lora_manager.py index 5b69de7..09979ab 100644 --- a/core/custom_nodes/lora_manager.py +++ b/core/custom_nodes/lora_manager.py @@ -1,9 +1,11 @@ """Backend adapter for ComfyUI-Lora-Manager loader nodes.""" import os -from typing import Any, Dict, List, Optional +from typing import Any, Callable, Dict, List, Optional +from ..contracts import CustomNodeMetadata, ModelReference, ResolvedModel from ..log_system import create_module_logger +from ..type_utils import to_bool from .base import CustomNodeModelAdapter log = create_module_logger(__name__) @@ -20,11 +22,11 @@ TEXT_WIDGET_INDEX = 1 def analyze_references( node: Dict[str, Any], - available_models: Optional[List[Dict[str, Any]]] = None, + available_models: Optional[List[ResolvedModel]] = None, *, is_active: bool, - get_widget_name_hint: Any, -) -> Optional[List[Dict[str, Any]]]: + get_widget_name_hint: Callable[[Dict[str, Any], int], str], +) -> Optional[List[ModelReference]]: """Extract LoRA references stored in Lora Manager's list widget.""" widgets_values = node.get("widgets_values", []) if len(widgets_values) < 3: @@ -36,11 +38,11 @@ def analyze_references( available_models if available_models is not None else get_model_files() ) lora_files = [ - model for model in all_loras if model.get("category") == "loras" + model for model in all_loras if model.category == "loras" ] - lora_lookup: Dict[str, List[Dict[str, Any]]] = {} + lora_lookup: Dict[str, List[ResolvedModel]] = {} for lora_file in lora_files: - filename = lora_file.get("filename", "") + filename = lora_file.filename if not filename: continue base_name = os.path.splitext(filename)[0] @@ -53,7 +55,7 @@ def analyze_references( node_id = node.get("id") node_type = node.get("type", "") node_title = str(node.get("title", "") or "").strip() - model_refs: List[Dict[str, Any]] = [] + model_refs: List[ModelReference] = [] for lora_item in lora_list: if not isinstance(lora_item, dict): continue @@ -65,7 +67,7 @@ def analyze_references( lora_exists = False lora_full_path = None if name in lora_lookup: - lora_full_path = lora_lookup[name][0].get("path") + lora_full_path = lora_lookup[name][0].path lora_exists = ( os.path.exists(lora_full_path) if lora_full_path else False ) @@ -74,7 +76,7 @@ def analyze_references( test_name = name + extension if test_name not in lora_lookup: continue - lora_full_path = lora_lookup[test_name][0].get("path") + lora_full_path = lora_lookup[test_name][0].path lora_exists = ( os.path.exists(lora_full_path) if lora_full_path else False ) @@ -83,27 +85,29 @@ def analyze_references( log.debug(f"Lora {name}: exists={lora_exists}, path={lora_full_path}") model_refs.append( - { - "node_id": node_id, - "node_type": node_type, - "widget_index": LORA_LIST_WIDGET_INDEX, - "widget_name": get_widget_name_hint( - node, LORA_LIST_WIDGET_INDEX - ), - "original_path": name, - "name": name, - "strength": float(lora_item.get("strength", 1.0)), - "active": lora_item.get("active", True), - "node_title": node_title, - "category": "loras", - "category_hints": ["loras"], - "folder_key_hints": ["loras"], - "full_path": lora_full_path, - "exists": lora_exists, - "is_urn": False, - "custom_node_adapter": ADAPTER_ID, - "connected": is_active, - } + ModelReference.from_mapping( + { + "node_id": node_id, + "node_type": node_type, + "widget_index": LORA_LIST_WIDGET_INDEX, + "widget_name": get_widget_name_hint( + node, LORA_LIST_WIDGET_INDEX + ), + "original_path": name, + "name": name, + "strength": float(lora_item.get("strength", 1.0)), + "active": to_bool(lora_item.get("active", True), True), + "node_title": node_title, + "category": "loras", + "category_hints": ["loras"], + "folder_key_hints": ["loras"], + "full_path": lora_full_path, + "exists": lora_exists, + "is_urn": False, + "custom_node_adapter": ADAPTER_ID, + "connected": bool(is_active), + } + ) ) return model_refs @@ -125,20 +129,16 @@ def has_potential_reference(node: Dict[str, Any]) -> bool: def update_model_path( node: Dict[str, Any], widget_index: int, - resolved_model: Optional[Dict[str, Any]], - mapping: Optional[Dict[str, Any]], + resolved_model: Optional[ResolvedModel], + metadata: CustomNodeMetadata, ) -> Optional[bool]: """Update one LoRA name in both the list and formatted text widgets.""" - mapping = mapping or {} - adapter_id = mapping.get("custom_node_adapter") - is_legacy_mapping = mapping.get("is_lora_v2") is True + adapter_id = metadata.adapter_id + is_legacy_mapping = metadata.is_legacy_lora_manager if adapter_id != ADAPTER_ID and not is_legacy_mapping: return None - original_name = ( - mapping.get("custom_node_original_identity") - or mapping.get("original_lora_name") - ) + original_name = metadata.original_identity if not original_name or widget_index != LORA_LIST_WIDGET_INDEX: return None @@ -153,9 +153,7 @@ def update_model_path( new_name = None if resolved_model: - new_name = resolved_model.get("filename") or resolved_model.get( - "name", "" - ) + new_name = resolved_model.filename if new_name and "." in new_name: new_name = new_name.rsplit(".", 1)[0] if not new_name: @@ -204,20 +202,20 @@ def update_model_path( return True -def should_skip_existing(reference: Dict[str, Any]) -> bool: +def should_skip_existing(reference: ModelReference) -> bool: """Existing list entries do not require matching or relinking.""" - return reference.get("exists") is True + return reference.exists is True def adapt_loaded_model( - reference: Dict[str, Any], + reference: ModelReference, model_name: str, strength: Any, ) -> tuple[str, Any]: """Use the list entry's display name and strength.""" return ( - reference.get("name", model_name), - reference.get("strength", strength), + reference.extra_value("name", model_name), + reference.extra_value("strength", strength), ) diff --git a/core/custom_nodes/registry.py b/core/custom_nodes/registry.py index 3421142..3db7c0e 100644 --- a/core/custom_nodes/registry.py +++ b/core/custom_nodes/registry.py @@ -1,7 +1,12 @@ """Registry and dispatch helpers for backend custom-node model adapters.""" -from typing import Any, Dict, List, Optional +from typing import Any, Callable, Dict, List, Optional +from ..contracts import ( + CustomNodeMetadata, + ModelReference, + ResolvedModel, +) from .base import CustomNodeModelAdapter from .lora_manager import ADAPTER as LORA_MANAGER_ADAPTER from .rgthree_power_lora_loader import ( @@ -35,16 +40,18 @@ def get_custom_node_model_adapter( def get_custom_node_adapter_for_reference( - reference: Dict[str, Any], + reference: Optional[ModelReference], ) -> Optional[CustomNodeModelAdapter]: """Resolve an adapter from normalized or legacy reference metadata.""" - adapter_id = reference.get("custom_node_adapter") + if reference is None: + return None + adapter_id = reference.extra_value("custom_node_adapter") if adapter_id in _ADAPTERS_BY_ID: return _ADAPTERS_BY_ID[adapter_id] - adapter = get_custom_node_model_adapter(reference.get("node_type")) + adapter = get_custom_node_model_adapter(reference.node_type) if adapter: return adapter - if reference.get("is_lora_v2"): + if CustomNodeMetadata.from_mapping(reference.extra).is_legacy_lora_manager: return LORA_MANAGER_ADAPTER return None @@ -70,38 +77,55 @@ def get_custom_node_widget_categories() -> Dict[str, Dict[int, str]]: def get_custom_node_resolution_metadata( - reference: Dict[str, Any], -) -> Dict[str, Any]: - """Normalize adapter metadata carried into workflow update mappings.""" + reference: Optional[ModelReference], +) -> CustomNodeMetadata: + """Normalize adapter metadata carried into workflow update resolutions.""" adapter = get_custom_node_adapter_for_reference(reference) - if not adapter: - return {} - original_identity = ( - reference.get("custom_node_original_identity") - or reference.get("original_lora_name") - or reference.get("name") - or reference.get("original_path") - ) - return { + if not adapter or reference is None: + return CustomNodeMetadata() + metadata = { "custom_node_adapter": adapter.adapter_id, - "custom_node_original_identity": original_identity, + "custom_node_original_identity": ( + reference.extra_value("custom_node_original_identity") + or reference.extra_value("original_lora_name") + or reference.extra_value("name") + or reference.original_path + ), } + legacy_v2_flag = reference.extra_value("is_lora_v2") + if legacy_v2_flag is not None: + metadata["is_lora_v2"] = legacy_v2_flag + legacy_flag = reference.extra_value("is_legacy_lora_manager") + if legacy_flag is not None: + metadata["is_legacy_lora_manager"] = legacy_flag + return CustomNodeMetadata.from_mapping(metadata) def analyze_custom_node_references( node: Dict[str, Any], - available_models: Optional[List[Dict[str, Any]]] = None, - **context: Any, -) -> Optional[List[Dict[str, Any]]]: + available_models: Optional[List[ResolvedModel]] = None, + *, + is_active: bool, + get_widget_name_hint: Callable[[Dict[str, Any], int], str], +) -> Optional[List[ModelReference]]: """Dispatch custom serialized model extraction when an adapter handles it.""" adapter = get_custom_node_model_adapter(node) if not adapter or not adapter.analyze_references: return None - return adapter.analyze_references( + raw_references = adapter.analyze_references( node, available_models, - **context, + is_active=is_active, + get_widget_name_hint=get_widget_name_hint, ) + if raw_references is None: + return None + return [ + reference + if isinstance(reference, ModelReference) + else ModelReference.from_mapping(reference) + for reference in raw_references + ] def custom_node_has_potential_model_reference( @@ -119,8 +143,8 @@ def custom_node_has_potential_model_reference( def update_custom_node_model_path( node: Dict[str, Any], widget_index: int, - resolved_model: Optional[Dict[str, Any]], - mapping: Optional[Dict[str, Any]], + resolved_model: Optional[ResolvedModel], + metadata: CustomNodeMetadata, ) -> Optional[bool]: """Dispatch a custom workflow update, returning None when not handled.""" adapter = get_custom_node_model_adapter(node) @@ -130,12 +154,12 @@ def update_custom_node_model_path( node, widget_index, resolved_model, - mapping, + metadata, ) def should_skip_existing_custom_node_reference( - reference: Dict[str, Any], + reference: ModelReference, ) -> bool: """Return whether an existing adapter reference needs no resolution.""" adapter = get_custom_node_adapter_for_reference(reference) @@ -147,7 +171,7 @@ def should_skip_existing_custom_node_reference( def adapt_custom_node_loaded_model( - reference: Dict[str, Any], + reference: ModelReference, model_name: str, strength: Any, ) -> tuple[str, Any]: diff --git a/core/download/directories.py b/core/download/directories.py index 7f680a1..0e65e26 100644 --- a/core/download/directories.py +++ b/core/download/directories.py @@ -3,6 +3,7 @@ import os from typing import Any, List, Optional +from ..path_utils import normalize_folder_path_values from .dependencies import require_download_dependencies @@ -90,7 +91,9 @@ def get_download_directory( def _get_folder_paths(folder_name: str) -> List[str]: """Return paths for one key without rejecting the other aliases.""" try: - return list(folder_paths.get_folder_paths(folder_name) or []) + return normalize_folder_path_values( + folder_paths.get_folder_paths(folder_name) + ) except Exception as exc: facade.log.debug( f"Could not get folder path for {folder_name}: {exc}" diff --git a/core/local_hash_matches.py b/core/local_hash_matches.py index 4f84356..bd3b3bd 100644 --- a/core/local_hash_matches.py +++ b/core/local_hash_matches.py @@ -1,8 +1,10 @@ """Shared helpers for enriching local model matches found by SHA-256.""" -from typing import Any, Callable, Dict, List, Optional +from typing import Callable, List, Optional -SearchLocalMatchesByHash = Callable[..., List[Dict[str, Any]]] +from .contracts import ModelMatch + +SearchLocalMatchesByHash = Callable[..., List[ModelMatch]] def collect_local_hash_matches_for_result( @@ -14,7 +16,7 @@ def collect_local_hash_matches_for_result( force_rescan: bool = False, source: str = "download_source", filename: str = "", -) -> List[Dict[str, Any]]: +) -> List[ModelMatch]: """Find local matches and add the shared hash-lookup result fields.""" if not sha256: return [] @@ -26,11 +28,10 @@ def collect_local_hash_matches_for_result( force_rescan=force_rescan, ) return [ - { - **match, - "hash_lookup_source": source, - "hash_lookup_filename": filename, - "hash_lookup_sha256": sha256, - } + match.with_extra( + hash_lookup_source=source, + hash_lookup_filename=filename, + hash_lookup_sha256=sha256, + ) for match in matches ] diff --git a/core/log_system/logger.py b/core/log_system/logger.py index b6127b3..d6f5f29 100644 --- a/core/log_system/logger.py +++ b/core/log_system/logger.py @@ -186,7 +186,16 @@ class DirectConsoleHandler(logging.Handler): stream = self.stream or _get_console_output_stream() if stream is None: return - stream.write(self.format(record) + self.terminator) + message = self.format(record) + self.terminator + try: + stream.write(message) + except UnicodeEncodeError: + encoding = getattr(stream, "encoding", None) or "utf-8" + safe_message = message.encode( + encoding, + errors="replace", + ).decode(encoding, errors="replace") + stream.write(safe_message) stream.flush() except Exception: self.handleError(record) diff --git a/core/matcher.py b/core/matcher.py index 9d60728..d7965db 100644 --- a/core/matcher.py +++ b/core/matcher.py @@ -8,9 +8,11 @@ import heapq import os import re from difflib import SequenceMatcher -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Callable, Dict, List, Optional, Tuple +from .contracts import ModelMatch, ResolvedModel, SearchResult from .log_system import create_module_logger +from .metadata_model_utils import normalize_models log = create_module_logger(__name__) @@ -92,10 +94,10 @@ def calculate_similarity_with_normalization(str1: str, str2: str) -> float: def find_matches( target_model: str, - candidate_models: List[Dict[str, str]], + candidate_models: List[ResolvedModel], threshold: float = 0.0, max_results: int = 10, -) -> List[Dict[str, any]]: +) -> List[ModelMatch]: """ Find similar models using fuzzy matching. @@ -106,13 +108,7 @@ def find_matches( max_results: Maximum number of results to return Returns: - List of match dictionaries sorted by similarity (highest first): - { - 'model': original model dict from candidates, - 'filename': model filename, - 'similarity': similarity score (0.0 to 1.0), - 'confidence': confidence percentage (0 to 100) - } + Typed local matches sorted by similarity (highest first). """ # Keep only the best N matches instead of collecting and sorting everything. best_matches = [] @@ -129,10 +125,12 @@ def find_matches( # Normalize target filename once for exact match comparisons target_norm = normalize_filename(target_filename) - for candidate in candidate_models: + normalized_candidates = normalize_models(candidate_models) + + for candidate in normalized_candidates: # Get filename from candidate (prefer 'filename' key, fallback to extracting from 'path' or 'relative_path') - candidate_filename = candidate.get("filename") - candidate_path = candidate.get("path", "") or candidate.get("relative_path", "") + candidate_filename = candidate.filename + candidate_path = candidate.path or candidate.relative_path # If no filename key, try to extract from path or relative_path if not candidate_filename and candidate_path: @@ -143,20 +141,16 @@ def find_matches( # Normalize candidate path separators based on current OS # This ensures paths with \ vs / separators are treated as identical - candidate_path_normalized = candidate.get("_match_path_norm") - if candidate_path_normalized is None: - candidate_path_normalized = ( - os.path.normpath(candidate_path) if candidate_path else "" - ) - candidate["_match_path_norm"] = candidate_path_normalized + candidate_path_normalized = ( + os.path.normpath(candidate_path) if candidate_path else "" + ) - candidate_relative_path = candidate.get("relative_path", "") - candidate_relative_path_normalized = candidate.get("_match_relative_path_norm") - if candidate_relative_path_normalized is None: - candidate_relative_path_normalized = ( - os.path.normpath(candidate_relative_path) if candidate_relative_path else "" - ) - candidate["_match_relative_path_norm"] = candidate_relative_path_normalized + candidate_relative_path = candidate.relative_path + candidate_relative_path_normalized = ( + os.path.normpath(candidate_relative_path) + if candidate_relative_path + else "" + ) # Check if normalized paths are identical (100% match) # This handles cases where paths differ only by separator (e.g., path/to/model vs path\to\model) @@ -176,12 +170,14 @@ def find_matches( if path_match: # Exact path match after normalization = 100% confidence similarity = 1.0 - match = { - "model": candidate, - "filename": candidate_filename, - "similarity": similarity, - "confidence": round(similarity * 100, 1), - } + match = ModelMatch.from_mapping( + { + "model": candidate, + "filename": candidate_filename, + "similarity": similarity, + "confidence": round(similarity * 100, 1), + } + ) if max_results <= 0: return [] entry = (similarity, match_counter, match) @@ -197,10 +193,7 @@ def find_matches( # First check for exact match (after normalization) - should be 100% # Only exact matches should get 100% confidence - candidate_norm = candidate.get("_match_filename_norm") - if candidate_norm is None: - candidate_norm = normalize_filename(candidate_filename) - candidate["_match_filename_norm"] = candidate_norm + candidate_norm = normalize_filename(candidate_filename) if target_norm == candidate_norm: similarity = 1.0 @@ -211,12 +204,14 @@ def find_matches( # Only include if above threshold if similarity >= threshold: - match = { - "model": candidate, - "filename": candidate_filename, - "similarity": similarity, - "confidence": confidence, - } + match = ModelMatch.from_mapping( + { + "model": candidate, + "filename": candidate_filename, + "similarity": similarity, + "confidence": confidence, + } + ) if max_results <= 0: return [] entry = (similarity, match_counter, match) @@ -226,7 +221,10 @@ def find_matches( elif similarity > best_matches[0][0]: heapq.heapreplace(best_matches, entry) - return [match for _, _, match in sorted(best_matches, key=lambda x: x[0], reverse=True)] + return [ + match + for _, _, match in sorted(best_matches, key=lambda x: x[0], reverse=True) + ] def normalize_base_model(value: str) -> str: @@ -577,14 +575,14 @@ def match_model_by_title_generic( model_id: int, title_query: str, model_name: str, - versions: list, + versions: list[Dict[str, Any]], base_model_context: Optional[str], - get_base_model_fn, - select_file_fn, - build_result_fn, - hydrate_version_fn=None, + get_base_model_fn: Callable[[Any], Optional[str]], + select_file_fn: Callable[[Any], Optional[Dict[str, Any]]], + build_result_fn: Callable[[Any, Any, float], Optional[SearchResult]], + hydrate_version_fn: Optional[Callable[[Any], Any]] = None, log_prefix: str = "ModelResolver", -) -> Optional[Dict[str, Any]]: +) -> Optional[SearchResult]: """Generic helper to resolve extensionless workflow values by model title across CivitAI and CivArchive.""" title_confidence = calculate_model_title_confidence(title_query, model_name) @@ -621,11 +619,20 @@ def match_model_by_title_generic( if not result: continue - result["confidence"] = title_confidence - result["title_confidence"] = title_confidence + if not isinstance(result, SearchResult): + raise TypeError( + "Model title matcher builders must return SearchResult" + ) + + result = result.with_updates(confidence=title_confidence).with_extra( + title_confidence=title_confidence + ) + result_version_id = result.version_id + result_filename = result.filename + result_base_model = result.base_model log.info( - f"{log_prefix} model-title match: query={title_query}, model_id={model_id}, model_name={model_name}, version_id={result.get('version_id')}, filename={result.get('filename')}, confidence={title_confidence}, base={result.get('base_model')}" + f"{log_prefix} model-title match: query={title_query}, model_id={model_id}, model_name={model_name}, version_id={result_version_id}, filename={result_filename}, confidence={title_confidence}, base={result_base_model}" ) return result diff --git a/core/metadata_audit.py b/core/metadata_audit.py index 1dadc41..66d8b3c 100644 --- a/core/metadata_audit.py +++ b/core/metadata_audit.py @@ -9,8 +9,9 @@ import os from concurrent.futures import ThreadPoolExecutor, as_completed from typing import Any, Dict, List, Optional, Tuple +from .contracts import ResolvedModel from .log_system import create_module_logger -from .metadata_model_utils import dedupe_models, is_model_file_path +from .metadata_model_utils import ModelRecord, dedupe_models, is_model_file_path from .metadata_utils import merge_counted_payload from .path_utils import ( find_metadata_sidecar_path, @@ -171,7 +172,7 @@ def _get_batch_size( return max(MIN_AUDIT_BATCH_SIZE, min(MAX_AUDIT_BATCH_SIZE, batch_size)) -def _make_batches(models: List[Dict[str, Any]], batch_size: int) -> List[List[Dict[str, Any]]]: +def _make_batches(models: List[ResolvedModel], batch_size: int) -> List[List[ResolvedModel]]: if not models: return [] safe_batch_size = max(1, int(batch_size or 1)) @@ -181,9 +182,9 @@ def _make_batches(models: List[Dict[str, Any]], batch_size: int) -> List[List[Di ] -def _audit_one_model(model: Dict[str, Any]) -> Dict[str, Any]: +def _audit_one_model(model: ResolvedModel) -> Dict[str, Any]: result = _empty_audit_counts() - model_path = str(model.get("path") or "").strip() + model_path = model.path.strip() if not model_path: return result @@ -254,9 +255,9 @@ def _audit_one_model(model: Dict[str, Any]) -> Dict[str, Any]: result["mismatches"].append( { "filename": model_filename, - "relative_path": model.get("relative_path") or model_filename, - "category": model.get("category") or "", - "base_directory": model.get("base_directory") or "", + "relative_path": model.relative_path or model_filename, + "category": model.category, + "base_directory": model.base_directory, "model_path": model_path, "metadata_path": metadata_path, "metadata_size": metadata_size, @@ -272,7 +273,7 @@ def _audit_one_model(model: Dict[str, Any]) -> Dict[str, Any]: return result -def _audit_model_batch(batch: List[Dict[str, Any]]) -> Dict[str, Any]: +def _audit_model_batch(batch: List[ResolvedModel]) -> Dict[str, Any]: result = _empty_audit_counts() for model in batch: merge_counted_payload( @@ -285,7 +286,7 @@ def _audit_model_batch(batch: List[Dict[str, Any]]) -> Dict[str, Any]: def audit_metadata_sizes( - models: Optional[List[Dict[str, Any]]] = None, + models: Optional[List[ModelRecord]] = None, *, force_rescan: bool = True, worker_count: Optional[int] = None, @@ -299,7 +300,7 @@ def audit_metadata_sizes( models = get_model_files(force_rescan=force_rescan) - audit_models = dedupe_models(models or []) + audit_models = dedupe_models(models) total_models = len(audit_models) resolved_worker_count, cpu_count = resolve_worker_count( total_models, diff --git a/core/metadata_builder.py b/core/metadata_builder.py index e1311ca..0921933 100644 --- a/core/metadata_builder.py +++ b/core/metadata_builder.py @@ -12,8 +12,9 @@ from collections.abc import Iterable from concurrent.futures import CancelledError, ThreadPoolExecutor, as_completed from typing import Any, Callable, Dict, List, Optional, Tuple +from .contracts import ResolvedModel from .log_system import create_module_logger -from .metadata_model_utils import dedupe_models, is_model_file_path +from .metadata_model_utils import ModelRecord, dedupe_models, is_model_file_path from .metadata_utils import build_sidecar_file_identity, merge_counted_payload from .path_utils import ( HashCalculationCancelled, @@ -89,9 +90,11 @@ def get_metadata_build_capabilities() -> Dict[str, Any]: } -def normalize_metadata_build_mode(value: Any) -> str: +def normalize_metadata_build_mode(value: Optional[str]) -> str: """Return a supported metadata builder source mode.""" - normalized = str(value or "").strip().lower() + if value is not None and not isinstance(value, str): + raise TypeError("Metadata build mode must be a string") + normalized = (value or "").strip().lower() if normalized in METADATA_BUILD_MODES: return normalized return METADATA_BUILD_MODE_CALCULATE_FRESH @@ -106,7 +109,7 @@ def _trusted_external_sha256(metadata: Dict[str, Any]) -> str: def _build_external_metadata_seed( model_path: str, - model: Dict[str, Any], + model: ResolvedModel, ) -> Tuple[Dict[str, Any], str, str]: external_path = find_external_metadata_sidecar_path(model_path) if not external_path: @@ -122,7 +125,7 @@ def _build_external_metadata_seed( seed = build_model_resolver_metadata( model_path, external_raw, - str(model.get("category") or ""), + model.category, ) if trusted_sha256: seed["sha256"] = trusted_sha256 @@ -238,7 +241,7 @@ def _model_type_for_category(category: str) -> str: def _build_local_metadata_payload( *, existing: Dict[str, Any], - model: Dict[str, Any], + model: ResolvedModel, model_path: str, metadata_path: str, file_size: int, @@ -249,7 +252,7 @@ def _build_local_metadata_payload( ) -> Tuple[Dict[str, Any], List[str]]: filename = get_filename_from_path(model_path) stem = os.path.splitext(filename)[0] - category = str(model.get("category") or "") + category = model.category now = time.time() payload = dict(existing or {}) changed_fields: List[str] = [] @@ -460,7 +463,7 @@ def _build_result_payload( def _build_metadata_history_item( - model: Dict[str, Any], + model: ResolvedModel, *, action: str, filename: str, @@ -479,8 +482,8 @@ def _build_metadata_history_item( "action": action, "filename": filename, "relative_path": relative_path, - "category": model.get("category") or "", - "base_directory": model.get("base_directory") or "", + "category": model.category, + "base_directory": model.base_directory, "model_path": model_path, "metadata_path": metadata_path, "size": file_size, @@ -498,19 +501,19 @@ def _build_metadata_history_item( return item -def _history_items_for_model(model: Dict[str, Any], result: Dict[str, Any]) -> List[Dict[str, Any]]: +def _history_items_for_model(model: ResolvedModel, result: Dict[str, Any]) -> List[Dict[str, Any]]: updated_items = result.get("updated") or [] if updated_items: return [dict(item) for item in updated_items if isinstance(item, dict)] - filename = get_filename_from_path(str(model.get("path") or "")) or str(model.get("filename") or "Model") - model_path = str(model.get("path") or "").strip() + filename = get_filename_from_path(model.path) or model.filename or "Model" + model_path = model.path.strip() metadata_path = get_model_resolver_sidecar_path(model_path) if model_path else "" base_item = { "filename": filename, - "relative_path": model.get("relative_path") or filename, - "category": model.get("category") or "", - "base_directory": model.get("base_directory") or "", + "relative_path": model.relative_path or filename, + "category": model.category, + "base_directory": model.base_directory, "model_path": model_path, "metadata_path": metadata_path, } @@ -538,7 +541,7 @@ def _history_items_for_model(model: Dict[str, Any], result: Dict[str, Any]) -> L def _build_missing_local_metadata_parallel( - model_items: List[Dict[str, Any]], + model_items: List[ResolvedModel], *, worker_count: int, cpu_count: int, @@ -599,10 +602,10 @@ def _build_missing_local_metadata_parallel( snapshot.update(data) _emit(progress_callback, snapshot) - def run_one_model(model: Dict[str, Any]) -> Dict[str, Any]: - model_path = str(model.get("path") or "").strip() + def run_one_model(model: ResolvedModel) -> Dict[str, Any]: + model_path = model.path.strip() filename = get_filename_from_path(model_path) - model_key = get_model_path_identity(model.get("path")) or model_path + model_key = get_model_path_identity(model.path) or model_path def child_progress(data: Dict[str, Any]) -> None: if not isinstance(data, dict): @@ -653,9 +656,9 @@ def _build_missing_local_metadata_parallel( try: for future in as_completed(futures): model = futures[future] - model_path = str(model.get("path") or "").strip() + model_path = model.path.strip() filename = get_filename_from_path(model_path) - model_key = get_model_path_identity(model.get("path")) or model_path + model_key = get_model_path_identity(model.path) or model_path try: item_result = future.result() except CancelledError: @@ -667,7 +670,7 @@ def _build_missing_local_metadata_parallel( "errors": [ { "filename": filename, - "relative_path": model.get("relative_path") or filename, + "relative_path": model.relative_path or filename, "model_path": model_path, "metadata_path": get_model_resolver_sidecar_path(model_path), "message": str(exc) or "Metadata build failed.", @@ -740,7 +743,7 @@ def _build_missing_local_metadata_parallel( def build_missing_local_metadata( - models: Optional[List[Dict[str, Any]]] = None, + models: Optional[List[ModelRecord]] = None, *, force_rescan: bool = True, worker_count: Optional[int] = None, @@ -762,10 +765,10 @@ def build_missing_local_metadata( models = get_model_files(force_rescan=force_rescan) - all_models = dedupe_models(models or []) + all_models = dedupe_models(models) model_items = [ model for model in all_models - if is_model_file_path(str(model.get("path") or "").strip()) + if is_model_file_path(model.path.strip()) ] total = len(model_items) resolved_worker_count, cpu_count = resolve_worker_count( @@ -861,9 +864,9 @@ def build_missing_local_metadata( include_updated_count=False, ) - model_path = str(model.get("path") or "").strip() + model_path = model.path.strip() filename = get_filename_from_path(model_path) - relative_path = model.get("relative_path") or filename + relative_path = model.relative_path or filename metadata_path, existing_metadata_path, metadata_exists = _select_metadata_path(model_path) scanned_models += 1 @@ -1154,8 +1157,8 @@ def build_missing_local_metadata( history_item = { **error_item, "action": "error", - "category": model.get("category") or "", - "base_directory": model.get("base_directory") or "", + "category": model.category, + "base_directory": model.base_directory, } history.append(history_item) _emit( diff --git a/core/metadata_model_utils.py b/core/metadata_model_utils.py index 90a7355..b38c3ce 100644 --- a/core/metadata_model_utils.py +++ b/core/metadata_model_utils.py @@ -1,8 +1,10 @@ """Shared helpers for collections of local model metadata entries.""" import os -from typing import Any, Dict, List +from collections.abc import Iterable, Mapping +from typing import Any, List +from .contracts import ResolvedModel from .path_utils import get_filename_from_path, get_model_path_identity from .type_utils import MODEL_EXTENSIONS @@ -19,14 +21,36 @@ def is_model_file_path(path: str) -> bool: return os.path.splitext(filename)[1].lower() in MODEL_EXTENSIONS -def dedupe_models(models: List[Dict[str, Any]]) -> List[Dict[str, Any]]: - """Keep the first valid model entry for every filesystem identity.""" - result: List[Dict[str, Any]] = [] - seen = set() - for model in models or []: - if not isinstance(model, dict): +ModelRecord = ResolvedModel | Mapping[str, Any] + + +def normalize_models(models: Iterable[ModelRecord] | None) -> List[ResolvedModel]: + """Normalize local model inputs once at the metadata boundary.""" + result: List[ResolvedModel] = [] + if models is None: + return result + if isinstance(models, (str, bytes, Mapping)) or not isinstance(models, Iterable): + raise TypeError("models must be an iterable of model records") + + for model in models: + if isinstance(model, ResolvedModel): + result.append(model) continue - identity = get_model_path_identity(model.get("path")) + if not isinstance(model, Mapping): + continue + try: + result.append(ResolvedModel.from_mapping(model)) + except (TypeError, ValueError): + continue + return result + + +def dedupe_models(models: Iterable[ModelRecord] | None) -> List[ResolvedModel]: + """Keep the first normalized local model for every filesystem identity.""" + result: List[ResolvedModel] = [] + seen = set() + for model in normalize_models(models): + identity = get_model_path_identity(model.path) if not identity or identity in seen: continue seen.add(identity) diff --git a/core/path_templates.py b/core/path_templates.py index f1ade23..2c208cf 100644 --- a/core/path_templates.py +++ b/core/path_templates.py @@ -6,6 +6,7 @@ import re from collections import Counter, defaultdict from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence +from .contracts import ResolvedModel from .settings import ( normalize_download_category, normalize_relative_subfolder, @@ -251,7 +252,7 @@ def _infer_category_template(records: List[Dict[str, Any]]) -> Dict[str, Any]: def infer_download_path_templates( - models: Iterable[Mapping[str, Any]], + models: Iterable[ResolvedModel], base_models_config: Any, ) -> Dict[str, Any]: """Return template and base-model mapping suggestions for local model folders.""" @@ -261,15 +262,24 @@ def infer_download_path_templates( mapping_votes: Dict[str, Counter] = defaultdict(Counter) seen = set() - for model in models or []: - category = normalize_download_category(str(model.get("category") or "")) + if models is None: + model_values = () + elif isinstance(models, (str, bytes)): + raise TypeError("models must be an iterable of ResolvedModel values") + else: + model_values = models + + for model in model_values: + if not isinstance(model, ResolvedModel): + raise TypeError("models must contain ResolvedModel values") + category = normalize_download_category(str(model.category or "")) if category not in TEMPLATE_CATEGORIES: continue - relative_path = str(model.get("relative_path") or model.get("filename") or "") + relative_path = str(model.relative_path or model.filename or "") model_key = ( category, - str(model.get("path") or ""), + str(model.path or ""), relative_path, ) if model_key in seen: diff --git a/core/path_utils.py b/core/path_utils.py index 6c10804..67724a3 100644 --- a/core/path_utils.py +++ b/core/path_utils.py @@ -116,14 +116,22 @@ def _get_folder_paths_module(folder_paths_module: Optional[Any] = None) -> Optio return None -def _coerce_folder_paths(value: Any) -> List[str]: +def normalize_string_values(value: Any) -> List[str]: + """Normalize an external string collection without iterating scalar text.""" if isinstance(value, str): - return [value] - if not isinstance(value, (list, tuple)): + return [value] if value else [] + if not isinstance(value, (list, tuple, set, frozenset)): return [] - if value and isinstance(value[0], (list, tuple, set)): - return [str(path) for path in value[0] if path] - return [str(path) for path in value if isinstance(path, str) and path] + return [item for item in value if isinstance(item, str) and item] + + +def normalize_folder_path_values(value: Any) -> List[str]: + """Normalize ComfyUI folder-path values across supported registry shapes.""" + if isinstance(value, (list, tuple)) and value: + first_item = value[0] + if isinstance(first_item, (list, tuple, set, frozenset)): + value = first_item + return normalize_string_values(value) def get_configured_model_roots( @@ -160,12 +168,12 @@ def get_configured_model_roots( try: get_folder_paths = getattr(fp, "get_folder_paths", None) if callable(get_folder_paths): - paths = list(get_folder_paths(name) or []) + paths = normalize_folder_path_values(get_folder_paths(name)) except Exception: paths = [] if not paths and isinstance(folder_names_and_paths, dict): - paths = _coerce_folder_paths(folder_names_and_paths.get(name)) + paths = normalize_folder_path_values(folder_names_and_paths.get(name)) for path in paths: if not path: @@ -229,14 +237,21 @@ def prefer_local_base_directory( def dedupe_local_base_directories( - paths: list, + paths: Optional[list | tuple] = None, preferred_directory: str = "", comfy_root: str = "", ) -> list: """Deduplicate a list of directory paths, resolving symlinks/junctions.""" + if paths is None: + path_values = () + elif isinstance(paths, (list, tuple)): + path_values = paths + else: + raise TypeError("paths must be a list or tuple") + by_identity = {} ordered_identities = [] - for path in paths or []: + for path in path_values: if not path or not os.path.isdir(path): continue path_abs = get_path_abs(path) diff --git a/core/request_utils.py b/core/request_utils.py index 06a75b0..1f5b1a3 100644 --- a/core/request_utils.py +++ b/core/request_utils.py @@ -15,6 +15,10 @@ def extract_request_sha256( value = "" for key in keys: value = payload.get(key) + if value is None or value == "": + continue + if not isinstance(value, str): + raise TypeError(f"Request {key} must be a string") if value: break return normalize_sha256(value) @@ -31,6 +35,200 @@ async def read_optional_object_payload(request: Any) -> dict[str, Any]: return payload if isinstance(payload, dict) else {} +def read_text_field( + payload: Mapping[str, Any], + field_name: str, + *, + default: str = "", + contract_name: str = "Request", +) -> str: + """Read a request text field without silently stringifying invalid values.""" + value = payload.get(field_name, default) + if value is None: + return default + if not isinstance(value, str): + raise TypeError(f"{contract_name} {field_name} must be a string") + return value.strip() + + +def read_optional_text_field( + payload: Mapping[str, Any], + field_name: str, + *, + default: str | None = None, + contract_name: str = "Request", +) -> str | None: + """Read an optional request text field without stringifying bad values.""" + value = payload.get(field_name, default) + if value is None: + return default + if not isinstance(value, str): + raise TypeError(f"{contract_name} {field_name} must be a string") + return value.strip() or default + + +def read_first_text_field( + payload: Mapping[str, Any], + field_names: Tuple[str, ...], + *, + default: str = "", + contract_name: str = "Request", +) -> str: + """Read the first non-empty text alias without bypassing invalid values.""" + for field_name in field_names: + if field_name not in payload: + continue + value = payload[field_name] + if value is None: + continue + if not isinstance(value, str): + raise TypeError(f"{contract_name} {field_name} must be a string") + value = value.strip() + if value: + return value + return default + + +def read_first_identifier_field( + payload: Mapping[str, Any], + field_names: Tuple[str, ...], + *, + default: int | str | None = None, + contract_name: str = "Request", +) -> int | str | None: + """Read the first non-empty integer/string identifier alias.""" + for field_name in field_names: + if field_name not in payload: + continue + value = payload[field_name] + if value is None or value == "": + continue + if isinstance(value, bool) or not isinstance(value, (int, str)): + raise TypeError( + f"{contract_name} {field_name} must be an integer or string" + ) + if isinstance(value, str): + value = value.strip() + if not value: + continue + return value + return default + + +def read_identifier_field( + payload: Mapping[str, Any], + field_name: str, + *, + default: int | str | None = None, + contract_name: str = "Request", +) -> int | str | None: + """Read an optional integer/string identifier without accepting booleans.""" + value = payload.get(field_name, default) + if value is None: + return default + if isinstance(value, bool) or not isinstance(value, (int, str)): + raise TypeError( + f"{contract_name} {field_name} must be an integer or string" + ) + if isinstance(value, str): + return value.strip() or default + return value + + +def coerce_integer_identifier( + value: int | str | None, + field_name: str, + *, + contract_name: str = "Request", +) -> int | None: + """Convert a validated identifier to an integer without hiding bad values.""" + if value is None: + return None + if isinstance(value, bool) or not isinstance(value, (int, str)): + raise TypeError(f"{contract_name} {field_name} must be an integer") + if isinstance(value, int): + return value + text = value.strip() + if not text: + return None + try: + return int(text) + except ValueError as exc: + raise ValueError( + f"{contract_name} {field_name} must be an integer" + ) from exc + + +def read_bool_field( + payload: Mapping[str, Any], + field_name: str, + *, + default: bool = False, + contract_name: str = "Request", +) -> bool: + """Read a request boolean using only the supported JSON representations.""" + value = payload.get(field_name, default) + if value is None: + return default + if isinstance(value, bool): + return value + if isinstance(value, (int, float)) and value in (0, 1): + return bool(value) + if isinstance(value, str): + normalized = value.strip().lower() + if normalized in {"true", "1", "yes", "on"}: + return True + if normalized in {"false", "0", "no", "off"}: + return False + raise TypeError(f"{contract_name} {field_name} must be a boolean") + + +def read_first_bool_field( + payload: Mapping[str, Any], + field_names: Tuple[str, ...], + *, + default: bool = False, + contract_name: str = "Request", +) -> bool: + """Read the first non-empty boolean alias without coercing bad values.""" + for field_name in field_names: + if field_name not in payload: + continue + value = payload[field_name] + if value is None or value == "": + continue + return read_bool_field( + payload, + field_name, + default=default, + contract_name=contract_name, + ) + return default + + +def read_int_field( + payload: Mapping[str, Any], + field_name: str, + *, + default: int | None = None, + contract_name: str = "Request", +) -> int | None: + """Read an optional integer without accepting booleans or float truncation.""" + value = payload.get(field_name, default) + if value is None or value == "": + return default + if isinstance(value, bool) or not isinstance(value, (int, str)): + raise TypeError(f"{contract_name} {field_name} must be an integer") + if isinstance(value, str): + try: + return int(value.strip()) + except ValueError as exc: + raise ValueError( + f"{contract_name} {field_name} must be an integer" + ) from exc + return value + + def validate_workflow_payload( workflow: Any, *, diff --git a/core/resolver.py b/core/resolver.py index 6596e47..ccc6f15 100644 --- a/core/resolver.py +++ b/core/resolver.py @@ -8,12 +8,20 @@ import json import os import re import threading +from dataclasses import replace from typing import Any, Callable, Dict, List, Optional, Tuple from urllib.parse import unquote +from .contracts import ( + MissingModel, + ModelMatch, + ModelReference, + Resolution, + ResolvedModel, + WorkflowAnalysisResult, +) from .custom_nodes import ( custom_node_has_potential_model_reference, - get_custom_node_resolution_metadata, should_skip_existing_custom_node_reference, ) from .log_system import create_module_logger @@ -22,6 +30,7 @@ log = create_module_logger(__name__) from .local_hash_matches import collect_local_hash_matches_for_result from .matcher import find_matches, strip_known_model_extension +from .metadata_model_utils import normalize_models from .scanner import get_model_files, invalidate_model_files_cache from .type_utils import MODEL_EXTENSIONS as _MODEL_EXTENSIONS from .type_utils import ( @@ -63,7 +72,7 @@ from .path_utils import ( # Imported from .matcher _LOCAL_HASH_MATCH_CACHE_LOCK = threading.Lock() -_LOCAL_HASH_MATCH_CACHE: Optional[Dict[str, List[Dict[str, Any]]]] = None +_LOCAL_HASH_MATCH_CACHE: Optional[Dict[str, List[ModelMatch]]] = None _ACTIVE_DOWNLOAD_STATUSES = {"starting", "downloading", "paused", "cancelling"} @@ -80,13 +89,6 @@ def invalidate_model_caches() -> None: invalidate_local_hash_match_cache() -def _clone_hash_match(match: Dict[str, Any]) -> Dict[str, Any]: - cloned = dict(match) - if isinstance(cloned.get("model"), dict): - cloned["model"] = dict(cloned["model"]) - return cloned - - def _get_active_downloads_by_path() -> Dict[str, Dict[str, Any]]: try: from .download.api import get_all_progress @@ -124,9 +126,9 @@ def _get_active_downloads_by_path() -> Dict[str, Dict[str, Any]]: def annotate_local_matches_with_download_state( - matches: List[Dict[str, Any]], + matches: List[ModelMatch], active_downloads_by_path: Optional[Dict[str, Dict[str, Any]]] = None, -) -> List[Dict[str, Any]]: +) -> List[ModelMatch]: active_downloads = ( active_downloads_by_path if active_downloads_by_path is not None @@ -135,18 +137,14 @@ def annotate_local_matches_with_download_state( if not active_downloads: return matches - enriched_matches: List[Dict[str, Any]] = [] + enriched_matches: List[ModelMatch] = [] for match in matches: - if not isinstance(match, dict): - enriched_matches.append(match) - continue - - model = match.get("model") if isinstance(match.get("model"), dict) else {} + model = match.model candidate_paths = [ - model.get("path"), - model.get("resolved_path"), - match.get("path"), - match.get("resolved_path"), + model.path, + model.extra_value("resolved_path"), + match.extra_value("path"), + match.extra_value("resolved_path"), ] download_info = None for candidate_path in candidate_paths: @@ -159,24 +157,21 @@ def annotate_local_matches_with_download_state( enriched_matches.append(match) continue - enriched_match = dict(match) - enriched_model = dict(model) download_fields = { **download_info, "is_downloading": True, "downloading": True, } - enriched_match.update(download_fields) - if enriched_model: - enriched_model.update(download_fields) - enriched_match["model"] = enriched_model - enriched_matches.append(enriched_match) + enriched_model = model.with_extra(**download_fields) + enriched_matches.append( + replace(match, model=enriched_model).with_extra(**download_fields) + ) return enriched_matches -def _is_local_hash_match_candidate(model: Dict[str, Any]) -> bool: - model_path = str(model.get("path") or "").strip() +def _is_local_hash_match_candidate(model: ResolvedModel) -> bool: + model_path = str(model.path or "").strip() if not model_path: return False @@ -192,16 +187,18 @@ def _is_local_hash_match_candidate(model: Dict[str, Any]) -> bool: def _build_local_hash_match_cache( - available_models: List[Dict[str, Any]], -) -> Dict[str, List[Dict[str, Any]]]: - index: Dict[str, List[Dict[str, Any]]] = {} + available_models: List[ResolvedModel], +) -> Dict[str, List[ModelMatch]]: + index: Dict[str, List[ModelMatch]] = {} seen_entries = set() - for model in available_models: + normalized_models = normalize_models(available_models) + + for model in normalized_models: if not _is_local_hash_match_candidate(model): continue - model_path = model.get("path", "") + model_path = model.path if not model_path: continue @@ -218,7 +215,7 @@ def _build_local_hash_match_cache( if not metadata_hashes: continue - model_filename = model.get("filename") or get_filename_from_path(model_path) + model_filename = model.filename or get_filename_from_path(model_path) for metadata_hash in metadata_hashes: normalized_hash = normalize_sha256(metadata_hash) if not normalized_hash: @@ -230,29 +227,32 @@ def _build_local_hash_match_cache( continue seen_entries.add(entry_key) - model_with_metadata = { - **model, - "sha256": normalized_hash, - "metadata_path": metadata_path, - } + model_with_metadata = model.with_extra( + sha256=normalized_hash, + metadata_path=metadata_path, + ) index.setdefault(normalized_hash, []).append( - { - "model": model_with_metadata, - "filename": model_filename, - "similarity": 1.0, - "confidence": 100.0, - "match_type": "hash", - "hash_match": True, - "hash_source": "metadata", - "sha256": normalized_hash, - "metadata_path": metadata_path, - } + ModelMatch.from_mapping( + { + "model": model_with_metadata, + "filename": model_filename, + "similarity": 1.0, + "confidence": 100.0, + "match_type": "hash", + "hash_match": True, + "hash_source": "metadata", + "sha256": normalized_hash, + "metadata_path": metadata_path, + } + ) ) return index -def _get_local_hash_match_cache(force_rescan: bool = False) -> Dict[str, List[Dict[str, Any]]]: +def _get_local_hash_match_cache( + force_rescan: bool = False, +) -> Dict[str, List[ModelMatch]]: global _LOCAL_HASH_MATCH_CACHE if force_rescan: invalidate_local_hash_match_cache() @@ -364,17 +364,17 @@ def extract_workflow_hash_metadata(workflow_json: Dict[str, Any]) -> Dict[str, D def get_workflow_hash_info_for_ref( - workflow_hashes: Dict[str, Dict[str, Any]], model_ref: Dict[str, Any] + workflow_hashes: Dict[str, Dict[str, Any]], model_ref: ModelReference ) -> Optional[Dict[str, Any]]: if not workflow_hashes: return None - original_path = str(model_ref.get("original_path") or "") + original_path = model_ref.original_path candidates = [ - f"{model_ref.get('node_id')}:{model_ref.get('widget_index')}", + f"{model_ref.node_id}:{model_ref.widget_index}", original_path, get_filename_from_path(original_path), - model_ref.get("filename") or "", - model_ref.get("name") or "", + model_ref.extra_value("filename", "") or "", + model_ref.extra_value("name", "") or "", ] for candidate in candidates: if candidate and candidate in workflow_hashes: @@ -452,16 +452,16 @@ def workflow_url_points_to_file(url: str, filename: str) -> bool: return filename in decoded_url or unquote(filename) in decoded_url -def _deduplicate_local_matches(matches: List[Dict[str, Any]]) -> List[Dict[str, Any]]: +def _deduplicate_local_matches(matches: List[ModelMatch]) -> List[ModelMatch]: """Keep one highest-confidence match for each local model path.""" seen_absolute_paths = {} deduplicated_matches = [] for match in matches: - model_dict = match["model"] - absolute_path = model_dict.get("path", "") + model = match.model + absolute_path = model.path path_identity = get_path_identity(absolute_path) if absolute_path else "" dedupe_key = path_identity or os.path.normcase( - model_dict.get("relative_path", "") or match.get("filename", "") + model.relative_path or match.filename ) if dedupe_key not in seen_absolute_paths: @@ -469,7 +469,7 @@ def _deduplicate_local_matches(matches: List[Dict[str, Any]]) -> List[Dict[str, deduplicated_matches.append(match) else: existing_match = seen_absolute_paths[dedupe_key] - if match["confidence"] > existing_match["confidence"]: + if match.confidence > existing_match.confidence: idx = deduplicated_matches.index(existing_match) deduplicated_matches[idx] = match seen_absolute_paths[dedupe_key] = match @@ -483,7 +483,7 @@ def search_local_matches( similarity_threshold: float = 0.0, max_matches_per_model: int = 10, force_rescan: bool = False, -) -> List[Dict[str, Any]]: +) -> List[ModelMatch]: """ Search local model files using the same matcher as workflow analysis. @@ -497,12 +497,13 @@ def search_local_matches( Deduplicated list of local matches sorted by similarity """ available_models = get_model_files(force_rescan=force_rescan) + available_models = normalize_models(available_models) candidates = available_models if category and category != "unknown": - candidates = [m for m in available_models if m.get("category") == category] + candidates = [m for m in available_models if m.category == category] candidates.extend( - [m for m in available_models if m.get("category") != category] + [m for m in available_models if m.category != category] ) matches = find_matches( @@ -521,9 +522,13 @@ def _collect_hashes_from_container(value: Any) -> List[str]: return [h] if h else [] -def _metadata_file_matches_model(file_info: Dict[str, Any], model: Dict[str, Any]) -> bool: - model_filename = str(model.get("filename") or get_filename_from_path(model.get("path", ""))).lower() - model_relative = str(model.get("relative_path") or "").replace("\\", "/").lower() +def _metadata_file_matches_model( + file_info: Dict[str, Any], model: ResolvedModel +) -> bool: + model_filename = str( + model.filename or get_filename_from_path(model.path) + ).lower() + model_relative = str(model.relative_path or "").replace("\\", "/").lower() model_stem = strip_known_model_extension(get_filename_from_path(model_filename)).lower() candidates = [ @@ -546,7 +551,7 @@ def _metadata_file_matches_model(file_info: Dict[str, Any], model: Dict[str, Any def _extract_model_sha256_from_metadata( - metadata: Dict[str, Any], model: Dict[str, Any] + metadata: Dict[str, Any], model: ResolvedModel ) -> List[str]: if not isinstance(metadata, dict): return [] @@ -580,7 +585,7 @@ def search_local_matches_by_hash( category: Optional[str] = None, max_matches: int = 20, force_rescan: bool = False, -) -> List[Dict[str, Any]]: +) -> List[ModelMatch]: """ Find local models whose sidecar .metadata.json contains the given SHA256. @@ -592,12 +597,12 @@ def search_local_matches_by_hash( return [] index = _get_local_hash_match_cache(force_rescan=force_rescan) - matches = [_clone_hash_match(match) for match in index.get(normalized_hash, [])] + matches = list(index.get(normalized_hash, [])) if category and category != "unknown": matches.sort( key=lambda match: 0 - if match.get("model", {}).get("category") == category + if match.model.category == category else 1 ) @@ -609,7 +614,7 @@ def search_local_matches_by_hash( def get_local_model_hash_metadata( model_path: str, - model: Optional[Dict[str, Any]] = None, + model: Optional[ResolvedModel] = None, ) -> Dict[str, Any]: """ Return SHA256 hashes already stored in sidecar metadata for a local model. @@ -629,12 +634,20 @@ def get_local_model_hash_metadata( file_size = os.path.getsize(normalized_path) except OSError: file_size = 0 - model_info: Dict[str, Any] = { - **(model if isinstance(model, dict) else {}), - "path": normalized_path, - } - model_info.setdefault("filename", get_filename_from_path(normalized_path)) - model_info.setdefault("relative_path", model_info.get("filename", "")) + filename = get_filename_from_path(normalized_path) + model_info = ( + model.with_updates( + path=normalized_path, + filename=model.filename or filename, + relative_path=model.relative_path or model.filename or filename, + ) + if model is not None + else ResolvedModel( + path=normalized_path, + filename=filename, + relative_path=filename, + ) + ) metadata_path = find_metadata_sidecar_path(normalized_path) last_hash_status = "" @@ -821,7 +834,7 @@ def analyze_and_find_matches( progress_callback: Optional[Callable[[Dict[str, Any]], None]] = None, force_rescan: bool = False, analysis_id: Optional[str] = None, -) -> Dict[str, Any]: +) -> WorkflowAnalysisResult: """ Main entry point: analyze workflow and find matches for missing models. @@ -833,32 +846,8 @@ def analyze_and_find_matches( analysis_id: Optional identifier used to correlate logs for one request Returns: - Dictionary with analysis results: - { - 'missing_models': [ - { - 'node_id': node ID, - 'node_type': node type, - 'widget_index': widget index, - 'original_path': original path from workflow, - 'category': model category, - 'workflow_url': URL from workflow if found, - 'workflow_directory': directory from workflow if found, - 'matches': [ - { - 'model': model dict from scanner, - 'filename': model filename, - 'similarity': similarity score (0.0-1.0), - 'confidence': confidence percentage (0-100) - }, - ... - ] - }, - ... - ], - 'total_missing': count of missing models, - 'total_models_analyzed': count of all models in workflow - } + Typed workflow analysis result. The HTTP route serializes it only after + applying route-level filtering and local download-source enrichment. """ if progress_callback: progress_callback( @@ -883,13 +872,7 @@ def analyze_and_find_matches( "total": 0, } ) - return { - "missing_models": [], - "resolved_models": [], - "total_resolved": 0, - "total_missing": 0, - "total_models_analyzed": 0, - } + return WorkflowAnalysisResult() # Extract URLs from workflow (node.properties.models + regex) analysis_context = f" (analysis_id={analysis_id})" if analysis_id else "" @@ -919,16 +902,23 @@ def analyze_and_find_matches( progress_callback=progress_callback, analysis_id=analysis_id, ) - available_models = inventory["available_models"] - all_model_refs = inventory["model_refs"] + available_models = inventory.available_models + all_model_refs = inventory.model_refs + available_models = normalize_models(available_models) + all_model_refs = [ + ref + if isinstance(ref, ModelReference) + else ModelReference.from_mapping(ref) + for ref in all_model_refs + ] available_models_by_category = {} for model in available_models: - model_category = model.get("category", "") + model_category = model.category if model_category not in available_models_by_category: available_models_by_category[model_category] = [] available_models_by_category[model_category].append(model) - ordered_candidates_cache: Dict[str, List[Dict[str, Any]]] = {} + ordered_candidates_cache: Dict[str, List[ResolvedModel]] = {} if progress_callback: progress_callback( @@ -943,35 +933,57 @@ def analyze_and_find_matches( # Identify missing models missing_models = identify_missing_models(all_model_refs, available_models) resolved_model_refs = [ - model_ref for model_ref in all_model_refs if model_ref.get("exists", False) + model_ref for model_ref in all_model_refs if model_ref.exists ] - # Enrich missing models with workflow URLs + # Enrich typed missing models with workflow URLs, hash metadata, and URN + # hints. Raw mappings are created only for helpers that still consume + # external workflow/custom-node metadata. + enriched_missing_models = [] for missing in missing_models: - original_path = missing.get("original_path", "") + updates: Dict[str, Any] = {} + original_path = missing.reference.original_path filename = get_filename_from_path(original_path) url_info = get_workflow_url_info_for_filename(workflow_urls, filename) if url_info: - missing["workflow_url"] = url_info.get("url", "") - missing["workflow_model_url"] = url_info.get("model_url", "") - missing["workflow_directory"] = url_info.get("directory", "") - missing["url_source"] = url_info.get("source", "") - hash_info = get_workflow_hash_info_for_ref(workflow_hashes, missing) + updates.update( + { + "workflow_url": url_info.get("url", ""), + "workflow_model_url": url_info.get("model_url", ""), + "workflow_directory": url_info.get("directory", ""), + "url_source": url_info.get("source", ""), + } + ) + hash_info = get_workflow_hash_info_for_ref( + workflow_hashes, + missing.reference, + ) if hash_info: - missing["workflow_sha256"] = normalize_sha256(hash_info.get("sha256")) - missing["hash_lookup_source"] = hash_info.get("source") or "workflow_metadata" + updates.update( + { + "workflow_sha256": normalize_sha256(hash_info.get("sha256")), + "hash_lookup_source": hash_info.get("source") + or "workflow_metadata", + } + ) - # Handle URNs: mark for async resolution by frontend - # No sync CivitAI calls here - frontend will fetch asynchronously - for missing in missing_models: - if missing.get("is_urn"): - missing["needs_urn_resolve"] = True - urn = missing.get("urn") + # Handle URNs: mark for async resolution by frontend. No sync CivitAI + # calls happen here; the frontend resolves them asynchronously. + if missing.extra_value("is_urn", False): + updates["needs_urn_resolve"] = True + urn = missing.extra_value("urn") if urn: - missing["urn_model_id"] = urn.get("model_id") - missing["urn_version_id"] = urn.get("version_id") - missing["urn_type"] = urn.get("type", "") + updates.update( + { + "urn_model_id": urn.get("model_id"), + "urn_version_id": urn.get("version_id"), + "urn_type": urn.get("type", ""), + } + ) + + enriched_missing_models.append(missing.with_extra(**updates)) + missing_models = enriched_missing_models total_matching_models = len(missing_models) if progress_callback: @@ -984,22 +996,24 @@ def analyze_and_find_matches( } ) - local_match_cache: Dict[Tuple[str, str], List[Dict[str, Any]]] = {} + local_match_cache: Dict[Tuple[str, str], List[ModelMatch]] = {} active_downloads_by_path = _get_active_downloads_by_path() - def get_match_target(model_ref: Dict[str, Any]) -> Optional[str]: + def get_match_target(model_ref: MissingModel) -> Optional[str]: + reference = model_ref.reference target_for_matching = ( - model_ref.get("original_path") - or model_ref.get("expected_filename") - or model_ref.get("name") - or model_ref.get("filename") - or model_ref.get("full_path") + reference.original_path + or model_ref.extra_value("expected_filename") + or model_ref.extra_value("name") + or model_ref.extra_value("filename") + or model_ref.extra_value("full_path") or "" ) # For URNs, prefer expected_filename for matching - if model_ref.get("is_urn") and model_ref.get("expected_filename"): - return model_ref["expected_filename"] + expected_filename = model_ref.extra_value("expected_filename") + if model_ref.extra_value("is_urn") and expected_filename: + return expected_filename if isinstance(target_for_matching, str) and target_for_matching.startswith( "urn:air:" @@ -1019,14 +1033,14 @@ def analyze_and_find_matches( return target_for_matching - def get_match_category(model_ref: Dict[str, Any]) -> str: - category = model_ref.get("category") + def get_match_category(model_ref: MissingModel) -> str: + category = model_ref.reference.category if not category or category == "unknown": - node_type = model_ref.get("node_type", "") + node_type = model_ref.reference.node_type category = NODE_TYPE_TO_CATEGORY_HINTS.get(node_type, "unknown") return category or "unknown" - def get_candidates_for_category(category: str) -> List[Dict[str, Any]]: + def get_candidates_for_category(category: str) -> List[ResolvedModel]: if not category or category == "unknown": return available_models @@ -1034,18 +1048,20 @@ def analyze_and_find_matches( if candidates is None: preferred = available_models_by_category.get(category, []) others = [ - m for m in available_models if m.get("category") != category + m for m in available_models if m.category != category ] candidates = preferred + others ordered_candidates_cache[category] = candidates return candidates - def find_local_matches_for_ref(model_ref: Dict[str, Any]) -> List[Dict[str, Any]]: + def find_local_matches_for_ref(model_ref: MissingModel) -> List[ModelMatch]: target_for_matching = get_match_target(model_ref) category = get_match_category(model_ref) - workflow_sha256 = normalize_sha256(model_ref.get("workflow_sha256")) - hash_matches = [] + workflow_sha256 = normalize_sha256( + model_ref.extra_value("workflow_sha256") + ) + hash_matches: List[ModelMatch] = [] if workflow_sha256: hash_matches = collect_local_hash_matches_for_result( workflow_sha256, @@ -1053,8 +1069,11 @@ def analyze_and_find_matches( category=category, max_matches=max_matches_per_model, force_rescan=False, - source=model_ref.get("hash_lookup_source") or "workflow_metadata", - filename=get_filename_from_path(model_ref.get("original_path") or ""), + source=model_ref.extra_value("hash_lookup_source") + or "workflow_metadata", + filename=get_filename_from_path( + model_ref.reference.original_path or "" + ), ) if not target_for_matching: @@ -1088,41 +1107,35 @@ def analyze_and_find_matches( "message": f"Analyzing model {index} of {total_missing}", "current": index, "total": total_matching_models, - "model_name": missing.get("name") - or missing.get("original_path", ""), + "model_name": missing.extra_value("name") + or missing.reference.original_path, } ) - name = missing.get("name") or missing.get("original_path", "") - if should_skip_existing_custom_node_reference(missing): + name = missing.extra_value("name") or missing.reference.original_path + if should_skip_existing_custom_node_reference(missing.reference): log.info( f"Skipping existing custom-node model reference: {name}" ) continue - missing_with_matches.append({ - **missing, - "matches": find_local_matches_for_ref(missing), - }) + missing_with_matches.append( + missing.with_matches( + find_local_matches_for_ref(missing) + ) + ) # Existing models already have an exact local path. Fuzzy-matching every # resolved reference against the full local model index is redundant and # makes small workflow edits unnecessarily expensive. - resolved_with_matches = [ - { - **resolved, - "matches": [], - } - for resolved in resolved_model_refs - ] - - result = { - "missing_models": missing_with_matches, - "resolved_models": resolved_with_matches, - "total_resolved": len(resolved_with_matches), - "total_missing": len(missing_with_matches), - "total_models_analyzed": len(all_model_refs), - } + resolved_references = tuple(resolved_model_refs) + result = WorkflowAnalysisResult( + missing_models=tuple(missing_with_matches), + resolved_models=resolved_references, + total_resolved=len(resolved_references), + total_missing=len(missing_with_matches), + total_models_analyzed=len(all_model_refs), + ) if progress_callback: progress_callback( @@ -1138,59 +1151,30 @@ def analyze_and_find_matches( def apply_resolution( - workflow_json: Dict[str, Any], resolutions: List[Dict[str, Any]] + workflow_json: Dict[str, Any], + resolutions: List[Resolution], ) -> Dict[str, Any]: """ Apply model resolutions to workflow. Args: workflow_json: Workflow JSON dictionary (will be modified) - resolutions: List of resolution dictionaries: - { - 'node_id': node ID, - 'widget_index': widget index, - 'resolved_path': absolute path to resolved model, - 'category': model category (optional), - 'resolved_model': model dict from scanner (optional), - 'nested_key': nested key for dict-type widgets (optional) - } + resolutions: Validated ``Resolution`` objects. Returns: Updated workflow JSON dictionary """ - # Prepare mappings for workflow_updater - mappings = [] + typed_resolutions = [] for resolution in resolutions: - mapping = { - "node_id": resolution.get("node_id"), - "widget_index": resolution.get("widget_index"), - "resolved_path": resolution.get("resolved_path"), - "category": resolution.get("category"), - "resolved_model": resolution.get("resolved_model"), - "subgraph_id": resolution.get( - "subgraph_id" - ), # Include subgraph_id for subgraph nodes - "is_top_level": resolution.get( - "is_top_level" - ), # True for top-level nodes, False for nodes in subgraph definitions - "nested_key": resolution.get( - "nested_key" - ), # For dict-type widget values - "promoted_widget_name": resolution.get("promoted_widget_name"), - **get_custom_node_resolution_metadata(resolution), - } + try: + resolution.validate() + except (TypeError, ValueError) as exc: + log.warning(f"Skipping invalid workflow resolution: {exc}") + continue - # If resolved_model provided, extract path if needed - if resolution.get("resolved_model"): - resolved_model = resolution["resolved_model"] - if "path" in resolved_model and not mapping.get("resolved_path"): - mapping["resolved_path"] = resolved_model["path"] - if "base_directory" in resolved_model: - mapping["base_directory"] = resolved_model["base_directory"] - - mappings.append(mapping) + typed_resolutions.append(resolution) # Update workflow - updated_workflow = update_workflow_nodes(workflow_json, mappings) + updated_workflow = update_workflow_nodes(workflow_json, typed_resolutions) return updated_workflow diff --git a/core/routes/base_models.py b/core/routes/base_models.py index 0df5be4..75b67a1 100644 --- a/core/routes/base_models.py +++ b/core/routes/base_models.py @@ -2,6 +2,7 @@ import asyncio +from ..request_utils import read_bool_field from ..sources.popular import ( get_base_models_config, get_base_models_status, @@ -23,7 +24,14 @@ def register_base_model_routes(routes, web, json_api_endpoint): @json_api_endpoint("base-models status") async def get_base_models_status_route(request): """Get local and optional remote base-model status.""" - check_remote = request.query.get("check_remote") == "1" + try: + check_remote = read_bool_field( + request.query, + "check_remote", + contract_name="Base models status request", + ) + except TypeError as exc: + return web.json_response({"error": str(exc)}, status=400) status = await asyncio.to_thread(get_base_models_status, check_remote) return web.json_response(status) diff --git a/core/routes/helpers.py b/core/routes/helpers.py index c716214..3430978 100644 --- a/core/routes/helpers.py +++ b/core/routes/helpers.py @@ -4,7 +4,7 @@ import asyncio import threading from functools import wraps -from ..request_utils import read_optional_object_payload +from ..request_utils import read_optional_object_payload, read_text_field def register_service_route( @@ -170,7 +170,11 @@ def create_route_helpers(web, logger, load_settings, hash_calculation_cancelled) payload = await read_optional_object_payload(request) if "aria2c_path" in payload: settings = dict(settings) - settings["aria2c_path"] = payload.get("aria2c_path", "") + settings["aria2c_path"] = read_text_field( + payload, + "aria2c_path", + contract_name="Aria2 settings request", + ) return settings return ( diff --git a/core/scanner.py b/core/scanner.py index 66972bc..50dbf5d 100644 --- a/core/scanner.py +++ b/core/scanner.py @@ -8,12 +8,19 @@ import os import time from typing import Dict, List, Optional, Tuple +from .contracts import ResolvedModel from .log_system import create_module_logger log = create_module_logger(__name__) -from .path_utils import get_filename_from_path, get_model_path_identity, get_path_identity +from .path_utils import ( + get_filename_from_path, + get_model_path_identity, + get_path_identity, + normalize_folder_path_values, + normalize_string_values, +) # Import folder_paths lazily - it may not be available until ComfyUI is initialized try: @@ -22,7 +29,7 @@ except ImportError: folder_paths = None log.warning("Model Resolver: folder_paths not available yet - will retry later") -_MODEL_FILES_CACHE: Optional[List[Dict[str, str]]] = None +_MODEL_FILES_CACHE: Optional[List[ResolvedModel]] = None _MODEL_FILES_CACHE_AT: float = 0.0 _MODEL_FILES_CACHE_TTL_SECONDS = 2.0 @@ -53,7 +60,7 @@ def get_model_directories() -> Dict[str, Tuple[List[str], set]]: def scan_directory( directory: str, extensions: set, category: str -) -> List[Dict[str, str]]: +) -> List[ResolvedModel]: """ Recursively scan a single directory for model files. @@ -63,16 +70,12 @@ def scan_directory( category: Model category name (e.g., 'checkpoints', 'loras') Returns: - List of dictionaries with model information: - { - 'filename': 'model.safetensors', - 'path': 'absolute/path/to/model.safetensors', - 'relative_path': 'subfolder/model.safetensors' or 'model.safetensors', - 'category': 'checkpoints', - 'base_directory': 'absolute/path/to/base' - } + Typed local model records with path, category and directory metadata. """ models = [] + if not isinstance(directory, str) or not isinstance(category, str): + return models + extensions = set(normalize_string_values(extensions)) if not os.path.exists(directory) or not os.path.isdir(directory): # log.debug(f"Directory does not exist or is not accessible: {directory}") @@ -120,13 +123,13 @@ def scan_directory( relative_path = get_filename_from_path(root) models.append( - { - "filename": get_filename_from_path(root), - "path": root, - "relative_path": relative_path, - "category": category, - "base_directory": base_directory, - } + ResolvedModel( + filename=get_filename_from_path(root), + path=root, + relative_path=relative_path, + category=category, + base_directory=base_directory, + ) ) for filename in files: @@ -152,13 +155,13 @@ def scan_directory( relative_path = filename models.append( - { - "filename": filename, - "path": full_path, - "relative_path": relative_path, - "category": category, - "base_directory": base_directory, - } + ResolvedModel( + filename=filename, + path=full_path, + relative_path=relative_path, + category=category, + base_directory=base_directory, + ) ) except (OSError, PermissionError) as e: log.warning(f"Error scanning directory {directory}: {e}") @@ -166,12 +169,12 @@ def scan_directory( return models -def scan_all_directories() -> List[Dict[str, str]]: +def scan_all_directories() -> List[ResolvedModel]: """ Scan all configured model directories and return list of available models. Returns: - List of dictionaries with model information (same format as scan_directory) + Typed local model records (same fields as scan_directory) """ all_models = [] directories = get_model_directories() @@ -188,15 +191,22 @@ def scan_all_directories() -> List[Dict[str, str]]: extensions = set() try: if isinstance(value, (list, tuple)): - if len(value) >= 2: - paths = value[0] or [] - raw_exts = value[1] + if len(value) >= 2 and ( + isinstance(value[0], (list, tuple, set, frozenset)) + or isinstance(value[1], (list, tuple, set, frozenset)) + ): + paths = normalize_folder_path_values(value[0]) + raw_exts = value[1] if len(value) >= 2 else [] else: - # Unexpected format; treat value as paths - paths = list(value) + # Unexpected registry format; retain only string paths. + paths = normalize_folder_path_values(value) raw_exts = [] elif isinstance(value, dict): - paths = value.get("paths") or value.get("path") or [] + paths = normalize_folder_path_values( + value.get("paths") + if "paths" in value + else value.get("path") + ) raw_exts = value.get("extensions") or [] else: # Unknown format; skip category @@ -206,10 +216,10 @@ def scan_all_directories() -> List[Dict[str, str]]: continue # Normalize extensions to a set[str] - if isinstance(raw_exts, (list, tuple, set)): - extensions = {str(e).lower() for e in raw_exts} - elif raw_exts: - extensions = {str(raw_exts).lower()} + extensions = { + extension.lower() + for extension in normalize_string_values(raw_exts) + } except Exception as e: log.warning(f"Error interpreting folder_paths entry for {category}: {e}") continue @@ -225,8 +235,8 @@ def scan_all_directories() -> List[Dict[str, str]]: models = scan_directory(directory_path, extensions, category) for model in models: - identity = get_model_path_identity(model.get("path")) - model_key = (model.get("category", category), identity) + identity = get_model_path_identity(model.path) + model_key = (model.category or category, identity) if identity and model_key in seen_models: continue if identity: @@ -252,7 +262,7 @@ def invalidate_model_files_cache() -> None: pass -def get_model_files(force_rescan: bool = False) -> List[Dict[str, str]]: +def get_model_files(force_rescan: bool = False) -> List[ResolvedModel]: """ Get list of all available model files with metadata. @@ -262,7 +272,7 @@ def get_model_files(force_rescan: bool = False) -> List[Dict[str, str]]: force_rescan: If True, bypass the short-lived cache and rescan directories Returns: - List of model dictionaries (same format as scan_directory) + Typed local model records (same fields as scan_directory) """ global _MODEL_FILES_CACHE, _MODEL_FILES_CACHE_AT @@ -309,10 +319,10 @@ def find_local_file_path(filename: str, category: Optional[str] = None) -> Optio available_models = get_model_files() for m in available_models: if ( - m.get("relative_path", "").endswith(filename) - or m.get("filename", "") == filename + m.relative_path.endswith(filename) + or m.filename == filename ): - file_path = m.get("path") + file_path = m.path break except Exception: pass diff --git a/core/services/civitai_search_service.py b/core/services/civitai_search_service.py index d4b7286..beac652 100644 --- a/core/services/civitai_search_service.py +++ b/core/services/civitai_search_service.py @@ -1,8 +1,16 @@ """CivitAI model search service.""" +from typing import Any, Dict + from .. import path_utils +from ..contracts import SearchResult from ..metadata_utils import build_metadata_source_payload -from ..request_utils import extract_request_sha256 +from ..request_utils import ( + extract_request_sha256, + read_bool_field, + read_optional_object_payload, + read_text_field, +) from ..routes.context import RouteContext from ..sources.common import collect_download_urls from .model_utils import CivitAISearchDependencies, ModelServiceDependencies @@ -31,27 +39,60 @@ class CivitAISearchService(ModelServiceDependencies): request_public_url = self.request_public_url resolve_civarchive_by_hash = self.resolve_civarchive_by_hash search_huggingface_for_file = self.search_huggingface_for_file - to_bool = self.to_bool web = self.web write_model_resolver_metadata = self.write_model_resolver_metadata - data = await request.json() - filename = data.get("filename", "") - category = data.get("category", "") - resolved_path = data.get("resolved_path", "") - local_only = to_bool(data.get("local_only"), False) - force_refresh = to_bool( - data.get("force_refresh") or data.get("force"), False - ) - provided_hash = extract_request_sha256( - data, - keys=("sha256", "hash", "file_hash"), - ) - hf_token = data.get("hf_token", "") - brave_search_api_key = data.get("brave_search_api_key", "") - hf_use_brave_fallback = to_bool( - data.get("hf_use_brave_fallback", True), - True, - ) + data = await read_optional_object_payload(request) + try: + filename = read_text_field( + data, + "filename", + contract_name="CivitAI search request", + ) + category = read_text_field( + data, + "category", + contract_name="CivitAI search request", + ) + resolved_path = read_text_field( + data, + "resolved_path", + contract_name="CivitAI search request", + ) + hf_token = read_text_field( + data, + "hf_token", + contract_name="CivitAI search request", + ) + brave_search_api_key = read_text_field( + data, + "brave_search_api_key", + contract_name="CivitAI search request", + ) + local_only = read_bool_field( + data, + "local_only", + contract_name="CivitAI search request", + ) + force_value = dict(data) + if force_value.get("force_refresh") is None: + force_value["force_refresh"] = force_value.get("force") + force_refresh = read_bool_field( + force_value, + "force_refresh", + contract_name="CivitAI search request", + ) + hf_use_brave_fallback = read_bool_field( + data, + "hf_use_brave_fallback", + default=True, + contract_name="CivitAI search request", + ) + provided_hash = extract_request_sha256( + data, + keys=("sha256", "hash", "file_hash"), + ) + except TypeError as exc: + return web.json_response({"error": str(exc)}, status=400) if not filename: return web.json_response( @@ -292,6 +333,18 @@ class CivitAISearchService(ModelServiceDependencies): return False return not require_filename or result_filename_matches(result) + def serialize_provider_result( + result: SearchResult | None, + ) -> Dict[str, Any]: + """Cross from typed provider results into metadata response data.""" + if isinstance(result, SearchResult): + return result.to_dict() + if result is not None: + raise TypeError( + f"{type(result).__name__} is not a typed provider result" + ) + return {} + def huggingface_page_url(result): try: from urllib.parse import quote as _quote @@ -310,7 +363,12 @@ class CivitAISearchService(ModelServiceDependencies): return f"https://huggingface.co/{repo_id}/blob/main/{hf_path}" def prepare_remote_result(result, source_name): - result = dict(result or {}) + if isinstance(result, SearchResult): + result = result.to_dict() + elif not isinstance(result, dict): + raise TypeError( + f"{type(result).__name__} is not a remote result payload" + ) source_name = str(source_name or result.get("source") or "").lower() if source_name == "huggingface": download_url = result.get("download_url") or result.get("url") @@ -690,11 +748,13 @@ class CivitAISearchService(ModelServiceDependencies): if provided_hash: try: - result = resolve_civarchive_by_hash( - provided_hash, - query=filename, - exact_only=False, - model_type=infer_model_type_from_category(category), + result = serialize_provider_result( + resolve_civarchive_by_hash( + provided_hash, + query=filename, + exact_only=False, + model_type=infer_model_type_from_category(category), + ) ) if result: if not extract_result_sha256(result): @@ -742,19 +802,21 @@ class CivitAISearchService(ModelServiceDependencies): for hf_attempt in hf_attempts: try: - result = search_huggingface_for_file( - filename, - token=hf_token or None, - exact_only=True, - brave_api_key=brave_search_api_key or None, - use_api_search=hf_attempt["use_api_search"], - use_comfy_org_fallback=hf_attempt[ - "use_comfy_org_fallback" - ], - use_brave_fallback=hf_attempt[ - "use_brave_fallback" - ], - force_refresh=force_refresh, + result = serialize_provider_result( + search_huggingface_for_file( + filename, + token=hf_token or None, + exact_only=True, + brave_api_key=brave_search_api_key or None, + use_api_search=hf_attempt["use_api_search"], + use_comfy_org_fallback=hf_attempt[ + "use_comfy_org_fallback" + ], + use_brave_fallback=hf_attempt[ + "use_brave_fallback" + ], + force_refresh=force_refresh, + ) ) if result and result_hash_matches( result, diff --git a/core/services/custom_url_service.py b/core/services/custom_url_service.py index 8cae865..453aea9 100644 --- a/core/services/custom_url_service.py +++ b/core/services/custom_url_service.py @@ -1,6 +1,12 @@ """Custom model URL service.""" +from ..contracts import SearchResult from ..local_hash_matches import collect_local_hash_matches_for_result +from ..request_utils import ( + read_first_text_field, + read_optional_object_payload, + read_text_field, +) from ..routes.context import RouteContext from .model_utils import CustomUrlDependencies, ModelServiceDependencies @@ -40,6 +46,10 @@ class CustomUrlService(ModelServiceDependencies): def _custom_result_timestamp(): return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) + def _as_search_result(value, source=""): + if isinstance(value, SearchResult): + return value + return None @@ -80,8 +90,11 @@ class CustomUrlService(ModelServiceDependencies): ) def _collect_custom_url_local_hash_matches(result, category): - source_key = str(result.get("source") or "custom").strip().lower() - sha256 = normalize_sha256(extract_sha256_from_metadata(result)) + result_mapping = result.to_dict() + source_key = str(result.source or "custom").strip().lower() + sha256 = normalize_sha256( + extract_sha256_from_metadata(result_mapping) + ) if not sha256: return [] @@ -92,8 +105,8 @@ class CustomUrlService(ModelServiceDependencies): category=category or None, max_matches=20, source=source_key, - filename=result.get("filename") - or result.get("path") + filename=result.filename + or result.extra_value("path") or "", ) except Exception as hash_error: @@ -102,8 +115,40 @@ class CustomUrlService(ModelServiceDependencies): ) return [] - data = await request.json() - raw_url = str(data.get("url") or data.get("custom_url") or "").strip() + data = await read_optional_object_payload(request) + try: + raw_url = read_first_text_field( + data, + ("url", "custom_url"), + contract_name="Custom URL request", + ) + category = read_text_field( + data, + "category", + contract_name="Custom URL request", + ) + filename = read_text_field( + data, + "filename", + contract_name="Custom URL request", + ) + original_path = read_text_field( + data, + "original_path", + contract_name="Custom URL request", + ) + civitai_key = read_text_field( + data, + "civitai_key", + contract_name="Custom URL request", + ) + hf_token = read_text_field( + data, + "hf_token", + contract_name="Custom URL request", + ) + except TypeError as exc: + return web.json_response({"error": str(exc)}, status=400) if not raw_url: return web.json_response( {"error": "URL is required"}, status=400 @@ -123,14 +168,11 @@ class CustomUrlService(ModelServiceDependencies): status=400, ) - category = data.get("category") or "" expected_filename = ( - data.get("filename") - or get_filename_from_path(data.get("original_path") or "") + filename + or get_filename_from_path(original_path) or "" ) - civitai_key = data.get("civitai_key", "") - hf_token = data.get("hf_token", "") result = None source = "" @@ -145,8 +187,8 @@ class CustomUrlService(ModelServiceDependencies): if civitai_parsed: source = "civitai" - model_id = civitai_parsed.get("model_id") - version_id = civitai_parsed.get("version_id") + model_id = civitai_parsed.model_id + version_id = civitai_parsed.version_id if model_id: details = await asyncio.to_thread( get_civitai_model_details, @@ -185,17 +227,17 @@ class CustomUrlService(ModelServiceDependencies): ) elif civarchive_parsed: source = "civarchive" - if civarchive_parsed.get("sha256"): + if civarchive_parsed.sha256: result = await asyncio.to_thread( resolve_civarchive_by_hash, - civarchive_parsed.get("sha256"), + civarchive_parsed.sha256, expected_filename, False, normalize_category_to_model_type(category), ) else: - model_id = civarchive_parsed.get("model_id") - version_id = civarchive_parsed.get("version_id") + model_id = civarchive_parsed.model_id + version_id = civarchive_parsed.version_id if model_id: result = await asyncio.to_thread( resolve_civarchive_model_version, @@ -217,11 +259,14 @@ class CustomUrlService(ModelServiceDependencies): expected_filename=expected_filename, ) if result: - result = dict(result) - result["source"] = "civarchive" - result["details_source"] = "civarchive" - result["match_type"] = "custom_url" - result["custom_url"] = True + result = _as_search_result(result, source="civarchive") + if result is not None: + result = result.with_updates( + source="civarchive", + details_source="civarchive", + match_type="custom_url", + custom_url=True, + ) else: from urllib.parse import urlparse @@ -258,32 +303,44 @@ class CustomUrlService(ModelServiceDependencies): status=400, ) - result = dict(result) - result["provided_url"] = normalized_url - result["url_source"] = "custom" - result["searched_at"] = result.get("searched_at") or _custom_result_timestamp() - result.setdefault("category", category) - if expected_filename and not result.get("filename"): - result["filename"] = expected_filename + result = _as_search_result(result, source=source) + if result is None: + return web.json_response( + {"error": "The URL resolved to an invalid model result"}, + status=400, + ) - if not (result.get("download_url") or result.get("url")): + result = result.with_extra( + provided_url=normalized_url, + url_source="custom", + searched_at=result.extra_value("searched_at") + or _custom_result_timestamp(), + category=result.extra_value("category") or category, + ) + if expected_filename and not result.filename: + result = result.with_updates(filename=expected_filename) + + if not (result.download_url or result.url): return web.json_response( {"error": "The URL resolved, but no download URL was found"}, status=400, ) - source = source or result.get("source") or "custom" + source = source or result.source or "custom" local_hash_matches = _collect_custom_url_local_hash_matches( result, category, ) + result_mapping = result.to_dict() response = { "success": True, "source": source, - "result": result, - "custom": [result], + "result": result_mapping, + "custom": [result_mapping], "searched_sources": ["custom"], - "local_hash_matches": local_hash_matches, + "local_hash_matches": [ + match.to_dict() for match in local_hash_matches + ], } - response[source] = result + response[source] = result_mapping return web.json_response(response) diff --git a/core/services/directory_service.py b/core/services/directory_service.py index 2f444a6..a258ffc 100644 --- a/core/services/directory_service.py +++ b/core/services/directory_service.py @@ -2,6 +2,7 @@ import os +from ..path_utils import normalize_folder_path_values, normalize_string_values from ..routes.context import RouteContext @@ -99,7 +100,11 @@ class DirectoryService: candidate_keys.append(candidate_key) paths = [] for candidate_key in candidate_keys: - paths.extend(folder_paths.get_folder_paths(candidate_key) or []) + paths.extend( + normalize_folder_path_values( + folder_paths.get_folder_paths(candidate_key) + ) + ) if folder_key == "ultralytics": normalized_ultralytics_paths = [] for path in paths: @@ -134,7 +139,9 @@ class DirectoryService: or raw_key in roots ): continue - raw_paths = folder_paths.get_folder_paths(raw_key) or [] + raw_paths = normalize_folder_path_values( + folder_paths.get_folder_paths(raw_key) + ) roots[raw_key] = self.dedupe_local_base_directories( raw_paths, comfy_root=comfy_root, @@ -220,7 +227,9 @@ class DirectoryService: raw_base_dirs = [] for folder_key in available_folder_keys: - for base_dir in folder_paths.get_folder_paths(folder_key) or []: + for base_dir in normalize_folder_path_values( + folder_paths.get_folder_paths(folder_key) + ): if not base_dir or not os.path.isdir(base_dir): continue raw_base_dirs.append(base_dir) @@ -239,7 +248,9 @@ class DirectoryService: return "" for folder_key in available_folder_keys: - filenames = folder_paths.get_filename_list(folder_key) or [] + filenames = normalize_string_values( + folder_paths.get_filename_list(folder_key) + ) for rel_path in filenames: if not isinstance(rel_path, str): continue diff --git a/core/services/download_service.py b/core/services/download_service.py index 506a7dc..f12b651 100644 --- a/core/services/download_service.py +++ b/core/services/download_service.py @@ -3,7 +3,15 @@ import os from urllib.parse import unquote, urlparse -from ..request_utils import read_optional_object_payload +from ..contracts import DownloadSpec +from ..request_utils import ( + coerce_integer_identifier, + read_bool_field, + read_first_identifier_field, + read_first_text_field, + read_optional_object_payload, + read_text_field, +) from ..routes.context import RouteContext @@ -18,7 +26,6 @@ class DownloadService: self.clear_completed_downloads_fn = context.require( "clear_completed_downloads" ) - self.first_non_empty = context.require("first_non_empty") self.get_all_progress_fn = context.require("get_all_progress") self.get_aria2_status_fn = context.require("get_aria2_status") self.get_civarchive_model_details = context.require( @@ -57,8 +64,6 @@ class DownloadService: self.start_aria2_daemon = context.require("start_aria2_daemon") self.start_background_download = context.require("start_background_download") self.stop_aria2_daemon = context.require("stop_aria2_daemon") - self.to_bool = context.require("to_bool") - self.to_int = context.require("to_int") self.validate_public_http_url = context.require("validate_public_http_url") self.web = context.require("web") @@ -68,23 +73,122 @@ class DownloadService: async def download_model(self, request): """Validate and start a background model download.""" - data = await request.json() - url = data.get("url", "") - filename = data.get("filename", "") + data = await read_optional_object_payload(request) + try: + url = read_text_field( + data, + "url", + contract_name="Download request", + ) + filename = read_text_field( + data, + "filename", + contract_name="Download request", + ) + category_value = read_text_field( + data, + "category", + default="checkpoints", + contract_name="Download request", + ) + subfolder = read_text_field( + data, + "subfolder", + contract_name="Download request", + ) + base_directory = read_text_field( + data, + "base_directory", + contract_name="Download request", + ) + hf_token = read_text_field( + data, + "hf_token", + contract_name="Download request", + ) + civitai_key = read_text_field( + data, + "civitai_key", + contract_name="Download request", + ) + civitai_session_token = read_text_field( + data, + "civitai_session_token", + contract_name="Download request", + ) + except TypeError as exc: + return self.web.json_response({"error": str(exc)}, status=400) + category = self.normalize_download_category( - data.get("category", "checkpoints") + category_value ) - subfolder = data.get("subfolder", "") - base_directory = data.get("base_directory", "") path_metadata = data.get("path_metadata", {}) - if not isinstance(path_metadata, dict): + if path_metadata is None: path_metadata = {} - download_metadata = data.get("download_metadata") or data.get( - "metadata", {} - ) - if not isinstance(download_metadata, dict): + elif not isinstance(path_metadata, dict): + return self.web.json_response( + {"error": "Download path metadata must be an object"}, + status=400, + ) + download_metadata = data.get("download_metadata") + if download_metadata is None: + download_metadata = data.get("metadata") + if download_metadata is None: download_metadata = {} + if not isinstance(download_metadata, dict): + return self.web.json_response( + {"error": "Download metadata must be an object"}, + status=400, + ) download_metadata = dict(download_metadata) + try: + metadata_source = read_first_text_field( + download_metadata, + ("details_source", "source"), + contract_name="Download metadata", + ) + if not metadata_source: + metadata_source = read_first_text_field( + path_metadata, + ("source",), + contract_name="Download path metadata", + ) + + raw_model_id = read_first_identifier_field( + download_metadata, + ("model_id", "modelId"), + contract_name="Download metadata", + ) + if raw_model_id is None: + raw_model_id = read_first_identifier_field( + path_metadata, + ("model_id",), + contract_name="Download path metadata", + ) + raw_version_id = read_first_identifier_field( + download_metadata, + ("version_id", "versionId"), + contract_name="Download metadata", + ) + if raw_version_id is None: + raw_version_id = read_first_identifier_field( + path_metadata, + ("version_id",), + contract_name="Download path metadata", + ) + + model_id = coerce_integer_identifier( + raw_model_id, + "model_id", + contract_name="Download metadata", + ) + version_id = coerce_integer_identifier( + raw_version_id, + "version_id", + contract_name="Download metadata", + ) + except (TypeError, ValueError) as exc: + return self.web.json_response({"error": str(exc)}, status=400) settings = self.load_resolver_settings() if not base_directory: base_directory = self.get_default_root_for_category(category, settings) @@ -119,7 +223,6 @@ class DownloadService: headers = {} if self.host_matches_domain(download_host, "huggingface.co"): - hf_token = data.get("hf_token", "") if hf_token: headers["Authorization"] = f"Bearer {hf_token}" elif self.host_matches_domain( @@ -127,12 +230,8 @@ class DownloadService: "civitai.com", "civitai.red", ): - civitai_key = data.get("civitai_key", "") if civitai_key and "token=" not in url: url += f"{'&' if '?' in url else '?'}token={civitai_key}" - civitai_session_token = str( - data.get("civitai_session_token", "") or "" - ).strip() if civitai_session_token: headers["Cookie"] = ( f"__Secure-civitai-token={civitai_session_token}" @@ -151,33 +250,9 @@ class DownloadService: download_metadata.setdefault("path_metadata", path_metadata) download_metadata.setdefault( "source", - self.first_non_empty( - download_metadata.get("details_source"), - path_metadata.get("source"), - inferred_source, - ), + metadata_source or inferred_source, ) - - model_id = self.to_int( - self.first_non_empty( - download_metadata.get("model_id"), - download_metadata.get("modelId"), - path_metadata.get("model_id"), - ) - ) - version_id = self.to_int( - self.first_non_empty( - download_metadata.get("version_id"), - download_metadata.get("versionId"), - path_metadata.get("version_id"), - ) - ) - source_name = str( - self.first_non_empty( - download_metadata.get("details_source"), - download_metadata.get("source"), - ) - ).lower() + source_name = (metadata_source or inferred_source).lower() try: if ( source_name == "civitai" @@ -188,7 +263,7 @@ class DownloadService: self.get_civitai_model_details, model_id, version_id, - data.get("civitai_key", ""), + civitai_key, ) if details: download_metadata["civitai_details"] = details @@ -227,15 +302,21 @@ class DownloadService: target_directory = "" target_path = "" - download_id = self.start_background_download( - url=url, - filename=filename, - category=category, - headers=headers if headers else None, - subfolder=subfolder, - base_directory=base_directory, - metadata=download_metadata, - ) + try: + download_spec = DownloadSpec( + url=url, + filename=filename, + category=category, + headers=headers, + subfolder=subfolder, + base_directory=base_directory, + metadata=download_metadata, + expected_sha256=download_metadata.get("sha256"), + ) + except (TypeError, ValueError) as exc: + return self.web.json_response({"error": str(exc)}, status=400) + + download_id = self.start_background_download(**download_spec.to_kwargs()) return self.web.json_response( { @@ -291,13 +372,19 @@ class DownloadService: async def aria2_status(self, request): """Report aria2 availability using optional request settings.""" - settings = await self.get_override_settings_from_request(request) + try: + settings = await self.get_override_settings_from_request(request) + except TypeError as exc: + return self.web.json_response({"error": str(exc)}, status=400) result = await self.asyncio.to_thread(self.get_aria2_status_fn, settings) return self.web.json_response(result) async def aria2_start(self, request): """Start the aria2 daemon without starting a download.""" - settings = await self.get_override_settings_from_request(request) + try: + settings = await self.get_override_settings_from_request(request) + except TypeError as exc: + return self.web.json_response({"error": str(exc)}, status=400) result = await self.asyncio.to_thread(self.start_aria2_daemon, settings) return self._json_result_with_success_status(result) @@ -309,7 +396,14 @@ class DownloadService: async def aria2_install(self, request): """Install aria2 and persist the selected download backend.""" payload = await read_optional_object_payload(request) - force = self.to_bool(payload.get("force"), False) + try: + force = read_bool_field( + payload, + "force", + contract_name="Aria2 install request", + ) + except TypeError as exc: + return self.web.json_response({"error": str(exc)}, status=400) try: install_result = await self.asyncio.to_thread( self.install_aria2_engine, diff --git a/core/services/hash_service.py b/core/services/hash_service.py index 9f5faf6..3ca9c9b 100644 --- a/core/services/hash_service.py +++ b/core/services/hash_service.py @@ -4,9 +4,18 @@ import os import time from .. import path_utils +from ..contracts import ModelReference, ResolvedModel from ..local_hash_matches import collect_local_hash_matches_for_result from ..progress import generate_progress_id -from ..request_utils import extract_request_sha256, validate_workflow_payload +from ..request_utils import ( + extract_request_sha256, + read_bool_field, + read_first_text_field, + read_int_field, + read_optional_object_payload, + read_text_field, + validate_workflow_payload, +) from ..routes.context import RouteContext @@ -59,8 +68,6 @@ class HashService: ) self.extension = context.require("self") self.time = context.get("time") or time - self.to_bool = context.require("to_bool") - self.to_int = context.require("to_int") self.web = context.require("web") self.write_json_atomic = context.require("write_json_atomic") @@ -84,16 +91,22 @@ class HashService: async def local_model_hashes(self, request): """Return SHA256 hashes already stored in local sidecar metadata.""" - data = await request.json() + data = await read_optional_object_payload(request) model = data.get("model") if isinstance(data.get("model"), dict) else {} - path = ( - data.get("path") - or data.get("file_path") - or data.get("resolved_path") - or model.get("path") - or model.get("resolved_path") - or "" - ) + try: + path = read_first_text_field( + data, + ("path", "file_path", "resolved_path"), + contract_name="Local model hashes request", + ) + if not path: + path = read_first_text_field( + model, + ("path", "resolved_path"), + contract_name="Local model hashes model", + ) + except TypeError as exc: + return self.web.json_response({"error": str(exc)}, status=400) if not path: return self.web.json_response( @@ -113,7 +126,14 @@ class HashService: async def get_model_preview(self, request): """Serve an adjacent model preview from a configured model directory.""" - model_path = str(request.query.get("path") or "").strip() + try: + model_path = read_text_field( + request.query, + "path", + contract_name="Model preview request", + ) + except TypeError as exc: + return self.web.Response(text=str(exc), status=400) is_preview_probe = str(getattr(request, "method", "")).upper() == "HEAD" if not model_path: return self.web.Response(text="path is required", status=400) @@ -153,7 +173,7 @@ class HashService: async def workflow_model_hashes(self, request): """Return hash metadata for existing local models used by a workflow.""" - data = await request.json() + data = await read_optional_object_payload(request) workflow_json, workflow_error = validate_workflow_payload( data.get("workflow"), none_is_missing=False, @@ -182,7 +202,16 @@ class HashService: self.get_workflow_model_inventory, workflow_json, ) - refs = inventory["model_refs"] + refs = [] + for raw_ref in inventory.model_refs: + try: + refs.append( + raw_ref + if isinstance(raw_ref, ModelReference) + else ModelReference.from_mapping(raw_ref) + ) + except (TypeError, ValueError): + continue by_node = {} by_path = {} @@ -190,17 +219,17 @@ class HashService: seen = set() for ref in refs: - if not isinstance(ref, dict) or not ref.get("exists"): + if not ref.exists: continue - full_path = str(ref.get("full_path") or "").strip() + full_path = ref.extra_value("full_path") or "" if not full_path: continue - model_info = { - "path": full_path, - "filename": self.get_filename_from_path(full_path), - "relative_path": ref.get("original_path") or "", - "category": ref.get("category") or "", - } + model_info = ResolvedModel( + path=full_path, + filename=self.get_filename_from_path(full_path), + relative_path=ref.original_path or "", + category=ref.category or "", + ) metadata = self.get_local_model_hash_metadata( full_path, model=model_info, @@ -210,15 +239,15 @@ class HashService: continue entry = { - "node_id": ref.get("node_id"), - "node_type": ref.get("node_type") or "", - "widget_index": ref.get("widget_index"), - "widget_name": ref.get("widget_name") or "", - "path": ref.get("original_path") or "", + "node_id": ref.node_id, + "node_type": ref.node_type or "", + "widget_index": ref.widget_index, + "widget_name": ref.extra_value("widget_name") or "", + "path": ref.original_path or "", "filename": self.get_filename_from_path( - ref.get("original_path") or full_path + ref.original_path or full_path ), - "category": ref.get("category") or "", + "category": ref.category or "", "sha256": sha256, "size": metadata.get("size") or 0, } @@ -253,30 +282,52 @@ class HashService: async def local_matches_by_hash(self, request): """Search local model metadata sidecars for a remote SHA256.""" - data = await request.json() - sha256 = extract_request_sha256( - data, - keys=("sha256", "hash", "SHA256"), - ) + data = await read_optional_object_payload(request) + try: + sha256 = extract_request_sha256( + data, + keys=("sha256", "hash", "SHA256"), + ) + except TypeError as exc: + return self.web.json_response({"error": str(exc)}, status=400) if not sha256: return self.web.json_response( {"error": "sha256 is required"}, status=400 ) - category = data.get("category", "") - source = str( - data.get("source") - or data.get("hash_lookup_source") - or "download_source" - ).strip().lower().replace("-", "_") - filename = ( - data.get("filename") - or data.get("path") - or data.get("model_name") - or "" - ) - max_matches = self.to_int(data.get("max_matches"), 20) - force_rescan = self.to_bool(data.get("force_rescan"), False) + try: + category = read_text_field( + data, + "category", + contract_name="Local hash matches request", + ) + source = read_first_text_field( + data, + ("source", "hash_lookup_source"), + default="download_source", + contract_name="Local hash matches request", + ).lower().replace("-", "_") + filename = read_first_text_field( + data, + ("filename", "path", "model_name"), + contract_name="Local hash matches request", + ) + except TypeError as exc: + return self.web.json_response({"error": str(exc)}, status=400) + try: + max_matches = read_int_field( + data, + "max_matches", + default=20, + contract_name="Local hash matches request", + ) + force_rescan = read_bool_field( + data, + "force_rescan", + contract_name="Local hash matches request", + ) + except (TypeError, ValueError) as exc: + return self.web.json_response({"error": str(exc)}, status=400) enriched_matches = collect_local_hash_matches_for_result( sha256, @@ -290,8 +341,10 @@ class HashService: return self.web.json_response( { "sha256": sha256, - "local_hash_matches": enriched_matches, - "matches": enriched_matches, + "local_hash_matches": [ + match.to_dict() for match in enriched_matches + ], + "matches": [match.to_dict() for match in enriched_matches], } ) @@ -309,7 +362,17 @@ class HashService: {"success": False, "error": "JSON body must be an object"}, status=400, ) - target_path = data.get("path", "") + try: + target_path = read_text_field( + data, + "path", + contract_name="Open containing folder request", + ) + except TypeError as exc: + return self.web.json_response( + {"success": False, "error": str(exc)}, + status=400, + ) try: normalized_path = self.normalize_file_manager_path(target_path) @@ -366,11 +429,10 @@ class HashService: def resolve_hash_file_request(self, data): """Validate and normalize a file path accepted by hash endpoints.""" - file_path = ( - data.get("file_path") - or data.get("resolved_path") - or data.get("path") - or "" + file_path = read_first_text_field( + data, + ("file_path", "resolved_path", "path"), + contract_name="Hash file request", ) if not file_path: return "", "file_path is required" @@ -387,7 +449,7 @@ class HashService: async def _get_validated_hash_file_request(self, request): """Read and validate the file request shared by hash endpoints.""" - data = await request.json() + data = await read_optional_object_payload(request) normalized_path, error = self.resolve_hash_file_request(data) if not error: return normalized_path, None diff --git a/core/services/loaded_models_service.py b/core/services/loaded_models_service.py index 7721efa..92a6475 100644 --- a/core/services/loaded_models_service.py +++ b/core/services/loaded_models_service.py @@ -1,6 +1,13 @@ """Loaded model inspection used by the HTTP route adapter.""" -from ..request_utils import validate_workflow_payload +from ..contracts import ModelReference +from ..metadata_model_utils import normalize_models +from ..path_utils import normalize_string_values +from ..request_utils import ( + read_first_text_field, + read_optional_object_payload, + validate_workflow_payload, +) from ..routes.context import RouteContext from ..workflow.traversal import iter_workflow_nodes_with_scope @@ -23,14 +30,19 @@ class LoadedModelsService: async def get_loaded_models(self, request): """Get all currently loaded models in the workflow.""" - data = await request.json() + data = await read_optional_object_payload(request) workflow_json, workflow_error = validate_workflow_payload( data.get("workflow"), empty_is_missing=True, ) - loaded_id = str( - data.get("loaded_id") or data.get("progress_id") or "" - ).strip() + try: + loaded_id = read_first_text_field( + data, + ("loaded_id", "progress_id"), + contract_name="Loaded models request", + ) + except TypeError as exc: + return self.web.json_response({"error": str(exc)}, status=400) if workflow_error: return self.web.json_response( @@ -80,8 +92,15 @@ class LoadedModelsService: workflow_json, progress_callback=update_workflow_analysis_progress, ) - available_models = inventory["available_models"] - all_model_refs = inventory["model_refs"] + available_models = inventory.available_models + all_model_refs = inventory.model_refs + available_models = normalize_models(available_models) + all_model_refs = [ + ref + if isinstance(ref, ModelReference) + else ModelReference.from_mapping(ref) + for ref in all_model_refs + ] # Create lookup for full paths by filename (with and without extension) path_by_filename = {} @@ -94,19 +113,17 @@ class LoadedModelsService: total=total_local_models, ) for index, model_info in enumerate(available_models, start=1): - rel_path = model_info.get("relative_path", "") + rel_path = model_info.relative_path if rel_path: filename = self.get_filename_from_path(rel_path) - path_by_filename[filename] = model_info.get("path") + path_by_filename[filename] = model_info.path # Also add without extension for matching (simple approach) if "." in filename: filename_no_ext = filename.rsplit(".", 1)[0] if filename_no_ext not in path_by_filename: - path_by_filename[filename_no_ext] = model_info.get( - "path" - ) + path_by_filename[filename_no_ext] = model_info.path # Add the full relative path as key too - path_by_filename[rel_path] = model_info.get("path") + path_by_filename[rel_path] = model_info.path if ( total_local_models @@ -147,7 +164,9 @@ class LoadedModelsService: total=len(folder_categories), ) try: - filenames = folder_paths.get_filename_list(category_name) + filenames = normalize_string_values( + folder_paths.get_filename_list(category_name) + ) for filename in filenames: full_path = folder_paths.get_full_path( category_name, filename @@ -192,20 +211,18 @@ class LoadedModelsService: }, ) - def node_matches_ref(node, ref): + def node_matches_ref(node, ref: ModelReference): if not isinstance(node, dict): return False - if str(node.get("id")) != str(ref.get("node_id")): + if str(node.get("id")) != str(ref.node_id): return False scope = get_node_scope(node) if scope["is_top_level"] != ( - ref.get("is_top_level") is not False + ref.is_top_level is not False ): return False if not scope["is_top_level"]: - return str(scope["subgraph_id"]) == str( - ref.get("subgraph_id") or "" - ) + return str(scope["subgraph_id"]) == str(ref.subgraph_id or "") return True # Collect all loaded models with their values @@ -232,18 +249,18 @@ class LoadedModelsService: total=total_refs, ) - original_path = ref.get("original_path", "") - node_id = ref.get("node_id") - widget_index = ref.get("widget_index") - node_type = ref.get("node_type", "") - category = ref.get("category", "unknown") + original_path = ref.original_path + node_id = ref.node_id + widget_index = ref.widget_index + node_type = ref.node_type + category = ref.category or "unknown" # Determine model name and strength model_name = self.get_filename_from_path(original_path) strength = None if ( - ref.get("strength") is None + ref.extra_value("strength") is None and category in {"lora", "loras"} ): for node in nodes: @@ -254,8 +271,8 @@ class LoadedModelsService: ) break - if ref.get("strength") is not None: - strength = ref.get("strength") + if ref.extra_value("strength") is not None: + strength = ref.extra_value("strength") model_name, strength = self.adapt_custom_node_loaded_model( ref, @@ -264,11 +281,11 @@ class LoadedModelsService: ) # Check if model exists locally - exists = ref.get("exists", False) + exists = ref.exists # If URN, resolve to display name - if ref.get("is_urn"): - urn = ref.get("urn", {}) + if ref.extra_value("is_urn"): + urn = ref.extra_value("urn") or {} # Use model name from URN as display name model_name = ( f"urn:{urn.get('type', 'model')}:{urn.get('model_id')}" @@ -284,34 +301,34 @@ class LoadedModelsService: "node_id": node_id, "widget_index": widget_index, "node_type": node_type, - "node_title": ref.get("node_title", ""), - "subgraph_id": ref.get("subgraph_id") or "", - "subgraph_name": ref.get("subgraph_name") or "", - "is_top_level": ref.get("is_top_level") is not False, - "locate_node_id": ref.get("locate_node_id"), - "locate_node_type": ref.get("locate_node_type", ""), - "locate_node_title": ref.get("locate_node_title", ""), + "node_title": ref.extra_value("node_title", ""), + "subgraph_id": ref.subgraph_id or "", + "subgraph_name": ref.extra_value("subgraph_name") or "", + "is_top_level": ref.is_top_level is not False, + "locate_node_id": ref.extra_value("locate_node_id"), + "locate_node_type": ref.extra_value("locate_node_type", ""), + "locate_node_title": ref.extra_value("locate_node_title", ""), "locate_subgraph_id": ( - ref.get("locate_subgraph_id") or "" + ref.extra_value("locate_subgraph_id") or "" ), "locate_subgraph_name": ( - ref.get("locate_subgraph_name") or "" + ref.extra_value("locate_subgraph_name") or "" ), "locate_is_top_level": ( - ref.get("locate_is_top_level") is not False + ref.extra_value("locate_is_top_level") is not False ), - "locate_via_promoted_widget": ref.get( + "locate_via_promoted_widget": ref.extra_value( "locate_via_promoted_widget", False ), "exists": exists, "strength": strength, "original_path": original_path, - "is_urn": ref.get("is_urn", False), - "custom_node_adapter": ref.get( + "is_urn": ref.extra_value("is_urn", False), + "custom_node_adapter": ref.extra_value( "custom_node_adapter" ), - "active": ref.get("active"), - "connected": ref.get("connected", True), + "active": ref.extra_value("active"), + "connected": ref.extra_value("connected", True), "resolved_path": ( path_by_filename.get(model_name) or path_by_filename.get(original_path) diff --git a/core/services/metadata_service.py b/core/services/metadata_service.py index 68015c3..02256c5 100644 --- a/core/services/metadata_service.py +++ b/core/services/metadata_service.py @@ -1,7 +1,12 @@ """Metadata audit and build orchestration used by HTTP route adapters.""" from ..progress import generate_progress_id -from ..request_utils import read_optional_object_payload +from ..request_utils import ( + read_bool_field, + read_int_field, + read_optional_object_payload, + read_optional_text_field, +) from ..routes.context import RouteContext @@ -26,8 +31,6 @@ class MetadataService: ) self.run_in_background_thread = context.require("run_in_background_thread") self.extension = context.require("self") - self.to_bool = context.require("to_bool") - self.to_int = context.require("to_int") self.web = context.require("web") @property @@ -37,13 +40,32 @@ class MetadataService: async def metadata_size_audit(self, request): """Check local metadata sidecars for stale file-size values.""" payload = await read_optional_object_payload(request) - - force_rescan = self.to_bool(payload.get("force_rescan"), True) + try: + force_rescan = read_bool_field( + payload, + "force_rescan", + default=True, + contract_name="Metadata size audit request", + ) + worker_count = read_int_field( + payload, + "worker_count", + default=None, + contract_name="Metadata size audit request", + ) + batch_size = read_int_field( + payload, + "batch_size", + default=None, + contract_name="Metadata size audit request", + ) + except (TypeError, ValueError) as exc: + return self.web.json_response({"error": str(exc)}, status=400) result = await self.asyncio.to_thread( self.audit_metadata_sizes_fn, force_rescan=force_rescan, - worker_count=payload.get("worker_count"), - batch_size=payload.get("batch_size"), + worker_count=worker_count, + batch_size=batch_size, ) return self.web.json_response(result) @@ -56,12 +78,28 @@ class MetadataService: async def metadata_build_start(self, request): """Start building missing local metadata sidecars in the background.""" payload = await read_optional_object_payload(request) - - force_rescan = self.to_bool(payload.get("force_rescan"), True) - worker_count = self.to_int(payload.get("worker_count"), 0) - metadata_mode = self.normalize_metadata_build_mode( - payload.get("metadata_mode") - ) + try: + force_rescan = read_bool_field( + payload, + "force_rescan", + default=True, + contract_name="Metadata build request", + ) + worker_count = read_int_field( + payload, + "worker_count", + default=0, + contract_name="Metadata build request", + ) + metadata_mode = self.normalize_metadata_build_mode( + read_optional_text_field( + payload, + "metadata_mode", + contract_name="Metadata build request", + ) + ) + except (TypeError, ValueError) as exc: + return self.web.json_response({"error": str(exc)}, status=400) self.metadata_builder_progress.cleanup() progress_id = generate_progress_id("metadata_build") self.metadata_builder_progress.update( diff --git a/core/services/model_details_service.py b/core/services/model_details_service.py index 274f900..e8c06fa 100644 --- a/core/services/model_details_service.py +++ b/core/services/model_details_service.py @@ -1,5 +1,11 @@ """Model details service.""" +from ..request_utils import ( + coerce_integer_identifier, + read_identifier_field, + read_optional_object_payload, + read_text_field, +) from ..routes.context import RouteContext from .model_utils import ModelDetailsDependencies, ModelServiceDependencies @@ -18,14 +24,45 @@ class ModelDetailsService(ModelServiceDependencies): get_civitai_model_details = self.get_civitai_model_details get_huggingface_model_details = self.get_huggingface_model_details web = self.web - data = await request.json() - source = str(data.get("source", "")).strip().lower() - model_id = data.get("model_id") - version_id = data.get("version_id") - civitai_key = data.get("civitai_key", "") - hf_token = data.get("hf_token", "") - file_path = data.get("file_path", "") - branch = data.get("branch", "") + data = await read_optional_object_payload(request) + try: + source = read_text_field( + data, + "source", + contract_name="Model details request", + ).lower() + model_id = read_identifier_field( + data, + "model_id", + contract_name="Model details request", + ) + version_id = read_identifier_field( + data, + "version_id", + contract_name="Model details request", + ) + civitai_key = read_text_field( + data, + "civitai_key", + contract_name="Model details request", + ) + hf_token = read_text_field( + data, + "hf_token", + contract_name="Model details request", + ) + file_path = read_text_field( + data, + "file_path", + contract_name="Model details request", + ) + branch = read_text_field( + data, + "branch", + contract_name="Model details request", + ) + except TypeError as exc: + return web.json_response({"error": str(exc)}, status=400) if not download_available: return web.json_response( @@ -41,27 +78,33 @@ class ModelDetailsService(ModelServiceDependencies): ) if source == "huggingface": - model_id = str(model_id or "").strip() - branch = str(branch or version_id or "main").strip() or "main" + if model_id is not None and not isinstance(model_id, str): + return web.json_response( + {"error": "Model details model_id must be a string for Hugging Face"}, + status=400, + ) + if version_id is not None and not isinstance(version_id, str): + return web.json_response( + {"error": "Model details version_id must be a string for Hugging Face"}, + status=400, + ) + model_id = model_id.strip() if isinstance(model_id, str) else "" + branch = branch or version_id or "main" version_id = branch else: try: - model_id = ( - int(model_id) - if model_id is not None and str(model_id).strip() - else None + model_id = coerce_integer_identifier( + model_id, + "model_id", + contract_name="Model details request", ) - except (TypeError, ValueError): - model_id = None - - try: - version_id = ( - int(version_id) - if version_id is not None and str(version_id).strip() - else None + version_id = coerce_integer_identifier( + version_id, + "version_id", + contract_name="Model details request", ) - except (TypeError, ValueError): - version_id = None + except (TypeError, ValueError) as exc: + return web.json_response({"error": str(exc)}, status=400) if not model_id: return web.json_response( diff --git a/core/services/model_utils.py b/core/services/model_utils.py index 7ea743d..93cf241 100644 --- a/core/services/model_utils.py +++ b/core/services/model_utils.py @@ -3,9 +3,11 @@ from dataclasses import dataclass, fields from typing import Any, Dict, Iterable, Optional, Protocol, Tuple +from ..contracts import ModelMatch, ProviderUrlReference, SearchResult from ..routes.context import RouteContext ModelResult = Optional[Dict[str, Any]] +SearchResultValue = Optional[SearchResult] class ExtractSha256Protocol(Protocol): @@ -47,7 +49,7 @@ class NormalizeSha256Protocol(Protocol): class ModelResultBuilderProtocol(Protocol): - def __call__(self, source: str, **fields: Any) -> Dict[str, Any]: ... + def __call__(self, source: str, **fields: Any) -> SearchResult: ... class ReadJsonProtocol(Protocol): @@ -76,7 +78,7 @@ class ResolveCivarchiveByHashProtocol(Protocol): query: str = "", exact_only: bool = False, model_type: Optional[str] = None, - ) -> ModelResult: ... + ) -> SearchResultValue: ... class SearchHuggingFaceProtocol(Protocol): @@ -91,11 +93,7 @@ class SearchHuggingFaceProtocol(Protocol): use_brave_fallback: bool = True, force_refresh: bool = False, progress_callback: Any = None, - ) -> ModelResult: ... - - -class ToBoolProtocol(Protocol): - def __call__(self, value: Any, default: bool = False) -> bool: ... + ) -> SearchResultValue: ... class WriteMetadataProtocol(Protocol): @@ -112,7 +110,7 @@ class WriteMetadataProtocol(Protocol): class BuildCivarchiveCustomResultProtocol(Protocol): def __call__( self, details: Dict[str, Any], expected_filename: str = "" - ) -> ModelResult: ... + ) -> SearchResultValue: ... class BuildCivitaiCustomResultProtocol(Protocol): @@ -121,7 +119,7 @@ class BuildCivitaiCustomResultProtocol(Protocol): details: Dict[str, Any], expected_filename: str = "", api_key: Optional[str] = None, - ) -> ModelResult: ... + ) -> SearchResultValue: ... class BuildHuggingFaceCustomResultProtocol(Protocol): @@ -130,7 +128,7 @@ class BuildHuggingFaceCustomResultProtocol(Protocol): url: str, expected_filename: str = "", token: Optional[str] = None, - ) -> ModelResult: ... + ) -> SearchResultValue: ... class GetCivarchiveModelDetailsProtocol(Protocol): @@ -170,7 +168,7 @@ class HostMatchesDomainProtocol(Protocol): class ParseProviderUrlProtocol(Protocol): - def __call__(self, url: str) -> Optional[Dict[str, Any]]: ... + def __call__(self, url: str) -> Optional[ProviderUrlReference]: ... class ResolveCivarchiveModelVersionProtocol(Protocol): @@ -181,7 +179,7 @@ class ResolveCivarchiveModelVersionProtocol(Protocol): query: str = "", exact_only: bool = False, prefer_page: bool = False, - ) -> ModelResult: ... + ) -> SearchResultValue: ... class ResolveCivitaiVersionProtocol(Protocol): @@ -190,7 +188,7 @@ class ResolveCivitaiVersionProtocol(Protocol): version_id: int, expected_filename: str = "", api_key: Optional[str] = None, - ) -> ModelResult: ... + ) -> SearchResultValue: ... class SearchLocalMatchesByHashProtocol(Protocol): @@ -200,7 +198,7 @@ class SearchLocalMatchesByHashProtocol(Protocol): category: Optional[str] = None, max_matches: int = 20, force_rescan: bool = False, - ) -> list[Dict[str, Any]]: ... + ) -> list[ModelMatch]: ... class ValidatePublicHttpUrlProtocol(Protocol): @@ -227,7 +225,6 @@ class CivitAISearchDependencies: request_public_url: RequestPublicUrlProtocol resolve_civarchive_by_hash: Optional[ResolveCivarchiveByHashProtocol] search_huggingface_for_file: Optional[SearchHuggingFaceProtocol] - to_bool: ToBoolProtocol web: Any write_model_resolver_metadata: Optional[WriteMetadataProtocol] @@ -265,7 +262,6 @@ class CivitAISearchDependencies: search_huggingface_for_file=context.get( "search_huggingface_for_file" ), - to_bool=context.require("to_bool"), web=context.require("web"), write_model_resolver_metadata=context.get( "write_model_resolver_metadata" diff --git a/core/services/scanner_service.py b/core/services/scanner_service.py index 2270208..b76fa5e 100644 --- a/core/services/scanner_service.py +++ b/core/services/scanner_service.py @@ -1,5 +1,7 @@ """Local model scanning orchestration used by HTTP route adapters.""" +from ..metadata_model_utils import normalize_models +from ..request_utils import read_first_bool_field from ..routes.context import RouteContext @@ -15,29 +17,44 @@ class ScannerService: self.invalidate_local_hash_match_cache = context.require( "invalidate_local_hash_match_cache" ) - self.to_bool = context.require("to_bool") self.web = context.require("web") async def get_models(self, request): """Return the locally available models, optionally forcing a rescan.""" - force_rescan = self.to_bool( - request.query.get("force") or request.query.get("force_rescan"), - False, - ) + try: + force_rescan = read_first_bool_field( + request.query, + ("force", "force_rescan"), + contract_name="Models request", + ) + except TypeError as exc: + return self.web.json_response({"error": str(exc)}, status=400) if force_rescan: self.invalidate_local_hash_match_cache() - models = self.get_model_files(force_rescan=force_rescan) - return self.web.json_response(models) + models = normalize_models( + self.get_model_files(force_rescan=force_rescan) + ) + return self.web.json_response( + [model.to_dict() for model in models] + ) async def get_path_template_suggestions(self, request): """Infer download path templates from locally available models.""" from ..sources.popular import get_base_models_config - force_rescan = request.query.get("force") == "1" + try: + force_rescan = read_first_bool_field( + request.query, + ("force", "force_rescan"), + contract_name="Path template suggestions request", + ) + except TypeError as exc: + return self.web.json_response({"error": str(exc)}, status=400) models = await self.asyncio.to_thread( self.get_model_files, force_rescan, ) + models = normalize_models(models) base_models_config = get_base_models_config() suggestions = await self.asyncio.to_thread( self.infer_download_path_templates, diff --git a/core/services/search_dependencies.py b/core/services/search_dependencies.py index 3ada55e..01733b5 100644 --- a/core/services/search_dependencies.py +++ b/core/services/search_dependencies.py @@ -3,6 +3,7 @@ from dataclasses import dataclass from typing import Any, Protocol +from ..contracts import ModelCatalogEntry, ModelMatch, SearchResult from ..routes.context import RouteContext @@ -22,16 +23,8 @@ class SearchTrackerProtocol(Protocol): def update(self, *args: Any, **kwargs: Any) -> Any: ... -class BoolParserProtocol(Protocol): - def __call__(self, value: Any, default: bool = False) -> bool: ... - - -class IntParserProtocol(Protocol): - def __call__(self, value: Any, default: int = 0) -> int: ... - - class ModelResultBuilderProtocol(Protocol): - def __call__(self, source: str, **fields: Any) -> dict[str, Any]: ... + def __call__(self, source: str, **fields: Any) -> SearchResult: ... class ClearCacheProtocol(Protocol): @@ -51,15 +44,19 @@ class FormatSizeProtocol(Protocol): class SearchProviderProtocol(Protocol): - def __call__(self, filename: str, **options: Any) -> Any: ... + def __call__( + self, + filename: str, + **options: Any, + ) -> SearchResult | list[SearchResult] | None: ... class PopularModelUrlProtocol(Protocol): - def __call__(self, filename: str) -> dict[str, Any] | None: ... + def __call__(self, filename: str) -> ModelCatalogEntry | None: ... class ModelListSearchProtocol(Protocol): - def __call__(self, filename: str) -> dict[str, Any] | None: ... + def __call__(self, filename: str) -> SearchResult | None: ... class LocalHashSearchProtocol(Protocol): @@ -70,25 +67,29 @@ class LocalHashSearchProtocol(Protocol): category: str | None, max_matches: int, force_rescan: bool, - ) -> list[dict[str, Any]]: ... + ) -> list[ModelMatch]: ... class ResolveUrnProtocol(Protocol): - def __call__(self, model_id: Any, version_id: Any) -> dict[str, Any] | None: ... + def __call__( + self, + model_id: int | str, + version_id: int | str, + ) -> dict[str, Any] | None: ... class ResolveArchiveVersionProtocol(Protocol): def __call__( self, - model_id: Any, - version_id: Any, + model_id: int | str, + version_id: int | str, *, query: str, - ) -> dict[str, Any] | None: ... + ) -> SearchResult | None: ... class GetDownloadUrlProtocol(Protocol): - def __call__(self, version_id: Any) -> str | None: ... + def __call__(self, version_id: int | str) -> str | None: ... @dataclass(frozen=True, slots=True) @@ -120,8 +121,6 @@ class SearchDependencies: search_local_matches_by_hash: LocalHashSearchProtocol search_lora_manager_archive_for_file: SearchProviderProtocol search_model_list: ModelListSearchProtocol - to_bool: BoolParserProtocol - to_int: IntParserProtocol web: Any @staticmethod @@ -184,7 +183,5 @@ class SearchDependencies: "search_lora_manager_archive_for_file" ), search_model_list=context.require("search_model_list"), - to_bool=context.require("to_bool"), - to_int=context.require("to_int"), web=context.require("web"), ) diff --git a/core/services/search_orchestrator.py b/core/services/search_orchestrator.py index 086cff3..0226851 100644 --- a/core/services/search_orchestrator.py +++ b/core/services/search_orchestrator.py @@ -1,7 +1,8 @@ """Request orchestration for model source searches.""" +from ..contracts import SearchResult from ..local_hash_matches import collect_local_hash_matches_for_result -from ..request_utils import extract_request_sha256 +from ..request_utils import read_optional_object_payload from ..routes.context import RouteContext from .search_cache import SearchResultCache from .search_dependencies import SearchDependencies @@ -53,8 +54,6 @@ class SearchOrchestrator: dependencies.search_lora_manager_archive_for_file ) self.search_model_list = dependencies.search_model_list - self.to_bool = dependencies.to_bool - self.to_int = dependencies.to_int self.web = dependencies.web self.search_cache = SearchResultCache(self.search_result_timestamps) self.provider_runner = SearchProviderRunner(self) @@ -83,6 +82,35 @@ class SearchOrchestrator: parts.append(f"{key}={formatted}") return " ".join(parts) + @staticmethod + def normalize_search_result(source_key, result): + """Validate that a provider returned the shared result contract.""" + if result is None or isinstance(result, SearchResult): + return result + if isinstance(result, list) and all( + isinstance(item, SearchResult) for item in result + ): + return result + raise TypeError( + f"Search source '{source_key}' returned an untyped result: " + f"{type(result).__name__}" + ) + + @staticmethod + def serialize_search_result(result): + """Serialize canonical results at the HTTP/cache boundary.""" + if result is None: + return None + if isinstance(result, SearchResult): + return result.to_dict() + if isinstance(result, list) and all( + isinstance(item, SearchResult) for item in result + ): + return [item.to_dict() for item in result] + raise TypeError( + f"Cannot serialize untyped search result: {type(result).__name__}" + ) + def normalize_result_extra(self, extra): if not extra: return {} @@ -100,6 +128,8 @@ class SearchOrchestrator: return normalized def format_result_details(self, result, extra=None): + if isinstance(result, SearchResult): + result = result.to_dict() if isinstance(result, list): return self.format_log_fields(count=len(result)) if not isinstance(result, dict): @@ -136,62 +166,27 @@ class SearchOrchestrator: async def search_sources(self, request): """Execute all requested sources and return a JSON response.""" try: - data = await request.json() - filename = str(data.get("filename", "") or "").strip() - sha256 = extract_request_sha256( - data, - keys=("sha256", "hash", "file_hash"), - ) - category = data.get("category", "") - base_model_context = data.get("base_model_context", "") - progress_id = str(data.get("progress_id") or "").strip() - progress_source = str(data.get("progress_source") or "").strip() - civitai_candidate_limit = self.to_int( - data.get("civitai_candidate_limit"), - 5, - ) - civitai_candidate_limit = max(1, min(civitai_candidate_limit, 20)) - civarchive_candidate_limit = self.to_int( - data.get("civarchive_candidate_limit"), - 10, - ) - civarchive_candidate_limit = max( - 1, - min(civarchive_candidate_limit, 30), - ) - - is_urn = self.to_bool(data.get("is_urn", False), False) - civitai_key = data.get("civitai_key", "") - civitai_session_token = data.get("civitai_session_token", "") - hf_token = data.get("hf_token", "") - brave_search_api_key = data.get("brave_search_api_key", "") - civitai_use_trpc_search = self.to_bool( - data.get("civitai_use_trpc_search", True), - True, - ) - civitai_use_api_search = self.to_bool( - data.get("civitai_use_api_search", True), - True, - ) - civitai_use_html_fallback = self.to_bool( - data.get("civitai_use_html_fallback", True), - True, - ) - hf_use_api_search = self.to_bool( - data.get("hf_use_api_search", True), - True, - ) - hf_use_comfy_org_fallback = self.to_bool( - data.get("hf_use_comfy_org_fallback", True), - True, - ) - hf_use_brave_fallback = self.to_bool( - data.get("hf_use_brave_fallback", True), - True, - ) - - model_id = data.get("model_id") - version_id = data.get("version_id") + data = await read_optional_object_payload(request) + try: + search_request = SearchRequest.from_mapping( + data, + ) + except (TypeError, ValueError) as exc: + return self.web.json_response( + {"error": str(exc)}, + status=400, + ) + filename = search_request.filename + sha256 = search_request.sha256 + category = search_request.category + base_model_context = search_request.base_model_context + progress_id = search_request.progress_id + progress_source = search_request.progress_source + model_id = search_request.model_id + version_id = search_request.version_id + is_urn = search_request.is_urn + normalized_sources = search_request.normalized_sources + force_search = search_request.force_search if not filename and not sha256 and not (is_urn and model_id and version_id): return self.web.json_response( { @@ -203,61 +198,6 @@ class SearchOrchestrator: status=400, ) - raw_sources = data.get("sources", ["all"]) - if isinstance(raw_sources, str): - raw_sources = [raw_sources] - elif not isinstance(raw_sources, list): - raw_sources = ["all"] - - normalized_sources = frozenset( - str(source).strip().lower() - for source in raw_sources - if str(source).strip() - ) - if not normalized_sources: - normalized_sources = frozenset({"all"}) - if "all" in normalized_sources: - normalized_sources = frozenset( - { - "local", - "huggingface", - "civitai", - "civarchive", - "lora_manager_archive", - } - ) - - if not progress_source: - progress_source = ( - next(iter(normalized_sources)) - if len(normalized_sources) == 1 - else "all" - ) - force_search = self.to_bool(data.get("force_search"), False) - search_request = SearchRequest( - data=data, - filename=filename, - category=category, - base_model_context=base_model_context, - progress_id=progress_id, - progress_source=progress_source, - civitai_candidate_limit=civitai_candidate_limit, - civarchive_candidate_limit=civarchive_candidate_limit, - is_urn=is_urn, - civitai_key=civitai_key, - civitai_session_token=civitai_session_token, - hf_token=hf_token, - brave_search_api_key=brave_search_api_key, - civitai_use_trpc_search=civitai_use_trpc_search, - civitai_use_api_search=civitai_use_api_search, - civitai_use_html_fallback=civitai_use_html_fallback, - hf_use_api_search=hf_use_api_search, - hf_use_comfy_org_fallback=hf_use_comfy_org_fallback, - hf_use_brave_fallback=hf_use_brave_fallback, - force_search=force_search, - normalized_sources=normalized_sources, - sha256=sha256, - ) raise_if_search_cancelled = ( self.provider_runner.raise_if_search_cancelled ) @@ -334,10 +274,17 @@ class SearchOrchestrator: def iter_result_items(result): if isinstance(result, list): for item in result: - if isinstance(item, dict): - yield item - elif isinstance(result, dict): - yield result + if not isinstance(item, SearchResult): + raise TypeError( + "Search result lists must contain SearchResult values" + ) + yield item.to_dict() + elif isinstance(result, SearchResult): + yield result.to_dict() + elif result is not None: + raise TypeError( + "Search result sources must contain SearchResult values" + ) def collect_local_hash_matches(payload): matches = [] @@ -405,12 +352,10 @@ class SearchOrchestrator: continue for match in hash_matches: - model_path = ( - match.get("model", {}).get("path") - or match.get("path") - or "" + model_path = match.model.path or str( + match.extra_value("path") or "" ) - path_key = model_path.lower() + path_key = model_path.casefold() if path_key and path_key in seen_match_paths: continue if path_key: @@ -451,14 +396,20 @@ class SearchOrchestrator: elif source_key == "source_status": results["source_status"].update(source_result or {}) elif source_result: - results[source_key] = source_result + results[source_key] = self.normalize_search_result( + source_key, + source_result, + ) if source_found: results["found"] = True raise_if_search_cancelled(search_request, progress_source) - results["local_hash_matches"] = collect_local_hash_matches(results) - if results["local_hash_matches"]: + local_hash_matches = collect_local_hash_matches(results) + if local_hash_matches: results["found"] = True + results["local_hash_matches"] = [ + match.to_dict() for match in local_hash_matches + ] raise_if_search_cancelled(search_request, progress_source) self.logger.info( @@ -473,6 +424,17 @@ class SearchOrchestrator: 100, status="completed", ) + for source_key in ( + "popular", + "model_list", + "huggingface", + "civitai", + "civarchive", + "lora_manager_archive", + ): + results[source_key] = self.serialize_search_result( + results[source_key] + ) self.search_cache.stamp_results( results, force_search=force_search, diff --git a/core/services/search_providers.py b/core/services/search_providers.py index ee63bac..bf6ff10 100644 --- a/core/services/search_providers.py +++ b/core/services/search_providers.py @@ -1,8 +1,17 @@ """Provider-specific search execution and progress handling.""" +from collections.abc import Mapping from dataclasses import dataclass from typing import Any +from ..contracts import SearchResult +from ..request_utils import ( + extract_request_sha256, + read_bool_field, + read_identifier_field, + read_int_field, + read_text_field, +) from ..sources.civarchive import build_civarchive_failure_status from ..sources.civitai import build_civitai_result_payload from ..type_utils import select_primary_model_file @@ -16,7 +25,6 @@ class SearchCancelled(BaseException): class SearchRequest: """Normalized values shared by the provider search tasks.""" - data: dict[str, Any] filename: str category: str base_model_context: str @@ -38,6 +46,240 @@ class SearchRequest: force_search: bool normalized_sources: frozenset[str] sha256: str = "" + model_id: int | str | None = None + version_id: int | str | None = None + + _TEXT_FIELDS = ( + "filename", + "category", + "base_model_context", + "progress_id", + "progress_source", + "civitai_key", + "civitai_session_token", + "hf_token", + "brave_search_api_key", + "sha256", + ) + _BOOLEAN_FIELDS = ( + "is_urn", + "civitai_use_trpc_search", + "civitai_use_api_search", + "civitai_use_html_fallback", + "hf_use_api_search", + "hf_use_comfy_org_fallback", + "hf_use_brave_fallback", + "force_search", + ) + + @classmethod + def from_mapping( + cls, + value: Mapping[str, Any], + ) -> "SearchRequest": + """Adapt one HTTP payload into the typed provider request.""" + if not isinstance(value, Mapping): + raise TypeError("Search request must be an object") + data = dict(value) + + filename = read_text_field( + data, + "filename", + contract_name="SearchRequest", + ) + sha256 = extract_request_sha256( + data, + keys=("sha256", "hash", "file_hash"), + ) + category = read_text_field( + data, + "category", + contract_name="SearchRequest", + ) + base_model_context = read_text_field( + data, + "base_model_context", + contract_name="SearchRequest", + ) + progress_id = read_text_field( + data, + "progress_id", + contract_name="SearchRequest", + ) + progress_source = read_text_field( + data, + "progress_source", + contract_name="SearchRequest", + ) + + civitai_candidate_limit = read_int_field( + data, + "civitai_candidate_limit", + default=5, + contract_name="SearchRequest", + ) + civarchive_candidate_limit = read_int_field( + data, + "civarchive_candidate_limit", + default=10, + contract_name="SearchRequest", + ) + civitai_candidate_limit = max(1, min(civitai_candidate_limit, 20)) + civarchive_candidate_limit = max(1, min(civarchive_candidate_limit, 30)) + + def read_request_bool(field_name: str, default: bool) -> bool: + return read_bool_field( + data, + field_name, + default=default, + contract_name="SearchRequest", + ) + + is_urn = read_request_bool("is_urn", False) + raw_sources = data.get("sources", ["all"]) + if isinstance(raw_sources, str): + raw_sources = [raw_sources] + elif raw_sources is None or not isinstance(raw_sources, list): + raise TypeError("SearchRequest sources must be a string or array") + if any(not isinstance(source, str) for source in raw_sources): + raise TypeError("SearchRequest sources must contain only strings") + + normalized_sources = frozenset( + source.strip().lower() + for source in raw_sources + if source.strip() + ) + if not normalized_sources: + normalized_sources = frozenset({"all"}) + if "all" in normalized_sources: + normalized_sources = frozenset( + { + "local", + "huggingface", + "civitai", + "civarchive", + "lora_manager_archive", + } + ) + + if not progress_source: + progress_source = ( + next(iter(normalized_sources)) + if len(normalized_sources) == 1 + else "all" + ) + + return cls( + filename=filename, + category=category, + base_model_context=base_model_context, + progress_id=progress_id, + progress_source=progress_source, + civitai_candidate_limit=civitai_candidate_limit, + civarchive_candidate_limit=civarchive_candidate_limit, + is_urn=is_urn, + civitai_key=read_text_field( + data, + "civitai_key", + contract_name="SearchRequest", + ), + civitai_session_token=read_text_field( + data, + "civitai_session_token", + contract_name="SearchRequest", + ), + hf_token=read_text_field( + data, + "hf_token", + contract_name="SearchRequest", + ), + brave_search_api_key=read_text_field( + data, + "brave_search_api_key", + contract_name="SearchRequest", + ), + civitai_use_trpc_search=read_request_bool( + "civitai_use_trpc_search", + True, + ), + civitai_use_api_search=read_request_bool( + "civitai_use_api_search", + True, + ), + civitai_use_html_fallback=read_request_bool( + "civitai_use_html_fallback", + True, + ), + hf_use_api_search=read_request_bool("hf_use_api_search", True), + hf_use_comfy_org_fallback=read_request_bool( + "hf_use_comfy_org_fallback", + True, + ), + hf_use_brave_fallback=read_request_bool( + "hf_use_brave_fallback", + True, + ), + force_search=read_request_bool("force_search", False), + normalized_sources=normalized_sources, + sha256=sha256, + model_id=read_identifier_field( + data, + "model_id", + contract_name="SearchRequest", + ), + version_id=read_identifier_field( + data, + "version_id", + contract_name="SearchRequest", + ), + ) + + def __post_init__(self) -> None: + """Validate the normalized request before provider execution.""" + for field_name in self._TEXT_FIELDS: + value = getattr(self, field_name) + if value is None: + value = "" + if not isinstance(value, str): + raise TypeError(f"SearchRequest {field_name} must be a string") + object.__setattr__(self, field_name, value) + + for field_name in self._BOOLEAN_FIELDS: + value = getattr(self, field_name) + if not isinstance(value, bool): + raise TypeError(f"SearchRequest {field_name} must be a boolean") + + for field_name in ( + "civitai_candidate_limit", + "civarchive_candidate_limit", + ): + value = getattr(self, field_name) + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError( + f"SearchRequest {field_name} must be a positive integer" + ) + if value < 1: + raise ValueError( + f"SearchRequest {field_name} must be a positive integer" + ) + + for field_name in ("model_id", "version_id"): + value = getattr(self, field_name) + if value is None: + continue + if isinstance(value, bool) or not isinstance(value, (int, str)): + raise TypeError( + f"SearchRequest {field_name} must be an integer or string" + ) + if isinstance(value, str): + object.__setattr__(self, field_name, value.strip() or None) + + if not isinstance(self.normalized_sources, frozenset): + raise TypeError("SearchRequest normalized_sources must be a frozenset") + if any(not isinstance(source, str) for source in self.normalized_sources): + raise TypeError( + "SearchRequest normalized_sources must contain only strings" + ) @property def search_local(self): @@ -171,19 +413,22 @@ class SearchProviderRunner: raise def mark_any_model_fallback(self, request: SearchRequest, result): + if result is None: + return None if isinstance(result, list): return [ self.mark_any_model_fallback(request, item) for item in result ] - if not isinstance(result, dict): - return result - - marked = dict(result) - marked["any_model_match"] = True - marked["base_model_fallback"] = True - marked["requested_base_model"] = request.base_model_context - return marked + if isinstance(result, SearchResult): + return result.with_extra( + any_model_match=True, + base_model_fallback=True, + requested_base_model=request.base_model_context, + ) + raise TypeError( + f"{type(result).__name__} is not a typed search result" + ) def execute_search_with_fallback( self, @@ -248,7 +493,6 @@ class SearchProviderRunner: return source_results, source_found popular_info = self.owner.get_popular_model_url(request.filename) - self.owner.log_search_result("popular", popular_info) self.owner.search_tracker.update( request.progress_id, "local", @@ -262,30 +506,44 @@ class SearchProviderRunner: model_list_result, { "confidence": ( - model_list_result.get("confidence") + model_list_result.confidence if model_list_result else None ) }, ) if popular_info: - popular_result = { - "source": "popular", - "filename": request.filename, - **popular_info, - } + popular_result = self.owner.build_model_result( + "popular", + filename=request.filename, + name=popular_info.name or request.filename, + type=popular_info.model_type, + url=( + popular_info.url + or popular_info.download_url + or "" + ), + download_url=popular_info.download_url, + size=popular_info.size, + match_type="exact", + directory=popular_info.directory, + **dict(popular_info.extra), + ) if ( model_list_result - and model_list_result.get("filename", "").lower() + and model_list_result.filename.lower() == request.filename.lower() - and model_list_result.get("size") + and model_list_result.size is not None ): - popular_result["size"] = model_list_result.get("size") + popular_result = popular_result.with_updates( + size=model_list_result.size + ) source_results["popular"] = popular_result + self.owner.log_search_result("popular", popular_result) source_found = True if model_list_result: - confidence = model_list_result.get("confidence", 0) + confidence = model_list_result.confidence if (request.is_urn and confidence >= 70) or not request.is_urn: source_results["model_list"] = model_list_result source_found = True @@ -340,8 +598,8 @@ class SearchProviderRunner: source_found = False if request.is_urn: - model_id = request.data.get("model_id") - version_id = request.data.get("version_id") + model_id = request.model_id + version_id = request.version_id if model_id and version_id: self.owner.search_tracker.update( @@ -418,20 +676,7 @@ class SearchProviderRunner: {"results_count": len(civitai_results)}, ) if civitai_results: - first_result = civitai_results[0] - source_results["civitai"] = self.owner.build_model_result( - "civitai", - model_id=first_result.get("model_id"), - version_id=first_result.get("version_id"), - name=first_result.get("name"), - filename=first_result.get("filename"), - type=first_result.get("type"), - download_url=first_result.get("download_url"), - url=first_result.get("url"), - size=first_result.get("size"), - base_model=first_result.get("base_model"), - tags=first_result.get("tags", []), - ) + source_results["civitai"] = civitai_results[0] source_found = True else: civitai_result = self.execute_search_with_fallback( @@ -478,8 +723,8 @@ class SearchProviderRunner: source_found = False if request.is_urn: - model_id = request.data.get("model_id") - version_id = request.data.get("version_id") + model_id = request.model_id + version_id = request.version_id if model_id and version_id: self.owner.search_tracker.update( request.progress_id, @@ -628,8 +873,8 @@ class SearchProviderRunner: or getattr(request, "sha256", "") or ( request.is_urn - and request.data.get("model_id") - and request.data.get("version_id") + and request.model_id + and request.version_id ) ): tasks.append( diff --git a/core/services/search_support_service.py b/core/services/search_support_service.py index b0372c7..876c399 100644 --- a/core/services/search_support_service.py +++ b/core/services/search_support_service.py @@ -1,5 +1,10 @@ """Search cache, credentials, and index operations for HTTP routes.""" +from ..request_utils import ( + read_bool_field, + read_optional_object_payload, + read_text_field, +) from ..routes.context import RouteContext @@ -56,10 +61,16 @@ class SearchSupportService: log_name, ): try: - data = await request.json() - value = data.get(payload_key, "") + data = await read_optional_object_payload(request) + value = read_text_field( + data, + payload_key, + contract_name=f"{log_name} request", + ) result = await self.asyncio.to_thread(check_func, value) return self.web.json_response(result) + except TypeError as error: + return self.web.json_response({"error": str(error)}, status=400) except Exception as error: self.extension.logger.exception(f"{log_name} check error: {error}") return self.web.json_response({"error": str(error)}, status=500) @@ -108,8 +119,15 @@ class SearchSupportService: async def refresh_author_index(self, request): """Refresh HuggingFace author fallback index.""" - data = await request.json() - hf_token = data.get("hf_token", "") + data = await read_optional_object_payload(request) + try: + hf_token = read_text_field( + data, + "hf_token", + contract_name="Author index refresh request", + ) + except TypeError as error: + return self.web.json_response({"error": str(error)}, status=400) result = await self.asyncio.to_thread( self.refresh_known_author_fallback_indexes, hf_token or None, @@ -119,11 +137,14 @@ class SearchSupportService: async def model_list_status(self, request): """Return local model-list status and optionally compare with GitHub.""" - check_remote = str(request.query.get("check_remote", "")).lower() in { - "1", - "true", - "yes", - } + try: + check_remote = read_bool_field( + request.query, + "check_remote", + contract_name="Model list status request", + ) + except TypeError as exc: + return self.web.json_response({"error": str(exc)}, status=400) return self.web.json_response( self.get_model_list_update_status(check_remote=check_remote) ) diff --git a/core/services/workflow_service.py b/core/services/workflow_service.py index 80f22a1..b47d4db 100644 --- a/core/services/workflow_service.py +++ b/core/services/workflow_service.py @@ -1,6 +1,13 @@ """Workflow analysis service.""" -from ..request_utils import validate_workflow_payload +from ..contracts import MissingModel, Resolution, WorkflowAnalysisResult +from ..custom_nodes import get_custom_node_resolution_metadata +from ..request_utils import ( + read_bool_field, + read_optional_object_payload, + read_text_field, + validate_workflow_payload, +) from ..routes.context import RouteContext @@ -25,7 +32,6 @@ class WorkflowService: self.should_skip_existing_custom_node_reference = context.require( "should_skip_existing_custom_node_reference" ) - self.to_bool = context.require("to_bool") self.web = context.require("web") async def analyze_workflow(self, request): @@ -39,15 +45,25 @@ class WorkflowService: invalidate_local_hash_match_cache = self.invalidate_local_hash_match_cache search_model_list = self.search_model_list should_skip_existing_custom_node_reference = self.should_skip_existing_custom_node_reference - to_bool = self.to_bool web = self.web try: - data = await request.json() + data = await read_optional_object_payload(request) workflow_json, workflow_error = validate_workflow_payload( data.get("workflow") ) - analysis_id = str(data.get("analysis_id") or "").strip() - force_rescan = to_bool(data.get("force_rescan"), False) + try: + analysis_id = read_text_field( + data, + "analysis_id", + contract_name="Workflow analysis request", + ) + force_rescan = read_bool_field( + data, + "force_rescan", + contract_name="Workflow analysis request", + ) + except TypeError as exc: + return web.json_response({"error": str(exc)}, status=400) if force_rescan: invalidate_local_hash_match_cache() @@ -72,7 +88,7 @@ class WorkflowService: self._update_analysis_progress(analysis_id, payload) # Analyze and find matches - result = await asyncio.to_thread( + analysis_result = await asyncio.to_thread( analyze_and_find_matches, workflow_json, 0.0, @@ -82,38 +98,36 @@ class WorkflowService: analysis_id=analysis_id, ) - missing_models = result.get("missing_models", []) - filtered_missing = [] - for missing in missing_models: - name = missing.get("name") or missing.get("original_path", "") - if should_skip_existing_custom_node_reference(missing): + filtered_missing: list[MissingModel] = [] + for missing in analysis_result.missing_models: + name = ( + missing.extra_value("name") + or missing.reference.original_path + ) + if should_skip_existing_custom_node_reference(missing.reference): self.logger.info( "Filtered existing custom-node model " f"reference: {name}" ) continue filtered_missing.append(missing) - result["missing_models"] = filtered_missing - result["total_missing"] = len(filtered_missing) # If download available, check for download sources only from LOCAL sources # (workflow_url, popular, model-list.json) - skip automatic online search # Online search is now only triggered on-demand via search button if download_available: - for missing in result.get("missing_models", []): + for index, missing in enumerate(filtered_missing): # Check if there's a 100% local match - matches = missing.get("matches", []) - has_perfect_match = any( - m.get("confidence", 0) == 100 for m in matches - ) + matches = missing.matches or () + has_perfect_match = any(match.confidence == 100 for match in matches) if not has_perfect_match: filename = get_filename_from_path( - missing.get("original_path", "") + missing.reference.original_path ) # 0. Check workflow URL first (highest priority - directly from workflow) - workflow_url = missing.get("workflow_url", "") + workflow_url = missing.extra_value("workflow_url", "") if workflow_url: # Determine source from URL if "huggingface.co" in workflow_url: @@ -126,21 +140,24 @@ class WorkflowService: # Try to get file size using cached remote helper file_size = fetch_remote_file_size_cached(workflow_url, timeout=5) - missing["download_source"] = { - "source": source, - "url": workflow_url, - "model_url": missing.get( - "workflow_model_url", workflow_url - ), - "filename": filename, - "directory": missing.get( - "workflow_directory", "" - ) - or missing.get("category", "checkpoints"), - "match_type": "exact", - "url_source": "workflow", - "size": file_size, - } + filtered_missing[index] = missing.with_extra( + download_source={ + "source": source, + "url": workflow_url, + "model_url": missing.extra_value( + "workflow_model_url", workflow_url + ), + "filename": filename, + "directory": missing.extra_value( + "workflow_directory", "" + ) + or missing.reference.category + or "checkpoints", + "match_type": "exact", + "url_source": "workflow", + "size": file_size, + } + ) continue # 1. Check popular models (always exact match) @@ -149,20 +166,27 @@ class WorkflowService: popular_model_list_result = search_model_list( filename, exact_only=True ) - missing["download_source"] = { - "source": "popular", - "url": popular_info.get("url"), - "filename": filename, - "type": popular_info.get("type"), - "directory": popular_info.get("directory"), - "size": ( - popular_model_list_result.get("size") + filtered_missing[index] = missing.with_extra( + download_source={ + "source": "popular", + "url": popular_info.url + or popular_info.download_url, + "filename": popular_info.filename or filename, + "type": popular_info.model_type, + "directory": popular_info.directory, + "size": ( + popular_model_list_result.size + if popular_model_list_result + and popular_model_list_result.size + is not None + else None + ) if popular_model_list_result - else None - ) - or popular_info.get("size"), - "match_type": "exact", - } + and popular_model_list_result.size is not None + else popular_info.size, + "match_type": "exact", + } + ) continue # 2. Check model list (ComfyUI Manager database) @@ -171,38 +195,46 @@ class WorkflowService: filename, exact_only=True ) if model_list_result: - missing["download_source"] = { - "source": "model_list", - "url": model_list_result.get("url"), - "filename": model_list_result.get("filename"), - "name": model_list_result.get("name"), - "type": model_list_result.get("type"), - "directory": model_list_result.get("directory"), - "size": model_list_result.get("size"), - "match_type": model_list_result.get( - "match_type" - ), - "confidence": model_list_result.get( - "confidence" - ), - } + filtered_missing[index] = missing.with_extra( + download_source={ + "source": "model_list", + "url": model_list_result.url, + "filename": model_list_result.filename, + "name": model_list_result.name, + "type": model_list_result.model_type, + "directory": model_list_result.extra_value( + "directory" + ), + "size": model_list_result.size, + "match_type": model_list_result.match_type, + "confidence": model_list_result.confidence, + } + ) continue # NOTE: Search for online sources (HuggingFace, CivitAI) is # now done on-demand via /model_resolver/search endpoint # when user clicks "Search Online" button, not automatically + result = WorkflowAnalysisResult( + missing_models=tuple(filtered_missing), + resolved_models=analysis_result.resolved_models, + total_resolved=analysis_result.total_resolved, + total_missing=len(filtered_missing), + total_models_analyzed=analysis_result.total_models_analyzed, + ) + if analysis_id: self.analysis_progress.update( analysis_id, status="completed", stage="completed", message="Analysis complete", - current=result.get("total_missing", 0), - total=result.get("total_missing", 0), + current=result.total_missing, + total=result.total_missing, ) - return web.json_response(result) + return web.json_response(result.to_dict()) except Exception as e: if "analysis_id" in locals() and analysis_id: self.analysis_progress.update( @@ -220,24 +252,42 @@ class WorkflowService: """Apply model resolution and return updated workflow.""" apply_resolution = self.apply_resolution web = self.web - data = await request.json() + data = await read_optional_object_payload(request) workflow_json, workflow_error = validate_workflow_payload( data.get("workflow"), empty_is_missing=True, require_object=False, ) - resolutions = data.get("resolutions", []) + raw_resolutions = data.get("resolutions", []) if workflow_error: return web.json_response( {"error": workflow_error}, status=400 ) - if not resolutions: + if not isinstance(raw_resolutions, list) or not raw_resolutions: return web.json_response( {"error": "Resolutions array is required"}, status=400 ) + try: + resolutions = [] + for raw_resolution in raw_resolutions: + if not isinstance(raw_resolution, dict): + raise ValueError("Each resolution must be an object") + resolution = Resolution.from_mapping(raw_resolution) + resolution.validate() + resolutions.append( + resolution.with_custom_node_metadata( + get_custom_node_resolution_metadata(resolution.reference) + ) + ) + except (TypeError, ValueError) as exc: + return web.json_response( + {"error": f"Invalid resolution: {exc}"}, + status=400, + ) + # Apply resolutions updated_workflow = apply_resolution(workflow_json, resolutions) @@ -249,12 +299,26 @@ class WorkflowService: """Search local model files by filename/path.""" invalidate_local_hash_match_cache = self.invalidate_local_hash_match_cache search_local_matches = self.search_local_matches - to_bool = self.to_bool web = self.web - data = await request.json() - filename = data.get("filename", "") - category = data.get("category", "") - force_rescan = to_bool(data.get("force_rescan"), False) + data = await read_optional_object_payload(request) + try: + filename = read_text_field( + data, + "filename", + contract_name="Local matches request", + ) + category = read_text_field( + data, + "category", + contract_name="Local matches request", + ) + force_rescan = read_bool_field( + data, + "force_rescan", + contract_name="Local matches request", + ) + except TypeError as exc: + return web.json_response({"error": str(exc)}, status=400) if force_rescan: invalidate_local_hash_match_cache() @@ -270,4 +334,6 @@ class WorkflowService: max_matches_per_model=10, force_rescan=force_rescan, ) - return web.json_response({"matches": matches}) + return web.json_response( + {"matches": [match.to_dict() for match in matches]} + ) diff --git a/core/sources/civarchive.py b/core/sources/civarchive.py index a5ee84a..48f6efb 100644 --- a/core/sources/civarchive.py +++ b/core/sources/civarchive.py @@ -13,6 +13,7 @@ from urllib.parse import urljoin, urlparse import requests +from ..contracts import ProviderUrlReference, SearchResult from ..log_system import create_module_logger from ..matcher import ( base_model_matches as _base_model_matches, @@ -588,7 +589,7 @@ def _search_page( return [result for result in results if isinstance(result, dict)] -def parse_civarchive_url(url: str) -> Optional[Dict[str, Any]]: +def parse_civarchive_url(url: str) -> Optional[ProviderUrlReference]: """Parse CivArchive URLs into sha256 or model/version identifiers.""" from ..type_utils import parse_provider_model_url if not url: @@ -604,7 +605,7 @@ def _extract_sha256(value: str) -> Optional[str]: if normalized: return normalized parsed = parse_civarchive_url(value) - return parsed.get("sha256") if parsed else None + return parsed.sha256 if parsed else None @@ -1024,29 +1025,34 @@ def get_civarchive_model_details( resolved = resolve_civarchive_model_version(model_id, version_id) if not resolved: return None + resolved_mapping = resolved.to_dict() fallback_version = { - "id": resolved.get("version_id"), - "name": resolved.get("version_name"), - "baseModel": resolved.get("base_model"), - "trainedWords": resolved.get("trained_words", []), - "images": resolved.get("images", []), + "id": resolved_mapping.get("version_id"), + "name": resolved_mapping.get("version_name"), + "baseModel": resolved_mapping.get("base_model"), + "trainedWords": resolved_mapping.get("trained_words", []), + "images": resolved_mapping.get("images", []), "files": [ { - "name": resolved.get("filename"), - "type": resolved.get("type"), - "downloadUrl": resolved.get("download_url"), - "sizeKB": (resolved.get("size") or 0) / 1024 if resolved.get("size") else None, + "name": resolved_mapping.get("filename"), + "type": resolved_mapping.get("type"), + "downloadUrl": resolved_mapping.get("download_url"), + "sizeKB": ( + (resolved_mapping.get("size") or 0) / 1024 + if resolved_mapping.get("size") + else None + ), "primary": True, } ], } context = { "id": model_id, - "name": resolved.get("name"), - "type": resolved.get("type"), - "tags": resolved.get("tags", []), - "creator": resolved.get("creator", {}), - "platform": resolved.get("platform"), + "name": resolved_mapping.get("name"), + "type": resolved_mapping.get("type"), + "tags": resolved_mapping.get("tags", []), + "creator": resolved_mapping.get("creator", {}), + "platform": resolved_mapping.get("platform"), } selected = _normalize_archive_version(fallback_version, context) return { @@ -1189,7 +1195,7 @@ def _build_civarchive_result( sha256: Optional[str] = None, hashes: Optional[Dict[str, Any]] = None, hash_value: Any = _RESULT_HASH_UNSET, -) -> Dict[str, Any]: +) -> SearchResult: extra_fields = { "civitai_model_id": civitai_model_id, "civitai_model_version_id": civitai_model_version_id, @@ -1231,7 +1237,7 @@ def _build_result_from_normalized_version( version: Dict[str, Any], file_info: Dict[str, Any], match_type: str, -) -> Optional[Dict[str, Any]]: +) -> Optional[SearchResult]: download_urls = _collect_archive_download_urls( file_info, **_NORMALIZED_DOWNLOAD_URL_OPTIONS, @@ -1313,7 +1319,7 @@ def _find_model_title_match_in_model_details( model_details: Dict[str, Any], title_query: str, base_model_context: Optional[str] = None, -) -> Optional[Dict[str, Any]]: +) -> Optional[SearchResult]: """For extensionless workflow values, resolve by CivArchive model page title.""" from ..matcher import match_model_by_title_generic @@ -1366,7 +1372,7 @@ def _build_result_from_model_details( model_details: Dict[str, Any], query: str = "", exact_only: bool = False, -) -> Optional[Dict[str, Any]]: +) -> Optional[SearchResult]: selected_version = model_details.get("selected_version") if not isinstance(selected_version, dict): return None @@ -1417,11 +1423,12 @@ def _build_result_from_model_details( hashes = selected_file.get("hashes") if isinstance(selected_file.get("hashes"), dict) else {} sha256 = extract_file_sha256(selected_file) - result["confidence"] = best_confidence - result["sha256"] = sha256 - result["hash"] = sha256 - result["hashes"] = hashes - return result + result = result.with_updates( + confidence=best_confidence, + sha256=sha256, + hashes=hashes, + ) + return result.with_extra(hash=sha256) def _extract_hash_page_files(html_text: str) -> List[Dict[str, Any]]: @@ -1580,7 +1587,7 @@ def _build_result_from_hash_page_file( file_info: Dict[str, Any], query: str, sha256: str = "", -) -> Optional[Dict[str, Any]]: +) -> Optional[SearchResult]: model_id = to_int(model_details.get("model_id") or file_info.get("model_id")) version_id = to_int(version.get("id") or file_info.get("model_version_id")) normalized_file = _normalize_archive_file(file_info, model_id, version_id) @@ -1599,13 +1606,15 @@ def _build_result_from_hash_page_file( if not result: return None - result["confidence"] = confidence + result = result.with_updates(confidence=confidence) if sha256: - result["sha256"] = sha256 - result["hash"] = sha256 - hashes = result.get("hashes") if isinstance(result.get("hashes"), dict) else {} + result = result.with_updates(sha256=sha256) + result = result.with_extra(hash=sha256) + hashes = dict(result.hashes or {}) if not hashes.get("SHA256") and not hashes.get("sha256"): - result["hashes"] = {**hashes, "SHA256": sha256} + result = result.with_updates( + hashes={**hashes, "SHA256": sha256} + ) return result @@ -1614,7 +1623,7 @@ def _build_result_from_payload( query: str = "", preferred_filename: str = "", exact_only: bool = False, -) -> Optional[Dict[str, Any]]: +) -> Optional[SearchResult]: context, version, top_files = _extract_model_context(payload) if not version: return None @@ -1728,7 +1737,7 @@ def resolve_civarchive_by_hash( query: str = "", exact_only: bool = False, model_type: Optional[str] = None, -) -> Optional[Dict[str, Any]]: +) -> Optional[SearchResult]: """Resolve a model by SHA256 hash through CivArchive.""" if not sha256: return None @@ -1760,7 +1769,7 @@ def resolve_civarchive_model_version( query: str = "", exact_only: bool = False, prefer_page: bool = False, -) -> Optional[Dict[str, Any]]: +) -> Optional[SearchResult]: """Resolve a model/version through CivArchive.""" if model_id is None: return None @@ -1796,9 +1805,9 @@ def _candidate_identity(candidate: Dict[str, Any]) -> Tuple[Any, Any, Any]: parsed = parse_civarchive_url(candidate.get("url", "")) if parsed: return ( - parsed.get("sha256"), - parsed.get("model_id"), - parsed.get("version_id"), + parsed.sha256, + parsed.model_id, + parsed.version_id, ) return (candidate.get("id"), candidate.get("url"), None) @@ -1816,7 +1825,7 @@ def _build_result_from_search_candidate( query: str, page_text: str = "", parsed: Optional[Dict[str, Any]] = None, -) -> Optional[Dict[str, Any]]: +) -> Optional[SearchResult]: candidate_name = _candidate_display_name(candidate) query_basename = get_filename_from_path(query or "").strip() filename = ( @@ -1881,7 +1890,7 @@ def _resolve_civarchive_model_link( exact_only: bool = False, base_model_context: Optional[str] = None, allow_model_title_match: bool = False, -) -> Optional[Dict[str, Any]]: +) -> Optional[SearchResult]: details = get_civarchive_model_details( model_id, version_id, @@ -1915,7 +1924,7 @@ def _resolve_hash_search_candidate_from_page( base_model_context: Optional[str] = None, allow_model_title_match: bool = False, expected_model_type: Optional[str] = None, -) -> Optional[Dict[str, Any]]: +) -> Optional[SearchResult]: page_text = _request_page_text(f"/sha256/{sha256.lower()}") hash_page_files = _extract_hash_page_files(page_text or "") model_links = _prefer_model_links_for_expected_type( @@ -1969,7 +1978,7 @@ def _resolve_hash_search_candidate_from_page( match_type="exact" if confidence == 100.0 else "similar", ) if result: - result["confidence"] = confidence + result = result.with_updates(confidence=confidence) selected_file = None if not result: @@ -1998,14 +2007,16 @@ def _resolve_hash_search_candidate_from_page( ) if result: - if exact_only and float(result.get("confidence") or 0) < 100.0: + if exact_only and result.confidence < 100.0: continue - result.setdefault("sha256", sha256) - result.setdefault("hash", sha256) - hashes = result.get("hashes") if isinstance(result.get("hashes"), dict) else {} + if not result.sha256: + result = result.with_updates(sha256=sha256) + if not result.extra_value("hash"): + result = result.with_extra(hash=sha256) + hashes = dict(result.hashes or {}) if sha256 and not hashes.get("SHA256") and not hashes.get("sha256"): hashes = {**hashes, "SHA256": sha256} - result["hashes"] = hashes + result = result.with_updates(hashes=hashes) return result return _build_result_from_search_candidate( @@ -2023,25 +2034,25 @@ def _resolve_search_candidate( base_model_context: Optional[str] = None, allow_model_title_match: bool = False, model_type: Optional[str] = None, -) -> Optional[Dict[str, Any]]: +) -> Optional[SearchResult]: parsed = parse_civarchive_url(candidate.get("url", "")) if not parsed: return None - if parsed.get("sha256"): + if parsed.sha256: result = _resolve_hash_search_candidate_from_page( candidate, - parsed["sha256"], + parsed.sha256, query=query, exact_only=exact_only, base_model_context=base_model_context, allow_model_title_match=allow_model_title_match, expected_model_type=model_type, ) - elif parsed.get("model_id"): + elif parsed.model_id: result = _resolve_civarchive_model_link( - parsed["model_id"], - parsed.get("version_id"), + parsed.model_id, + parsed.version_id, query=query, exact_only=exact_only, base_model_context=base_model_context, @@ -2050,27 +2061,29 @@ def _resolve_search_candidate( else: result = None - if result and not result.get("confidence"): - result["confidence"] = calculate_archived_model_confidence( + if result and not result.confidence: + result = result.with_updates( + confidence=calculate_archived_model_confidence( query, candidate.get("name", ""), "", - result.get("filename", ""), + result.filename, + ) ) return result -def _is_hash_verified_exact_match(result: Optional[Dict[str, Any]], confidence: float) -> bool: +def _is_hash_verified_exact_match(result: Optional[SearchResult], confidence: float) -> bool: if not result or confidence < 100.0: return False - return bool(result.get("sha256") or result.get("hash")) + return bool(result.sha256 or result.extra_value("hash")) def search_civarchive( query: str, model_type: Optional[str] = None, limit: int = DEFAULT_CIVARCHIVE_CANDIDATE_LIMIT, -) -> List[Dict[str, Any]]: +) -> List[SearchResult]: """ Search CivArchive by model name or filename. """ @@ -2104,7 +2117,7 @@ def search_civarchive( if not resolved: continue - if resolved.get("confidence", 0) < 40: + if resolved.confidence < 40: continue results.append(resolved) if len(results) >= detail_limit: @@ -2112,8 +2125,8 @@ def search_civarchive( results.sort( key=lambda item: ( - item.get("confidence", 0), - 1 if item.get("match_type") == "exact" else 0, + item.confidence, + 1 if item.match_type == "exact" else 0, ), reverse=True, ) @@ -2129,7 +2142,7 @@ def search_civarchive_for_file( limit: int = DEFAULT_CIVARCHIVE_CANDIDATE_LIMIT, progress_callback: Optional[Callable[[Dict[str, Any]], None]] = None, sha256: Optional[str] = None, -) -> Optional[Dict[str, Any]]: +) -> Optional[SearchResult]: """ Search CivArchive for the best downloadable match for a model filename. """ @@ -2264,30 +2277,37 @@ def search_civarchive_for_file( if not resolved: continue - if resolved.get("match_type") == "model_title": + if resolved.match_type == "model_title": + title_confidence = resolved.extra_value("title_confidence") confidence = float( - resolved.get("title_confidence") - or resolved.get("confidence") + title_confidence + or resolved.confidence or calculate_model_title_confidence( normalized_filename, - resolved.get("name", ""), + resolved.name, ) ) else: confidence = calculate_archived_model_confidence( normalized_filename, - resolved.get("name", ""), - resolved.get("version_name", ""), - resolved.get("filename", ""), + resolved.name, + resolved.version_name, + resolved.filename, ) if exact_only and confidence < 100.0: continue - resolved["confidence"] = confidence - resolved["match_type"] = "exact" if confidence == 100.0 else "similar" + resolved = resolved.with_updates( + confidence=confidence, + match_type="exact" if confidence == 100.0 else "similar", + ) should_update, rank, base_model_matches = should_update_best_match( - confidence, resolved.get("base_model"), base_model_context, best_rank, exact_only + confidence, + resolved.base_model, + base_model_context, + best_rank, + exact_only, ) if exact_only and confidence < 100.0: continue @@ -2320,7 +2340,7 @@ def search_civarchive_for_file( if best_match and ( not base_model_context - or _base_model_matches(best_match.get("base_model"), base_model_context) + or _base_model_matches(best_match.base_model, base_model_context) ): break except CivArchiveSearchError: @@ -2341,7 +2361,7 @@ def search_civarchive_for_file( if ( best_match and base_model_context - and not _base_model_matches(best_match.get("base_model"), base_model_context) + and not _base_model_matches(best_match.base_model, base_model_context) and not _is_hash_verified_exact_match(best_match, best_confidence) ): best_match = None @@ -2363,7 +2383,7 @@ def search_civarchive_for_file( def build_civarchive_custom_result( details: Dict[str, Any], expected_filename: str = "", -) -> Optional[Dict[str, Any]]: +) -> Optional[SearchResult]: if not isinstance(details, dict): return None diff --git a/core/sources/civitai.py b/core/sources/civitai.py index c8bd623..d9da5c0 100644 --- a/core/sources/civitai.py +++ b/core/sources/civitai.py @@ -12,6 +12,7 @@ from urllib.parse import quote, urlencode import requests +from ..contracts import ProviderUrlReference, SearchResult from ..log_system import create_module_logger from ..matcher import ( base_model_matches as _base_model_matches, @@ -203,7 +204,7 @@ def _build_civitai_result_from_version( file_info: Dict[str, Any], tags: Optional[List[str]] = None, match_type: str = "exact", -) -> Dict[str, Any]: +) -> SearchResult: """Normalize CivitAI model/version/file data into the search result format.""" version_id = version.get("id") size = extract_file_size(file_info) @@ -238,7 +239,7 @@ def _find_model_title_match_in_model( title_query: str, api_key: Optional[str] = None, base_model_context: Optional[str] = None, -) -> Optional[Dict[str, Any]]: +) -> Optional[SearchResult]: """For extensionless workflow values, resolve by CivitAI model page title.""" from ..matcher import match_model_by_title_generic @@ -266,7 +267,7 @@ def _find_model_title_match_in_model( match_type="model_title", ) if res: - res["version_name"] = v.get("name", "") + res = res.with_updates(version_name=v.get("name", "")) return res return match_model_by_title_generic( @@ -733,7 +734,7 @@ def _find_civitai_file_in_model( exact_only: bool = False, preferred_version_id: Optional[int] = None, base_model_context: Optional[str] = None, -) -> Optional[Dict[str, Any]]: +) -> Optional[SearchResult]: """Load one CivitAI model and search all its versions for the requested file.""" filename_lower = filename.lower() filename_base = os.path.splitext(filename_lower)[0] @@ -743,7 +744,7 @@ def _find_civitai_file_in_model( def build_result_from_resolved_version( resolved: Dict[str, Any], version_id: int - ) -> Dict[str, Any]: + ) -> SearchResult: expected_filename = resolved.get("expected_filename", "") primary_file = select_primary_model_file( resolved.get("files") or [], @@ -751,8 +752,8 @@ def _find_civitai_file_in_model( fallback_to_first=True, ) or {} - return { - "source": "civitai", + return build_model_result( + "civitai", **build_civitai_result_payload( model_id=model_id, version_id=version_id, @@ -769,7 +770,8 @@ def _find_civitai_file_in_model( filename, expected_filename ), ), - } + normalize_hashes=True, + ) def resolved_version_matches(resolved: Dict[str, Any]) -> bool: expected_filename = str(resolved.get("expected_filename", "")).lower() @@ -849,9 +851,10 @@ def _find_civitai_file_in_model( best_resolved_confidence = ranking_score best_resolved_result = build_result_from_resolved_version( resolved, version_id + ).with_updates( + match_type="similar", + confidence=confidence, ) - best_resolved_result["match_type"] = "similar" - best_resolved_result["confidence"] = confidence files = version.get("files", []) file_names = [f.get("name", "") for f in files if isinstance(f, dict)] @@ -861,7 +864,7 @@ def _find_civitai_file_in_model( if best_resolved_result and best_resolved_confidence >= 50.0: log.info( - f"CivitAI best probable match: model_id={model_id}, version_id={best_resolved_result.get('version_id')}, filename={best_resolved_result.get('filename')}, confidence={best_resolved_confidence}" + f"CivitAI best probable match: model_id={model_id}, version_id={best_resolved_result.version_id}, filename={best_resolved_result.filename}, confidence={best_resolved_confidence}" ) return best_resolved_result @@ -878,15 +881,17 @@ def _find_civitai_file_in_model( tags=data.get("tags", []), match_type=match["match_type"], ) - result["confidence"] = match.get( - "confidence", - calculate_filename_confidence(filename, result.get("filename", "")), + result = result.with_updates( + confidence=match.get( + "confidence", + calculate_filename_confidence(filename, result.filename), + ) ) if match["match_type"] == "similar": log.info( - f"CivitAI version-list probable match: model_id={model_id}, version_id={result.get('version_id')}, filename={result.get('filename')}, confidence={result['confidence']}" + f"CivitAI version-list probable match: model_id={model_id}, version_id={result.version_id}, filename={result.filename}, confidence={result.confidence}" ) - if _base_model_matches(result.get("base_model"), base_model_context): + if _base_model_matches(result.base_model, base_model_context): return result if not base_model_context: return result @@ -895,7 +900,7 @@ def _find_civitai_file_in_model( -def parse_civitai_url(url: str) -> Optional[Dict[str, Any]]: +def parse_civitai_url(url: str) -> Optional[ProviderUrlReference]: """ Parse a CivitAI URL to extract model/version info. """ @@ -924,7 +929,7 @@ def search_civitai_for_file( use_html_fallback: bool = True, progress_callback: Optional[Callable[[Dict[str, Any]], None]] = None, sha256: Optional[str] = None, -) -> Optional[Dict[str, Any]]: +) -> Optional[SearchResult]: """ Search CivitAI for a specific model file. Returns the first model that actually has this exact filename. @@ -936,7 +941,7 @@ def search_civitai_for_file( If False, also try partial matching (for local file resolution). Returns: - Dict with download info if found, None otherwise + A typed SearchResult if a matching model is found, otherwise None. """ global _search_cache @@ -1146,7 +1151,7 @@ def search_civitai_for_file( base_model_context=base_model_context, ) if result: - confidence = float(result.get("confidence") or 0.0) + confidence = result.confidence if confidence >= 100.0: _search_cache[cache_key] = result _report_progress( @@ -1154,12 +1159,12 @@ def search_civitai_for_file( "found", "Found exact CivitAI match", 92, - model_id=result.get("model_id"), - version_id=result.get("version_id"), + model_id=result.model_id, + version_id=result.version_id, confidence=confidence, ) log.info( - f"Found exact CivitAI match for {filename}: model_id={result.get('model_id')}, version_id={result.get('version_id')}, candidate_limit={candidate_limit}" + f"Found exact CivitAI match for {filename}: model_id={result.model_id}, version_id={result.version_id}, candidate_limit={candidate_limit}" ) return result if confidence > best_confidence: @@ -1173,12 +1178,12 @@ def search_civitai_for_file( "found", "Found CivitAI match", 92, - model_id=best_result.get("model_id"), - version_id=best_result.get("version_id"), + model_id=best_result.model_id, + version_id=best_result.version_id, confidence=best_confidence, ) log.info( - f"Found best CivitAI match for {filename}: model_id={best_result.get('model_id')}, version_id={best_result.get('version_id')}, confidence={best_confidence}, candidate_limit={candidate_limit}" + f"Found best CivitAI match for {filename}: model_id={best_result.model_id}, version_id={best_result.version_id}, confidence={best_confidence}, candidate_limit={candidate_limit}" ) return best_result @@ -1217,7 +1222,7 @@ def search_civitai( model_type: Optional[str] = None, limit: int = 10, api_key: Optional[str] = None, -) -> List[Dict[str, Any]]: +) -> List[SearchResult]: """ Search CivitAI for models (general search). Returns models that might be relevant. @@ -1245,7 +1250,7 @@ def search_civitai( for model in data.get("items", []): model_id = model.get("id") model_name = model.get("name", "") - model_type = model.get("type", "") + civitai_model_type = model.get("type", "") versions = model.get("modelVersions", []) if versions: @@ -1258,24 +1263,24 @@ def search_civitai( prefer_first_marked=True, ) - result = { - "source": "civitai", - "model_id": model_id, - "version_id": version_id, - "name": model_name, - "type": model_type, - "url": f"https://civitai.com/models/{model_id}", - "download_url": get_civitai_download_url(version_id, api_key), - "downloads": model.get("stats", {}).get("downloadCount", 0), - "base_model": latest.get("baseModel"), - "tags": model.get("tags", []), - } - - if primary_file: - result["filename"] = primary_file.get("name", "") - result["size"] = extract_file_size(primary_file) or 0 - - results.append(result) + results.append( + build_model_result( + "civitai", + model_id=model_id, + version_id=version_id, + name=model_name, + type=civitai_model_type, + filename=(primary_file or {}).get("name") or model_name, + url=f"https://civitai.com/models/{model_id}", + download_url=get_civitai_download_url(version_id, api_key), + size=(extract_file_size(primary_file) or 0) + if primary_file + else None, + base_model=latest.get("baseModel"), + tags=model.get("tags", []), + downloads=model.get("stats", {}).get("downloadCount", 0), + ) + ) except Exception as e: log.error(f"CivitAI search error: {e}") @@ -2337,7 +2342,7 @@ def build_civitai_custom_result( details: Dict[str, Any], expected_filename: str = "", api_key: Optional[str] = None, -) -> Optional[Dict[str, Any]]: +) -> Optional[SearchResult]: if not isinstance(details, dict): return None @@ -2394,7 +2399,7 @@ def resolve_civitai_version_custom_result( version_id: int, expected_filename: str = "", api_key: Optional[str] = None, -) -> Optional[Dict[str, Any]]: +) -> Optional[SearchResult]: if not version_id: return None diff --git a/core/sources/huggingface.py b/core/sources/huggingface.py index 2fbe968..21a3f04 100644 --- a/core/sources/huggingface.py +++ b/core/sources/huggingface.py @@ -13,6 +13,7 @@ from concurrent.futures import ThreadPoolExecutor, as_completed from typing import Any, Callable, Dict, List, Optional from urllib.parse import quote, unquote, urlparse +from ..contracts import HuggingFaceFileReference, SearchResult from ..log_system import create_module_logger from ..progress import get_progress_reporter @@ -443,7 +444,7 @@ def refresh_known_author_fallback_indexes( return status -def parse_huggingface_url(url: str) -> Optional[Dict[str, str]]: +def parse_huggingface_url(url: str) -> Optional[HuggingFaceFileReference]: """ Parse a HuggingFace URL to extract repo and filename. @@ -453,12 +454,15 @@ def parse_huggingface_url(url: str) -> Optional[Dict[str, str]]: - hf://user/repo/file.safetensors Returns: - Dictionary with 'repo' and 'filename' keys, or None if not HF URL + Typed repository-file identity, or None if not an HF URL """ if url.startswith("hf://"): parts = url[5:].split("/", 2) if len(parts) >= 3: - return {"repo": f"{parts[0]}/{parts[1]}", "filename": parts[2]} + return HuggingFaceFileReference( + repo=f"{parts[0]}/{parts[1]}", + filename=parts[2], + ) return None parsed = urlparse(url) @@ -467,11 +471,11 @@ def parse_huggingface_url(url: str) -> Optional[Dict[str, str]]: match = re.match(r"^/([^/]+/[^/]+)/(resolve|blob)/([^/]+)/(.+)$", parsed.path) if match: - return { - "repo": match.group(1), - "branch": match.group(3), - "filename": match.group(4), - } + return HuggingFaceFileReference( + repo=match.group(1), + branch=match.group(3), + filename=match.group(4), + ) return None @@ -700,7 +704,7 @@ def _build_huggingface_result( file_info: Dict[str, Any], match_type: str, headers: Optional[Dict[str, str]] = None, -) -> Dict[str, Any]: +) -> SearchResult: sha256 = _extract_huggingface_file_sha256(file_info) download_url = get_huggingface_download_url(repo_id, file_path) size = resolve_file_size( @@ -765,10 +769,10 @@ def get_huggingface_file_sha256( if not parsed: return "" - repo_id = str(parsed.get("repo") or "").strip() - branch = str(parsed.get("branch") or "main").strip() or "main" + repo_id = parsed.repo + branch = parsed.branch target_path = _normalize_huggingface_path( - unquote(str(parsed.get("filename") or "")) + unquote(parsed.filename) ) if not repo_id or not target_path: return "" @@ -787,7 +791,7 @@ def _find_matching_file_in_repo( filename: str, exact_only: bool = False, headers: Optional[Dict[str, str]] = None, -) -> Optional[Dict[str, Any]]: +) -> Optional[SearchResult]: filename_lower = filename.lower() filename_base = os.path.splitext(filename_lower)[0] partial_match = None @@ -845,7 +849,7 @@ def _find_matching_file_in_author_index( exact_only: bool = False, headers: Optional[Dict[str, str]] = None, sha256: Optional[str] = None, -) -> Optional[Dict[str, Any]]: +) -> Optional[SearchResult]: files = index.get("files") or [] requested_sha256 = normalize_sha256(sha256) @@ -988,9 +992,9 @@ def _search_brave_for_huggingface_candidates( parsed = parse_huggingface_url(candidate_url) if parsed: - repo_id = parsed.get("repo", "") - file_path = parsed.get("filename", "") - branch = parsed.get("branch", "main") + repo_id = parsed.repo + file_path = parsed.filename + branch = parsed.branch else: parsed_url = urlparse(candidate_url) if "huggingface.co" not in parsed_url.netloc: @@ -1032,7 +1036,7 @@ def search_huggingface_for_file( force_refresh: bool = False, progress_callback: Optional[Callable[[Dict[str, Any]], None]] = None, sha256: Optional[str] = None, -) -> Optional[Dict[str, Any]]: +) -> Optional[SearchResult]: """ Search HuggingFace for a specific model file. Returns the first repo that contains a matching file. @@ -1044,7 +1048,7 @@ def search_huggingface_for_file( If False, also try partial matching (for local file resolution). Returns: - Dict with url, repo, filename if found, None otherwise + A typed SearchResult if a matching file is found, otherwise None. """ global _search_cache @@ -1069,14 +1073,14 @@ def search_huggingface_for_file( return _search_cache[cache_key] try: - def result_matches_requested_hash(result: Optional[Dict[str, Any]]) -> bool: + def result_matches_requested_hash(result: Optional[SearchResult]) -> bool: if not requested_sha256: return True return normalize_sha256( - result.get("sha256") - or result.get("hash") - or (result.get("hashes") or {}).get("SHA256") - or (result.get("hashes") or {}).get("sha256") + result.sha256 + or result.extra_value("hash") + or (result.hashes or {}).get("SHA256") + or (result.hashes or {}).get("sha256") ) == requested_sha256 headers = {} @@ -1087,7 +1091,7 @@ def search_huggingface_for_file( def find_author_index_result( author: str, - ) -> tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]]]: + ) -> tuple[Optional[Dict[str, Any]], Optional[SearchResult]]: index = _get_author_index( author, headers={}, @@ -1155,13 +1159,13 @@ def search_huggingface_for_file( "found", "Found HuggingFace SHA-256 match", 92, - repo=result.get("repo_id"), - match_type=result.get("match_type"), + repo=result.extra_value("repo_id"), + match_type=result.match_type, ) log.info( f"HuggingFace found source=author_index " - f"match=hash file={result['filename']} " - f"repo={result['repo_id']} path={result['path']}" + f"match=hash file={result.filename} " + f"repo={result.extra_value('repo_id')} path={result.extra_value('path')}" ) return result @@ -1250,10 +1254,10 @@ def search_huggingface_for_file( "Found HuggingFace match", 92, repo=repo_id, - match_type=result.get("match_type"), + match_type=result.match_type, ) log.info( - f"HuggingFace found match={result['match_type']} file={filename} repo={repo_id} path={result['path']}" + f"HuggingFace found match={result.match_type} file={filename} repo={repo_id} path={result.extra_value('path')}" ) return result @@ -1295,11 +1299,11 @@ def search_huggingface_for_file( "found", "Found HuggingFace fallback match", 92, - repo=result.get("repo_id"), - match_type=result.get("match_type"), + repo=result.extra_value("repo_id"), + match_type=result.match_type, ) log.info( - f"HuggingFace found source=author_index match={result['match_type']} file={filename} repo={result['repo_id']} path={result['path']}" + f"HuggingFace found source=author_index match={result.match_type} file={filename} repo={result.extra_value('repo_id')} path={result.extra_value('path')}" ) return result @@ -1356,10 +1360,10 @@ def search_huggingface_for_file( "Found HuggingFace fallback match", 92, repo=repo_id, - match_type=result.get("match_type"), + match_type=result.match_type, ) log.info( - f"HuggingFace found source=author_fallback match={result['match_type']} file={filename} repo={repo_id} path={result['path']}" + f"HuggingFace found source=author_fallback match={result.match_type} file={filename} repo={repo_id} path={result.extra_value('path')}" ) return result @@ -1411,7 +1415,7 @@ def search_huggingface_for_file( ) if ( result - and result.get("filename", "").lower() == filename.lower() + and result.filename.lower() == filename.lower() and result_matches_requested_hash(result) ): _search_cache[cache_key] = result @@ -1421,10 +1425,10 @@ def search_huggingface_for_file( "Found HuggingFace Brave match", 92, repo=repo_id, - match_type=result.get("match_type"), + match_type=result.match_type, ) log.info( - f"HuggingFace found source=brave_fallback match=exact file={filename} repo={repo_id} path={result['path']}" + f"HuggingFace found source=brave_fallback match=exact file={filename} repo={repo_id} path={result.extra_value('path')}" ) return result @@ -1461,7 +1465,7 @@ def search_huggingface( model_type: Optional[str] = None, limit: int = 10, token: Optional[str] = None, -) -> List[Dict[str, Any]]: +) -> List[SearchResult]: """ Search HuggingFace Hub for models (general search). Returns repos that might be relevant, not guaranteed to have exact file. @@ -1486,14 +1490,14 @@ def search_huggingface( repo_id = model.get("id", "") results.append( - { - "source": "huggingface", - "repo": repo_id, - "name": model.get("modelId", repo_id), - "downloads": model.get("downloads", 0), - "likes": model.get("likes", 0), - "url": f"https://huggingface.co/{repo_id}", - } + build_model_result( + "huggingface", + name=model.get("modelId", repo_id), + url=f"https://huggingface.co/{repo_id}", + repo=repo_id, + downloads=model.get("downloads", 0), + likes=model.get("likes", 0), + ) ) except Exception as e: @@ -1506,16 +1510,16 @@ def build_huggingface_custom_result( url: str, expected_filename: str = "", token: Optional[str] = None, -) -> Optional[Dict[str, Any]]: +) -> Optional[SearchResult]: from urllib.parse import quote, unquote parsed = parse_huggingface_url(url) if not parsed: return None - repo_id = parsed.get("repo") or "" - branch = parsed.get("branch") or "main" - file_path = unquote(parsed.get("filename") or "") + repo_id = parsed.repo + branch = parsed.branch + file_path = unquote(parsed.filename) if not repo_id or not file_path: return None diff --git a/core/sources/lora_manager_archive.py b/core/sources/lora_manager_archive.py index e002ca1..fb3d80e 100644 --- a/core/sources/lora_manager_archive.py +++ b/core/sources/lora_manager_archive.py @@ -11,6 +11,7 @@ import sqlite3 import time from typing import Any, Callable, Dict, List, Optional +from ..contracts import SearchResult from ..log_system import create_module_logger from ..matcher import ( base_model_matches as _base_model_matches, @@ -386,7 +387,7 @@ def _load_full_rows_for_versions( def _build_result_from_row( conn: sqlite3.Connection, row: sqlite3.Row, query: str -) -> Optional[Dict[str, Any]]: +) -> Optional[SearchResult]: try: model_data = json.loads(row["model_data"] or "{}") except Exception: @@ -469,7 +470,7 @@ def search_lora_manager_archive( base_model_context: Optional[str] = None, limit: int = 10, progress_callback: Optional[Callable[[Dict[str, Any]], None]] = None, -) -> List[Dict[str, Any]]: +) -> List[SearchResult]: """ Search the local comfyui-lora-manager CivitAI archive by model name. @@ -613,19 +614,19 @@ def search_lora_manager_archive( if not full_row: continue result = _build_result_from_row(conn, full_row, normalized_query) - if not result or result["confidence"] < 40: + if not result or result.confidence < 40: continue if base_model_context and not _base_model_matches( - result.get("base_model"), base_model_context + result.base_model, base_model_context ): continue results.append(result) results.sort( key=lambda item: ( - _base_model_score(item.get("base_model"), base_model_context), - item.get("confidence", 0), - 1 if item.get("match_type") == "exact" else 0, + _base_model_score(item.base_model, base_model_context), + item.confidence, + 1 if item.match_type == "exact" else 0, ), reverse=True, ) @@ -666,7 +667,7 @@ def search_lora_manager_archive_for_file( limit: int = 10, progress_callback: Optional[Callable[[Dict[str, Any]], None]] = None, sha256: Optional[str] = None, -) -> Optional[Dict[str, Any]]: +) -> Optional[SearchResult]: """ Search the archive for the best model/version match for a filename. """ @@ -696,16 +697,16 @@ def search_lora_manager_archive_for_file( matching_size = extract_file_size(matching_file) if matching_size is None: matching_size = matching_file.get("size") - result.update( - { - "filename": matching_file.get("name") or result.get("filename", ""), - "download_url": matching_file.get("downloadUrl") or result.get("download_url"), - "size": matching_size if matching_size is not None else result.get("size"), - "sha256": requested_sha256, - "hashes": matching_file.get("hashes") or {"SHA256": requested_sha256}, - "match_type": "hash", - "confidence": 100.0, - } + result = result.with_updates( + filename=matching_file.get("name") or result.filename, + download_url=matching_file.get("downloadUrl") + or result.download_url, + size=matching_size if matching_size is not None else result.size, + sha256=requested_sha256, + hashes=matching_file.get("hashes") + or {"SHA256": requested_sha256}, + match_type="hash", + confidence=100.0, ) _report_progress( progress_callback, @@ -774,19 +775,29 @@ def search_lora_manager_archive_for_file( candidate_index=candidate_index, candidate_count=total_candidates, ) - candidate_filename = candidate.get("filename", "") - confidence = calculate_archived_model_confidence(filename, candidate.get("name", ""), candidate.get("version_name", ""), candidate_filename) + candidate_filename = candidate.filename + confidence = calculate_archived_model_confidence( + filename, + candidate.name, + candidate.version_name, + candidate_filename, + ) should_update, rank, base_model_matches = should_update_best_match( - confidence, candidate.get("base_model"), base_model_context, best_rank, exact_only + confidence, + candidate.base_model, + base_model_context, + best_rank, + exact_only, ) if base_model_context and not base_model_matches: continue if should_update: best_rank = rank best_confidence = confidence - best_match = dict(candidate) - best_match["confidence"] = confidence - best_match["match_type"] = "exact" if confidence == 100.0 else "similar" + best_match = candidate.with_updates( + confidence=confidence, + match_type="exact" if confidence == 100.0 else "similar", + ) if confidence == 100.0 and base_model_matches: _report_progress( progress_callback, diff --git a/core/sources/model_list.py b/core/sources/model_list.py index 6334e29..30614a3 100644 --- a/core/sources/model_list.py +++ b/core/sources/model_list.py @@ -8,8 +8,10 @@ import hashlib import os from typing import Any, Dict, List, Optional +from ..contracts import SearchResult from ..log_system import create_module_logger from ..matcher import calculate_similarity, normalize_filename +from ..type_utils import build_model_result log = create_module_logger(__name__) @@ -177,21 +179,28 @@ def update_model_list_from_remote() -> Dict[str, Any]: -def _build_model_list_result(model: Dict[str, Any]) -> Dict[str, Any]: - return { - "source": "model_list", - "filename": model.get("filename", ""), - "url": model.get("url", ""), - "name": model.get("name", ""), - "type": model.get("type", ""), - "directory": model.get("save_path", "checkpoints"), - "size": model.get("size", ""), - } +def _build_model_list_result( + model: Dict[str, Any], + *, + match_type: str = "similar", + confidence: float = 0.0, +) -> SearchResult: + return build_model_result( + "model_list", + filename=model.get("filename", ""), + name=model.get("name", ""), + url=model.get("url", ""), + type=model.get("type", ""), + directory=model.get("save_path", "checkpoints"), + size=model.get("size"), + match_type=match_type, + confidence=confidence, + ) def search_model_list( filename: str, exact_only: bool = False -) -> Optional[Dict[str, Any]]: +) -> Optional[SearchResult]: """ Search model-list.json for a model by filename. Uses exact match first, then fuzzy matching (unless exact_only=True). @@ -202,7 +211,7 @@ def search_model_list( If False, also try fuzzy matching (for local file resolution). Returns: - Dict with url, filename, type, etc. if found, None otherwise + Typed search result with URL, filename, type, etc. if found. """ models = _load_model_list() if not models: @@ -218,10 +227,7 @@ def search_model_list( if model_filename.lower() == filename_lower: url = model.get("url", "") if url: - return { - **_build_model_list_result(model), - "match_type": "exact", - } + return _build_model_list_result(model, match_type="exact") # If exact_only is True, don't try fuzzy matching - prevents confusing # users with wrong model suggestions for downloads @@ -242,11 +248,11 @@ def search_model_list( url = model.get("url", "") if url: score = calculate_similarity(filename_norm, normalize_filename(model_filename)) - return { - **_build_model_list_result(model), - "match_type": "fuzzy", - "confidence": round(score * 100, 1), - } + return _build_model_list_result( + model, + match_type="fuzzy", + confidence=round(score * 100, 1), + ) # 3. Try normalized similarity matching on all models (fallback) best_match = None @@ -264,16 +270,19 @@ def search_model_list( url = model.get("url", "") if url: best_score = score - best_match = { - **_build_model_list_result(model), - "match_type": "similar", - "confidence": round(score * 100, 1), - } + best_match = _build_model_list_result( + model, + match_type="similar", + confidence=round(score * 100, 1), + ) return best_match -def search_model_list_multiple(filename: str, limit: int = 5) -> List[Dict[str, Any]]: +def search_model_list_multiple( + filename: str, + limit: int = 5, +) -> List[SearchResult]: """ Search model-list.json and return multiple fuzzy matches. @@ -303,14 +312,14 @@ def search_model_list_multiple(filename: str, limit: int = 5) -> List[Dict[str, url = model.get("url", "") if url: results.append( - { - **_build_model_list_result(model), - "confidence": round(score * 100, 1), - } + _build_model_list_result( + model, + confidence=round(score * 100, 1), + ) ) # Sort by confidence descending - results.sort(key=lambda x: x["confidence"], reverse=True) + results.sort(key=lambda x: x.confidence, reverse=True) return results[:limit] diff --git a/core/sources/popular.py b/core/sources/popular.py index 887c803..4671fc6 100644 --- a/core/sources/popular.py +++ b/core/sources/popular.py @@ -8,6 +8,7 @@ import os import re from typing import Any, Dict, List, Optional, Set +from ..contracts import ModelCatalogEntry, SearchResult from ..log_system import create_module_logger log = create_module_logger(__name__) @@ -16,7 +17,7 @@ log = create_module_logger(__name__) from ..catalog_manager import CatalogManager from ..network_utils import request_source_json from ..path_utils import METADATA_DIR, read_json_safe -from ..type_utils import normalize_alphanumeric_key +from ..type_utils import build_model_result, normalize_alphanumeric_key POPULAR_MODELS_FILE = os.path.join(METADATA_DIR, "popular-models.json") MODEL_ALIASES_FILE = os.path.join(METADATA_DIR, "model-aliases.json") @@ -129,7 +130,7 @@ def load_base_model_aliases() -> Dict[str, List[str]]: return _base_models_aliases_cache -def get_popular_model_url(filename: str) -> Optional[Dict[str, Any]]: +def get_popular_model_url(filename: str) -> Optional[ModelCatalogEntry]: """ Look up a model filename in the popular models database. @@ -137,19 +138,30 @@ def get_popular_model_url(filename: str) -> Optional[Dict[str, Any]]: filename: Model filename to look up Returns: - Dictionary with url, type, directory if found, None otherwise + Typed catalog entry with URL, type and directory if found """ models = _load_popular_models() # Direct lookup if filename in models: - return models[filename].copy() + try: + return ModelCatalogEntry.from_mapping( + models[filename], + filename=filename, + ) + except (TypeError, ValueError) as exc: + log.warning(f"Invalid popular model entry for {filename}: {exc}") + return None # Try lowercase filename_lower = filename.lower() for name, info in models.items(): if name.lower() == filename_lower: - return info.copy() + try: + return ModelCatalogEntry.from_mapping(info, filename=name) + except (TypeError, ValueError) as exc: + log.warning(f"Invalid popular model entry for {name}: {exc}") + return None # Try aliases aliases = _load_model_aliases() @@ -158,14 +170,21 @@ def get_popular_model_url(filename: str) -> Optional[Dict[str, Any]]: filename in alias_list or filename_lower in [a.lower() for a in alias_list] ) and canonical in models: - result = models[canonical].copy() - result["canonical_name"] = canonical - return result + try: + return ModelCatalogEntry.from_mapping( + models[canonical], + filename=canonical, + ).with_extra(canonical_name=canonical) + except (TypeError, ValueError) as exc: + log.warning( + f"Invalid popular model entry for {canonical}: {exc}" + ) + return None return None -def search_popular_models(query: str, limit: int = 10) -> List[Dict[str, Any]]: +def search_popular_models(query: str, limit: int = 10) -> List[SearchResult]: """ Search popular models database by filename pattern. @@ -174,7 +193,7 @@ def search_popular_models(query: str, limit: int = 10) -> List[Dict[str, Any]]: limit: Maximum results to return Returns: - List of matching models with url info + Typed search results with URL information """ models = _load_popular_models() query_lower = query.lower() @@ -182,9 +201,29 @@ def search_popular_models(query: str, limit: int = 10) -> List[Dict[str, Any]]: results = [] for name, info in models.items(): if query_lower in name.lower(): - result = info.copy() - result["filename"] = name - results.append(result) + try: + catalog_entry = ModelCatalogEntry.from_mapping( + info, + filename=name, + ) + except (TypeError, ValueError) as exc: + log.warning(f"Invalid popular model entry for {name}: {exc}") + continue + results.append( + build_model_result( + "popular", + filename=catalog_entry.filename or name, + name=catalog_entry.name or name, + url=catalog_entry.url or catalog_entry.download_url or "", + download_url=catalog_entry.download_url, + type=catalog_entry.model_type, + size=catalog_entry.size, + match_type=catalog_entry.match_type or "similar", + confidence=catalog_entry.confidence, + directory=catalog_entry.directory, + **dict(catalog_entry.extra), + ) + ) if len(results) >= limit: break diff --git a/core/type_utils.py b/core/type_utils.py index 3373e56..28f328e 100644 --- a/core/type_utils.py +++ b/core/type_utils.py @@ -8,6 +8,8 @@ import math import re from typing import Any, Dict, List +from .contracts import ProviderUrlReference, SearchResult + def as_dict(value: Any) -> Dict[str, Any]: """ @@ -706,7 +708,10 @@ def extract_response_file_size(response: Any) -> Optional[int]: return None -def parse_civitai_model_path(path: str, query_string: str) -> Optional[Dict[str, int]]: +def parse_civitai_model_path( + path: str, + query_string: str, +) -> Optional[ProviderUrlReference]: """ Common helper to parse `/models/{model_id}?modelVersionId={version_id}` URLs for CivitAI and CivArchive. @@ -716,15 +721,17 @@ def parse_civitai_model_path(path: str, query_string: str) -> Optional[Dict[str, model_match = re.search(r"/models/(\d+)", path) if model_match: - result = {"model_id": int(model_match.group(1))} + model_id = int(model_match.group(1)) query = parse_qs(query_string) version_id_list = query.get("modelVersionId") if version_id_list: try: - result["version_id"] = int(version_id_list[0]) + version_id = int(version_id_list[0]) except (ValueError, TypeError): - pass - return result + version_id = None + else: + version_id = None + return ProviderUrlReference(model_id=model_id, version_id=version_id) return None @@ -887,8 +894,10 @@ def normalize_sha256(value: Any) -> str: """Return a normalized SHA256 hex string or an empty string.""" if value is None: return "" + if not isinstance(value, str): + return "" - text = str(value).strip() + text = value.strip() for prefix in ("sha256:", "sha256="): if text.lower().startswith(prefix): text = text[len(prefix):].strip() @@ -899,19 +908,19 @@ def normalize_sha256(value: Any) -> str: return text.lower() if SHA256_PATTERN.match(text) else "" -def extract_file_sha256(file_info: Any) -> Any: +def extract_file_sha256(file_info: Any) -> str: """Extract a provider file's SHA256 value using the common field precedence.""" if not isinstance(file_info, dict): return "" hashes = file_info.get("hashes") if isinstance(file_info.get("hashes"), dict) else {} - return ( - file_info.get("sha256") - or file_info.get("hash") - or hashes.get("SHA256") - or hashes.get("sha256") - or "" + value = first_non_empty( + file_info.get("sha256"), + file_info.get("hash"), + hashes.get("SHA256"), + hashes.get("sha256"), ) + return value if isinstance(value, str) else "" def normalize_hashes_dict(hashes: Optional[Dict[str, Any]]) -> Dict[str, str]: @@ -1310,7 +1319,7 @@ def normalize_model_file_info( ) -> Dict[str, Any]: """Normalize a single model file dictionary to a unified schema.""" hashes = raw.get("hashes") if isinstance(raw.get("hashes"), dict) else {} - sha256 = raw.get("sha256") or hashes.get("SHA256") or hashes.get("sha256") or "" + sha256 = extract_file_sha256(raw) size = raw.get("size") if size is None and raw.get("sizeKB") is not None: @@ -1331,16 +1340,20 @@ def normalize_model_file_info( "sha256": sha256, "hashes": hashes, "metadata": raw.get("metadata") if isinstance(raw.get("metadata"), dict) else {}, - "model_id": model_id or raw.get("modelId"), - "version_id": version_id or raw.get("modelVersionId"), + "model_id": model_id if model_id is not None else raw.get("modelId"), + "version_id": ( + version_id + if version_id is not None + else raw.get("modelVersionId") + ), } def build_model_result( source: str, *, - model_id: Any = None, - version_id: Any = None, + model_id: int | str | None = None, + version_id: int | str | None = None, name: str = "", version_name: str = "", type: str = "", @@ -1362,79 +1375,46 @@ def build_model_result( custom_url: bool = False, result_mode: str = "search", **extra: Any, -) -> Dict[str, Any]: +) -> SearchResult: """Build the shared model-result contract for search and custom URL results.""" if normalize_hashes: hashes = normalize_hashes_dict(hashes) if not sha256 and "sha256" in hashes: sha256 = hashes["sha256"] - if result_mode not in {"search", "custom_url", "compact_custom_url"}: - raise ValueError(f"Unsupported model result mode: {result_mode}") - - if result_mode == "compact_custom_url": - result = { + result = SearchResult.from_mapping( + { "source": source, - "details_source": details_source or source, + "model_id": model_id, + "version_id": version_id, "name": name, + "version_name": version_name, + "type": type, "filename": filename, "url": url, - "version_url": version_url or url, "download_url": download_url, + "size": size, + "base_model": base_model, + "tags": tags, "match_type": match_type, + "confidence": confidence, + "sha256": sha256, + "hashes": hashes, + "trained_words": trained_words, + "images": images, + "details_source": details_source, + "version_url": version_url, "custom_url": custom_url, + "result_mode": result_mode, } - if model_id is not None: - result["model_id"] = model_id - if version_id is not None: - result["version_id"] = version_id - result.update(extra) - return result - - result = { - "source": source, - "model_id": model_id, - "version_id": version_id, - "name": name, - "version_name": version_name, - "type": type, - "filename": filename, - "url": url, - "download_url": download_url, - "size": size, - "base_model": base_model, - "tags": tags or [], - "match_type": match_type, - "confidence": confidence, - "sha256": sha256, - "hashes": hashes or {}, - "trained_words": trained_words or [], - "images": images or [], - } - - if result_mode == "custom_url": - result.pop("confidence", None) - result.update( - { - "details_source": details_source or source, - "version_url": version_url or url, - "match_type": "custom_url", - "custom_url": custom_url, - } - ) - else: - if details_source is not None: - result["details_source"] = details_source - if version_url is not None: - result["version_url"] = version_url - if custom_url: - result["custom_url"] = True - - result.update(extra) - return result + ) + return result.with_extra(**extra) -def parse_provider_model_url(url: str, allowed_domains: list[str]) -> Optional[Dict[str, Any]]: +def parse_provider_model_url( + url: str, + allowed_domains: list[str], +) -> Optional[ProviderUrlReference]: """Parse a provider model URL (e.g. CivitAI or mirror site) to extract model/version info.""" import re from urllib.parse import urlparse @@ -1457,13 +1437,13 @@ def parse_provider_model_url(url: str, allowed_domains: list[str]) -> Optional[D # Standard CivArchive sha256 mirrors pattern sha_match = re.search(r"/sha256/([a-fA-F0-9]{64})", parsed.path) if sha_match: - return {"sha256": sha_match.group(1).lower()} + return ProviderUrlReference(sha256=sha_match.group(1).lower()) # Standard CivitAI download pattern if "/api/download/models/" in parsed.path: match = re.search(r"/api/download/models/(\d+)", parsed.path) if match: - return {"version_id": int(match.group(1))} + return ProviderUrlReference(version_id=int(match.group(1))) return parse_civitai_model_path(parsed.path, parsed.query) diff --git a/core/workflow/analysis.py b/core/workflow/analysis.py index f833e2d..2294243 100644 --- a/core/workflow/analysis.py +++ b/core/workflow/analysis.py @@ -2,8 +2,9 @@ import hashlib import json -from typing import Any, Callable, Dict, List, Optional +from typing import Any, Callable, Dict, List, Optional, Sequence +from ..contracts import MissingModel, ModelReference, ResolvedModel from ..log_system import create_module_logger from . import references from . import subgraphs as subgraph_utils @@ -203,7 +204,7 @@ def _get_promoted_widget_context_cache_key( def analyze_workflow_models( workflow_json: Dict[str, Any], - available_models: Optional[List[Dict[str, Any]]] = None, + available_models: Optional[List[ResolvedModel]] = None, progress_callback: Optional[Callable[[Dict[str, Any]], None]] = None, *, previous_node_cache: Optional[ @@ -214,7 +215,7 @@ def analyze_workflow_models( ] = None, analysis_stats: Optional[Dict[str, int]] = None, analysis_context: Optional[str] = None, -) -> List[Dict[str, Any]]: +) -> List[ModelReference]: """ Extract all model references from a workflow, including nested subgraphs. @@ -224,10 +225,10 @@ def analyze_workflow_models( progress_callback: Optional callback called while workflow nodes are scanned Returns: - List of model reference dictionaries (same format as get_node_model_info) - Each dict includes 'subgraph_id' if the model is in a subgraph + Typed model references. Each reference includes ``subgraph_id`` when + the model is in a subgraph. """ - all_model_refs = [] + all_model_refs: List[ModelReference] = [] reused_nodes = 0 analyzed_nodes = 0 log_context = f" ({analysis_context})" if analysis_context else "" @@ -315,7 +316,12 @@ def analyze_workflow_models( and cached_node.get("fingerprint") == node_fingerprint ) if reused: - base_model_refs = cached_node.get("refs", []) + base_model_refs = [ + ref + if isinstance(ref, ModelReference) + else ModelReference.from_mapping(ref) + for ref in cached_node.get("refs", []) + ] reused_nodes += 1 else: base_model_refs = references.get_node_model_info( @@ -323,10 +329,9 @@ def analyze_workflow_models( available_models=available_models, ) analyzed_nodes += 1 - # Keep cached node references as raw extraction results. Subgraph - # instance context is workflow-level state and must be reapplied - # when a cached node is reused. - model_refs = [dict(ref) for ref in base_model_refs] + # Subgraph instance context is workflow-level state and must be + # reapplied when a cached node is reused. + model_refs = list(base_model_refs) node_type = node.get("type", "") # Check if node type is a subgraph UUID @@ -340,26 +345,32 @@ def analyze_workflow_models( # For top-level subgraph instance nodes, subgraph_path is None # This distinguishes them from nodes within subgraph definitions if subgraph_id: - for ref in model_refs: + for index, ref in enumerate(model_refs): context = promoted_widget_contexts.get( "instance_widgets", {} - ).get((str(node.get("id")), ref.get("widget_index"))) + ).get((str(node.get("id")), ref.widget_index)) if context: promoted_instance_slots.add( - (str(node.get("id")), ref.get("widget_index")) + (str(node.get("id")), ref.widget_index) ) - subgraph_utils._apply_instance_promoted_widget_context( - ref, context, available_models + model_refs[index] = ( + subgraph_utils._apply_instance_promoted_widget_context( + ref, + context, + available_models, + ) ) - for ref in model_refs: - ref["subgraph_id"] = subgraph_id - ref["subgraph_name"] = subgraph_name - ref["subgraph_path"] = None - ref["is_top_level"] = True + for index, ref in enumerate(model_refs): + model_refs[index] = ref.with_updates( + subgraph_id=subgraph_id, + subgraph_name=subgraph_name, + subgraph_path=None, + is_top_level=True, + ) if node_cache_out is not None: node_cache_out[node_cache_key] = { "fingerprint": node_fingerprint, - "refs": base_model_refs, + "refs": list(base_model_refs), } report_node_progress(node, reused=reused) all_model_refs.extend(model_refs) @@ -404,7 +415,12 @@ def analyze_workflow_models( and cached_node.get("fingerprint") == node_fingerprint ) if reused: - base_model_refs = cached_node.get("refs", []) + base_model_refs = [ + ref + if isinstance(ref, ModelReference) + else ModelReference.from_mapping(ref) + for ref in cached_node.get("refs", []) + ] reused_nodes += 1 reused_subgraph_nodes += 1 else: @@ -414,18 +430,19 @@ def analyze_workflow_models( ) analyzed_nodes += 1 analyzed_subgraph_nodes += 1 - model_refs = [] + model_refs: List[ModelReference] = [] for base_ref in base_model_refs: - scoped_ref = dict(base_ref) - scoped_ref["subgraph_id"] = subgraph_id - scoped_ref["subgraph_name"] = subgraph_name - scoped_ref["subgraph_path"] = [ - "definitions", - "subgraphs", - subgraph_id, - "nodes", - ] - scoped_ref["is_top_level"] = False + scoped_ref = base_ref.with_updates( + subgraph_id=subgraph_id, + subgraph_name=subgraph_name, + subgraph_path=[ + "definitions", + "subgraphs", + subgraph_id, + "nodes", + ], + is_top_level=False, + ) promoted_refs = subgraph_utils._promote_model_reference_to_instances( scoped_ref, promoted_widget_contexts, @@ -437,15 +454,14 @@ def analyze_workflow_models( continue ref = scoped_ref - subgraph_utils._apply_promoted_widget_locator( - ref, - promoted_widget_contexts, + ref = subgraph_utils._apply_promoted_widget_locator( + ref, promoted_widget_contexts ) model_refs.append(ref) if node_cache_out is not None: node_cache_out[node_cache_key] = { "fingerprint": node_fingerprint, - "refs": base_model_refs, + "refs": list(base_model_refs), } report_node_progress( node, @@ -484,8 +500,9 @@ def analyze_workflow_models( return all_model_refs def identify_missing_models( - workflow_models: List[Dict[str, Any]], available_models: List[Dict[str, str]] = None -) -> List[Dict[str, Any]]: + workflow_models: Sequence[ModelReference], + available_models: Optional[Sequence[ResolvedModel]] = None, +) -> List[MissingModel]: """ Identify which models from the workflow are missing. Deduplicates by category and filename - same model file only appears once even if @@ -496,48 +513,57 @@ def identify_missing_models( available_models: Optional list of available models (if None, checks via folder_paths) Returns: - List of missing model references (deduplicated by category and filename). - Each entry has 'all_node_refs' containing all node references for that model. + Typed missing model references (deduplicated by category and filename). + Each entry retains 'all_node_refs' for HTTP compatibility. """ # Group missing models by category and filename to deduplicate workflow # references without merging unrelated model folders that happen to share a name. - missing_by_model: Dict[tuple[str, str], Dict[str, Any]] = {} + missing_by_model: Dict[tuple[str, str], MissingModel] = {} - for model_ref in workflow_models: + for raw_model_ref in workflow_models: + model_ref = ( + raw_model_ref + if isinstance(raw_model_ref, ModelReference) + else ModelReference.from_mapping(raw_model_ref) + ) # If exists is False, it's missing - if not model_ref.get("exists", False): - filename = model_ref.get("original_path", "") - category = model_ref.get("category", "") + if not model_ref.exists: + filename = model_ref.original_path + category = model_ref.category group_key = (str(category or ""), str(filename or "")) if group_key not in missing_by_model: # First occurrence - use this as the primary entry - missing_by_model[group_key] = { - **model_ref, - "reference_count": 1, - "all_node_refs": [ - model_ref.copy() - ], # Track all nodes needing this model - } + missing_by_model[group_key] = MissingModel( + reference=model_ref, + all_node_refs=(model_ref,), + extra=dict(model_ref.extra), + ) else: # Duplicate - just add to the node refs list existing = missing_by_model[group_key] - existing["all_node_refs"].append(model_ref.copy()) - existing["reference_count"] = len(existing["all_node_refs"]) - if model_ref.get("auto_download_capable"): - existing["auto_download_capable"] = True - if model_ref.get("auto_download_candidate"): - existing["auto_download_candidate"] = True - if model_ref.get("input_choice_matches_value"): - existing["input_choice_matches_value"] = True - existing_source = str(existing.get("input_choice_source") or "").lower() - model_source = str(model_ref.get("input_choice_source") or "").lower() + updates = {} + for key in ( + "auto_download_capable", + "auto_download_candidate", + "input_choice_matches_value", + ): + if model_ref.extra_value(key): + updates[key] = True + existing_source = str( + existing.extra_value("input_choice_source") or "" + ).lower() + model_source = str( + model_ref.extra_value("input_choice_source") or "" + ).lower() if existing_source != "hybrid" and model_source in { "static", "hybrid", "workflow_schema", }: - existing["input_choice_source"] = model_source + updates["input_choice_source"] = model_source + missing_by_model[group_key] = existing.with_references( + (*existing.all_node_refs, model_ref) + ).with_extra(**updates) - # Return deduplicated list return list(missing_by_model.values()) diff --git a/core/workflow/inventory.py b/core/workflow/inventory.py index 49a27fc..b890c46 100644 --- a/core/workflow/inventory.py +++ b/core/workflow/inventory.py @@ -5,6 +5,7 @@ import time from collections import OrderedDict from typing import Any, Callable, Dict, List, Optional +from ..contracts import ModelReference, WorkflowModelInventory from ..log_system import create_module_logger from . import analysis @@ -17,18 +18,18 @@ _WORKFLOW_MODEL_INVENTORY_CACHE: OrderedDict[str, Dict[str, Any]] = OrderedDict( def _build_workflow_node_cache( workflow_json: Dict[str, Any], - model_refs: List[Dict[str, Any]], + model_refs: List[ModelReference], ) -> Dict[tuple[str, str, str], Dict[str, Any]]: fingerprints = analysis._get_workflow_node_fingerprints(workflow_json) - refs_by_node: Dict[tuple[str, str, str], List[Dict[str, Any]]] = { + refs_by_node: Dict[tuple[str, str, str], List[ModelReference]] = { key: [] for key in fingerprints } for ref in model_refs: - is_top_level = ref.get("is_top_level") is not False + is_top_level = ref.is_top_level is not False key = analysis._get_workflow_scope_key( - ref.get("node_id"), - subgraph_id=ref.get("subgraph_id"), + ref.node_id, + subgraph_id=ref.subgraph_id, is_top_level=is_top_level, ) if key in refs_by_node: @@ -71,16 +72,26 @@ def get_workflow_model_inventory( force_rescan: bool = False, progress_callback: Optional[Callable[[Dict[str, Any]], None]] = None, analysis_id: Optional[str] = None, -) -> Dict[str, List[Dict[str, Any]]]: +) -> WorkflowModelInventory: """ Return the shared base model inventory for a workflow. Missing Models and Loaded Models use this snapshot so the same unchanged workflow is not parsed separately by both endpoints. Callers must treat the - returned lists as read-only. + returned tuples as read-only. """ from ..scanner import get_model_files + def normalize_model_refs(values: List[Any]) -> List[ModelReference]: + if not isinstance(values, (list, tuple)): + raise TypeError("workflow model references must be an array") + return [ + value + if isinstance(value, ModelReference) + else ModelReference.from_mapping(value) + for value in values + ] + cache_key = _get_workflow_model_inventory_cache_key(workflow_json) analysis_context = _get_analysis_log_context(cache_key, analysis_id) promoted_context_key = analysis._get_promoted_widget_context_cache_key(workflow_json) @@ -124,7 +135,7 @@ def get_workflow_model_inventory( previous_inventory = entry if cached is not None: - model_refs = cached["model_refs"] + model_refs = normalize_model_refs(cached["model_refs"]) if progress_callback: progress_callback( { @@ -138,15 +149,15 @@ def get_workflow_model_inventory( log.debug( f"Reusing shared workflow model analysis ({analysis_context})" ) - return { - "available_models": cached["available_models"], - "model_refs": model_refs, - } + return WorkflowModelInventory( + available_models=tuple(cached["available_models"]), + model_refs=tuple(model_refs), + ) node_cache = {} analysis_stats = {} if previous_inventory is not None: - available_models = previous_inventory["available_models"] + available_models = list(previous_inventory["available_models"]) model_refs = analysis.analyze_workflow_models( workflow_json, available_models=available_models, @@ -156,6 +167,7 @@ def get_workflow_model_inventory( analysis_stats=analysis_stats, analysis_context=analysis_context, ) + model_refs = normalize_model_refs(model_refs) else: available_models = get_model_files(force_rescan=force_rescan) model_refs = analysis.analyze_workflow_models( @@ -164,6 +176,7 @@ def get_workflow_model_inventory( progress_callback=progress_callback, analysis_context=analysis_context, ) + model_refs = normalize_model_refs(model_refs) node_cache = _build_workflow_node_cache(workflow_json, model_refs) analysis_stats = { "total_nodes": len(node_cache), @@ -196,7 +209,7 @@ def get_workflow_model_inventory( ): _WORKFLOW_MODEL_INVENTORY_CACHE.popitem(last=False) - return { - "available_models": available_models, - "model_refs": model_refs, - } + return WorkflowModelInventory( + available_models=tuple(available_models), + model_refs=tuple(model_refs), + ) diff --git a/core/workflow/references.py b/core/workflow/references.py index cdfb59f..3d327eb 100644 --- a/core/workflow/references.py +++ b/core/workflow/references.py @@ -4,11 +4,14 @@ import os import re from typing import Any, Dict, List, Optional +from ..contracts import ModelReference, ResolvedModel from ..custom_nodes import ( analyze_custom_node_references, get_custom_node_model_adapter, ) from ..log_system import create_module_logger +from ..metadata_model_utils import normalize_models +from ..path_utils import normalize_string_values from ..type_utils import MODEL_EXTENSIONS, URN_TYPE_MAP, normalize_download_category from . import dynamic_widgets from .widgets import ( @@ -209,7 +212,7 @@ def _normalize_model_path_for_lookup(value: str) -> str: def _resolve_from_available_models( filename: str, categories: Optional[List[str]], - available_models: Optional[List[Dict[str, Any]]], + available_models: Optional[List[ResolvedModel]], ) -> Optional[tuple[str, str]]: """Resolve using scanner data with case-sensitive relative path matching.""" if not available_models: @@ -223,11 +226,13 @@ def _resolve_from_available_models( if requested_is_absolute: requested_key = _normalize_model_path_for_lookup(os.path.abspath(filename)) + normalized_models = normalize_models(available_models) + if categories is None: category_order = [] seen_categories = set() - for model in available_models: - category = model.get("category") + for model in normalized_models: + category = model.category if category and category not in seen_categories: seen_categories.add(category) category_order.append(category) @@ -239,8 +244,8 @@ def _resolve_from_available_models( for category in category_order: normalized_category = normalize_download_category(category) - for model in available_models: - model_category = model.get("category") + for model in normalized_models: + model_category = model.category if ( model_category != category and normalize_download_category(model_category) != normalized_category @@ -248,14 +253,14 @@ def _resolve_from_available_models( continue if requested_is_absolute: - model_path = model.get("path") or "" + model_path = model.path or "" else: - model_path = model.get("relative_path") or model.get("filename") or "" + model_path = model.relative_path or model.filename or "" if _normalize_model_path_for_lookup(model_path) != requested_key: continue - full_path = model.get("path") + full_path = model.path if full_path and os.path.exists(full_path): return (category, full_path) @@ -272,7 +277,9 @@ def _category_has_exact_filename(category: str, filename: str) -> Optional[bool] return False try: - available_filenames = folder_paths.get_filename_list(category) or [] + available_filenames = normalize_string_values( + folder_paths.get_filename_list(category) + ) except Exception: return None @@ -286,7 +293,7 @@ def _category_has_exact_filename(category: str, filename: str) -> Optional[bool] def try_resolve_model_path( value: str, categories: List[str] = None, - available_models: Optional[List[Dict[str, Any]]] = None, + available_models: Optional[List[ResolvedModel]] = None, ) -> Optional[tuple[str, str]]: """Resolve a workflow model value using scanner data or ComfyUI paths.""" if not isinstance(value, str) or not value.strip(): @@ -336,8 +343,8 @@ def try_resolve_model_path( return None def get_node_model_info( - node: Dict[str, Any], available_models: Optional[List[Dict[str, Any]]] = None -) -> List[Dict[str, Any]]: + node: Dict[str, Any], available_models: Optional[List[ResolvedModel]] = None +) -> List[ModelReference]: """ Extract model references from a single node. @@ -348,16 +355,8 @@ def get_node_model_info( node: Node dictionary from workflow JSON Returns: - List of model reference dictionaries: - { - 'node_id': node id, - 'node_type': node type, - 'widget_index': index in widgets_values, - 'original_path': original path from workflow, - 'category': model category (if found), - 'exists': True if model exists, - 'connected': True if node has any connected inputs/outputs - } + Typed model references. Provider/custom-node-specific fields are kept + in ``ModelReference.extra``. """ model_refs = [] node_id = node.get("id") @@ -380,7 +379,7 @@ def get_node_model_info( is_active = is_connected and not is_bypassed if not widgets_values: - return model_refs + return [] custom_model_refs = analyze_custom_node_references( node, @@ -389,7 +388,10 @@ def get_node_model_info( get_widget_name_hint=get_widget_name_hint, ) if custom_model_refs is not None: - return custom_model_refs + return [ + ref if isinstance(ref, ModelReference) else ModelReference.from_mapping(ref) + for ref in custom_model_refs + ] # For each widget value, check if it looks like a model file or URN for idx, value in enumerate(widgets_values): @@ -646,4 +648,7 @@ def get_node_model_info( } ) - return model_refs + return [ + ref if isinstance(ref, ModelReference) else ModelReference.from_mapping(ref) + for ref in model_refs + ] diff --git a/core/workflow/subgraphs.py b/core/workflow/subgraphs.py index d333761..25aeedd 100644 --- a/core/workflow/subgraphs.py +++ b/core/workflow/subgraphs.py @@ -3,6 +3,7 @@ import os from typing import Any, Dict, List, Optional +from ..contracts import ModelReference, ResolvedModel from . import references from .widgets import ( _has_widget_input, @@ -330,46 +331,50 @@ def _build_promoted_widget_contexts( def _apply_instance_promoted_widget_context( - ref: Dict[str, Any], + ref: ModelReference, context: Optional[Dict[str, Any]], - available_models: Optional[List[Dict[str, Any]]], -) -> None: + available_models: Optional[List[ResolvedModel]], +) -> ModelReference: if not context: - return + return ref - ref["promoted_widget_name"] = context.get("proxy_widget_name", "") - ref["promoted_inner_node_id"] = context.get("node_id") - ref["promoted_inner_node_type"] = context.get("node_type", "") - ref["promoted_inner_node_title"] = context.get("node_title", "") - ref["promoted_inner_widget_index"] = context.get("widget_index") - ref["promoted_inner_widget_name"] = context.get("widget_name", "") + updates = { + "promoted_widget_name": context.get("proxy_widget_name", ""), + "promoted_inner_node_id": context.get("node_id"), + "promoted_inner_node_type": context.get("node_type", ""), + "promoted_inner_node_title": context.get("node_title", ""), + "promoted_inner_widget_index": context.get("widget_index"), + "promoted_inner_widget_name": context.get("widget_name", ""), + } promoted_value = context.get("promoted_value") - nested_key = ref.get("nested_key") + nested_key = ref.extra_value("nested_key") if nested_key and isinstance(promoted_value, dict): promoted_value = promoted_value.get(nested_key) if _has_promoted_value(promoted_value): - ref["original_path"] = str(promoted_value).strip() - ref["is_urn"] = bool(references.URN_REGEX.match(ref["original_path"])) + original_path = str(promoted_value).strip() + updates["original_path"] = original_path + updates["is_urn"] = bool(references.URN_REGEX.match(original_path)) category = context.get("category") - original_path = ref.get("original_path", "") + original_path = updates.get("original_path", ref.original_path) if not category or not original_path: - return + return ref.with_updates(**updates) resolved = references.try_resolve_model_path( original_path, [category], available_models=available_models, ) - ref["category"] = category + updates["category"] = category if resolved: _, full_path = resolved - ref["full_path"] = full_path - ref["exists"] = os.path.exists(full_path) + updates["full_path"] = full_path + updates["exists"] = os.path.exists(full_path) else: - ref["full_path"] = None - ref["exists"] = False + updates["full_path"] = None + updates["exists"] = False + return ref.with_updates(**updates) def _select_promoted_locator( @@ -385,37 +390,37 @@ def _select_promoted_locator( def _get_promoted_widget_locators( - ref: Dict[str, Any], contexts: Dict[str, Dict[Any, Any]] + ref: ModelReference, contexts: Dict[str, Dict[Any, Any]] ) -> List[Dict[str, Any]]: - subgraph_id = str(ref.get("subgraph_id") or "") + subgraph_id = str(ref.subgraph_id or "") if not subgraph_id: return [] - exact_key = (subgraph_id, str(ref.get("node_id")), ref.get("widget_index")) + exact_key = (subgraph_id, str(ref.node_id), ref.widget_index) locators = contexts.get("inner_widgets", {}).get(exact_key, []) - if not locators and ref.get("widget_name"): + if not locators and ref.extra_value("widget_name"): name_key = ( subgraph_id, - str(ref.get("node_id")), - normalize_widget_name(ref.get("widget_name")), + str(ref.node_id), + normalize_widget_name(ref.extra_value("widget_name")), ) locators = contexts.get("inner_widget_names", {}).get(name_key, []) return locators def _promote_model_reference_to_instances( - ref: Dict[str, Any], + ref: ModelReference, contexts: Dict[str, Dict[Any, Any]], - available_models: Optional[List[Dict[str, Any]]] = None, + available_models: Optional[List[ResolvedModel]] = None, existing_instance_slots: Optional[set] = None, -) -> Optional[List[Dict[str, Any]]]: +) -> Optional[List[ModelReference]]: locators = _get_promoted_widget_locators(ref, contexts) if not locators: return None if existing_instance_slots is None: existing_instance_slots = set() - promoted_refs = [] + promoted_refs: List[ModelReference] = [] for locator in locators: instance_slot = ( str(locator.get("node_id")), @@ -427,46 +432,43 @@ def _promote_model_reference_to_instances( continue value = locator.get("promoted_value") - nested_key = ref.get("nested_key") + nested_key = ref.extra_value("nested_key") if nested_key and isinstance(value, dict): value = value.get(nested_key) if not _has_promoted_value(value): continue value_str = str(value).strip() - promoted = dict(ref) - promoted.update( - { - "node_id": locator.get("node_id"), - "node_type": locator.get("node_type", ""), - "node_title": locator.get("node_title", ""), - "widget_index": locator.get("proxy_widget_index"), - "widget_name": locator.get("proxy_widget_name", ""), - "original_path": value_str, - "subgraph_id": locator.get("subgraph_id", ""), - "subgraph_name": locator.get("subgraph_name", ""), - "subgraph_path": None, - "is_top_level": True, - "promoted_widget_name": locator.get("proxy_widget_name", ""), - "promoted_inner_node_id": ref.get("node_id"), - "promoted_inner_node_type": ref.get("node_type", ""), - "promoted_inner_node_title": ref.get("node_title", ""), - "promoted_inner_widget_index": ref.get("widget_index"), - "promoted_inner_widget_name": ref.get("widget_name", ""), - "locate_node_id": locator.get("node_id"), - "locate_node_type": locator.get("node_type", ""), - "locate_node_title": locator.get("node_title", ""), - "locate_subgraph_id": "", - "locate_subgraph_name": locator.get("subgraph_name", ""), - "locate_is_top_level": True, - "locate_via_promoted_widget": True, - "is_urn": bool(references.URN_REGEX.match(value_str)), - } + promoted = ref.with_updates( + node_id=locator.get("node_id"), + node_type=locator.get("node_type", ""), + widget_index=locator.get("proxy_widget_index"), + original_path=value_str, + subgraph_id=locator.get("subgraph_id", ""), + is_top_level=True, + node_title=locator.get("node_title", ""), + widget_name=locator.get("proxy_widget_name", ""), + subgraph_name=locator.get("subgraph_name", ""), + subgraph_path=None, + promoted_widget_name=locator.get("proxy_widget_name", ""), + promoted_inner_node_id=ref.node_id, + promoted_inner_node_type=ref.node_type, + promoted_inner_node_title=ref.extra_value("node_title", ""), + promoted_inner_widget_index=ref.widget_index, + promoted_inner_widget_name=ref.extra_value("widget_name", ""), + locate_node_id=locator.get("node_id"), + locate_node_type=locator.get("node_type", ""), + locate_node_title=locator.get("node_title", ""), + locate_subgraph_id="", + locate_subgraph_name=locator.get("subgraph_name", ""), + locate_is_top_level=True, + locate_via_promoted_widget=True, + is_urn=bool(references.URN_REGEX.match(value_str)), ) - category_hints = promoted.get("category_hints") or [] - if not category_hints and promoted.get("category"): - category_hints = [promoted["category"]] + category_hints = promoted.extra_value("category_hints") or [] + if not category_hints and promoted.category: + category_hints = [promoted.category] resolved = references.try_resolve_model_path( value_str, category_hints or None, @@ -474,12 +476,13 @@ def _promote_model_reference_to_instances( ) if resolved: category, full_path = resolved - promoted["category"] = category - promoted["full_path"] = full_path - promoted["exists"] = os.path.exists(full_path) + promoted = promoted.with_updates( + category=category, + full_path=full_path, + exists=os.path.exists(full_path), + ) else: - promoted["full_path"] = None - promoted["exists"] = False + promoted = promoted.with_updates(full_path=None, exists=False) promoted_refs.append(promoted) existing_instance_slots.add(instance_slot) @@ -488,17 +491,19 @@ def _promote_model_reference_to_instances( def _apply_promoted_widget_locator( - ref: Dict[str, Any], contexts: Dict[str, Dict[Any, Any]] -) -> None: + ref: ModelReference, contexts: Dict[str, Dict[Any, Any]] +) -> ModelReference: locators = _get_promoted_widget_locators(ref, contexts) - locator = _select_promoted_locator(locators, ref.get("original_path", "")) + locator = _select_promoted_locator(locators, ref.original_path) if not locator: - return + return ref - ref["locate_node_id"] = locator.get("node_id") - ref["locate_node_type"] = locator.get("node_type", "") - ref["locate_node_title"] = locator.get("node_title", "") - ref["locate_subgraph_id"] = "" - ref["locate_subgraph_name"] = locator.get("subgraph_name", "") - ref["locate_is_top_level"] = True - ref["locate_via_promoted_widget"] = True + return ref.with_updates( + locate_node_id=locator.get("node_id"), + locate_node_type=locator.get("node_type", ""), + locate_node_title=locator.get("node_title", ""), + locate_subgraph_id="", + locate_subgraph_name=locator.get("subgraph_name", ""), + locate_is_top_level=True, + locate_via_promoted_widget=True, + ) diff --git a/core/workflow_updater.py b/core/workflow_updater.py index 80ee4dd..f01c738 100644 --- a/core/workflow_updater.py +++ b/core/workflow_updater.py @@ -7,11 +7,17 @@ Updates workflow JSON by replacing model paths in nodes. import os from typing import Any, Dict, List, Optional +from .contracts import Resolution, ResolvedModel from .custom_nodes import update_custom_node_model_path from .log_system import create_module_logger log = create_module_logger(__name__) -from .path_utils import get_filename_from_path, is_path_within +from .path_utils import ( + get_filename_from_path, + is_path_within, + normalize_folder_path_values, + normalize_string_values, +) def convert_to_relative_path( @@ -45,7 +51,9 @@ def convert_to_relative_path( # Get all available filenames for this category # This returns paths with OS-native separators (backslashes on Windows) - available_filenames = folder_paths.get_filename_list(category) + available_filenames = normalize_string_values( + folder_paths.get_filename_list(category) + ) # Try to find a matching entry in ComfyUI's list # Compare by finding the file that resolves to our absolute path @@ -83,31 +91,32 @@ def convert_to_relative_path( def get_base_directory_for_model( - model_dict: Dict[str, str], category: str + model: ResolvedModel, category: str ) -> Optional[str]: """ Get the base directory for a model based on its metadata. Args: - model_dict: Model dictionary with 'base_directory' or 'path' key + model: Typed local model with optional base directory and path category: Model category Returns: Base directory path if found, None otherwise """ - # Try to get base_directory from model dict - if "base_directory" in model_dict: - return model_dict["base_directory"] + if model.base_directory: + return model.base_directory # If we have the full path, try to find the category base directory - if "path" in model_dict: - full_path = model_dict["path"] + if model.path: + full_path = model.path # Import here to avoid circular dependency import folder_paths # Try to get category directories if category in folder_paths.folder_names_and_paths: - category_paths = folder_paths.get_folder_paths(category) + category_paths = normalize_folder_path_values( + folder_paths.get_folder_paths(category) + ) # Find which base directory this path belongs to for base_dir in category_paths: if is_path_within(full_path, base_dir): @@ -116,37 +125,35 @@ def get_base_directory_for_model( return None -def update_model_path( - workflow: Dict[str, Any], - node_id: int, - widget_index: int, - resolved_path: str, - category: str = None, - base_directory: str = None, - resolved_model: Dict[str, Any] = None, - subgraph_id: str = None, - is_top_level: bool = None, - mapping: Dict[str, Any] = None, -) -> bool: +def update_model_path(workflow: Dict[str, Any], resolution: Resolution) -> bool: """ Update a single model path in a workflow node, supporting both top-level and subgraph nodes. Args: workflow: Workflow JSON dictionary - node_id: ID of the node to update - widget_index: Widget index to update - resolved_path: Absolute path to resolved model - category: Model category - base_directory: Base directory for category - resolved_model: Model dict from scanner - subgraph_id: Subgraph ID (if node is in subgraph) - is_top_level: Whether node is in top-level (not subgraph definition) - mapping: Full mapping dictionary for standard or custom node updates + resolution: Validated typed workflow resolution Returns: True if update was successful, False otherwise """ + try: + resolution.validate() + except ValueError as exc: + log.warning(f"Invalid workflow resolution: {exc}") + return False + node = None + node_id = resolution.node_id + widget_index = resolution.widget_index + resolved_path = resolution.resolved_path or ( + resolution.resolved_model.path if resolution.resolved_model else "" + ) + category = resolution.category + base_directory = resolution.base_directory + subgraph_id = resolution.subgraph_id + is_top_level = resolution.is_top_level + resolved_model_contract = resolution.resolved_model + metadata = resolution.custom_node_metadata # Determine if this is a top-level node or inside a subgraph definition # - If is_top_level is True, it's a top-level node (even if it's a subgraph instance) @@ -211,8 +218,7 @@ def update_model_path( if widget_index >= len(widgets_values): can_extend_promoted_input = bool( is_top_level is True - and mapping - and mapping.get("promoted_widget_name") + and resolution.promoted_widget_name ) if can_extend_promoted_input: widgets_values.extend([None] * (widget_index + 1 - len(widgets_values))) @@ -225,14 +231,14 @@ def update_model_path( return False # Get category from resolved_model if not provided - if not category and resolved_model: - category = resolved_model.get("category") + if not category and resolved_model_contract: + category = resolved_model_contract.category custom_update_result = update_custom_node_model_path( node, widget_index, - resolved_model, - mapping, + resolved_model_contract, + metadata, ) if custom_update_result is not None: return custom_update_result @@ -243,8 +249,8 @@ def update_model_path( if os.path.isabs(resolved_path): # Use category from resolved_model for path conversion effective_category = category - if resolved_model: - effective_category = resolved_model.get("category", category) + if resolved_model_contract and resolved_model_contract.category: + effective_category = resolved_model_contract.category relative_path = convert_to_relative_path( resolved_path, effective_category, base_directory @@ -254,7 +260,7 @@ def update_model_path( # Update the widget value # Handle model references stored inside dictionary widget values. - nested_key = mapping.get("nested_key") if mapping else None + nested_key = resolution.nested_key if nested_key and isinstance(widgets_values[widget_index], dict): widgets_values[widget_index][nested_key] = relative_path log.debug( @@ -267,63 +273,36 @@ def update_model_path( def update_workflow_nodes( - workflow: Dict[str, Any], mappings: List[Dict[str, Any]] + workflow: Dict[str, Any], resolutions: List[Resolution] ) -> Dict[str, Any]: """ Apply multiple model path changes to a workflow. Args: workflow: Workflow JSON dictionary (will be modified in place) - mappings: List of mapping dictionaries: - { - 'node_id': node ID, - 'widget_index': widget index, - 'resolved_path': absolute path to resolved model, - 'category': model category (optional), - 'base_directory': base directory for category (optional), - 'resolved_model': model dict from scanner (optional, for base_directory) - } + resolutions: Validated typed workflow resolutions. Returns: Updated workflow dictionary (same reference, modified in place) """ updated_count = 0 - for mapping in mappings: - node_id = mapping.get("node_id") - widget_index = mapping.get("widget_index") - resolved_path = mapping.get("resolved_path") - - if not all([node_id is not None, widget_index is not None, resolved_path]): - log.warning(f"Invalid mapping: {mapping}") + for resolution in resolutions: + try: + resolution.validate() + except ValueError as exc: + log.warning(f"Skipping invalid workflow resolution: {exc}") continue - # Try to get base_directory from resolved_model if provided - base_directory = mapping.get("base_directory") - if not base_directory and "resolved_model" in mapping: - resolved_model = mapping["resolved_model"] - category = mapping.get("category", "") - base_directory = get_base_directory_for_model(resolved_model, category) + # Try to get base_directory from the typed local model if not explicit. + if not resolution.base_directory and resolution.resolved_model: + base_directory = get_base_directory_for_model( + resolution.resolved_model, + resolution.category or "", + ) + resolution = resolution.with_base_directory(base_directory) - category = mapping.get("category") - resolved_model = mapping.get("resolved_model") - subgraph_id = mapping.get("subgraph_id") - is_top_level = mapping.get( - "is_top_level" - ) # True for top-level nodes, False for nodes in subgraph definitions - - success = update_model_path( - workflow, - node_id, - widget_index, - resolved_path, - category, - base_directory, - resolved_model, - subgraph_id, - is_top_level, - mapping, - ) + success = update_model_path(workflow, resolution) if success: updated_count += 1 diff --git a/tests/test_contracts.py b/tests/test_contracts.py new file mode 100644 index 0000000..411e0dd --- /dev/null +++ b/tests/test_contracts.py @@ -0,0 +1,1450 @@ +import pytest + +from core.contracts import ( + CustomNodeMetadata, + DownloadSpec, + HuggingFaceFileReference, + MissingModel, + ModelCatalogEntry, + ModelMatch, + ModelReference, + ProviderUrlReference, + Resolution, + ResolvedModel, + SearchResult, + WorkflowAnalysisResult, + WorkflowModelInventory, +) +from core.custom_nodes.base import CustomNodeModelAdapter +from core.custom_nodes.registry import ( + get_custom_node_adapter_for_reference, + get_custom_node_resolution_metadata, +) +from core.resolver import apply_resolution +from core.services.search_orchestrator import SearchOrchestrator + + +def test_search_result_normalizes_provider_aliases_and_preserves_extra_fields(): + result = SearchResult.from_mapping( + { + "source": "huggingface", + "name": "model.safetensors", + "downloadUrl": "https://example.test/model.safetensors", + "sizeKB": 2, + "hashes": {"SHA256": "abc123"}, + "repository": "author/model", + } + ) + + serialized = result.to_dict() + + assert result.source == "huggingface" + assert result.filename == "model.safetensors" + assert result.download_url == "https://example.test/model.safetensors" + assert result.size == 2048 + assert result.sha256 == "abc123" + assert serialized["repository"] == "author/model" + assert serialized["type"] == "" + + +def test_provider_url_reference_validates_and_serializes_identity(): + reference = ProviderUrlReference.from_mapping( + {"modelId": "123", "modelVersionId": 456} + ) + + assert reference.model_id == 123 + assert reference.version_id == 456 + assert reference.to_dict() == {"model_id": 123, "version_id": 456} + + with pytest.raises(ValueError, match="model_id must be an integer"): + ProviderUrlReference(model_id="not-a-number") + + +def test_huggingface_file_reference_requires_typed_repository_path(): + reference = HuggingFaceFileReference( + repo="org/repo", + filename="models/model.safetensors", + ) + + assert reference.to_dict() == { + "repo": "org/repo", + "branch": "main", + "filename": "models/model.safetensors", + } + with pytest.raises(TypeError, match="repo must be a string"): + HuggingFaceFileReference(repo=123, filename="model.safetensors") + + +def test_model_catalog_entry_validates_source_fields_and_preserves_extra(): + entry = ModelCatalogEntry.from_mapping( + { + "url": "https://example.test/model", + "download_url": "https://example.test/download", + "type": "checkpoints", + "directory": "checkpoints", + "size": 2048, + "canonical_name": "model", + }, + filename="model.safetensors", + ) + + assert entry.filename == "model.safetensors" + assert entry.model_type == "checkpoints" + assert entry.extra_value("canonical_name") == "model" + assert entry.to_dict()["download_url"] == "https://example.test/download" + + with pytest.raises(TypeError, match="ModelCatalogEntry url"): + ModelCatalogEntry.from_mapping( + {"url": 123}, + filename="model.safetensors", + ) + + +def test_search_result_normalizes_legacy_identity_and_file_aliases(): + result = SearchResult.from_mapping( + { + "source": "civitai", + "modelName": "Example model", + "modelVersionId": 456, + "fileName": "example.safetensors", + "versionUrl": "https://example.test/models/123?modelVersionId=456", + "SHA256": "abc123", + } + ) + + assert result.name == "Example model" + assert result.version_id == 456 + assert result.filename == "example.safetensors" + assert result.version_url.endswith("modelVersionId=456") + assert result.sha256 == "abc123" + assert "modelName" not in result.to_dict() + + +def test_search_result_preserves_explicit_empty_provider_hash(): + result = SearchResult.from_mapping( + {"source": "huggingface", "sha256": ""} + ) + + assert result.sha256 == "" + assert result.to_dict()["sha256"] == "" + + +def test_search_result_mapping_source_override_is_strict(): + with pytest.raises(TypeError, match="SearchResult source"): + SearchResult.from_mapping({"name": "model"}, source=123) + + +def test_search_result_normalizes_human_readable_size_to_bytes(): + result = SearchResult(source="model_list", size="12 MB") + + assert result.size == 12 * 1024**2 + assert result.to_dict()["size"] == 12 * 1024**2 + + +@pytest.mark.parametrize( + "field_name", + [ + "name", + "version_name", + "model_type", + "filename", + "url", + "download_url", + "base_model", + "sha256", + "details_source", + "version_url", + "match_type", + ], +) +def test_search_result_rejects_non_string_canonical_text_fields(field_name): + with pytest.raises(TypeError, match=f"SearchResult {field_name}"): + SearchResult(source="test", **{field_name: 123}) + + with pytest.raises(TypeError, match="SearchResult download_url"): + SearchResult.from_mapping( + {"source": "test", "download_url": 123} + ) + + +@pytest.mark.parametrize( + ("field_name", "value"), + [ + ("tags", 0), + ("trained_words", "trigger"), + ("images", {}), + ("hashes", []), + ("extra", []), + ], +) +def test_search_result_rejects_non_collection_canonical_fields(field_name, value): + with pytest.raises(TypeError, match=f"SearchResult {field_name}"): + SearchResult(source="test", **{field_name: value}) + + +@pytest.mark.parametrize( + ("field_name", "value"), + [("tags", 0), ("trained_words", "trigger"), ("images", {}), ("hashes", [])], +) +def test_search_result_mapping_adapter_rejects_non_collection_fields( + field_name, + value, +): + with pytest.raises(TypeError, match=f"SearchResult {field_name}"): + SearchResult.from_mapping({"source": "test", field_name: value}) + + +@pytest.mark.parametrize( + ("payload", "field_name"), + [ + ({"source": "test", "name": 123}, "name"), + ({"source": "test", "sha256": 123}, "sha256"), + ({"source": "test", "sizeKB": "invalid"}, "size"), + ], +) +def test_search_result_mapping_adapter_rejects_invalid_provider_scalars( + payload, + field_name, +): + with pytest.raises((TypeError, ValueError), match=field_name): + SearchResult.from_mapping(payload) + + +def test_search_result_mapping_adapter_rejects_invalid_provider_source(): + with pytest.raises(TypeError, match="SearchResult source"): + SearchResult.from_mapping({"source": {"name": "civitai"}}) + + +def test_search_result_rejects_malformed_numeric_fields(): + with pytest.raises(ValueError, match="SearchResult confidence"): + SearchResult(source="test", confidence="not-a-number") + with pytest.raises(ValueError, match="SearchResult size"): + SearchResult(source="test", size="not-a-size") + + +def test_search_result_round_trip_preserves_custom_url_modes(): + compact = SearchResult( + source="civitai", + name="Compact result", + filename="compact.safetensors", + url="https://example.test/models/1", + result_mode="compact_custom_url", + custom_url=True, + ) + full = SearchResult( + source="civitai", + name="Full result", + filename="full.safetensors", + url="https://example.test/models/2", + result_mode="custom_url", + custom_url=True, + ) + + assert SearchResult.from_mapping(compact.to_dict()).result_mode == ( + "compact_custom_url" + ) + assert SearchResult.from_mapping(full.to_dict()).result_mode == "custom_url" + + +def test_search_result_rejects_unknown_result_modes(): + with pytest.raises(ValueError, match="Unsupported model result mode"): + SearchResult(source="civitai", result_mode="unknown") + + +def test_contract_adapters_reject_malformed_payloads(): + with pytest.raises(TypeError, match="SearchResult must be an object"): + SearchResult.from_mapping(None) + with pytest.raises(ValueError, match="model identity"): + ModelReference.from_mapping({}) + with pytest.raises(ValueError, match="path or filename"): + ResolvedModel.from_mapping({}) + with pytest.raises(TypeError, match="model object"): + ModelMatch.from_mapping({"model": None}) + + +def test_extra_fields_cannot_override_canonical_contract_fields(): + reference = ModelReference( + node_id=1, + original_path="canonical.safetensors", + category="checkpoints", + extra={"category": "wrong"}, + ) + model = ResolvedModel( + path="canonical.safetensors", + filename="canonical.safetensors", + extra={"path": "wrong.safetensors"}, + ) + result = SearchResult( + source="civitai", + filename="canonical.safetensors", + extra={"filename": "wrong.safetensors"}, + ) + + assert reference.to_dict()["category"] == "checkpoints" + assert model.to_dict()["path"] == "canonical.safetensors" + assert result.to_dict()["filename"] == "canonical.safetensors" + + +def test_extra_fields_cannot_create_empty_canonical_values(): + model = ResolvedModel(filename="model.safetensors").with_extra( + path="wrong.safetensors", + category="wrong", + provider="test", + ) + reference = ModelReference( + node_id=1, + original_path="model.safetensors", + extra={"subgraph_id": "wrong", "provider": "test"}, + ) + + serialized_model = model.to_dict() + serialized_reference = reference.to_dict() + + assert "path" not in serialized_model + assert "category" not in serialized_model + assert serialized_model["provider"] == "test" + assert "subgraph_id" not in serialized_reference + assert serialized_reference["provider"] == "test" + + +def test_resolved_model_aliases_cannot_leak_through_extra_or_raw_payloads(): + model = ResolvedModel( + path="canonical.safetensors", + filename="canonical.safetensors", + raw={ + "resolved_path": "wrong.safetensors", + "name": "wrong.safetensors", + }, + ).with_extra( + resolved_path="another-wrong.safetensors", + name="another-wrong.safetensors", + provider="scanner", + ) + + serialized = model.to_dict() + + assert serialized["path"] == "canonical.safetensors" + assert serialized["filename"] == "canonical.safetensors" + assert "resolved_path" not in serialized + assert "name" not in serialized + assert serialized["provider"] == "scanner" + + +def test_search_result_extra_cannot_restore_omitted_mode_fields(): + custom = SearchResult( + source="civitai", + filename="custom.safetensors", + result_mode="custom_url", + custom_url=True, + ).with_extra( + confidence=99, + result_mode="search", + filename="wrong.safetensors", + provider="civitai", + ) + compact = SearchResult( + source="civitai", + filename="compact.safetensors", + result_mode="compact_custom_url", + custom_url=True, + ).with_extra( + confidence=99, + size=123, + type="wrong", + provider="civitai", + ) + + custom_payload = custom.to_dict() + compact_payload = compact.to_dict() + + assert "confidence" not in custom_payload + assert "result_mode" not in custom_payload + assert custom_payload["filename"] == "custom.safetensors" + assert custom_payload["provider"] == "civitai" + assert "confidence" not in compact_payload + assert "size" not in compact_payload + assert "type" not in compact_payload + assert compact_payload["filename"] == "compact.safetensors" + + +def test_missing_model_extra_cannot_inject_matches_when_unset(): + missing = MissingModel.from_mapping( + { + "node_id": 1, + "original_path": "missing.safetensors", + } + ).with_extra( + matches=[{"bad": True}], + reference_count=99, + provider="workflow", + ) + + serialized = missing.to_dict() + + assert "matches" not in serialized + assert serialized["reference_count"] == 1 + assert serialized["provider"] == "workflow" + + +def test_missing_model_requires_consistent_reference_collection(): + with pytest.raises(TypeError, match="all_node_refs"): + MissingModel.from_mapping( + { + "node_id": 1, + "original_path": "missing.safetensors", + "all_node_refs": [ + {"node_id": 1, "original_path": "missing.safetensors"}, + None, + ], + } + ) + + with pytest.raises(ValueError, match="reference_count"): + MissingModel.from_mapping( + { + "node_id": 1, + "original_path": "missing.safetensors", + "all_node_refs": [ + {"node_id": 1, "original_path": "missing.safetensors"} + ], + "reference_count": 2, + } + ) + + +def test_missing_model_rejects_malformed_matches_in_mapping(): + with pytest.raises(TypeError, match="MissingModel matches"): + MissingModel.from_mapping( + { + "node_id": 1, + "original_path": "missing.safetensors", + "matches": [None], + } + ) + + with pytest.raises(TypeError, match="MissingModel matches"): + MissingModel.from_mapping( + { + "node_id": 1, + "original_path": "missing.safetensors", + "matches": {"bad": True}, + } + ) + + +def test_missing_model_derives_direct_reference_count(): + reference = ModelReference(node_id=1, original_path="missing.safetensors") + missing = MissingModel( + reference=reference, + all_node_refs=(reference, reference), + ) + + assert missing.reference_count == 2 + assert len(missing.all_node_refs) == 2 + + +@pytest.mark.parametrize("value", ["1", 1.5, True]) +def test_missing_model_rejects_non_integer_reference_count(value): + reference = ModelReference(node_id=1, original_path="missing.safetensors") + + with pytest.raises(TypeError, match="reference_count"): + MissingModel(reference=reference, reference_count=value) + + with pytest.raises(TypeError, match="reference_count"): + MissingModel.from_mapping( + { + "node_id": 1, + "original_path": "missing.safetensors", + "reference_count": value, + } + ) + + +def test_missing_model_rejects_non_positive_reference_count(): + reference = ModelReference(node_id=1, original_path="missing.safetensors") + + with pytest.raises(ValueError, match="reference_count"): + MissingModel(reference=reference, reference_count=0) + + with pytest.raises(ValueError, match="reference_count"): + MissingModel.from_mapping( + { + "node_id": 1, + "original_path": "missing.safetensors", + "reference_count": -1, + } + ) + + +def test_contract_boolean_adapters_do_not_use_string_truthiness(): + reference = ModelReference.from_mapping( + { + "node_id": 1, + "original_path": "missing.safetensors", + "exists": "false", + "is_top_level": "false", + } + ) + result = SearchResult.from_mapping( + { + "source": "civitai", + "custom_url": "false", + } + ) + + assert reference.exists is False + assert reference.is_top_level is False + assert result.custom_url is False + assert result.result_mode == "search" + + metadata = CustomNodeMetadata.from_mapping( + {"is_legacy_lora_manager": "false"} + ) + assert metadata.is_legacy_lora_manager is False + + legacy_reference = ModelReference.from_mapping( + { + "node_id": 2, + "original_path": "legacy.safetensors", + "is_lora_v2": "false", + } + ) + assert get_custom_node_adapter_for_reference(legacy_reference) is None + assert ( + get_custom_node_resolution_metadata(legacy_reference).is_legacy_lora_manager + is False + ) + + legacy_v2_reference = ModelReference.from_mapping( + { + "node_id": 4, + "original_path": "legacy-v2.safetensors", + "is_lora_v2": "true", + } + ) + assert get_custom_node_adapter_for_reference(legacy_v2_reference) is not None + assert ( + get_custom_node_resolution_metadata(legacy_v2_reference).is_legacy_lora_manager + is True + ) + + workflow_flags = ModelReference.from_mapping( + { + "node_id": 3, + "original_path": "workflow.safetensors", + "active": "false", + "auto_download_capable": "false", + "auto_download_candidate": "false", + "connected": "false", + "input_choice_matches_value": "false", + "is_legacy_lora_manager": "false", + "is_lora_v2": "false", + "is_urn": "false", + "locate_is_top_level": "false", + "locate_via_promoted_widget": "false", + } + ) + for flag in ( + "active", + "auto_download_capable", + "auto_download_candidate", + "connected", + "input_choice_matches_value", + "is_legacy_lora_manager", + "is_lora_v2", + "is_urn", + "locate_is_top_level", + "locate_via_promoted_widget", + ): + assert workflow_flags.extra_value(flag) is False + + +def test_custom_node_metadata_preserves_canonical_legacy_flag(): + reference = ModelReference.from_mapping( + { + "node_id": 2, + "original_path": "old-lora", + "is_legacy_lora_manager": "true", + } + ) + + metadata = get_custom_node_resolution_metadata(reference) + + assert metadata.adapter_id == "lora-manager" + assert metadata.is_legacy_lora_manager is True + + +def test_direct_contract_constructors_normalize_nested_values(): + reference = ModelReference( + node_id=1, + original_path="model.safetensors", + exists="false", + is_top_level="true", + ) + assert reference.exists is False + assert reference.is_top_level is True + + match = ModelMatch( + model={"path": "models/model.safetensors"}, + similarity="0.25", + ) + assert isinstance(match.model, ResolvedModel) + assert match.similarity == 0.25 + + missing = MissingModel( + reference={"node_id": 1, "original_path": "missing.safetensors"}, + all_node_refs=[ + {"node_id": 1, "original_path": "missing.safetensors"} + ], + matches=[ + {"model": {"path": "models/model.safetensors"}} + ], + ) + assert isinstance(missing.reference, ModelReference) + assert isinstance(missing.all_node_refs[0], ModelReference) + assert isinstance(missing.matches[0], ModelMatch) + + analysis = WorkflowAnalysisResult( + missing_models=[ + {"node_id": 1, "original_path": "missing.safetensors"} + ], + resolved_models=[ + {"node_id": 2, "original_path": "loaded.safetensors"} + ], + total_missing=1, + total_resolved=1, + total_models_analyzed=2, + ) + assert isinstance(analysis.missing_models[0], MissingModel) + assert isinstance(analysis.resolved_models[0], ModelReference) + + resolution = Resolution( + node_id=1, + widget_index=0, + resolved_path="new.safetensors", + resolved_model={"path": "new.safetensors"}, + reference={"node_id": 1, "original_path": "old.safetensors"}, + is_top_level="false", + ) + assert isinstance(resolution.resolved_model, ResolvedModel) + assert isinstance(resolution.reference, ModelReference) + assert resolution.is_top_level is False + + with pytest.raises(TypeError, match="ModelMatch model"): + ModelMatch(model=None) + with pytest.raises(TypeError, match="MissingModel reference"): + MissingModel(reference=None) + + +@pytest.mark.parametrize("value", ["maybe", 2, [], {}]) +def test_boolean_contracts_reject_unknown_values(value): + with pytest.raises(ValueError, match="exists"): + ModelReference.from_mapping( + { + "node_id": 1, + "original_path": "model.safetensors", + "exists": value, + } + ) + + +@pytest.mark.parametrize("value", [123, [], {}]) +def test_model_reference_known_text_extras_reject_non_text_values(value): + with pytest.raises(TypeError, match="ModelReference full_path"): + ModelReference.from_mapping( + { + "node_id": 1, + "original_path": "model.safetensors", + "full_path": value, + } + ) + + with pytest.raises(ValueError, match="is_top_level"): + ModelReference.from_mapping( + { + "node_id": 1, + "original_path": "model.safetensors", + "is_top_level": value, + } + ) + + with pytest.raises(ValueError, match="custom_url"): + SearchResult.from_mapping( + {"source": "test", "custom_url": value} + ) + + with pytest.raises(ValueError, match="is_legacy_lora_manager"): + CustomNodeMetadata.from_mapping( + {"is_legacy_lora_manager": value} + ) + + with pytest.raises(ValueError, match="is_top_level"): + Resolution.from_mapping( + { + "node_id": 1, + "widget_index": 0, + "original_path": "missing.safetensors", + "resolved_path": "model.safetensors", + "is_top_level": value, + } + ) + + +def test_canonical_boolean_contracts_reject_explicit_null(): + assert ( + ModelReference.from_mapping( + {"node_id": 1, "original_path": "model.safetensors"} + ).exists + is False + ) + assert SearchResult.from_mapping({"source": "test"}).custom_url is False + assert ( + CustomNodeMetadata.from_mapping({}).is_legacy_lora_manager is False + ) + + with pytest.raises(ValueError, match="exists"): + ModelReference.from_mapping( + { + "node_id": 1, + "original_path": "model.safetensors", + "exists": None, + } + ) + + with pytest.raises(ValueError, match="custom_url"): + SearchResult.from_mapping({"source": "test", "custom_url": None}) + + with pytest.raises(ValueError, match="is_legacy_lora_manager"): + CustomNodeMetadata.from_mapping( + {"is_legacy_lora_manager": None} + ) + + +@pytest.mark.parametrize("value", [None, "maybe", 2, [], {}]) +def test_model_reference_boolean_extras_reject_invalid_values(value): + with pytest.raises(ValueError, match="ModelReference active"): + ModelReference.from_mapping( + { + "node_id": 1, + "original_path": "model.safetensors", + "active": value, + } + ) + + +def test_lora_manager_adapter_normalizes_external_boolean_values(): + from core.custom_nodes.lora_manager import analyze_references + + references = analyze_references( + { + "id": 1, + "type": "LoraLoaderV2", + "widgets_values": [ + "", + "", + [{"name": "character", "active": None, "strength": 1.0}], + ], + }, + available_models=[], + is_active=False, + get_widget_name_hint=lambda _node, _index: "lora_name", + ) + + assert references[0].extra_value("active") is True + assert references[0].extra_value("connected") is False + + +@pytest.mark.parametrize( + ("factory", "payload", "message"), + [ + ( + ModelReference.from_mapping, + {"node_id": 1, "original_path": 123}, + "ModelReference original_path", + ), + ( + ResolvedModel.from_mapping, + {"path": 123}, + "ResolvedModel path", + ), + ( + ResolvedModel.from_mapping, + {"path": 0, "resolved_path": "valid.safetensors"}, + "ResolvedModel path", + ), + ], +) +def test_model_contracts_reject_non_string_text_fields(factory, payload, message): + with pytest.raises(TypeError, match=message): + factory(payload) + + +@pytest.mark.parametrize( + ("factory", "payload", "message"), + [ + ( + ModelReference.from_mapping, + {"node_id": True, "original_path": "model.safetensors"}, + "ModelReference node_id", + ), + ( + ModelReference.from_mapping, + { + "node_id": 1, + "widget_index": "0", + "original_path": "model.safetensors", + }, + "ModelReference widget_index", + ), + ( + SearchResult.from_mapping, + {"source": "test", "model_id": False}, + "SearchResult model_id", + ), + ( + Resolution.from_mapping, + { + "node_id": True, + "widget_index": 0, + "resolved_path": "model.safetensors", + }, + "Resolution node_id", + ), + ], +) +def test_identifier_contracts_reject_invalid_types(factory, payload, message): + with pytest.raises(TypeError, match=message): + factory(payload) + + +def test_identifier_contracts_normalize_string_ids(): + reference = ModelReference.from_mapping( + { + "node_id": " 12 ", + "widget_index": 0, + "original_path": "model.safetensors", + } + ) + result = SearchResult.from_mapping( + {"source": "test", "model_id": " 12 ", "version_id": " 34 "} + ) + + assert reference.node_id == "12" + assert result.model_id == "12" + assert result.version_id == "34" + + +@pytest.mark.parametrize( + ("factory", "payload", "message"), + [ + ( + ModelMatch.from_mapping, + { + "model": {"path": "models/model.safetensors"}, + "filename": 123, + }, + "ModelMatch filename", + ), + ( + CustomNodeMetadata.from_mapping, + {"adapter_id": 123}, + "CustomNodeMetadata adapter_id", + ), + ( + CustomNodeMetadata.from_mapping, + {"original_identity": 123}, + "CustomNodeMetadata original_identity", + ), + ( + Resolution.from_mapping, + { + "node_id": 1, + "widget_index": 0, + "resolved_path": 123, + }, + "Resolution resolved_path", + ), + ( + Resolution.from_mapping, + { + "node_id": 1, + "widget_index": 0, + "resolved_path": "model.safetensors", + "category": 123, + }, + "Resolution category", + ), + ], +) +def test_nested_contracts_reject_non_string_canonical_fields( + factory, + payload, + message, +): + with pytest.raises((TypeError, ValueError), match=message): + factory(payload) + + +def test_contract_collections_reject_falsey_non_arrays(): + reference = ModelReference(node_id=1, original_path="model.safetensors") + + with pytest.raises(TypeError, match="all_node_refs"): + MissingModel(reference=reference, all_node_refs=0) + with pytest.raises(TypeError, match="missing_models"): + WorkflowAnalysisResult(missing_models=0) + with pytest.raises(TypeError, match="available_models"): + WorkflowModelInventory(available_models=0) + + +@pytest.mark.parametrize("field_name", ["similarity", "confidence"]) +def test_model_match_rejects_malformed_numeric_fields(field_name): + with pytest.raises(ValueError, match=field_name): + ModelMatch.from_mapping( + { + "model": {"path": "models/model.safetensors"}, + field_name: "not-a-number", + } + ) + + +def test_model_contracts_are_not_mapping_compatibility_wrappers(): + reference = ModelReference(node_id=1, original_path="model.safetensors") + model = ResolvedModel(path="models/model.safetensors") + + with pytest.raises(TypeError): + reference["original_path"] + with pytest.raises(TypeError): + model["path"] + + +def test_contract_collections_are_immutable_and_thaw_at_json_boundary(): + result = SearchResult( + source="civitai", + tags=["tag-a"], + hashes={"SHA256": "abc123"}, + extra={"nested": {"values": ["one"]}}, + ) + + assert result.tags == ("tag-a",) + serialized = result.to_dict() + assert serialized["tags"] == ["tag-a"] + assert serialized["hashes"] == {"SHA256": "abc123"} + assert serialized["nested"]["values"] == ["one"] + assert result.extra_value("nested")["values"] == ("one",) + + with pytest.raises(AttributeError): + result.tags.append("tag-b") + with pytest.raises(TypeError): + result.hashes["MD5"] = "hash" + with pytest.raises(TypeError): + result.extra_value("nested")["values"] = ("two",) + + model = ResolvedModel.from_mapping( + { + "path": "models/model.safetensors", + "extra": {"nested": {"values": ["one"]}}, + } + ) + assert model.to_dict()["extra"]["nested"]["values"] == ["one"] + with pytest.raises(TypeError): + model.raw["path"] = "other.safetensors" + + +def test_workflow_inventory_and_custom_node_metadata_are_typed_and_immutable(): + model = ResolvedModel(path="models/model.safetensors") + reference = ModelReference( + node_id=1, + original_path="model.safetensors", + category="checkpoints", + ) + inventory = WorkflowModelInventory( + available_models=[model], + model_refs=[reference], + ) + assert inventory.available_models == (model,) + assert inventory.model_refs == (reference,) + with pytest.raises(AttributeError): + inventory.available_models.append(model) + + metadata = get_custom_node_resolution_metadata( + ModelReference.from_mapping( + { + "node_id": 2, + "original_path": "old-lora", + "custom_node_adapter": "lora-manager", + "custom_node_original_identity": "old-lora", + } + ) + ) + assert isinstance(metadata, CustomNodeMetadata) + assert metadata.adapter_id == "lora-manager" + assert metadata.original_identity == "old-lora" + + adapter = CustomNodeModelAdapter( + adapter_id="test", + node_types=["TestLoader"], + widget_categories={0: "loras"}, + ) + assert adapter.node_types == ("TestLoader",) + assert adapter.widget_categories == {0: "loras"} + with pytest.raises(TypeError): + adapter.widget_categories[1] = "checkpoints" + + with pytest.raises(TypeError, match="node_types"): + CustomNodeModelAdapter( + adapter_id="test", + node_types=0, + widget_categories={}, + ) + with pytest.raises(TypeError, match="widget_categories"): + CustomNodeModelAdapter( + adapter_id="test", + node_types=("TestLoader",), + widget_categories=[], + ) + with pytest.raises(TypeError, match="update_model_path"): + CustomNodeModelAdapter( + adapter_id="test", + node_types=("TestLoader",), + widget_categories={}, + update_model_path="not-callable", + ) + + +def test_search_orchestrator_normalizes_and_serializes_typed_result(): + typed = SearchResult( + source="civitai", + name="Model", + filename="model.safetensors", + download_url="https://example.test/model", + size=4096, + ) + result = SearchOrchestrator.normalize_search_result( + "civitai", + typed, + ) + + assert result is typed + assert result.source == "civitai" + assert SearchOrchestrator.serialize_search_result(result)["size"] == 4096 + assert SearchOrchestrator.serialize_search_result(result)["download_url"] == ( + "https://example.test/model" + ) + + with pytest.raises(TypeError, match="untyped result"): + SearchOrchestrator.normalize_search_result( + "model_list", + {"filename": "legacy.safetensors"}, + ) + + +def test_missing_model_and_match_use_typed_contracts(): + missing = MissingModel.from_mapping( + { + "node_id": 4, + "node_type": "CheckpointLoaderSimple", + "widget_index": 0, + "original_path": "missing.safetensors", + "category": "checkpoints", + "exists": False, + "all_node_refs": [ + { + "node_id": 4, + "node_type": "CheckpointLoaderSimple", + "widget_index": 0, + "original_path": "missing.safetensors", + "category": "checkpoints", + "exists": False, + } + ], + "matches": [ + { + "model": { + "path": "models/checkpoint.safetensors", + "filename": "checkpoint.safetensors", + "category": "checkpoints", + }, + "filename": "checkpoint.safetensors", + "similarity": 0.9, + "confidence": 90, + } + ], + } + ) + + assert isinstance(missing.reference, ModelReference) + assert isinstance(missing.matches[0], ModelMatch) + assert isinstance(missing.matches[0].model, ResolvedModel) + assert missing.all_node_refs[0].node_id == 4 + assert missing.matches[0].model.path == "models/checkpoint.safetensors" + assert missing.to_dict()["matches"][0]["model"]["path"] == ( + "models/checkpoint.safetensors" + ) + + +def test_model_match_keeps_typed_nested_resolved_model(): + model = ResolvedModel( + path="models/checkpoint.safetensors", + filename="checkpoint.safetensors", + ) + + match = ModelMatch.from_mapping( + { + "model": model, + "filename": "checkpoint.safetensors", + "similarity": 0.9, + } + ) + + assert match.model is model + assert match.model.path == "models/checkpoint.safetensors" + + +def test_workflow_analysis_result_serializes_typed_models(): + missing = MissingModel.from_mapping( + { + "node_id": 1, + "widget_index": 0, + "original_path": "missing.safetensors", + "category": "checkpoints", + "exists": False, + } + ) + result = WorkflowAnalysisResult( + missing_models=(missing,), + resolved_models=( + ModelReference( + node_id=2, + widget_index=0, + original_path="loaded.safetensors", + category="checkpoints", + exists=True, + ), + ), + total_missing=1, + total_resolved=1, + total_models_analyzed=2, + ).to_dict() + + assert result["total_missing"] == 1 + assert result["missing_models"][0]["original_path"] == "missing.safetensors" + assert result["resolved_models"][0]["matches"] == [] + + +@pytest.mark.parametrize( + "field_name, value, error_type", + [ + ("total_missing", "1", TypeError), + ("total_resolved", 1.5, TypeError), + ("total_models_analyzed", -1, ValueError), + ], +) +def test_workflow_analysis_result_requires_non_negative_integer_counts( + field_name, value, error_type +): + with pytest.raises(error_type, match=field_name): + WorkflowAnalysisResult(**{field_name: value}) + + +def test_workflow_analysis_result_counts_match_model_collections(): + missing = MissingModel.from_mapping( + {"node_id": 1, "original_path": "missing.safetensors"} + ) + resolved = ModelReference(node_id=2, original_path="loaded.safetensors") + + with pytest.raises(ValueError, match="total_missing"): + WorkflowAnalysisResult( + missing_models=(missing,), + total_missing=0, + ) + + with pytest.raises(ValueError, match="total_resolved"): + WorkflowAnalysisResult( + resolved_models=(resolved,), + total_resolved=0, + ) + + with pytest.raises(ValueError, match="total_models_analyzed"): + WorkflowAnalysisResult( + missing_models=(missing,), + total_missing=1, + total_models_analyzed=0, + ) + + +def test_download_spec_normalizes_mapping_and_keeps_expected_hash_in_metadata(): + spec = DownloadSpec.from_mapping( + { + "download_url": " https://example.test/model.safetensors ", + "name": " model.safetensors ", + "category": "checkpoints", + "expected_sha256": " abc123 ", + "headers": {"Authorization": "Bearer token"}, + } + ) + + kwargs = spec.to_kwargs() + + assert spec.url == "https://example.test/model.safetensors" + assert spec.filename == "model.safetensors" + assert spec.expected_sha256 == "abc123" + assert kwargs["headers"] == {"Authorization": "Bearer token"} + assert kwargs["metadata"]["sha256"] == "abc123" + + +def test_download_spec_uses_sha256_alias_when_expected_hash_is_blank(): + spec = DownloadSpec.from_mapping( + { + "url": "https://example.test/model.safetensors", + "filename": "model.safetensors", + "category": "checkpoints", + "expected_sha256": " ", + "sha256": " alias-hash ", + } + ) + + assert spec.expected_sha256 == "alias-hash" + assert spec.to_kwargs()["metadata"]["sha256"] == "alias-hash" + + +@pytest.mark.parametrize( + "field_name, value", + [ + ("url", 123), + ("filename", 456), + ("category", 789), + ("subfolder", 123), + ("base_directory", 456), + ], +) +def test_download_spec_rejects_non_string_text_fields(field_name, value): + payload = { + "url": "https://example.test/model.safetensors", + "filename": "model.safetensors", + "category": "checkpoints", + "subfolder": "", + "base_directory": "", + } + payload[field_name] = value + + with pytest.raises(TypeError, match=rf"DownloadSpec {field_name}"): + DownloadSpec(**payload) + + +@pytest.mark.parametrize( + "field_name, value, alias_key", + [ + ("url", 0, "download_url"), + ("filename", False, "name"), + ("category", 0, None), + ("subfolder", 0, None), + ("base_directory", False, None), + ], +) +def test_download_spec_mapping_rejects_falsey_non_string_text_fields( + field_name, value, alias_key +): + payload = { + "url": "https://example.test/model.safetensors", + "filename": "model.safetensors", + "category": "checkpoints", + } + payload[field_name] = value + if alias_key: + payload[alias_key] = ( + "https://example.test/model.safetensors" + if alias_key == "download_url" + else "model.safetensors" + ) + + with pytest.raises(TypeError, match=rf"DownloadSpec {field_name}"): + DownloadSpec.from_mapping(payload) + + +def test_download_spec_rejects_non_string_headers(): + with pytest.raises(TypeError, match="DownloadSpec headers"): + DownloadSpec( + url="https://example.test/model.safetensors", + filename="model.safetensors", + category="checkpoints", + headers={"X-Test": 123}, + ) + + with pytest.raises(TypeError, match="DownloadSpec headers"): + DownloadSpec( + url="https://example.test/model.safetensors", + filename="model.safetensors", + category="checkpoints", + headers=["invalid"], + ) + + with pytest.raises(TypeError, match="DownloadSpec headers"): + DownloadSpec.from_mapping( + { + "url": "https://example.test/model.safetensors", + "filename": "model.safetensors", + "category": "checkpoints", + "headers": [], + } + ) + + +def test_download_spec_rejects_non_mapping_metadata(): + with pytest.raises(TypeError, match="DownloadSpec metadata"): + DownloadSpec( + url="https://example.test/model.safetensors", + filename="model.safetensors", + category="checkpoints", + metadata=[], + ) + + with pytest.raises(TypeError, match="DownloadSpec metadata"): + DownloadSpec.from_mapping( + { + "url": "https://example.test/model.safetensors", + "filename": "model.safetensors", + "category": "checkpoints", + "metadata": [], + } + ) + + +@pytest.mark.parametrize("value", [0, 123, True, [], {}]) +def test_download_spec_rejects_non_string_expected_sha256(value): + with pytest.raises(TypeError, match="DownloadSpec expected_sha256"): + DownloadSpec( + url="https://example.test/model.safetensors", + filename="model.safetensors", + category="checkpoints", + expected_sha256=value, + ) + + with pytest.raises(TypeError, match="DownloadSpec expected_sha256"): + DownloadSpec.from_mapping( + { + "url": "https://example.test/model.safetensors", + "filename": "model.safetensors", + "category": "checkpoints", + "expected_sha256": value, + } + ) + + +@pytest.mark.parametrize( + "payload, message", + [ + ({"filename": "model.safetensors", "category": "checkpoints"}, "URL"), + ({"url": "https://example.test/model.safetensors", "category": "checkpoints"}, "filename"), + ({"url": "https://example.test/model.safetensors", "filename": "model.safetensors"}, "category"), + ], +) +def test_download_spec_requires_core_fields(payload, message): + with pytest.raises(ValueError, match=message): + DownloadSpec.from_mapping(payload) + + +def test_resolution_uses_model_reference_and_preserves_workflow_context(): + resolution = Resolution.from_mapping( + { + "node_id": 7, + "widget_index": 1, + "original_path": "old.safetensors", + "resolved_path": "new.safetensors", + "category": "checkpoints", + "subgraph_id": "group-1", + "is_top_level": False, + "nested_key": "model", + } + ) + + resolution.validate() + resolution = resolution.with_custom_node_metadata( + {"custom_node_adapter": "test"} + ) + + assert isinstance(resolution.reference, ModelReference) + assert resolution.reference.original_path == "old.safetensors" + assert resolution.node_id == 7 + assert resolution.widget_index == 1 + assert resolution.subgraph_id == "group-1" + assert resolution.nested_key == "model" + assert isinstance(resolution.custom_node_metadata, CustomNodeMetadata) + assert resolution.custom_node_metadata.adapter_id == "test" + + +def test_resolution_accepts_path_from_resolved_model(): + resolution = Resolution.from_mapping( + { + "node_id": 1, + "widget_index": 0, + "resolved_model": {"path": "models/checkpoint.safetensors"}, + } + ) + + resolution.validate() + + +@pytest.mark.parametrize( + "payload", + [ + { + "node_id": 1, + "widget_index": "0", + "resolved_path": "new.safetensors", + }, + { + "node_id": 1, + "widget_index": -1, + "resolved_path": "new.safetensors", + }, + { + "node_id": 1, + "widget_index": 0, + "resolved_path": 123, + }, + { + "node_id": 1, + "widget_index": 0, + "resolved_model": {"path": 123}, + }, + ], +) +def test_resolution_rejects_invalid_field_types(payload): + with pytest.raises((TypeError, ValueError)): + Resolution.from_mapping(payload).validate() + + +def test_apply_resolution_accepts_typed_resolution(): + workflow = { + "nodes": [ + { + "id": 1, + "type": "CheckpointLoaderSimple", + "widgets_values": ["old.safetensors"], + } + ] + } + + resolution = Resolution.from_mapping( + { + "node_id": 1, + "widget_index": 0, + "resolved_path": "new.safetensors", + "category": "checkpoints", + "resolved_model": { + "path": "new.safetensors", + "base_directory": "models", + }, + } + ) + updated = apply_resolution(workflow, [resolution]) + + assert updated["nodes"][0]["widgets_values"] == ["new.safetensors"] + + +def test_apply_resolution_skips_invalid_payload_without_raising(): + workflow = {"nodes": []} + + updated = apply_resolution(workflow, [Resolution()]) + + assert updated == workflow diff --git a/tests/test_download_routes.py b/tests/test_download_routes.py index 410ccf2..ceeac4a 100644 --- a/tests/test_download_routes.py +++ b/tests/test_download_routes.py @@ -1,7 +1,7 @@ import asyncio import json from types import SimpleNamespace -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest from aiohttp import web @@ -203,6 +203,29 @@ async def test_download_route_rejects_missing_unsafe_and_unsupported_inputs(): } +@pytest.mark.asyncio +@pytest.mark.parametrize( + "payload", + [ + {"url": 123}, + {"url": "https://example.com/model.safetensors", "filename": 123}, + { + "url": "https://example.com/model.safetensors", + "civitai_session_token": {"token": "invalid"}, + }, + ], +) +async def test_download_route_rejects_malformed_text_fields(payload): + handlers, values = _build_download_routes() + handler = handlers[("POST", "/model_resolver/download")] + + response = await handler(_request(payload)) + + assert response.status == 400 + assert json.loads(response.text)["error"].startswith("Download request") + values["start_background_download"].assert_not_called() + + @pytest.mark.asyncio async def test_download_route_supports_huggingface_headers_and_optional_inputs(): handlers, values = _build_download_routes() @@ -212,8 +235,7 @@ async def test_download_route_supports_huggingface_headers_and_optional_inputs() _request( { "url": "https://huggingface.co/org/repo/resolve/main/model.safetensors", - "path_metadata": ["invalid"], - "metadata": ["invalid"], + "metadata": {}, "base_directory": r"C:\custom", "hf_token": "hf-token", } @@ -230,6 +252,106 @@ async def test_download_route_supports_huggingface_headers_and_optional_inputs() values["get_default_root_for_category"].assert_not_called() +@pytest.mark.asyncio +async def test_download_route_rejects_invalid_download_metadata(): + handlers, values = _build_download_routes() + handler = handlers[("POST", "/model_resolver/download")] + + response = await handler( + _request( + { + "url": "https://example.com/model.safetensors", + "metadata": [], + } + ) + ) + + assert response.status == 400 + assert json.loads(response.text) == { + "error": "Download metadata must be an object" + } + values["start_background_download"].assert_not_called() + + +@pytest.mark.asyncio +async def test_download_route_rejects_invalid_path_metadata_and_identifiers(): + handlers, values = _build_download_routes() + handler = handlers[("POST", "/model_resolver/download")] + + response = await handler( + _request( + { + "url": "https://example.com/model.safetensors", + "path_metadata": ["invalid"], + } + ) + ) + + assert response.status == 400 + assert json.loads(response.text) == { + "error": "Download path metadata must be an object" + } + + response = await handler( + _request( + { + "url": "https://example.com/model.safetensors", + "download_metadata": {"model_id": {"value": 123}}, + } + ) + ) + + assert response.status == 400 + assert json.loads(response.text) == { + "error": "Download metadata model_id must be an integer or string" + } + values["start_background_download"].assert_not_called() + + +@pytest.mark.asyncio +async def test_download_route_rejects_invalid_expected_sha256(): + handlers, values = _build_download_routes() + handler = handlers[("POST", "/model_resolver/download")] + + response = await handler( + _request( + { + "url": "https://example.com/model.safetensors", + "metadata": {"sha256": 123}, + } + ) + ) + + assert response.status == 400 + assert json.loads(response.text) == { + "error": "DownloadSpec expected_sha256 must be a string" + } + values["start_background_download"].assert_not_called() + + +@pytest.mark.asyncio +async def test_download_route_handles_download_spec_type_errors(): + handlers, _values = _build_download_routes() + handler = handlers[("POST", "/model_resolver/download")] + + with patch( + "core.services.download_service.DownloadSpec", + side_effect=TypeError("invalid download specification"), + ): + response = await handler( + _request( + { + "url": "https://example.com/model.safetensors", + } + ) + ) + + assert response.status == 400 + assert json.loads(response.text) == { + "error": "invalid download specification" + } + + @pytest.mark.asyncio async def test_download_route_derives_encoded_filename_from_url_path(): handlers, values = _build_download_routes() @@ -331,6 +453,20 @@ async def test_aria2_install_handles_invalid_and_non_mapping_payloads(): assert values["install_aria2_engine"].call_args_list[-1].args == (False,) +@pytest.mark.asyncio +async def test_aria2_install_rejects_malformed_force_flag(): + handlers, values = _build_download_routes() + handler = handlers[("POST", "/model_resolver/aria2/install")] + + response = await handler(_request({"force": {"enabled": True}})) + + assert response.status == 400 + assert json.loads(response.text) == { + "error": "Aria2 install request force must be a boolean" + } + values["install_aria2_engine"].assert_not_called() + + @pytest.mark.asyncio async def test_download_control_routes_preserve_progress_and_status_codes(): handlers, values = _build_download_routes() diff --git a/tests/test_hash_routes.py b/tests/test_hash_routes.py index e523302..baf4467 100644 --- a/tests/test_hash_routes.py +++ b/tests/test_hash_routes.py @@ -8,6 +8,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from aiohttp import web +from core.contracts import ModelMatch, ModelReference, WorkflowModelInventory from core.path_utils import HashCalculationCancelled, normalize_absolute_path from core.routes.context import RouteContext from core.routes.hashes import register_hash_routes @@ -83,7 +84,7 @@ def _build_hash_routes(overrides=None): f"{path}.modelresolver.json" ), "get_workflow_model_inventory": MagicMock( - return_value={"model_refs": []} + return_value=WorkflowModelInventory() ), "is_path_in_configured_model_roots": MagicMock(return_value=True), "json_api_endpoint": lambda _name: lambda handler: handler, @@ -165,6 +166,9 @@ async def test_hash_routes_validate_local_metadata_and_preview_paths(): response = await preview_handler(SimpleNamespace(query={})) assert response.status == 400 + response = await preview_handler(SimpleNamespace(query={"path": 123})) + assert response.status == 400 + with tempfile.TemporaryDirectory() as temp_dir: model_path = os.path.join(temp_dir, "model.safetensors") preview_path = os.path.join(temp_dir, "model.png") @@ -267,7 +271,13 @@ async def test_workflow_model_hashes_skips_invalid_refs_and_deduplicates_models( handlers, values = _build_hash_routes( { "get_workflow_model_inventory": MagicMock( - return_value={"model_refs": refs} + return_value=WorkflowModelInventory( + model_refs=tuple( + ModelReference.from_mapping(ref) + for ref in refs + if ref.get("full_path") + ) + ) ), "get_local_model_hash_metadata": get_metadata, } @@ -287,7 +297,11 @@ async def test_workflow_model_hashes_skips_invalid_refs_and_deduplicates_models( @pytest.mark.asyncio async def test_local_matches_by_hash_normalizes_and_enriches_results(): - matches = [{"path": r"C:\models\local.safetensors"}] + matches = [ + ModelMatch.from_mapping( + {"model": {"path": r"C:\models\local.safetensors"}} + ) + ] search_matches = MagicMock(return_value=matches) handlers, values = _build_hash_routes( {"search_local_matches_by_hash": search_matches} @@ -310,14 +324,12 @@ async def test_local_matches_by_hash_normalizes_and_enriches_results(): body = json.loads(response.text) assert body["sha256"] == sha256 - assert body["local_hash_matches"] == [ - { - "path": r"C:\models\local.safetensors", - "hash_lookup_source": "civit_ai", - "hash_lookup_filename": "remote.safetensors", - "hash_lookup_sha256": sha256, - } - ] + assert len(body["local_hash_matches"]) == 1 + match = body["local_hash_matches"][0] + assert match["model"]["path"] == r"C:\models\local.safetensors" + assert match["hash_lookup_source"] == "civit_ai" + assert match["hash_lookup_filename"] == "remote.safetensors" + assert match["hash_lookup_sha256"] == sha256 values["search_local_matches_by_hash"].assert_called_once_with( sha256, category="checkpoints", @@ -326,6 +338,22 @@ async def test_local_matches_by_hash_normalizes_and_enriches_results(): ) +@pytest.mark.asyncio +async def test_local_matches_by_hash_rejects_malformed_limit(): + handlers, values = _build_hash_routes() + handler = handlers[("POST", "/model_resolver/local-matches-by-hash")] + + response = await handler( + _request({"sha256": "a" * 64, "max_matches": {"value": 20}}) + ) + + assert response.status == 400 + assert json.loads(response.text) == { + "error": "Local hash matches request max_matches must be an integer" + } + values["search_local_matches_by_hash"].assert_not_called() + + @pytest.mark.asyncio async def test_open_containing_folder_maps_file_manager_errors(): handlers, values = _build_hash_routes() diff --git a/tests/test_huggingface.py b/tests/test_huggingface.py index 6b47999..5acaa4e 100644 --- a/tests/test_huggingface.py +++ b/tests/test_huggingface.py @@ -1,12 +1,13 @@ import json -import re import tempfile import time import unittest from pathlib import Path from unittest.mock import patch +from core.contracts import SearchResult from core.matcher import build_filename_search_queries +from core.sources.common import probe_remote_file_size from core.sources.huggingface import ( HF_AUTHOR_FALLBACKS, _build_author_index_from_models, @@ -23,14 +24,39 @@ from core.sources.huggingface import ( get_huggingface_model_details, get_known_author_fallback_indexes_status, parse_huggingface_url, + search_huggingface, search_huggingface_for_file, ) -from core.sources.common import probe_remote_file_size from core.type_utils import extract_file_size +def _result_dict(result): + assert isinstance(result, SearchResult) + return result.to_dict() + + class HuggingFaceSourceTests(unittest.TestCase): + @patch("core.sources.huggingface.execute_provider_json_request") + def test_general_search_returns_typed_results(self, request_json): + request_json.return_value = [ + { + "id": "Comfy-Org/example", + "modelId": "Comfy-Org/example", + "downloads": 12, + "likes": 3, + } + ] + + results = search_huggingface("example", limit=1) + + self.assertEqual(1, len(results)) + self.assertIsInstance(results[0], SearchResult) + self.assertEqual("Comfy-Org/example", results[0].name) + self.assertEqual("Comfy-Org/example", results[0].extra_value("repo")) + self.assertEqual(12, results[0].extra_value("downloads")) + self.assertEqual("https://huggingface.co/Comfy-Org/example", results[0].url) + def test_known_author_fallbacks_include_comfy_org_and_kijai(self): self.assertEqual(["Comfy-Org", "Kijai"], HF_AUTHOR_FALLBACKS) @@ -142,6 +168,7 @@ class HuggingFaceSourceTests(unittest.TestCase): ) self.assertIsNotNone(result) + result = _result_dict(result) self.assertEqual("hash", result["match_type"]) self.assertEqual(sha256, result["sha256"]) self.assertEqual("models/example.safetensors", result["path"]) @@ -180,6 +207,7 @@ class HuggingFaceSourceTests(unittest.TestCase): ) self.assertIsNotNone(result) + result = _result_dict(result) self.assertEqual("hash", result["match_type"]) self.assertEqual(sha256, result["sha256"]) get_index.assert_called_once() @@ -434,6 +462,7 @@ class HuggingFaceSourceTests(unittest.TestCase): ) self.assertIsNotNone(result) + result = _result_dict(result) self.assertEqual("hash", result["match_type"]) fetch_index.assert_called_once_with("Comfy-Org", headers={}) @@ -482,6 +511,7 @@ class HuggingFaceSourceTests(unittest.TestCase): ) self.assertIsNotNone(result) + result = _result_dict(result) self.assertEqual("hash", result["match_type"]) fetch_index.assert_not_called() @@ -521,16 +551,16 @@ class HuggingFaceSourceTests(unittest.TestCase): url = "https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors" result = parse_huggingface_url(url) self.assertIsNotNone(result) - self.assertEqual(result["repo"], "runwayml/stable-diffusion-v1-5") - self.assertEqual(result["branch"], "main") - self.assertEqual(result["filename"], "v1-5-pruned-emaonly.safetensors") + self.assertEqual(result.repo, "runwayml/stable-diffusion-v1-5") + self.assertEqual(result.branch, "main") + self.assertEqual(result.filename, "v1-5-pruned-emaonly.safetensors") def test_parse_huggingface_url_valid_hf_protocol(self): url = "hf://stabilityai/stable-diffusion-xl-base-1.0/sd_xl_base_1.0.safetensors" result = parse_huggingface_url(url) self.assertIsNotNone(result) - self.assertEqual(result["repo"], "stabilityai/stable-diffusion-xl-base-1.0") - self.assertEqual(result["filename"], "sd_xl_base_1.0.safetensors") + self.assertEqual(result.repo, "stabilityai/stable-diffusion-xl-base-1.0") + self.assertEqual(result.filename, "sd_xl_base_1.0.safetensors") def test_parse_huggingface_url_invalid(self): self.assertIsNone(parse_huggingface_url("https://example.com/not-hf")) @@ -633,6 +663,7 @@ class HuggingFaceSourceTests(unittest.TestCase): result = build_huggingface_custom_result(url) self.assertIsNotNone(result) + result = _result_dict(result) self.assertEqual(sha256, result["sha256"]) self.assertEqual({"SHA256": sha256}, result["hashes"]) self.assertEqual(123, result["size"]) @@ -641,14 +672,17 @@ class HuggingFaceSourceTests(unittest.TestCase): def test_custom_result_leaves_hash_empty_when_huggingface_does_not_provide_it(self): url = "https://huggingface.co/user/repo/blob/main/model.safetensors" - with patch("core.sources.huggingface._get_repo_tree", return_value=None): - with patch( + with ( + patch("core.sources.huggingface._get_repo_tree", return_value=None), + patch( "core.sources.common.fetch_remote_file_size_cached", return_value=456, - ): - result = build_huggingface_custom_result(url) + ), + ): + result = build_huggingface_custom_result(url) self.assertIsNotNone(result) + result = _result_dict(result) self.assertEqual("", result["sha256"]) self.assertEqual({}, result["hashes"]) self.assertEqual(456, result["size"]) @@ -676,8 +710,9 @@ class HuggingFaceSourceTests(unittest.TestCase): }, {"path": "diffusion_models/README.md", "size": 10}, ] - with patch("core.sources.huggingface._get_repo_tree", return_value=tree): - with patch( + with ( + patch("core.sources.huggingface._get_repo_tree", return_value=tree), + patch( "core.sources.huggingface.execute_provider_json_request", return_value={ "modelId": "Comfy-Org/example", @@ -686,11 +721,12 @@ class HuggingFaceSourceTests(unittest.TestCase): "likes": 4, "tags": ["comfyui"], }, - ): - details = get_huggingface_model_details( - "Comfy-Org/example", - "diffusion_models/model_fp8_scaled.safetensors", - ) + ), + ): + details = get_huggingface_model_details( + "Comfy-Org/example", + "diffusion_models/model_fp8_scaled.safetensors", + ) self.assertIsNotNone(details) self.assertEqual("huggingface", details["source"]) @@ -710,19 +746,21 @@ class HuggingFaceSourceTests(unittest.TestCase): def test_model_details_passes_huggingface_token_to_tree_and_metadata(self): tree = [{"path": "model.safetensors", "size": 100}] - with patch( - "core.sources.huggingface._get_repo_tree", - return_value=tree, - ) as mock_tree: - with patch( + with ( + patch( + "core.sources.huggingface._get_repo_tree", + return_value=tree, + ) as mock_tree, + patch( "core.sources.huggingface.execute_provider_json_request", return_value={}, - ) as mock_request: - details = get_huggingface_model_details( - "private/repo", - "model.safetensors", - token="hf_secret", - ) + ) as mock_request, + ): + details = get_huggingface_model_details( + "private/repo", + "model.safetensors", + token="hf_secret", + ) self.assertIsNotNone(details) expected_headers = {"Authorization": "Bearer hf_secret"} diff --git a/tests/test_loaded_models_routes.py b/tests/test_loaded_models_routes.py index 95eace3..d6ef2a0 100644 --- a/tests/test_loaded_models_routes.py +++ b/tests/test_loaded_models_routes.py @@ -7,6 +7,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from aiohttp import web +from core.contracts import WorkflowModelInventory from core.path_utils import get_filename_from_path from core.routes.context import RouteContext from core.routes.loaded_models import register_loaded_model_routes @@ -35,7 +36,17 @@ def _json_api_endpoint(_name, **_kwargs): return lambda handler: handler +def _as_inventory(inventory=None): + if not isinstance(inventory, WorkflowModelInventory): + inventory = WorkflowModelInventory( + available_models=(inventory or {}).get("available_models", []), + model_refs=(inventory or {}).get("model_refs", []), + ) + return inventory + + def _build_routes(inventory=None): + inventory = _as_inventory(inventory) routes = _Routes() extension = SimpleNamespace( loaded_progress=MagicMock(), @@ -54,7 +65,6 @@ def _build_routes(inventory=None): ), "get_workflow_model_inventory": MagicMock( return_value=inventory - or {"available_models": [], "model_refs": []} ), "json_api_endpoint": _json_api_endpoint, "routes": routes, @@ -247,7 +257,7 @@ async def test_loaded_models_progress_counts_top_level_and_subgraph_nodes(): "total": 5, } ) - return {"available_models": [], "model_refs": []} + return WorkflowModelInventory() values["get_workflow_model_inventory"].side_effect = inventory_with_progress workflow = { @@ -415,7 +425,7 @@ async def test_loaded_models_route_uses_adapter_name_and_original_path_fallback( def inventory_with_progress(_workflow, progress_callback): progress_callback(progress_payload) - return inventory + return _as_inventory(inventory) values["get_workflow_model_inventory"].side_effect = inventory_with_progress diff --git a/tests/test_local_hash_matches.py b/tests/test_local_hash_matches.py index 09f0d33..d318a09 100644 --- a/tests/test_local_hash_matches.py +++ b/tests/test_local_hash_matches.py @@ -1,11 +1,21 @@ from unittest.mock import Mock +from core.contracts import ModelMatch from core.local_hash_matches import collect_local_hash_matches_for_result def test_collect_local_hash_matches_forwards_lookup_options_and_enriches_results(): sha256 = "a" * 64 - lookup = Mock(return_value=[{"path": r"C:\models\local.safetensors"}]) + lookup = Mock( + return_value=[ + ModelMatch.from_mapping( + { + "model": {"path": r"C:\models\local.safetensors"}, + "filename": "local.safetensors", + } + ) + ] + ) matches = collect_local_hash_matches_for_result( sha256, @@ -23,14 +33,11 @@ def test_collect_local_hash_matches_forwards_lookup_options_and_enriches_results max_matches=7, force_rescan=True, ) - assert matches == [ - { - "path": r"C:\models\local.safetensors", - "hash_lookup_source": "civitai", - "hash_lookup_filename": "remote.safetensors", - "hash_lookup_sha256": sha256, - } - ] + assert len(matches) == 1 + assert matches[0].model.path == r"C:\models\local.safetensors" + assert matches[0].extra_value("hash_lookup_source") == "civitai" + assert matches[0].extra_value("hash_lookup_filename") == "remote.safetensors" + assert matches[0].extra_value("hash_lookup_sha256") == sha256 def test_collect_local_hash_matches_skips_empty_hash_without_lookup(): diff --git a/tests/test_log_system.py b/tests/test_log_system.py new file mode 100644 index 0000000..f81f647 --- /dev/null +++ b/tests/test_log_system.py @@ -0,0 +1,26 @@ +import io +import logging + +from core.log_system.logger import DirectConsoleHandler + + +def test_console_handler_replaces_characters_unsupported_by_windows_encoding(): + buffer = io.BytesIO() + stream = io.TextIOWrapper(buffer, encoding="cp1250", errors="strict") + handler = DirectConsoleHandler(stream) + handler.setFormatter(logging.Formatter("%(message)s")) + + handler.emit( + logging.LogRecord( + "test", + logging.INFO, + __file__, + 1, + "Download complete ✓", + (), + None, + ) + ) + stream.flush() + + assert "Download complete ?" in buffer.getvalue().decode("cp1250") diff --git a/tests/test_lora_manager_archive.py b/tests/test_lora_manager_archive.py index e25c1d6..1727b15 100644 --- a/tests/test_lora_manager_archive.py +++ b/tests/test_lora_manager_archive.py @@ -1,17 +1,20 @@ -import unittest -import sqlite3 import json +import sqlite3 +import unittest from unittest.mock import patch + +from core.contracts import SearchResult from core.sources.lora_manager_archive import ( - _normalize_model_type, - _extract_search_tokens, - _query_candidate_rows, - _load_full_rows_for_versions, _build_result_from_row, + _extract_search_tokens, _find_hash_match_row, + _load_full_rows_for_versions, + _normalize_model_type, + _query_candidate_rows, search_lora_manager_archive_for_file, ) + class LoraManagerArchiveTests(unittest.TestCase): def test_normalize_model_type(self): @@ -146,12 +149,12 @@ class LoraManagerArchiveTests(unittest.TestCase): # Test building result from row result = _build_result_from_row(conn, row, "ghost_checkpoint") self.assertIsNotNone(result) - self.assertEqual(result["filename"], "ghost_v1.safetensors") - self.assertEqual(result["name"], "ghost_checkpoint") - self.assertEqual(result["trained_words"], ["ghost_word"]) - self.assertEqual(result["size"], 3 * 1024) - self.assertEqual(result["sha256"], archive_hash) - self.assertEqual(result["hashes"], {"SHA256": archive_hash}) + self.assertEqual(result.filename, "ghost_v1.safetensors") + self.assertEqual(result.name, "ghost_checkpoint") + self.assertEqual(list(result.trained_words), ["ghost_word"]) + self.assertEqual(result.size, 3 * 1024) + self.assertEqual(result.sha256, archive_hash) + self.assertEqual(dict(result.hashes), {"SHA256": archive_hash}) hash_match = _find_hash_match_row(conn, archive_hash) self.assertIsNotNone(hash_match) @@ -168,6 +171,27 @@ class LoraManagerArchiveTests(unittest.TestCase): ) self.assertIsNotNone(hash_result) - self.assertEqual(hash_result["match_type"], "hash") - self.assertEqual(hash_result["confidence"], 100.0) - self.assertEqual(hash_result["sha256"], archive_hash) + self.assertEqual(hash_result.match_type, "hash") + self.assertEqual(hash_result.confidence, 100.0) + self.assertEqual(hash_result.sha256, archive_hash) + + def test_filename_search_keeps_typed_candidates(self): + filename = "typed-lora-archive-model.safetensors" + candidate = SearchResult( + source="lora_manager_archive", + name="typed-lora-archive-model", + version_name="v1", + filename=filename, + download_url="https://example.com/typed-lora-archive-model.safetensors", + ) + + with patch( + "core.sources.lora_manager_archive.search_lora_manager_archive", + return_value=[candidate], + ): + result = search_lora_manager_archive_for_file(filename) + + self.assertIsInstance(result, SearchResult) + self.assertEqual(filename, result.filename) + self.assertEqual(100.0, result.confidence) + self.assertEqual("exact", result.match_type) diff --git a/tests/test_matcher.py b/tests/test_matcher.py index bbfdb33..3472611 100644 --- a/tests/test_matcher.py +++ b/tests/test_matcher.py @@ -1,6 +1,7 @@ import unittest from unittest.mock import patch +from core.contracts import SearchResult from core.matcher import ( build_filename_search_queries, calculate_archived_model_confidence, @@ -12,6 +13,7 @@ from core.matcher import ( ) from core.progress import report_progress from core.sources.model_list import search_model_list, search_model_list_multiple +from core.sources.popular import search_popular_models from core.type_utils import ( MODEL_CONTAINER_SUFFIXES, MODEL_VARIANT_SUFFIXES, @@ -47,8 +49,8 @@ class MatcherTests(unittest.TestCase): ] matches = find_matches("sd_xl_base_1.0.safetensors", candidates) self.assertEqual(2, len(matches)) - self.assertEqual("sd_xl_base_1.0.safetensors", matches[0]["filename"]) - self.assertEqual(1.0, matches[0]["similarity"]) + self.assertEqual("sd_xl_base_1.0.safetensors", matches[0].filename) + self.assertEqual(1.0, matches[0].similarity) # With higher threshold, only the best match should be returned matches_filtered = find_matches("sd_xl_base_1.0.safetensors", candidates, threshold=0.7) @@ -165,9 +167,21 @@ class MatcherTests(unittest.TestCase): "qwen3vl-4b-heretic_int8.safetensors", "qwen_image_fp8_e4m3fn.safetensors", ], - [match["filename"] for match in matches], + [match.filename for match in matches], ) + def test_find_matches_does_not_mutate_candidate_mappings(self): + candidate = {"filename": "model.safetensors", "category": "checkpoints"} + + matches = find_matches("model.safetensors", [candidate]) + + self.assertEqual(1, len(matches)) + self.assertEqual( + {"filename": "model.safetensors", "category": "checkpoints"}, + candidate, + ) + self.assertEqual("checkpoints", matches[0].model.category) + def test_search_model_list_unification(self): try: search_model_list("nonexistent_model_name_xyz.safetensors", exact_only=True) @@ -200,25 +214,43 @@ class MatcherTests(unittest.TestCase): fuzzy = search_model_list("target-model-extra.safetensors") multiple = search_model_list_multiple("exact", limit=1) - self.assertEqual( - exact, - { - "source": "model_list", - "filename": "exact.safetensors", - "url": "https://example.test/exact", - "name": "Exact model", - "type": "checkpoints", - "directory": "checkpoints", - "size": "12 MB", - "match_type": "exact", + self.assertIsInstance(exact, SearchResult) + self.assertEqual(exact.source, "model_list") + self.assertEqual(exact.filename, "exact.safetensors") + self.assertEqual(exact.url, "https://example.test/exact") + self.assertEqual(exact.name, "Exact model") + self.assertEqual(exact.model_type, "checkpoints") + self.assertEqual(exact.extra_value("directory"), "checkpoints") + self.assertEqual(exact.size, 12 * 1024**2) + self.assertEqual(exact.match_type, "exact") + self.assertEqual(fuzzy.filename, "target-model.safetensors") + self.assertEqual(fuzzy.match_type, "fuzzy") + self.assertGreater(fuzzy.confidence, 0) + self.assertEqual(multiple[0].filename, "exact.safetensors") + self.assertGreater(multiple[0].confidence, 0) + self.assertEqual(multiple[0].match_type, "similar") + + def test_popular_search_returns_typed_search_results(self): + with patch( + "core.sources.popular._load_popular_models", + return_value={ + "popular-model.safetensors": { + "url": "https://example.test/popular", + "type": "checkpoints", + "directory": "checkpoints", + "canonical_name": "popular-model", + } }, - ) - self.assertEqual(fuzzy["filename"], "target-model.safetensors") - self.assertEqual(fuzzy["match_type"], "fuzzy") - self.assertIn("confidence", fuzzy) - self.assertEqual(multiple[0]["filename"], "exact.safetensors") - self.assertIn("confidence", multiple[0]) - self.assertNotIn("match_type", multiple[0]) + ): + results = search_popular_models("popular-model", limit=1) + + self.assertEqual(1, len(results)) + self.assertIsInstance(results[0], SearchResult) + self.assertEqual("popular-model.safetensors", results[0].filename) + self.assertEqual("https://example.test/popular", results[0].url) + self.assertEqual("checkpoints", results[0].model_type) + self.assertEqual("checkpoints", results[0].extra_value("directory")) + self.assertEqual("popular-model", results[0].extra_value("canonical_name")) def test_build_filename_search_queries(self): # Verify that common precisions are extracted correctly diff --git a/tests/test_metadata_builder.py b/tests/test_metadata_builder.py index 2aa3b7a..3026055 100644 --- a/tests/test_metadata_builder.py +++ b/tests/test_metadata_builder.py @@ -497,6 +497,10 @@ class MetadataBuilderTests(unittest.TestCase): normalize_metadata_build_mode("unsupported"), ) + def test_invalid_metadata_mode_type_is_rejected(self): + with self.assertRaises(TypeError): + normalize_metadata_build_mode(123) + def test_exact_sidecar_remains_canonical_after_sibling_is_removed(self): with tempfile.TemporaryDirectory() as tmpdir: bin_path = self._write_model(tmpdir, "adapter.bin", b"bin") diff --git a/tests/test_metadata_model_helpers.py b/tests/test_metadata_model_helpers.py index 86d5f9e..79f6c1b 100644 --- a/tests/test_metadata_model_helpers.py +++ b/tests/test_metadata_model_helpers.py @@ -3,8 +3,13 @@ import tempfile import unittest from unittest.mock import patch +from core.contracts import ResolvedModel from core.metadata_audit import _model_key -from core.metadata_model_utils import dedupe_models, is_model_file_path +from core.metadata_model_utils import ( + dedupe_models, + is_model_file_path, + normalize_models, +) from core.path_utils import get_model_path_identity from core.resolver import _build_local_hash_match_cache from core.scanner import scan_directory @@ -51,7 +56,16 @@ class MetadataModelHelperTests(unittest.TestCase): get_model_path_identity(duplicate["path"]), ) self.assertEqual("", get_model_path_identity(empty["path"])) - self.assertEqual([first], dedupe_models(models)) + deduped = dedupe_models(models) + self.assertEqual(1, len(deduped)) + self.assertIsInstance(deduped[0], ResolvedModel) + self.assertEqual(model_path, deduped[0].path) + + def test_normalize_models_rejects_invalid_collection_payloads(self): + with self.assertRaises(TypeError): + normalize_models(0) + with self.assertRaises(TypeError): + normalize_models({"path": "model.safetensors"}) def test_scanner_and_metadata_identity_share_path_fallback(self): model_path = os.path.join("models", "model.safetensors") @@ -82,7 +96,7 @@ class MetadataModelHelperTests(unittest.TestCase): ): models = scan_directory("models", {".safetensors"}, "checkpoints") - self.assertEqual(["model.safetensors"], [model["filename"] for model in models]) + self.assertEqual(["model.safetensors"], [model.filename for model in models]) with patch("core.resolver.os.path.isdir", return_value=False), patch( "core.resolver.find_metadata_sidecar_path", @@ -99,7 +113,7 @@ class MetadataModelHelperTests(unittest.TestCase): ) self.assertEqual(1, len(index["a" * 64])) - self.assertEqual(model_path, index["a" * 64][0]["model"]["path"]) + self.assertEqual(model_path, index["a" * 64][0].model.path) def test_model_path_identity_rejects_empty_and_whitespace_paths(self): self.assertEqual("", get_model_path_identity(None)) diff --git a/tests/test_path_security.py b/tests/test_path_security.py index 1bb4fe0..c1ed38a 100644 --- a/tests/test_path_security.py +++ b/tests/test_path_security.py @@ -33,10 +33,13 @@ from core.network_utils import ( validate_public_http_url, ) from core.path_utils import ( + dedupe_local_base_directories, get_safe_metadata_sidecar_path, get_safe_model_resolver_sidecar_path, is_path_in_configured_model_roots, is_path_within, + normalize_folder_path_values, + normalize_string_values, ) @@ -52,6 +55,33 @@ class DummyFolderPaths: class PathSecurityTests(unittest.TestCase): + def test_external_string_collections_do_not_iterate_scalar_text(self): + self.assertEqual( + [r"C:\\models\\checkpoints"], + normalize_folder_path_values(r"C:\\models\\checkpoints"), + ) + self.assertEqual( + [r"C:\\models\\checkpoints"], + normalize_folder_path_values( + (r"C:\\models\\checkpoints", {".safetensors"}) + ), + ) + self.assertEqual( + ["model.safetensors"], + normalize_string_values("model.safetensors"), + ) + self.assertEqual( + ["model.safetensors"], + normalize_string_values([None, "model.safetensors", 123]), + ) + + def test_dedupe_local_base_directories_rejects_scalar_payloads(self): + self.assertEqual([], dedupe_local_base_directories(None)) + with self.assertRaises(TypeError): + dedupe_local_base_directories(0) + with self.assertRaises(TypeError): + dedupe_local_base_directories("models") + def test_is_path_within_rejects_sibling_directory_with_shared_prefix(self): with tempfile.TemporaryDirectory() as temp_root: model_root = os.path.join(temp_root, "models") diff --git a/tests/test_path_templates.py b/tests/test_path_templates.py index ef9fbbd..af8d9ef 100644 --- a/tests/test_path_templates.py +++ b/tests/test_path_templates.py @@ -1,9 +1,9 @@ import unittest +from core.contracts import ResolvedModel from core.path_templates import infer_download_path_templates from core.settings import calculate_template_subfolder - BASE_MODELS = { "base_models": [ { @@ -42,26 +42,36 @@ BASE_MODELS = { } +def _typed_models(items): + return [ResolvedModel.from_mapping(item) for item in items] + + class PathTemplateInferenceTests(unittest.TestCase): + def test_rejects_non_model_collections(self): + with self.assertRaises(TypeError): + infer_download_path_templates(0, BASE_MODELS) + def test_detects_base_model_and_first_tag_layout(self): result = infer_download_path_templates( - [ - { - "category": "loras", - "relative_path": r"ponyxl\style\model-a.safetensors", - "path": r"C:\models\loras\ponyxl\style\model-a.safetensors", - }, - { - "category": "lora", - "relative_path": r"ponyxl\character\model-b.safetensors", - "path": r"C:\models\loras\ponyxl\character\model-b.safetensors", - }, - { - "category": "loras", - "relative_path": r"sdxl\style\model-c.safetensors", - "path": r"C:\models\loras\sdxl\style\model-c.safetensors", - }, - ], + _typed_models( + [ + { + "category": "loras", + "relative_path": r"ponyxl\style\model-a.safetensors", + "path": r"C:\models\loras\ponyxl\style\model-a.safetensors", + }, + { + "category": "lora", + "relative_path": r"ponyxl\character\model-b.safetensors", + "path": r"C:\models\loras\ponyxl\character\model-b.safetensors", + }, + { + "category": "loras", + "relative_path": r"sdxl\style\model-c.safetensors", + "path": r"C:\models\loras\sdxl\style\model-c.safetensors", + }, + ] + ), BASE_MODELS, ) @@ -72,18 +82,20 @@ class PathTemplateInferenceTests(unittest.TestCase): def test_detects_flat_layout(self): result = infer_download_path_templates( - [ - { - "category": "vae", - "relative_path": "vae-a.safetensors", - "path": r"C:\models\vae\vae-a.safetensors", - }, - { - "category": "vae", - "relative_path": "vae-b.safetensors", - "path": r"C:\models\vae\vae-b.safetensors", - }, - ], + _typed_models( + [ + { + "category": "vae", + "relative_path": "vae-a.safetensors", + "path": r"C:\models\vae\vae-a.safetensors", + }, + { + "category": "vae", + "relative_path": "vae-b.safetensors", + "path": r"C:\models\vae\vae-b.safetensors", + }, + ] + ), BASE_MODELS, ) @@ -92,28 +104,30 @@ class PathTemplateInferenceTests(unittest.TestCase): def test_detects_nested_checkpoint_base_model_mappings(self): result = infer_download_path_templates( - [ - { - "category": "checkpoints", - "relative_path": r"SDXL\Pony\pony-realism.safetensors", - "path": r"C:\models\checkpoints\SDXL\Pony\pony-realism.safetensors", - }, - { - "category": "checkpoints", - "relative_path": r"SDXL\Illustrious\wai.safetensors", - "path": r"C:\models\checkpoints\SDXL\Illustrious\wai.safetensors", - }, - { - "category": "checkpoints", - "relative_path": r"SDXL\Illustrious\realistic\real.safetensors", - "path": r"C:\models\checkpoints\SDXL\Illustrious\realistic\real.safetensors", - }, - { - "category": "checkpoints", - "relative_path": r"SDXL\SDXL\juggernaut.safetensors", - "path": r"C:\models\checkpoints\SDXL\SDXL\juggernaut.safetensors", - }, - ], + _typed_models( + [ + { + "category": "checkpoints", + "relative_path": r"SDXL\Pony\pony-realism.safetensors", + "path": r"C:\models\checkpoints\SDXL\Pony\pony-realism.safetensors", + }, + { + "category": "checkpoints", + "relative_path": r"SDXL\Illustrious\wai.safetensors", + "path": r"C:\models\checkpoints\SDXL\Illustrious\wai.safetensors", + }, + { + "category": "checkpoints", + "relative_path": r"SDXL\Illustrious\realistic\real.safetensors", + "path": r"C:\models\checkpoints\SDXL\Illustrious\realistic\real.safetensors", + }, + { + "category": "checkpoints", + "relative_path": r"SDXL\SDXL\juggernaut.safetensors", + "path": r"C:\models\checkpoints\SDXL\SDXL\juggernaut.safetensors", + }, + ] + ), BASE_MODELS, ) @@ -142,18 +156,20 @@ class PathTemplateInferenceTests(unittest.TestCase): def test_detect_does_not_map_krea2_to_flux_krea_folder(self): result = infer_download_path_templates( - [ - { - "category": "loras", - "relative_path": r"FLUX\KREA\concept\snofs_krea_v1.safetensors", - "path": r"C:\models\loras\FLUX\KREA\concept\snofs_krea_v1.safetensors", - }, - { - "category": "loras", - "relative_path": r"FLUX\KREA\style\another_flux_krea.safetensors", - "path": r"C:\models\loras\FLUX\KREA\style\another_flux_krea.safetensors", - }, - ], + _typed_models( + [ + { + "category": "loras", + "relative_path": r"FLUX\KREA\concept\snofs_krea_v1.safetensors", + "path": r"C:\models\loras\FLUX\KREA\concept\snofs_krea_v1.safetensors", + }, + { + "category": "loras", + "relative_path": r"FLUX\KREA\style\another_flux_krea.safetensors", + "path": r"C:\models\loras\FLUX\KREA\style\another_flux_krea.safetensors", + }, + ] + ), BASE_MODELS, ) @@ -165,18 +181,20 @@ class PathTemplateInferenceTests(unittest.TestCase): def test_detect_does_not_map_anima_to_wan_animate_folder(self): result = infer_download_path_templates( - [ - { - "category": "checkpoints", - "relative_path": r"WAN\WAN2.2ANIMATE\model-a.safetensors", - "path": r"C:\models\checkpoints\WAN\WAN2.2ANIMATE\model-a.safetensors", - }, - { - "category": "checkpoints", - "relative_path": r"WAN\WAN2.2ANIMATE\model-b.safetensors", - "path": r"C:\models\checkpoints\WAN\WAN2.2ANIMATE\model-b.safetensors", - }, - ], + _typed_models( + [ + { + "category": "checkpoints", + "relative_path": r"WAN\WAN2.2ANIMATE\model-a.safetensors", + "path": r"C:\models\checkpoints\WAN\WAN2.2ANIMATE\model-a.safetensors", + }, + { + "category": "checkpoints", + "relative_path": r"WAN\WAN2.2ANIMATE\model-b.safetensors", + "path": r"C:\models\checkpoints\WAN\WAN2.2ANIMATE\model-b.safetensors", + }, + ] + ), BASE_MODELS, ) @@ -184,18 +202,20 @@ class PathTemplateInferenceTests(unittest.TestCase): def test_detect_accepts_version_and_plural_alias_suffixes(self): result = infer_download_path_templates( - [ - { - "category": "diffusion_models", - "relative_path": r"IDEOGRAM4\model-a.safetensors", - "path": r"C:\models\diffusion_models\IDEOGRAM4\model-a.safetensors", - }, - { - "category": "upscale_models", - "relative_path": r"ESRGAN\Upscalers\photo\model-a.pth", - "path": r"C:\models\upscale_models\ESRGAN\Upscalers\photo\model-a.pth", - }, - ], + _typed_models( + [ + { + "category": "diffusion_models", + "relative_path": r"IDEOGRAM4\model-a.safetensors", + "path": r"C:\models\diffusion_models\IDEOGRAM4\model-a.safetensors", + }, + { + "category": "upscale_models", + "relative_path": r"ESRGAN\Upscalers\photo\model-a.pth", + "path": r"C:\models\upscale_models\ESRGAN\Upscalers\photo\model-a.pth", + }, + ] + ), BASE_MODELS, ) diff --git a/tests/test_refactoring_targets.py b/tests/test_refactoring_targets.py index e16f19d..f1e558b 100644 --- a/tests/test_refactoring_targets.py +++ b/tests/test_refactoring_targets.py @@ -7,6 +7,8 @@ import unittest from typing import ClassVar from unittest.mock import AsyncMock, MagicMock, patch +from core.contracts import MissingModel, WorkflowAnalysisResult, WorkflowModelInventory + # Make sure parent package directory is in sys.path parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) if parent_dir not in sys.path: @@ -58,6 +60,9 @@ mock_routes.post = post_decorator # Import the module so that routes register node_mod = importlib.import_module("comfyui-model-resolver") +TypedSearchResult = importlib.import_module( + "comfyui-model-resolver.core.contracts" +).SearchResult # Force setup_routes to run to populate routes_registered mock registry node_mod.extension.routes_setup = False @@ -506,12 +511,12 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase): if function_name == "build_huggingface_custom_result": self.assertEqual(args, (url, "model.safetensors", None)) self.assertEqual(kwargs, {}) - return { - "source": "huggingface", - "filename": "model.safetensors", - "url": url, - "download_url": url, - } + return TypedSearchResult( + source="huggingface", + filename="model.safetensors", + url=url, + download_url=url, + ) self.fail(f"Unexpected threaded function: {function_name}") with ( @@ -551,14 +556,27 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase): self.assertEqual(getattr(func, "__name__", ""), "analyze_and_find_matches") self.assertEqual(args[:3], (workflow, 0.0, 10)) self.assertEqual(kwargs["force_rescan"], False) - return { - "missing_models": [ - { - "name": "already-matched.safetensors", - "matches": [{"confidence": 100}], - } - ] - } + return WorkflowAnalysisResult( + missing_models=( + MissingModel.from_mapping( + { + "name": "already-matched.safetensors", + "matches": [ + { + "model": { + "path": "models/already-matched.safetensors", + "filename": "already-matched.safetensors", + "category": "checkpoints", + }, + "confidence": 100, + } + ], + } + ), + ), + total_missing=1, + total_models_analyzed=1, + ) with ( patch("asyncio.to_thread", side_effect=fake_to_thread), @@ -602,7 +620,14 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase): request = AsyncMock() request.json.return_value = { "workflow": workflow, - "resolutions": [{"node_id": 1, "widget_index": 0}], + "resolutions": [ + { + "node_id": 1, + "widget_index": 0, + "resolved_path": "resolved.safetensors", + "category": "checkpoints", + } + ], } with patch("aiohttp.web.json_response") as mock_json_response: @@ -650,6 +675,22 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase): {"loaded_models": [], "total": 0} ) + async def test_loaded_models_route_rejects_non_text_progress_id(self): + post_handler = routes_registered[("POST", "/model_resolver/loaded")] + request = AsyncMock() + request.json.return_value = { + "workflow": {"nodes": []}, + "loaded_id": 123, + } + + with patch("aiohttp.web.json_response") as mock_json_response: + await post_handler(request) + + mock_json_response.assert_called_once_with( + {"error": "Loaded models request loaded_id must be a string"}, + status=400, + ) + async def test_loaded_models_route_rejects_non_object_workflow(self): post_handler = routes_registered[("POST", "/model_resolver/loaded")] request = AsyncMock() @@ -796,6 +837,38 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase): status=400, ) + async def test_model_details_route_rejects_non_numeric_provider_ids(self): + post_handler = routes_registered[("POST", "/model_resolver/model-details")] + request = AsyncMock() + request.json.return_value = { + "source": "civitai", + "model_id": "not-a-number", + } + + with patch("aiohttp.web.json_response") as mock_json_response: + await post_handler(request) + + mock_json_response.assert_called_once_with( + {"error": "Model details request model_id must be an integer"}, + status=400, + ) + + async def test_model_details_route_keeps_huggingface_ids_textual(self): + post_handler = routes_registered[("POST", "/model_resolver/model-details")] + request = AsyncMock() + request.json.return_value = { + "source": "huggingface", + "model_id": 123, + } + + with patch("aiohttp.web.json_response") as mock_json_response: + await post_handler(request) + + mock_json_response.assert_called_once_with( + {"error": "Model details model_id must be a string for Hugging Face"}, + status=400, + ) + async def test_civitai_search_route_rejects_missing_filename(self): post_handler = routes_registered[("POST", "/model_resolver/civitai-search")] request = AsyncMock() @@ -912,7 +985,7 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase): if function_name == "get_workflow_model_inventory": self.assertEqual((workflow,), args) self.assertEqual({}, kwargs) - return {"available_models": [], "model_refs": []} + return WorkflowModelInventory() self.fail(f"Unexpected threaded function: {function_name}") with ( @@ -1238,7 +1311,6 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase): "request_public_url", "resolve_civarchive_by_hash", "search_huggingface_for_file", - "to_bool", "web", "write_model_resolver_metadata", }, @@ -1356,7 +1428,9 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase): response = MagicMock(status_code=200) request_public_url = MagicMock(return_value=(response, mirror_url, {})) - resolve_civarchive_by_hash = MagicMock(return_value=result) + resolve_civarchive_by_hash = MagicMock( + return_value=TypedSearchResult.from_mapping(result, source="civarchive") + ) service = self._build_model_service_for_unit_test( service_module="civitai_search_service", service_class="CivitAISearchService", @@ -1389,8 +1463,12 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase): } try: - with patch( - "core.sources.civitai.get_model_info_by_hash", + civitai_module = importlib.import_module( + "comfyui-model-resolver.core.sources.civitai" + ) + with patch.object( + civitai_module, + "get_model_info_by_hash", return_value=None, ): api_response = await service.civitai_search(request) @@ -1404,6 +1482,100 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase): payload = json.loads(api_response.text) self.assertEqual(payload["download_url"], mirror_url) + async def test_civitai_service_accepts_typed_huggingface_result(self): + expected_hash = "c" * 64 + hf_result = TypedSearchResult( + source="huggingface", + name="Typed HuggingFace model", + filename="model.safetensors", + url=( + "https://huggingface.co/example/repo/blob/main/model.safetensors" + ), + download_url=( + "https://huggingface.co/example/repo/resolve/main/model.safetensors" + ), + sha256=expected_hash, + extra={ + "repo_id": "example/repo", + "path": "model.safetensors", + "from_metadata": True, + }, + ) + + def to_bool(value, default=False): + return default if value is None else bool(value) + + search_huggingface_for_file = MagicMock(return_value=hf_result) + service = self._build_model_service_for_unit_test( + service_module="civitai_search_service", + service_class="CivitAISearchService", + download_available=True, + extract_sha256_from_metadata=lambda value: value.get("sha256", ""), + get_filename_from_path=lambda value: str(value or "").replace( + "\\", "/" + ).rsplit("/", 1)[-1], + normalize_category_to_model_type=lambda value: "Checkpoint", + normalize_sha256=lambda value: str(value or ""), + resolve_civarchive_by_hash=MagicMock(return_value=None), + search_huggingface_for_file=search_huggingface_for_file, + to_bool=to_bool, + ) + request = AsyncMock() + with tempfile.NamedTemporaryFile( + suffix=".safetensors", delete=False + ) as handle: + model_path = handle.name + request.json.return_value = { + "filename": "model.safetensors", + "category": "checkpoints", + "resolved_path": model_path, + "sha256": expected_hash, + } + + civitai_module = importlib.import_module( + "comfyui-model-resolver.core.sources.civitai" + ) + try: + with patch.object( + civitai_module, + "get_model_info_by_hash", + return_value=None, + ): + response = await service.civitai_search(request) + finally: + os.unlink(model_path) + + self.assertEqual(response.status, 200) + payload = json.loads(response.text) + self.assertEqual("huggingface", payload["source"]) + self.assertEqual(expected_hash, payload["sha256"]) + self.assertEqual( + hf_result.download_url, + payload["download_url"], + ) + self.assertEqual( + "https://huggingface.co/example/repo/blob/main/model.safetensors", + payload["url"], + ) + search_huggingface_for_file.assert_called_once() + + async def test_civitai_service_rejects_malformed_text_fields(self): + service = self._build_model_service_for_unit_test( + service_module="civitai_search_service", + service_class="CivitAISearchService", + download_available=True, + ) + request = AsyncMock() + request.json.return_value = {"filename": 123} + + response = await service.civitai_search(request) + + self.assertEqual(response.status, 400) + self.assertEqual( + json.loads(response.text)["error"], + "CivitAI search request filename must be a string", + ) + async def test_civitai_service_preserves_remote_sidecar_payload_fields(self): expected_hash = "b" * 64 remote_result = { @@ -1667,6 +1839,22 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase): self.assertEqual(response.status, 400) self.assertEqual(json.loads(response.text)["error"], "Unsafe URL") + async def test_custom_url_service_rejects_malformed_text_fields(self): + service = self._build_model_service_for_unit_test( + service_module="custom_url_service", + service_class="CustomUrlService", + ) + request = AsyncMock() + request.json.return_value = {"url": 123} + + response = await service.custom_url(request) + + self.assertEqual(response.status, 400) + self.assertEqual( + json.loads(response.text)["error"], + "Custom URL request url must be a string", + ) + async def test_model_details_service_returns_unavailable_when_providers_are_missing(self): service = self._build_model_service_for_unit_test( service_module="model_details_service", @@ -1703,3 +1891,23 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase): json.loads(response.text)["error"], "Model details not found", ) + + async def test_model_details_service_rejects_malformed_identifiers(self): + service = self._build_model_service_for_unit_test( + service_module="model_details_service", + service_class="ModelDetailsService", + download_available=True, + ) + request = AsyncMock() + request.json.return_value = { + "source": "civitai", + "model_id": True, + } + + response = await service.model_details(request) + + self.assertEqual(response.status, 400) + self.assertEqual( + json.loads(response.text)["error"], + "Model details request model_id must be an integer or string", + ) diff --git a/tests/test_refactoring_unification.py b/tests/test_refactoring_unification.py index f80dde5..5623a6e 100644 --- a/tests/test_refactoring_unification.py +++ b/tests/test_refactoring_unification.py @@ -9,6 +9,7 @@ parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) if parent_dir not in sys.path: sys.path.insert(0, parent_dir) +from core.contracts import SearchResult from core.matcher import MODEL_TITLE_MATCH_THRESHOLD from core.type_utils import ( build_model_result, @@ -30,13 +31,17 @@ class TestRefactoringUnification(unittest.TestCase): from core import resolver calls = MagicMock() - with patch.object(resolver, "invalidate_model_files_cache") as invalidate_inventory: - with patch.object( + with ( + patch.object( + resolver, "invalidate_model_files_cache" + ) as invalidate_inventory, + patch.object( resolver, "invalidate_local_hash_match_cache" - ) as invalidate_hash_matches: - calls.attach_mock(invalidate_inventory, "inventory") - calls.attach_mock(invalidate_hash_matches, "hash_matches") - resolver.invalidate_model_caches() + ) as invalidate_hash_matches, + ): + calls.attach_mock(invalidate_inventory, "inventory") + calls.attach_mock(invalidate_hash_matches, "hash_matches") + resolver.invalidate_model_caches() calls.assert_has_calls([call.inventory(), call.hash_matches()]) invalidate_inventory.assert_called_once_with() @@ -74,12 +79,14 @@ class TestRefactoringUnification(unittest.TestCase): filename="test.safetensors", extra_field="hello_world", ) - self.assertEqual(res["source"], "civitai") - self.assertEqual(res["model_id"], 100) - self.assertEqual(res["version_id"], 200) - self.assertEqual(res["name"], "Test Model") - self.assertEqual(res["filename"], "test.safetensors") - self.assertEqual(res["extra_field"], "hello_world") + self.assertIsInstance(res, SearchResult) + serialized = res.to_dict() + self.assertEqual(serialized["source"], "civitai") + self.assertEqual(serialized["model_id"], 100) + self.assertEqual(serialized["version_id"], 200) + self.assertEqual(serialized["name"], "Test Model") + self.assertEqual(serialized["filename"], "test.safetensors") + self.assertEqual(serialized["extra_field"], "hello_world") def test_build_model_result_custom_url_contract(self): res = build_model_result( @@ -97,10 +104,11 @@ class TestRefactoringUnification(unittest.TestCase): result_mode="custom_url", ) - self.assertEqual(res["details_source"], "civitai") - self.assertEqual(res["version_url"], res["url"]) - self.assertTrue(res["custom_url"]) - self.assertNotIn("confidence", res) + serialized = res.to_dict() + self.assertEqual(serialized["details_source"], "civitai") + self.assertEqual(serialized["version_url"], serialized["url"]) + self.assertTrue(serialized["custom_url"]) + self.assertNotIn("confidence", serialized) def test_build_model_result_compact_custom_url_contract(self): res = build_model_result( @@ -126,7 +134,7 @@ class TestRefactoringUnification(unittest.TestCase): "download_url": "https://civitai.com/api/download/models/200", "match_type": "custom_url", "custom_url": True, - }, res) + }, res.to_dict()) @patch("requests.get") def test_request_page_text_success(self, mock_get): @@ -409,11 +417,12 @@ class TestRefactoringUnification(unittest.TestCase): hashes={"sha256": valid_sha.lower()}, ) self.assertEqual(expected, res) - self.assertEqual(res["source"], "civitai") - self.assertEqual(res["model_id"], 123) - self.assertEqual(res["version_id"], 456) - self.assertEqual(res["sha256"], valid_sha.lower()) - self.assertEqual(res["hashes"]["sha256"], valid_sha.lower()) + serialized = res.to_dict() + self.assertEqual(serialized["source"], "civitai") + self.assertEqual(serialized["model_id"], 123) + self.assertEqual(serialized["version_id"], 456) + self.assertEqual(serialized["sha256"], valid_sha.lower()) + self.assertEqual(serialized["hashes"]["sha256"], valid_sha.lower()) def test_job_progress_tracker_update_from_payload(self): from core.progress import JobProgressTracker diff --git a/tests/test_request_utils.py b/tests/test_request_utils.py index d71ce28..57f24b6 100644 --- a/tests/test_request_utils.py +++ b/tests/test_request_utils.py @@ -3,7 +3,19 @@ from unittest.mock import AsyncMock import pytest -from core.request_utils import extract_request_sha256, read_optional_object_payload +from core.request_utils import ( + coerce_integer_identifier, + extract_request_sha256, + read_bool_field, + read_first_bool_field, + read_first_identifier_field, + read_first_text_field, + read_identifier_field, + read_int_field, + read_optional_object_payload, + read_optional_text_field, + read_text_field, +) def test_extract_request_sha256_preserves_alias_precedence_and_normalization(): @@ -26,6 +38,101 @@ def test_extract_request_sha256_preserves_alias_precedence_and_normalization(): keys=("sha256", "hash", "SHA256"), ) == sha256 assert extract_request_sha256({}, keys=("sha256", "hash")) == "" + with pytest.raises(TypeError, match="Request sha256 must be a string"): + extract_request_sha256( + {"sha256": 123}, + keys=("sha256", "hash"), + ) + + +def test_read_text_field_preserves_optional_null_and_rejects_scalars(): + assert read_text_field({}, "name", default="fallback") == "fallback" + assert read_text_field({"name": None}, "name", default="fallback") == ( + "fallback" + ) + assert read_text_field({"name": " value "}, "name") == "value" + with pytest.raises(TypeError, match="Request name must be a string"): + read_text_field({"name": 123}, "name") + + +def test_read_optional_text_field_preserves_missing_value_and_validates(): + assert read_optional_text_field({}, "name") is None + assert read_optional_text_field({"name": " value "}, "name") == "value" + assert read_optional_text_field({"name": ""}, "name") is None + with pytest.raises(TypeError, match="Request name must be a string"): + read_optional_text_field({"name": 123}, "name") + + +def test_read_first_text_field_preserves_alias_order_and_rejects_scalars(): + assert read_first_text_field( + {"url": "", "custom_url": " https://example.test/model "}, + ("url", "custom_url"), + ) == "https://example.test/model" + with pytest.raises(TypeError, match="Request url must be a string"): + read_first_text_field( + {"url": 123, "custom_url": "https://example.test/model"}, + ("url", "custom_url"), + ) + + +def test_read_identifier_field_accepts_only_string_or_integer_identifiers(): + assert read_identifier_field({"id": " 12 "}, "id") == "12" + assert read_identifier_field({"id": 12}, "id") == 12 + assert read_identifier_field({}, "id") is None + with pytest.raises(TypeError, match="Request id must be an integer or string"): + read_identifier_field({"id": True}, "id") + + +def test_coerce_integer_identifier_rejects_non_numeric_values(): + assert coerce_integer_identifier(" 12 ", "id") == 12 + assert coerce_integer_identifier(34, "id") == 34 + assert coerce_integer_identifier("", "id") is None + with pytest.raises(ValueError, match="Request id must be an integer"): + coerce_integer_identifier("repo/name", "id") + + +def test_read_first_identifier_field_preserves_alias_order_and_rejects_scalars(): + assert read_first_identifier_field( + {"model_id": "", "modelId": " 12 "}, + ("model_id", "modelId"), + ) == "12" + with pytest.raises( + TypeError, + match="Request model_id must be an integer or string", + ): + read_first_identifier_field( + {"model_id": {"value": 12}, "modelId": 12}, + ("model_id", "modelId"), + ) + + +def test_read_bool_field_accepts_supported_forms_and_rejects_scalars(): + assert read_bool_field({"value": "true"}, "value") is True + assert read_bool_field({"value": 0}, "value") is False + assert read_bool_field({}, "value", default=True) is True + with pytest.raises(TypeError, match="Request value must be a boolean"): + read_bool_field({"value": {"enabled": True}}, "value") + + +def test_read_first_bool_field_preserves_alias_order_and_rejects_scalars(): + assert read_first_bool_field( + {"force": "", "force_rescan": "true"}, + ("force", "force_rescan"), + ) is True + with pytest.raises(TypeError, match="Request force must be a boolean"): + read_first_bool_field( + {"force": {"enabled": True}, "force_rescan": True}, + ("force", "force_rescan"), + ) + + +def test_read_int_field_rejects_boolean_and_invalid_numeric_values(): + assert read_int_field({"value": " 12 "}, "value") == 12 + assert read_int_field({}, "value", default=5) == 5 + with pytest.raises(TypeError, match="Request value must be an integer"): + read_int_field({"value": True}, "value") + with pytest.raises(ValueError, match="Request value must be an integer"): + read_int_field({"value": "not-a-number"}, "value") @pytest.mark.asyncio diff --git a/tests/test_resolver_robustness.py b/tests/test_resolver_robustness.py index 54e13b8..fb1c6f7 100644 --- a/tests/test_resolver_robustness.py +++ b/tests/test_resolver_robustness.py @@ -18,6 +18,7 @@ from core.download.api import ( from core.download.api import context as downloader_context from core.download.metadata import read_completed_metadata_sha256 from core.download.state import download_lock, download_progress +from core.contracts import Resolution, ResolvedModel from core.path_utils import get_metadata_sidecar_path, get_path_key from core.workflow_updater import ( convert_to_relative_path, @@ -179,12 +180,14 @@ class ModelResolverRobustnessTests(unittest.TestCase): # Act success = update_model_path( workflow=workflow, - node_id=15, - widget_index=2, - resolved_path="new_lora_name.safetensors", - category="loras", - resolved_model=resolved_model, - mapping=mapping + resolution=Resolution( + node_id=15, + widget_index=2, + resolved_path="new_lora_name.safetensors", + category="loras", + resolved_model=ResolvedModel.from_mapping(resolved_model), + custom_node_metadata=mapping, + ), ) # Assert @@ -229,12 +232,14 @@ class ModelResolverRobustnessTests(unittest.TestCase): # Act success = update_model_path( workflow=workflow, - node_id=99, - widget_index=0, - resolved_path="new_model.safetensors", - category="checkpoints", - subgraph_id="subgraph_a", - is_top_level=False + resolution=Resolution( + node_id=99, + widget_index=0, + resolved_path="new_model.safetensors", + category="checkpoints", + subgraph_id="subgraph_a", + is_top_level=False, + ), ) # Assert @@ -582,12 +587,14 @@ class ModelResolverRobustnessTests(unittest.TestCase): # Act success = update_model_path( workflow=workflow, - node_id=1, - widget_index=0, - resolved_path="model.safetensors", - category="checkpoints", - subgraph_id="nonexistent_subgraph", - is_top_level=False + resolution=Resolution( + node_id=1, + widget_index=0, + resolved_path="model.safetensors", + category="checkpoints", + subgraph_id="nonexistent_subgraph", + is_top_level=False, + ), ) # Assert diff --git a/tests/test_scanner.py b/tests/test_scanner.py index 28ae217..140238c 100644 --- a/tests/test_scanner.py +++ b/tests/test_scanner.py @@ -1,11 +1,34 @@ import unittest from unittest.mock import patch -from core.scanner import find_local_file_path, scan_directory +from core.contracts import ResolvedModel +from core.scanner import find_local_file_path, scan_all_directories, scan_directory class ScannerTests(unittest.TestCase): + @patch("core.scanner.get_model_directories") + @patch("core.scanner.scan_directory") + def test_scan_all_directories_normalizes_scalar_registry_paths( + self, mock_scan_directory, mock_get_model_directories + ): + mock_get_model_directories.return_value = { + "checkpoints": (r"C:\\models\\checkpoints", {".safetensors"}) + } + model = ResolvedModel( + path=r"C:\\models\\checkpoints\\model.safetensors", + filename="model.safetensors", + category="checkpoints", + ) + mock_scan_directory.return_value = [model] + + assert scan_all_directories() == [model] + mock_scan_directory.assert_called_once_with( + r"C:\\models\\checkpoints", + {".safetensors"}, + "checkpoints", + ) + def test_model_path_identity_empty_and_valid(self): from core.path_utils import get_model_path_identity @@ -37,17 +60,18 @@ class ScannerTests(unittest.TestCase): # Verify that we found the two models and ignored readme.txt self.assertEqual(len(models), 2) + self.assertTrue(all(isinstance(model, ResolvedModel) for model in models)) # First model check - m1 = next(m for m in models if m["filename"] == "root_model.safetensors") - self.assertEqual(m1["relative_path"].replace("\\", "/"), "root_model.safetensors") - self.assertEqual(m1["category"], "checkpoints") + m1 = next(m for m in models if m.filename == "root_model.safetensors") + self.assertEqual(m1.relative_path.replace("\\", "/"), "root_model.safetensors") + self.assertEqual(m1.category, "checkpoints") import os - self.assertEqual(m1["base_directory"], os.path.abspath("/models/checkpoints")) + self.assertEqual(m1.base_directory, os.path.abspath("/models/checkpoints")) # Second model check - m2 = next(m for m in models if m["filename"] == "nested_model.ckpt") - self.assertEqual(m2["relative_path"].replace("\\", "/"), "subdir/nested_model.ckpt") + m2 = next(m for m in models if m.filename == "nested_model.ckpt") + self.assertEqual(m2.relative_path.replace("\\", "/"), "subdir/nested_model.ckpt") @patch("os.walk") @patch("os.path.exists") @@ -69,7 +93,7 @@ class ScannerTests(unittest.TestCase): "/models/diffusion_models", {".gguf"}, "model_gguf" ) - self.assertEqual([model["filename"] for model in models], ["quantized.gguf"]) + self.assertEqual([model.filename for model in models], ["quantized.gguf"]) @patch("core.scanner.folder_paths") @patch("core.scanner.get_model_files") @@ -89,7 +113,10 @@ class ScannerTests(unittest.TestCase): # 3. Test found via local scanner fallback mock_folder_paths.get_full_path.return_value = None mock_get_files.return_value = [ - {"filename": "fallback.safetensors", "path": "/models/other/fallback.safetensors"} + ResolvedModel( + filename="fallback.safetensors", + path="/models/other/fallback.safetensors", + ) ] self.assertEqual( find_local_file_path("fallback.safetensors", "checkpoints"), diff --git a/tests/test_scanner_routes.py b/tests/test_scanner_routes.py index b290864..615a6d9 100644 --- a/tests/test_scanner_routes.py +++ b/tests/test_scanner_routes.py @@ -1,11 +1,12 @@ import asyncio import json from types import SimpleNamespace -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest from aiohttp import web +from core.contracts import ResolvedModel from core.path_utils import is_path_within from core.routes.context import RouteContext from core.routes.directories import register_directory_routes @@ -113,6 +114,15 @@ def _request(query=None): return SimpleNamespace(query=query or {}, match_info={}) +def _body_request(payload): + return SimpleNamespace( + can_read_body=True, + json=AsyncMock(return_value=payload), + query={}, + match_info={}, + ) + + @pytest.mark.asyncio async def test_models_route_forwards_force_rescan_and_invalidates_hash_cache(): handlers, values = _build_metadata_routes() @@ -125,6 +135,20 @@ async def test_models_route_forwards_force_rescan_and_invalidates_hash_cache(): values["get_model_files"].assert_called_once_with(force_rescan=True) +@pytest.mark.asyncio +async def test_models_route_rejects_malformed_force_flag(): + handlers, values = _build_metadata_routes() + handler = handlers[("GET", "/model_resolver/models")] + + response = await handler(_request({"force": {"enabled": True}})) + + assert response.status == 400 + assert json.loads(response.text) == { + "error": "Models request force must be a boolean" + } + values["get_model_files"].assert_not_called() + + @pytest.mark.asyncio async def test_path_suggestions_route_scans_and_infers_in_background(): handlers, values = _build_directory_routes() @@ -140,6 +164,36 @@ async def test_path_suggestions_route_scans_and_infers_in_background(): assert json.loads(response.text) == [{"category": "checkpoints"}] values["get_model_files"].assert_called_once_with(True) values["infer_download_path_templates"].assert_called_once_with( - [{"filename": "model.safetensors"}], + [ResolvedModel(filename="model.safetensors")], base_models, ) + + +@pytest.mark.asyncio +async def test_metadata_size_audit_rejects_malformed_numeric_options(): + handlers, values = _build_metadata_routes() + handler = handlers[("POST", "/model_resolver/metadata-size-audit")] + + response = await handler(_body_request({"worker_count": {"value": 4}})) + + assert response.status == 400 + assert json.loads(response.text) == { + "error": "Metadata size audit request worker_count must be an integer" + } + values["audit_metadata_sizes"].assert_not_called() + + +@pytest.mark.asyncio +async def test_metadata_build_start_rejects_malformed_options(): + handlers, values = _build_metadata_routes() + handler = handlers[("POST", "/model_resolver/metadata-build/start")] + + response = await handler( + _body_request({"force_rescan": [], "worker_count": 2}) + ) + + assert response.status == 400 + assert json.loads(response.text) == { + "error": "Metadata build request force_rescan must be a boolean" + } + values["run_in_background_thread"].assert_not_called() diff --git a/tests/test_search_orchestration.py b/tests/test_search_orchestration.py index 3ad1993..27d5b57 100644 --- a/tests/test_search_orchestration.py +++ b/tests/test_search_orchestration.py @@ -63,6 +63,7 @@ civitai_sources = importlib.import_module("comfyui-model-resolver.core.sources.c huggingface_sources = importlib.import_module("comfyui-model-resolver.core.sources.huggingface") civarchive_sources = importlib.import_module("comfyui-model-resolver.core.sources.civarchive") resolver_core = importlib.import_module("comfyui-model-resolver.core.resolver") +typed_contracts = importlib.import_module("comfyui-model-resolver.core.contracts") @contextmanager @@ -88,30 +89,36 @@ class LocalMatchDeduplicationTests(unittest.TestCase): def test_search_local_matches_keeps_highest_confidence_for_same_path(self): shared_path = os.path.join(os.getcwd(), "models", "shared.safetensors") matches = [ - { - "model": { - "path": shared_path, - "relative_path": "models/shared.safetensors", - }, - "filename": "shared.safetensors", - "confidence": 0.42, - }, - { - "model": { - "path": shared_path, - "relative_path": "models/shared.safetensors", - }, - "filename": "shared.safetensors", - "confidence": 0.91, - }, - { - "model": { - "path": os.path.join(os.getcwd(), "models", "other.safetensors"), - "relative_path": "models/other.safetensors", - }, - "filename": "other.safetensors", - "confidence": 0.73, - }, + resolver_core.ModelMatch.from_mapping( + { + "model": { + "path": shared_path, + "relative_path": "models/shared.safetensors", + }, + "filename": "shared.safetensors", + "confidence": 0.42, + } + ), + resolver_core.ModelMatch.from_mapping( + { + "model": { + "path": shared_path, + "relative_path": "models/shared.safetensors", + }, + "filename": "shared.safetensors", + "confidence": 0.91, + } + ), + resolver_core.ModelMatch.from_mapping( + { + "model": { + "path": os.path.join(os.getcwd(), "models", "other.safetensors"), + "relative_path": "models/other.safetensors", + }, + "filename": "other.safetensors", + "confidence": 0.73, + } + ), ] with ( @@ -126,7 +133,7 @@ class LocalMatchDeduplicationTests(unittest.TestCase): result = resolver_core.search_local_matches("shared.safetensors") self.assertEqual( - [match["confidence"] for match in result], + [match.confidence for match in result], [0.91, 0.73], ) @@ -157,8 +164,15 @@ class SearchOrchestrationTests(unittest.IsolatedAsyncioTestCase): async def test_search_sources_orchestrator_success(self): # Setup mocks - self.mock_civitai_search.return_value = {"name": "CivitAI Match", "confidence": 100.0} - self.mock_hf_search.return_value = {"name": "HF Match"} + self.mock_civitai_search.return_value = typed_contracts.SearchResult( + source="civitai", + name="CivitAI Match", + confidence=100.0, + ) + self.mock_hf_search.return_value = typed_contracts.SearchResult( + source="huggingface", + name="HF Match", + ) # Construct a mock request payload = { @@ -270,7 +284,11 @@ class SearchOrchestrationTests(unittest.IsolatedAsyncioTestCase): def side_effect(filename, base_model_context=None, **kwargs): if base_model_context is not None: return None - return {"name": "Fallback Match", "confidence": 90.0} + return typed_contracts.SearchResult( + source="civitai", + name="Fallback Match", + confidence=90.0, + ) self.mock_civitai_search.side_effect = side_effect @@ -300,17 +318,16 @@ class SearchOrchestrationTests(unittest.IsolatedAsyncioTestCase): "resolved_path": model_path, "sha256": file_hash, } - civarchive_result = { - "source": "civarchive", - "filename": os.path.basename(model_path), - "model_name": "Archive Match", - "name": "Archive Match", - "sha256": file_hash, - "hashes": {"SHA256": file_hash}, - "download_url": "https://civarchive.com/api/download/models/123.safetensors", - "url": "https://civarchive.com/models/123", - "version_url": "https://civarchive.com/models/123?modelVersionId=456", - } + civarchive_result = typed_contracts.SearchResult( + source="civarchive", + filename=os.path.basename(model_path), + name="Archive Match", + sha256=file_hash, + hashes={"SHA256": file_hash}, + download_url="https://civarchive.com/api/download/models/123.safetensors", + url="https://civarchive.com/models/123", + version_url="https://civarchive.com/models/123?modelVersionId=456", + ) mock_response = MagicMock() mock_response.status_code = 200 @@ -477,14 +494,14 @@ class SearchOrchestrationTests(unittest.IsolatedAsyncioTestCase): with open(metadata_path, "w", encoding="utf-8") as metadata_file: json.dump({"hashes": {"SHA256": file_hash}}, metadata_file) - self.mock_civitai_search.return_value = { - "source": "civitai", - "name": "Remote Match", - "filename": "remote.safetensors", - "download_url": "https://example.test/remote.safetensors", - "confidence": 52.0, - "hashes": {"SHA256": file_hash}, - } + self.mock_civitai_search.return_value = typed_contracts.SearchResult( + source="civitai", + name="Remote Match", + filename="remote.safetensors", + download_url="https://example.test/remote.safetensors", + confidence=52.0, + hashes={"SHA256": file_hash}, + ) mock_request = MagicMock() mock_request.json = AsyncMock(return_value={ "filename": "remote.safetensors", @@ -569,4 +586,4 @@ class SearchOrchestrationTests(unittest.IsolatedAsyncioTestCase): category="checkpoints", ) - self.assertEqual([match["filename"] for match in matches], ["cached.safetensors"]) + self.assertEqual([match.filename for match in matches], ["cached.safetensors"]) diff --git a/tests/test_search_service_components.py b/tests/test_search_service_components.py index 2e1b9ed..5873341 100644 --- a/tests/test_search_service_components.py +++ b/tests/test_search_service_components.py @@ -4,11 +4,17 @@ from unittest.mock import AsyncMock, MagicMock import pytest +from core.contracts import ModelCatalogEntry, ModelMatch, SearchResult from core.routes.context import RouteContext from core.services.search_cache import SearchResultCache from core.services.search_dependencies import SearchDependencies from core.services.search_orchestrator import SearchOrchestrator -from core.services.search_providers import SearchCancelled, SearchProviderRunner +from core.services.search_providers import ( + SearchCancelled, + SearchProviderRunner, + SearchRequest, +) +from core.type_utils import build_model_result def _request(**overrides): @@ -21,9 +27,128 @@ def _request(**overrides): "force_search": False, } values.update(overrides) + request_data = values.get("data") + if isinstance(request_data, dict): + values.setdefault("model_id", request_data.get("model_id")) + values.setdefault("version_id", request_data.get("version_id")) return SimpleNamespace(**values) +def _typed_search_request(**overrides): + values = { + "filename": "model.safetensors", + "category": "checkpoints", + "base_model_context": "SDXL", + "progress_id": "search-1", + "progress_source": "civitai", + "civitai_candidate_limit": 5, + "civarchive_candidate_limit": 10, + "is_urn": False, + "civitai_key": "", + "civitai_session_token": "", + "hf_token": "", + "brave_search_api_key": "", + "civitai_use_trpc_search": True, + "civitai_use_api_search": True, + "civitai_use_html_fallback": True, + "hf_use_api_search": True, + "hf_use_comfy_org_fallback": True, + "hf_use_brave_fallback": True, + "force_search": False, + "normalized_sources": frozenset({"civitai"}), + "sha256": "", + } + values.update(overrides) + return SearchRequest(**values) + + +def test_search_request_from_mapping_adapts_http_payload_once(): + request = SearchRequest.from_mapping( + { + "filename": " model.safetensors ", + "sources": ["CivitAI"], + "civitai_candidate_limit": 100, + "is_urn": "true", + "model_id": " 12 ", + "version_id": 34, + "sha256": "A" * 64, + }, + ) + + assert request.filename == "model.safetensors" + assert request.normalized_sources == frozenset({"civitai"}) + assert request.civitai_candidate_limit == 20 + assert request.is_urn is True + assert request.model_id == "12" + assert request.version_id == 34 + assert request.sha256 == "a" * 64 + + +def test_search_request_from_mapping_rejects_non_object_payloads(): + with pytest.raises(TypeError, match="Search request must be an object"): + SearchRequest.from_mapping( + [], + ) + + +@pytest.mark.parametrize( + ("payload", "message"), + [ + ({"filename": 123}, "SearchRequest filename must be a string"), + ( + {"filename": "model.safetensors", "category": 123}, + "SearchRequest category must be a string", + ), + ( + {"filename": "model.safetensors", "sources": {"civitai": True}}, + "SearchRequest sources must be a string or array", + ), + ( + {"filename": "model.safetensors", "sources": ["civitai", 1]}, + "SearchRequest sources must contain only strings", + ), + ( + {"filename": "model.safetensors", "sources": None}, + "SearchRequest sources must be a string or array", + ), + ], +) +def test_search_request_from_mapping_rejects_malformed_http_fields(payload, message): + with pytest.raises(TypeError, match=message): + SearchRequest.from_mapping( + payload, + ) + + +@pytest.mark.parametrize( + ("field_name", "value"), + [ + ("category", 123), + ("civitai_key", 123), + ("is_urn", "false"), + ("force_search", {"enabled": True}), + ("civitai_candidate_limit", "5"), + ("normalized_sources", {"civitai"}), + ("model_id", 1.5), + ("version_id", False), + ], +) +def test_search_request_rejects_invalid_normalized_field_types(field_name, value): + with pytest.raises((TypeError, ValueError), match=field_name): + _typed_search_request(**{field_name: value}) + + +@pytest.mark.parametrize( + "field_name", + ["civitai_candidate_limit", "civarchive_candidate_limit"], +) +def test_search_request_from_mapping_rejects_invalid_limit_types(field_name): + with pytest.raises((TypeError, ValueError), match=field_name): + SearchRequest.from_mapping( + {"filename": "model.safetensors", field_name: "invalid"}, + ) + + def _build_search_orchestrator(): tracker = MagicMock() tracker.is_cancelled.return_value = False @@ -33,19 +158,6 @@ def _build_search_orchestrator(): logger=MagicMock(), ) - def to_bool(value, default=False): - if value is None: - return default - if isinstance(value, str): - return value.strip().lower() in {"1", "true", "yes", "on"} - return bool(value) - - def to_int(value, default=0): - try: - return int(value) - except (TypeError, ValueError): - return default - def extract_sha256(metadata): return metadata.get("sha256") or ( metadata.get("hashes") or {} @@ -55,10 +167,7 @@ def _build_search_orchestrator(): "self": extension, "asyncio": asyncio, "CivArchiveSearchError": Exception, - "build_model_result": lambda source, **fields: { - "source": source, - **fields, - }, + "build_model_result": build_model_result, "clear_civarchive_search_cache": MagicMock(), "clear_civitai_search_cache": MagicMock(), "clear_huggingface_search_cache": MagicMock(), @@ -78,8 +187,6 @@ def _build_search_orchestrator(): "search_local_matches_by_hash": MagicMock(return_value=[]), "search_lora_manager_archive_for_file": MagicMock(), "search_model_list": MagicMock(), - "to_bool": to_bool, - "to_int": to_int, "web": SimpleNamespace( json_response=lambda payload, status=200: SimpleNamespace( payload=payload, @@ -200,7 +307,7 @@ def test_search_provider_runner_retries_without_base_model_context(): calls.append((base_model_context, progress_callback)) if len(calls) == 1: return None - return {"name": "Fallback model"} + return SearchResult(source="civitai", name="Fallback model") result = runner.execute_search_with_fallback( _request(), @@ -210,9 +317,9 @@ def test_search_provider_runner_retries_without_base_model_context(): ) assert [call[0] for call in calls] == ["SDXL", None] - assert result["any_model_match"] is True - assert result["base_model_fallback"] is True - assert result["requested_base_model"] == "SDXL" + assert result.extra_value("any_model_match") is True + assert result.extra_value("base_model_fallback") is True + assert result.extra_value("requested_base_model") == "SDXL" def test_search_provider_runner_searches_local_sources(): @@ -220,15 +327,20 @@ def test_search_provider_runner_searches_local_sources(): logger=MagicMock(), search_tracker=MagicMock(), format_log_fields=MagicMock(return_value="file=model.safetensors"), + build_model_result=build_model_result, get_popular_model_url=MagicMock( - return_value={"download_url": "https://example.test/model"} + return_value=ModelCatalogEntry( + filename="model.safetensors", + download_url="https://example.test/model", + ) ), search_model_list=MagicMock( - return_value={ - "filename": "model.safetensors", - "size": 123, - "confidence": 95, - } + return_value=SearchResult( + source="model_list", + filename="model.safetensors", + size=123, + confidence=95, + ) ), log_search_result=MagicMock(), ) @@ -238,9 +350,9 @@ def test_search_provider_runner_searches_local_sources(): result, found = runner.search_local_sources(_request(is_urn=False)) assert found is True - assert result["popular"]["size"] == 123 - assert result["popular"]["download_url"] == "https://example.test/model" - assert result["model_list"]["confidence"] == 95 + assert result["popular"].size == 123 + assert result["popular"].download_url == "https://example.test/model" + assert result["model_list"].confidence == 95 def test_search_provider_runner_scales_provider_progress(): @@ -358,7 +470,7 @@ def test_search_provider_runner_resolves_civitai_urn(): } ), get_civitai_download_url=MagicMock(return_value="https://example.test/download"), - build_model_result=MagicMock(side_effect=lambda source, **fields: {"source": source, **fields}), + build_model_result=MagicMock(side_effect=build_model_result), ) owner.search_tracker.is_cancelled.return_value = False owner.log_search_result = MagicMock() @@ -373,20 +485,20 @@ def test_search_provider_runner_resolves_civitai_urn(): ) assert found is True - assert result["civitai"]["model_id"] == 10 - assert result["civitai"]["version_id"] == 20 - assert result["civitai"]["name"] == "URN Model" - assert result["civitai"]["version_name"] == "v1" - assert result["civitai"]["filename"] == "urn.safetensors" - assert result["civitai"]["type"] == "checkpoints" - assert result["civitai"]["download_url"] == "https://example.test/download" - assert result["civitai"]["url"] == "https://civitai.com/models/10?modelVersionId=20" - assert result["civitai"]["size"] is None - assert result["civitai"]["base_model"] == "SDXL" - assert result["civitai"]["match_type"] == "exact" - assert result["civitai"]["confidence"] == 100.0 - assert result["civitai"]["sha256"] == "a" * 64 - assert result["civitai"]["hashes"] == {} + assert result["civitai"].model_id == 10 + assert result["civitai"].version_id == 20 + assert result["civitai"].name == "URN Model" + assert result["civitai"].version_name == "v1" + assert result["civitai"].filename == "urn.safetensors" + assert result["civitai"].model_type == "checkpoints" + assert result["civitai"].download_url == "https://example.test/download" + assert result["civitai"].url == "https://civitai.com/models/10?modelVersionId=20" + assert result["civitai"].size is None + assert result["civitai"].base_model == "SDXL" + assert result["civitai"].match_type == "exact" + assert result["civitai"].confidence == 100.0 + assert result["civitai"].sha256 == "a" * 64 + assert result["civitai"].hashes == {} def test_search_provider_runner_keeps_first_file_when_expected_name_is_missing(): @@ -413,7 +525,7 @@ def test_search_provider_runner_keeps_first_file_when_expected_name_is_missing() } ), get_civitai_download_url=MagicMock(return_value="https://example.test/download"), - build_model_result=MagicMock(side_effect=lambda source, **fields: {"source": source, **fields}), + build_model_result=MagicMock(side_effect=build_model_result), ) owner.search_tracker.is_cancelled.return_value = False owner.log_search_result = MagicMock() @@ -428,7 +540,7 @@ def test_search_provider_runner_keeps_first_file_when_expected_name_is_missing() ) assert found is True - assert result["civitai"]["sha256"] == "a" * 64 + assert result["civitai"].sha256 == "a" * 64 def test_search_provider_runner_falls_back_for_civitai_urn_without_ids(): @@ -438,18 +550,19 @@ def test_search_provider_runner_falls_back_for_civitai_urn_without_ids(): format_log_fields=MagicMock(return_value="file=model.safetensors"), search_civitai=MagicMock( return_value=[ - { - "model_id": 30, - "version_id": 40, - "name": "Fallback model", - "filename": "model.safetensors", - "type": "Checkpoint", - "download_url": "https://example.test/fallback", - "url": "https://example.test/model/30", - "size": 456, - "base_model": "SDXL", - "tags": ["fallback"], - } + SearchResult( + source="civitai", + model_id=30, + version_id=40, + name="Fallback model", + filename="model.safetensors", + model_type="Checkpoint", + download_url="https://example.test/fallback", + url="https://example.test/model/30", + size=456, + base_model="SDXL", + tags=["fallback"], + ) ] ), build_model_result=MagicMock( @@ -465,8 +578,9 @@ def test_search_provider_runner_falls_back_for_civitai_urn_without_ids(): ) assert found is True - assert result["civitai"]["model_id"] == 30 - assert result["civitai"]["download_url"] == "https://example.test/fallback" + assert isinstance(result["civitai"], SearchResult) + assert result["civitai"].model_id == 30 + assert result["civitai"].download_url == "https://example.test/fallback" owner.search_civitai.assert_called_once_with( "model.safetensors", model_type="checkpoints", @@ -480,7 +594,11 @@ def test_search_provider_runner_resolves_civarchive_urn(): search_tracker=MagicMock(), format_log_fields=MagicMock(return_value="ids=10@20"), resolve_civarchive_model_version=MagicMock( - return_value={"name": "Archived model", "version_id": 20} + return_value=SearchResult( + source="civarchive", + name="Archived model", + version_id=20, + ) ), log_search_result=MagicMock(), ) @@ -495,7 +613,8 @@ def test_search_provider_runner_resolves_civarchive_urn(): ) assert found is True - assert result["civarchive"] == {"name": "Archived model", "version_id": 20} + assert result["civarchive"].name == "Archived model" + assert result["civarchive"].version_id == 20 owner.resolve_civarchive_model_version.assert_called_once_with( 10, 20, @@ -509,7 +628,7 @@ def test_search_provider_runner_searches_lora_manager_archive(): search_tracker=MagicMock(), format_log_fields=MagicMock(return_value="file=model.safetensors"), search_lora_manager_archive_for_file=MagicMock( - return_value={"name": "Archived LoRA"} + return_value=SearchResult(source="lora_manager_archive", name="Archived LoRA") ), log_search_result=MagicMock(), ) @@ -521,7 +640,7 @@ def test_search_provider_runner_searches_lora_manager_archive(): ) assert found is True - assert result == {"lora_manager_archive": {"name": "Archived LoRA"}} + assert result["lora_manager_archive"].name == "Archived LoRA" owner.search_lora_manager_archive_for_file.assert_called_once() @@ -710,7 +829,7 @@ async def test_search_orchestrator_returns_cancelled_response(): @pytest.mark.asyncio -async def test_search_orchestrator_returns_error_response_for_unexpected_failure(): +async def test_search_orchestrator_rejects_invalid_request_payload(): orchestrator, dependencies = _build_search_orchestrator() request = SimpleNamespace( json=AsyncMock(side_effect=ValueError("invalid request")) @@ -718,12 +837,32 @@ async def test_search_orchestrator_returns_error_response_for_unexpected_failure response = await orchestrator.search_sources(request) - assert response.status == 500 - assert response.payload == {"error": "invalid request"} - error_update = dependencies["self"].search_tracker.update.call_args - assert error_update.args[:5] == ("", "", "error", "invalid request", 100) - assert error_update.kwargs == {"status": "error"} - dependencies["self"].logger.exception.assert_called_once() + assert response.status == 400 + assert response.payload == { + "error": "Filename is required for non-URN, or model_id+version_id for URN" + } + dependencies["self"].search_tracker.update.assert_not_called() + dependencies["self"].logger.exception.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "payload", + [ + {"filename": "model.safetensors", "category": 123}, + {"filename": "model.safetensors", "sources": {"civitai": True}}, + ], +) +async def test_search_orchestrator_returns_bad_request_for_contract_errors(payload): + orchestrator, dependencies = _build_search_orchestrator() + request = SimpleNamespace(json=AsyncMock(return_value=payload)) + + response = await orchestrator.search_sources(request) + + assert response.status == 400 + assert response.payload["error"] + dependencies["self"].search_tracker.update.assert_not_called() + dependencies["self"].logger.exception.assert_not_called() @pytest.mark.asyncio @@ -735,8 +874,18 @@ async def test_search_orchestrator_deduplicates_local_hash_matches_and_keeps_err def lookup_by_hash(sha256, **kwargs): if sha256 == first_hash: return [ - {"model": {"path": r"C:\\Models\\Local.safetensors"}}, - {"path": r"c:\\models\\local.safetensors"}, + ModelMatch.from_mapping( + { + "model": {"path": r"C:\\Models\\Local.safetensors"}, + "filename": "Local.safetensors", + } + ), + ModelMatch.from_mapping( + { + "model": {"path": r"c:\\models\\local.safetensors"}, + "filename": "local.safetensors", + } + ), ] raise RuntimeError("local index unavailable") @@ -744,12 +893,20 @@ async def test_search_orchestrator_deduplicates_local_hash_matches_and_keeps_err orchestrator.provider_runner = _StaticSearchRunner( { "civitai": [ - { - "filename": "remote.safetensors", - "hashes": {"SHA256": first_hash}, - }, - {"filename": "without-hash.safetensors"}, - {"filename": "failing.safetensors", "sha256": failing_hash}, + SearchResult( + source="civitai", + filename="remote.safetensors", + hashes={"SHA256": first_hash}, + ), + SearchResult( + source="civitai", + filename="without-hash.safetensors", + ), + SearchResult( + source="civitai", + filename="failing.safetensors", + sha256=failing_hash, + ), ], "source_errors": {"civarchive": "provider unavailable"}, }, diff --git a/tests/test_source_internals.py b/tests/test_source_internals.py index 7813306..fd57f74 100644 --- a/tests/test_source_internals.py +++ b/tests/test_source_internals.py @@ -18,6 +18,8 @@ from unittest.mock import MagicMock, patch import requests +from core.contracts import ResolvedModel, SearchResult + # --------------------------------------------------------------------------- # civarchive internals # --------------------------------------------------------------------------- @@ -35,7 +37,10 @@ from core.sources.civarchive import ( _resolve_file_size_bytes, _transform_file_entry, build_civarchive_custom_result, + get_civarchive_model_details, parse_civarchive_url, + search_civarchive, + search_civarchive_for_file, ) # --------------------------------------------------------------------------- @@ -62,8 +67,8 @@ from core.sources.civitai import ( get_model_info_by_hash, get_model_info_for_file, parse_civitai_url, - resolve_urn, resolve_civitai_version_custom_result, + resolve_urn, search_civitai, search_civitai_for_file, ) @@ -174,6 +179,8 @@ class CivitaiResultBuilderTests(unittest.TestCase): tags=["character"], ) + self.assertIsInstance(result, SearchResult) + result = result.to_dict() self.assertEqual("civitai", result["source"]) self.assertEqual(123, result["model_id"]) self.assertEqual(456, result["version_id"]) @@ -196,7 +203,10 @@ class CivitaiCustomVersionLookupTests(unittest.TestCase): ) as get_details, patch( "core.sources.civitai.build_civitai_custom_result", - return_value={"source": "civitai", "version_id": 456}, + return_value=SearchResult( + source="civitai", + version_id=456, + ), ), ): result = resolve_civitai_version_custom_result( @@ -205,7 +215,9 @@ class CivitaiCustomVersionLookupTests(unittest.TestCase): api_key="secret", ) - self.assertEqual({"source": "civitai", "version_id": 456}, result) + self.assertIsInstance(result, SearchResult) + self.assertEqual("civitai", result.source) + self.assertEqual(456, result.version_id) request_json.assert_called_once_with( "CivitAI custom URL version lookup", "https://civitai.com/api/v1/model-versions/456", @@ -242,6 +254,113 @@ class CivitaiCustomVersionLookupTests(unittest.TestCase): class CivarchiveResultBuilderTests(unittest.TestCase): + def test_name_search_sorts_typed_results(self): + lower_result = SearchResult( + source="civarchive", + name="typed-civarchive-lower", + filename="typed-civarchive-lower.safetensors", + confidence=55.0, + match_type="similar", + ) + higher_result = SearchResult( + source="civarchive", + name="typed-civarchive-higher", + filename="typed-civarchive-higher.safetensors", + confidence=90.0, + match_type="exact", + ) + + with ( + patch( + "core.sources.civarchive._search_page", + return_value=[ + {"id": 1, "url": "/models/1", "name": "lower"}, + {"id": 2, "url": "/models/2", "name": "higher"}, + ], + ), + patch( + "core.sources.civarchive._resolve_search_candidate", + side_effect=[lower_result, higher_result], + ), + ): + results = search_civarchive( + "typed-civarchive-sort-regression", + limit=2, + ) + + self.assertEqual( + [higher_result, lower_result], + results, + ) + + def test_details_fallback_serializes_typed_version_result(self): + resolved = SearchResult( + source="civarchive", + model_id=123, + version_id=456, + name="Typed fallback model", + version_name="v1", + model_type="LORA", + filename="typed-fallback.safetensors", + download_url="https://civarchive.com/download/456", + base_model="SDXL", + ) + + with ( + patch( + "core.sources.civarchive._request_model_payload", + return_value=None, + ), + patch( + "core.sources.civarchive.resolve_civarchive_model_version", + return_value=resolved, + ), + patch( + "core.sources.civarchive._normalize_archive_version", + return_value={ + "id": 456, + "url": "https://civarchive.com/models/123?modelVersionId=456", + }, + ), + ): + result = get_civarchive_model_details(123, version_id=456) + + self.assertEqual("civarchive", result["source"]) + self.assertEqual(456, result["version_id"]) + self.assertEqual("Typed fallback model", result["name"]) + + def test_filename_search_keeps_typed_resolved_result(self): + filename = "typed-civarchive-model.safetensors" + resolved = SearchResult( + source="civarchive", + name="typed-civarchive-model", + version_name="v1", + filename=filename, + download_url="https://civarchive.com/api/download/models/456", + ) + + with ( + patch( + "core.sources.civarchive._search_page", + return_value=[ + { + "name": "typed-civarchive-model", + "url": "/models/123?modelVersionId=456", + } + ], + ), + patch( + "core.sources.civarchive._resolve_search_candidate", + return_value=resolved, + ), + ): + result = search_civarchive_for_file(filename, limit=1) + + self.assertIsInstance(result, SearchResult) + self.assertEqual(filename, result.filename) + self.assertEqual(100.0, result.confidence) + self.assertEqual("exact", result.match_type) + def test_normalized_version_builder_preserves_search_result_contract(self): result = _build_result_from_normalized_version( model_details={ @@ -268,6 +387,8 @@ class CivarchiveResultBuilderTests(unittest.TestCase): match_type="similar", ) + self.assertIsInstance(result, SearchResult) + result = result.to_dict() self.assertEqual("civarchive", result["source"]) self.assertEqual(123, result["model_id"]) self.assertEqual(456, result["version_id"]) @@ -310,6 +431,8 @@ class CivarchiveResultBuilderTests(unittest.TestCase): preferred_filename="archive.safetensors", ) + self.assertIsInstance(result, SearchResult) + result = result.to_dict() self.assertEqual("civarchive", result["source"]) self.assertEqual(123, result["model_id"]) self.assertEqual(456, result["version_id"]) @@ -448,8 +571,9 @@ class CivitaiPrimaryFileSelectionTests(unittest.TestCase): results = search_civitai("example", limit=1) - self.assertEqual("first.safetensors", results[0]["filename"]) - self.assertEqual(1024, results[0]["size"]) + self.assertIsInstance(results[0], SearchResult) + self.assertEqual("first.safetensors", results[0].filename) + self.assertEqual(1024, results[0].size) clear_search_cache() @@ -542,7 +666,7 @@ class CustomUrlResultBuilderTests(unittest.TestCase): "hashes": {"SHA256": "abc123"}, "match_type": "custom_url", "custom_url": True, - }, result) + }, result.to_dict()) def test_civarchive_custom_result_preserves_provider_fields(self): result = build_civarchive_custom_result({ @@ -606,7 +730,7 @@ class CustomUrlResultBuilderTests(unittest.TestCase): "platform": "civitai", "match_type": "custom_url", "custom_url": True, - }, result) + }, result.to_dict()) class CivitaiModelDetailsTests(unittest.TestCase): @@ -1153,13 +1277,13 @@ class ParseCivarchiveUrlTests(unittest.TestCase): def test_model_url_parsed(self): result = parse_civarchive_url("https://civarchive.com/models/123") self.assertIsNotNone(result) - self.assertEqual(result.get("model_id"), 123) + self.assertEqual(result.model_id, 123) def test_model_version_url_parsed(self): result = parse_civarchive_url("https://civarchive.com/models/123?modelVersionId=456") self.assertIsNotNone(result) - self.assertEqual(result.get("model_id"), 123) - self.assertEqual(result.get("version_id"), 456) + self.assertEqual(result.model_id, 123) + self.assertEqual(result.version_id, 456) def test_non_civarchive_url_returns_none(self): self.assertIsNone(parse_civarchive_url("https://huggingface.co/user/repo")) @@ -1339,24 +1463,24 @@ class ParseCivitaiUrlTests(unittest.TestCase): def test_standard_model_url(self): result = parse_civitai_url("https://civitai.com/models/123456") self.assertIsNotNone(result) - self.assertEqual(result.get("model_id"), 123456) + self.assertEqual(result.model_id, 123456) def test_model_url_with_version(self): result = parse_civitai_url("https://civitai.com/models/123456?modelVersionId=789") self.assertIsNotNone(result) - self.assertEqual(result.get("model_id"), 123456) - self.assertEqual(result.get("version_id"), 789) + self.assertEqual(result.model_id, 123456) + self.assertEqual(result.version_id, 789) def test_civitai_red_model_url_with_version(self): result = parse_civitai_url("https://civitai.red/models/123456?modelVersionId=789") self.assertIsNotNone(result) - self.assertEqual(result.get("model_id"), 123456) - self.assertEqual(result.get("version_id"), 789) + self.assertEqual(result.model_id, 123456) + self.assertEqual(result.version_id, 789) def test_civitai_red_download_url(self): result = parse_civitai_url("https://civitai.red/api/download/models/789") self.assertIsNotNone(result) - self.assertEqual(result.get("version_id"), 789) + self.assertEqual(result.version_id, 789) def test_non_civitai_url_returns_none(self): self.assertIsNone(parse_civitai_url("https://civarchive.com/models/123")) @@ -1539,13 +1663,13 @@ class CivitaiTrpcSearchTests(unittest.TestCase): def test_public_api_candidates_used_when_trpc_disabled_and_html_empty(self): clear_search_cache() - expected = { - "source": "civitai", - "model_id": 2676616, - "version_id": 3005748, - "filename": "sickOllie_v1.safetensors", - "confidence": 100.0, - } + expected = SearchResult( + source="civitai", + model_id=2676616, + version_id=3005748, + filename="sickOllie_v1.safetensors", + confidence=100.0, + ) with patch( "core.sources.civitai._search_civitai_red_candidates", return_value=[], @@ -1564,7 +1688,12 @@ class CivitaiTrpcSearchTests(unittest.TestCase): use_html_fallback=True, ) - self.assertEqual(result, expected) + self.assertIsInstance(result, SearchResult) + self.assertEqual(expected.source, result.source) + self.assertEqual(expected.model_id, result.model_id) + self.assertEqual(expected.version_id, result.version_id) + self.assertEqual(expected.filename, result.filename) + self.assertEqual(expected.confidence, result.confidence) html_search.assert_called_once() api_search.assert_called_once() clear_search_cache() @@ -1906,12 +2035,12 @@ class ConvertToRelativePathTests(unittest.TestCase): class GetBaseDirectoryForModelTests(unittest.TestCase): def test_returns_base_directory_if_provided(self): - model = {"base_directory": "/some/path/models"} + model = ResolvedModel(base_directory="/some/path/models") result = get_base_directory_for_model(model, "checkpoints") self.assertEqual(result, "/some/path/models") def test_returns_none_without_path_or_base_directory(self): - result = get_base_directory_for_model({}, "checkpoints") + result = get_base_directory_for_model(ResolvedModel(), "checkpoints") self.assertIsNone(result) diff --git a/tests/test_title_matching_hashing.py b/tests/test_title_matching_hashing.py index b48f876..cfd80ac 100644 --- a/tests/test_title_matching_hashing.py +++ b/tests/test_title_matching_hashing.py @@ -1,15 +1,16 @@ -import sys import os +import sys import unittest -from unittest.mock import MagicMock, patch, AsyncMock +from unittest.mock import MagicMock parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) if parent_dir not in sys.path: sys.path.insert(0, parent_dir) -from core.sources.civitai import _find_model_title_match_in_model -from core.sources.civarchive import _find_model_title_match_in_model_details from core.path_utils import calculate_file_sha256 +from core.sources.civarchive import _find_model_title_match_in_model_details +from core.sources.civitai import _find_model_title_match_in_model + class TestTitleMatchingAndHashing(unittest.IsolatedAsyncioTestCase): @@ -44,11 +45,11 @@ class TestTitleMatchingAndHashing(unittest.IsolatedAsyncioTestCase): base_model_context="SDXL 1.0" ) self.assertIsNotNone(result) - self.assertEqual(result["model_id"], 1) - self.assertEqual(result["version_id"], 123) - self.assertEqual(result["filename"], "sdxl_base.safetensors") - self.assertEqual(result["base_model"], "SDXL 1.0") - self.assertIn("anime", result["tags"]) + self.assertEqual(result.model_id, 1) + self.assertEqual(result.version_id, 123) + self.assertEqual(result.filename, "sdxl_base.safetensors") + self.assertEqual(result.base_model, "SDXL 1.0") + self.assertIn("anime", result.tags) def test_civitai_find_model_title_match_rejected_confidence(self): model_data = { @@ -92,10 +93,10 @@ class TestTitleMatchingAndHashing(unittest.IsolatedAsyncioTestCase): base_model_context="SD 1.5" ) self.assertIsNotNone(result) - self.assertEqual(result["model_id"], 10) - self.assertEqual(result["version_id"], 999) - self.assertEqual(result["filename"], "v1-5-pruned.safetensors") - self.assertEqual(result["base_model"], "SD 1.5") + self.assertEqual(result.model_id, 10) + self.assertEqual(result.version_id, 999) + self.assertEqual(result.filename, "v1-5-pruned.safetensors") + self.assertEqual(result.base_model, "SD 1.5") async def test_calculate_file_sha256_with_mock_progress(self): import tempfile diff --git a/tests/test_unified_helpers.py b/tests/test_unified_helpers.py index 3846ec5..d8619e0 100644 --- a/tests/test_unified_helpers.py +++ b/tests/test_unified_helpers.py @@ -57,6 +57,7 @@ class UnifiedHelpersTests(unittest.TestCase): self.assertEqual("", extract_file_sha256(None)) self.assertEqual("", extract_file_sha256({"hashes": []})) self.assertEqual("fallback", extract_file_sha256({"sha256": "", "hash": "fallback"})) + self.assertEqual("", extract_file_sha256({"sha256": 123})) def test_extract_sha256_from_metadata_uses_first_valid_source(self): preferred = "a" * 64 @@ -1038,6 +1039,7 @@ class UnifiedHelpersTests(unittest.TestCase): self.assertEqual(normalize_sha256("SHA256:0X" + h.upper()), h) self.assertEqual(normalize_sha256("invalid_hash"), "") self.assertEqual(normalize_sha256(None), "") + self.assertEqual(normalize_sha256(123), "") def test_unique_ordered_strings(self): self.assertEqual(unique_ordered_strings(["a", "b", "a", "c", "b"]), ["a", "b", "c"]) @@ -1405,11 +1407,14 @@ class UnifiedHelpersTests(unittest.TestCase): # Valid url download format res = parse_provider_model_url("https://civitai.com/api/download/models/12345", allowed) - self.assertEqual(res, {"version_id": 12345}) + self.assertEqual(res.to_dict(), {"version_id": 12345}) # Valid model page format res = parse_provider_model_url("https://civitai.red/models/9876?modelVersionId=54321", allowed) - self.assertEqual(res, {"model_id": 9876, "version_id": 54321}) + self.assertEqual( + res.to_dict(), + {"model_id": 9876, "version_id": 54321}, + ) # Non-matching hostname res = parse_provider_model_url("https://example.com/api/download/models/12345", allowed) @@ -1417,7 +1422,9 @@ class UnifiedHelpersTests(unittest.TestCase): # Mirror sha256 path pattern res = parse_provider_model_url("https://civarchive.com/sha256/a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2", ["civarchive.com"]) - self.assertEqual(res, {"sha256": "a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2"}) + self.assertIsNotNone(res) + self.assertEqual(len(res.sha256), 64) + self.assertEqual(res.to_dict(), {"sha256": res.sha256}) def test_resolve_model_category(self): from core.type_utils import resolve_model_category diff --git a/tests/test_workflow_analyzer.py b/tests/test_workflow_analyzer.py index d87c6dd..7966995 100644 --- a/tests/test_workflow_analyzer.py +++ b/tests/test_workflow_analyzer.py @@ -5,6 +5,13 @@ from types import SimpleNamespace from unittest.mock import ANY, patch from core import resolver as resolver_core +from core.contracts import ( + ModelMatch, + ModelReference, + ResolvedModel, + WorkflowAnalysisResult, + WorkflowModelInventory, +) from core.scanner import invalidate_model_files_cache, scan_directory from core.workflow import analysis, dynamic_widgets, references from core.workflow.analysis import analyze_workflow_models, identify_missing_models @@ -54,10 +61,10 @@ class WorkflowAnalyzerCaseSensitivityTests(unittest.TestCase): ) missing = identify_missing_models(refs, available_models) - self.assertFalse(refs[0]["exists"]) + self.assertFalse(refs[0].exists) self.assertEqual(1, len(missing)) self.assertEqual( - r"Qwen\qwen_3_4b.safetensors", missing[0]["original_path"] + r"Qwen\qwen_3_4b.safetensors", missing[0].reference.original_path ) def test_exact_folder_case_is_not_reported_missing(self): @@ -84,7 +91,7 @@ class WorkflowAnalyzerCaseSensitivityTests(unittest.TestCase): ) missing = identify_missing_models(refs, available_models) - self.assertTrue(refs[0]["exists"]) + self.assertTrue(refs[0].exists) self.assertEqual([], missing) @@ -348,10 +355,10 @@ class WorkflowAnalyzerCategoryHintTests(unittest.TestCase): refs = analyze_workflow_models(workflow, available_models=[]) self.assertEqual(1, len(refs)) - self.assertEqual(3, refs[0]["widget_index"]) - self.assertEqual("start_ckpt_name", refs[0]["widget_name"]) - self.assertEqual(r"LTXV\model.safetensors", refs[0]["original_path"]) - self.assertEqual("checkpoints", refs[0]["category"]) + self.assertEqual(3, refs[0].widget_index) + self.assertEqual("start_ckpt_name", refs[0].extra_value("widget_name")) + self.assertEqual(r"LTXV\model.safetensors", refs[0].original_path) + self.assertEqual("checkpoints", refs[0].category) def test_implicit_seed_control_keeps_following_model_aligned(self): dynamic_hints = { @@ -415,14 +422,14 @@ class WorkflowAnalyzerCategoryHintTests(unittest.TestCase): refs = analyze_workflow_models(workflow, available_models=[]) self.assertEqual(1, len(refs)) - self.assertEqual(2, refs[0]["widget_index"]) - self.assertEqual("clip_name", refs[0]["widget_name"]) + self.assertEqual(2, refs[0].widget_index) + self.assertEqual("clip_name", refs[0].extra_value("widget_name")) self.assertEqual( "clip-vit-large-patch14.safetensors", - refs[0]["original_path"], + refs[0].original_path, ) - self.assertEqual("clip_vision", refs[0]["category"]) - self.assertFalse(refs[0]["exists"]) + self.assertEqual("clip_vision", refs[0].category) + self.assertFalse(refs[0].exists) def test_non_model_static_choices_in_hybrid_model_dropdowns_are_skipped(self): cases = [ @@ -546,7 +553,7 @@ class WorkflowAnalyzerCategoryHintTests(unittest.TestCase): refs = analyze_workflow_models(workflow, available_models=[]) self.assertEqual(1, len(refs)) - self.assertEqual("diffusion_models", refs[0]["category"]) + self.assertEqual("diffusion_models", refs[0].category) def test_show_anything_cached_model_text_is_not_a_model_reference(self): with tempfile.TemporaryDirectory() as tmpdir: @@ -605,7 +612,7 @@ class WorkflowAnalyzerCategoryHintTests(unittest.TestCase): refs = analyze_workflow_models(workflow, available_models=[]) self.assertEqual(1, len(refs)) - self.assertEqual("custom_model.safetensors", refs[0]["original_path"]) + self.assertEqual("custom_model.safetensors", refs[0].original_path) def test_set_node_model_slot_name_does_not_turn_constant_into_model(self): workflow = { @@ -740,7 +747,7 @@ class WorkflowAnalyzerCategoryHintTests(unittest.TestCase): refs = analyze_workflow_models(workflow, available_models=[]) self.assertEqual(1, len(refs)) - self.assertEqual("upscale_models", refs[0]["category"]) + self.assertEqual("upscale_models", refs[0].category) def test_upscale_output_hint_does_not_mark_every_text_widget_as_model(self): workflow = { @@ -758,8 +765,8 @@ class WorkflowAnalyzerCategoryHintTests(unittest.TestCase): refs = analyze_workflow_models(workflow, available_models=[]) self.assertEqual(1, len(refs)) - self.assertEqual("4x_NMKD-Siax_200k.pth", refs[0]["original_path"]) - self.assertEqual("upscale_models", refs[0]["category"]) + self.assertEqual("4x_NMKD-Siax_200k.pth", refs[0].original_path) + self.assertEqual("upscale_models", refs[0].category) def test_impact_sam_loader_uses_sams_category(self): workflow = { @@ -778,8 +785,8 @@ class WorkflowAnalyzerCategoryHintTests(unittest.TestCase): refs = analyze_workflow_models(workflow, available_models=[]) self.assertEqual(1, len(refs)) - self.assertEqual("sam_vit_b_01ec64.pth", refs[0]["original_path"]) - self.assertEqual("sams", refs[0]["category"]) + self.assertEqual("sam_vit_b_01ec64.pth", refs[0].original_path) + self.assertEqual("sams", refs[0].category) def test_impact_sam_loader_esam_option_is_not_missing_model(self): workflow = { @@ -819,8 +826,8 @@ class WorkflowAnalyzerCategoryHintTests(unittest.TestCase): refs = analyze_workflow_models(workflow, available_models=[]) self.assertEqual(1, len(refs)) - self.assertEqual("segm/person_yolov8m-seg.pt", refs[0]["original_path"]) - self.assertEqual("ultralytics", refs[0]["category"]) + self.assertEqual("segm/person_yolov8m-seg.pt", refs[0].original_path) + self.assertEqual("ultralytics", refs[0].category) def test_core_extra_loader_categories_match_comfyui_folder_paths(self): cases = [ @@ -855,7 +862,7 @@ class WorkflowAnalyzerCategoryHintTests(unittest.TestCase): ) self.assertEqual(1, len(refs)) - self.assertEqual(expected_category, refs[0]["category"]) + self.assertEqual(expected_category, refs[0].category) def test_multi_clip_loader_indexes_are_text_encoders(self): refs = analyze_workflow_models( @@ -878,7 +885,7 @@ class WorkflowAnalyzerCategoryHintTests(unittest.TestCase): ) self.assertEqual(4, len(refs)) - self.assertEqual({"text_encoders"}, {ref["category"] for ref in refs}) + self.assertEqual({"text_encoders"}, {ref.category for ref in refs}) def test_gguf_unet_loader_keeps_diffusion_badge_and_raw_folder_hint(self): for node_type in ( @@ -910,9 +917,15 @@ class WorkflowAnalyzerCategoryHintTests(unittest.TestCase): ) self.assertEqual(1, len(refs)) - self.assertEqual("diffusion_models", refs[0]["category"]) - self.assertEqual(["diffusion_models"], refs[0]["category_hints"]) - self.assertEqual(["model_gguf"], refs[0]["folder_key_hints"]) + self.assertEqual("diffusion_models", refs[0].category) + self.assertEqual( + ("diffusion_models",), + refs[0].extra_value("category_hints"), + ) + self.assertEqual( + ("model_gguf",), + refs[0].extra_value("folder_key_hints"), + ) def test_gguf_loader_resolves_scanner_category_alias(self): with tempfile.TemporaryDirectory() as tmpdir: @@ -953,9 +966,9 @@ class WorkflowAnalyzerCategoryHintTests(unittest.TestCase): ) self.assertEqual(1, len(refs)) - self.assertTrue(refs[0]["exists"]) - self.assertEqual(model_path, refs[0]["full_path"]) - self.assertEqual("diffusion_models", refs[0]["category"]) + self.assertTrue(refs[0].exists) + self.assertEqual(model_path, refs[0].extra_value("full_path")) + self.assertEqual("diffusion_models", refs[0].category) class ScannerFolderModelTests(unittest.TestCase): @@ -971,8 +984,8 @@ class ScannerFolderModelTests(unittest.TestCase): models = scan_directory(tmpdir, {"folder"}, "diffusers") self.assertEqual(1, len(models)) - self.assertEqual("wan_diffusers", models[0]["relative_path"]) - self.assertEqual(model_dir, models[0]["path"]) + self.assertEqual("wan_diffusers", models[0].relative_path) + self.assertEqual(model_dir, models[0].path) class WorkflowCategoryHintsTests(unittest.TestCase): @@ -1007,8 +1020,8 @@ class WorkflowCategoryHintsTests(unittest.TestCase): refs = analyze_workflow_models(workflow, available_models=[]) self.assertEqual(1, len(refs)) - self.assertEqual("model.safetensors", refs[0]["original_path"]) - self.assertEqual("checkpoints", refs[0]["category"]) + self.assertEqual("model.safetensors", refs[0].original_path) + self.assertEqual("checkpoints", refs[0].category) def test_node_type_to_category_hints_is_populated(self): from core.workflow.widgets import NODE_TYPE_TO_CATEGORY_HINTS @@ -1186,10 +1199,12 @@ class WorkflowCustomNodeAdapterTests(unittest.TestCase): refs = analyze_workflow_models(workflow, available_models=[]) self.assertEqual(1, len(refs)) - self.assertEqual("lora-manager", refs[0]["custom_node_adapter"]) - self.assertEqual("missing_style", refs[0]["original_path"]) - self.assertEqual(0.65, refs[0]["strength"]) - self.assertTrue(refs[0]["active"]) + self.assertEqual( + "lora-manager", refs[0].extra_value("custom_node_adapter") + ) + self.assertEqual("missing_style", refs[0].original_path) + self.assertEqual(0.65, refs[0].extra_value("strength")) + self.assertTrue(refs[0].extra_value("active")) class WorkflowMissingReferenceGroupingTests(unittest.TestCase): @@ -1219,17 +1234,17 @@ class WorkflowMissingReferenceGroupingTests(unittest.TestCase): self.assertEqual(2, len(refs)) self.assertEqual( ["rgthree-power-lora-loader", "rgthree-power-lora-loader"], - [ref["custom_node_adapter"] for ref in refs], + [ref.extra_value("custom_node_adapter") for ref in refs], ) self.assertEqual(1, len(missing)) - self.assertEqual(2, missing[0]["reference_count"]) + self.assertEqual(2, missing[0].reference_count) self.assertEqual( [401, 402], - [ref["node_id"] for ref in missing[0]["all_node_refs"]], + [ref.node_id for ref in missing[0].all_node_refs], ) self.assertEqual( ["lora", "lora"], - [ref["nested_key"] for ref in missing[0]["all_node_refs"]], + [ref.extra["nested_key"] for ref in missing[0].all_node_refs], ) def test_same_filename_different_categories_are_not_merged(self): @@ -1255,7 +1270,7 @@ class WorkflowMissingReferenceGroupingTests(unittest.TestCase): self.assertEqual(2, len(missing)) self.assertEqual( {"checkpoints", "loras"}, - {item["category"] for item in missing}, + {item.reference.category for item in missing}, ) @@ -1263,22 +1278,26 @@ class WorkflowResolverMatchingTests(unittest.TestCase): def test_missing_model_matches_keep_highest_confidence_for_same_path(self): shared_path = os.path.join(os.getcwd(), "models", "shared.safetensors") matches = [ - { - "model": { - "path": shared_path, - "relative_path": "models/shared.safetensors", - }, - "filename": "shared.safetensors", - "confidence": 0.35, - }, - { - "model": { - "path": shared_path, - "relative_path": "models/shared.safetensors", - }, - "filename": "shared.safetensors", - "confidence": 0.88, - }, + ModelMatch.from_mapping( + { + "model": { + "path": shared_path, + "relative_path": "models/shared.safetensors", + }, + "filename": "shared.safetensors", + "confidence": 0.35, + } + ), + ModelMatch.from_mapping( + { + "model": { + "path": shared_path, + "relative_path": "models/shared.safetensors", + }, + "filename": "shared.safetensors", + "confidence": 0.88, + } + ), ] workflow_ref = { "node_id": 1, @@ -1293,10 +1312,9 @@ class WorkflowResolverMatchingTests(unittest.TestCase): patch.object( resolver_core, "get_workflow_model_inventory", - return_value={ - "available_models": [], - "model_refs": [workflow_ref], - }, + return_value=WorkflowModelInventory( + model_refs=(ModelReference.from_mapping(workflow_ref),) + ), ), patch.object(resolver_core, "find_matches", return_value=matches), patch.object(resolver_core, "_get_active_downloads_by_path", return_value={}), @@ -1305,8 +1323,13 @@ class WorkflowResolverMatchingTests(unittest.TestCase): _workflow_with_model("shared.safetensors") ) + self.assertIsInstance(result, WorkflowAnalysisResult) + serialized = result.to_dict() self.assertEqual( - [match["confidence"] for match in result["missing_models"][0]["matches"]], + [ + match["confidence"] + for match in serialized["missing_models"][0]["matches"] + ], [0.88], ) @@ -1335,10 +1358,13 @@ class WorkflowResolverMatchingTests(unittest.TestCase): patch.object( resolver_core, "get_workflow_model_inventory", - return_value={ - "available_models": available_models, - "model_refs": [resolved_ref], - }, + return_value=WorkflowModelInventory( + available_models=tuple( + ResolvedModel.from_mapping(model) + for model in available_models + ), + model_refs=(ModelReference.from_mapping(resolved_ref),), + ), ), patch.object(resolver_core, "find_matches") as find_matches, ): @@ -1347,9 +1373,10 @@ class WorkflowResolverMatchingTests(unittest.TestCase): progress_callback=progress.append, ) + self.assertIsInstance(result, WorkflowAnalysisResult) find_matches.assert_not_called() - self.assertEqual(1, result["total_resolved"]) - self.assertEqual([], result["resolved_models"][0]["matches"]) + self.assertEqual(1, result.total_resolved) + self.assertEqual([], result.to_dict()["resolved_models"][0]["matches"]) self.assertFalse( any( "Analyzing resolved model" in str(update.get("message") or "") @@ -1384,8 +1411,15 @@ class WorkflowModelInventoryCacheTests(unittest.TestCase): first = get_workflow_model_inventory(workflow) second = get_workflow_model_inventory({"nodes": workflow["nodes"]}) - self.assertIs(first["available_models"], available_models) - self.assertIs(second["model_refs"], model_refs) + self.assertEqual( + ("shared.safetensors",), + tuple(model.filename for model in first.available_models), + ) + self.assertIsInstance(second.model_refs[0], ModelReference) + self.assertEqual( + [ref.to_dict() for ref in first.model_refs], + [ref.to_dict() for ref in second.model_refs], + ) get_models.assert_called_once_with(force_rescan=False) analyze_models.assert_called_once_with( workflow, @@ -1503,7 +1537,7 @@ class WorkflowModelInventoryCacheTests(unittest.TestCase): self.assertEqual(3, get_node_info.call_count) self.assertEqual( {"second.safetensors", "unchanged.safetensors"}, - {ref["original_path"] for ref in result["model_refs"]}, + {ref.original_path for ref in result.model_refs}, ) def test_changed_loader_keeps_previous_result_position(self): @@ -1555,7 +1589,7 @@ class WorkflowModelInventoryCacheTests(unittest.TestCase): "middle-updated.safetensors", "last.safetensors", ], - [ref["original_path"] for ref in result["model_refs"]], + [ref.original_path for ref in result.model_refs], ) def test_new_loader_does_not_reanalyze_existing_loader(self): diff --git a/tests/test_workflow_components.py b/tests/test_workflow_components.py index 0b7eb2a..a676e7f 100644 --- a/tests/test_workflow_components.py +++ b/tests/test_workflow_components.py @@ -2,6 +2,7 @@ from types import SimpleNamespace from unittest.mock import patch from core import resolver as resolver_core +from core.contracts import ModelReference, Resolution from core.workflow import analysis, references, subgraphs, traversal from core.workflow_updater import update_workflow_nodes @@ -274,7 +275,9 @@ def test_apply_instance_promoted_context_resolves_existing_model(tmp_path): "path": str(model_path), } ] - reference = {"original_path": "promoted.safetensors", "exists": False} + reference = ModelReference.from_mapping( + {"original_path": "promoted.safetensors", "exists": False} + ) context = { "proxy_widget_name": "ckpt_name", "node_id": 10, @@ -285,16 +288,16 @@ def test_apply_instance_promoted_context_resolves_existing_model(tmp_path): "category": "checkpoints", } - subgraphs._apply_instance_promoted_widget_context( + reference = subgraphs._apply_instance_promoted_widget_context( reference, context, available_models, ) - assert reference["promoted_inner_node_id"] == 10 - assert reference["category"] == "checkpoints" - assert reference["full_path"] == str(model_path) - assert reference["exists"] is True + assert reference.extra_value("promoted_inner_node_id") == 10 + assert reference.category == "checkpoints" + assert reference.extra_value("full_path") == str(model_path) + assert reference.exists is True def test_apply_promoted_locator_supports_name_fallback_and_value_selection(): @@ -307,20 +310,22 @@ def test_apply_promoted_locator_supports_name_fallback_and_value_selection(): ] }, } - reference = { - "subgraph_id": "subgraph-1", - "node_id": 10, - "widget_index": 99, - "widget_name": "ckpt_name", - "original_path": "second", - } + reference = ModelReference.from_mapping( + { + "subgraph_id": "subgraph-1", + "node_id": 10, + "widget_index": 99, + "widget_name": "ckpt_name", + "original_path": "second", + } + ) - subgraphs._apply_promoted_widget_locator(reference, contexts) + reference = subgraphs._apply_promoted_widget_locator(reference, contexts) - assert reference["locate_node_id"] == 2 - assert reference["locate_subgraph_name"] == "Two" - assert reference["locate_is_top_level"] is True - assert reference["locate_via_promoted_widget"] is True + assert reference.extra_value("locate_node_id") == 2 + assert reference.extra_value("locate_subgraph_name") == "Two" + assert reference.extra_value("locate_is_top_level") is True + assert reference.extra_value("locate_via_promoted_widget") is True def test_subgraph_input_reference_targets_instance_for_update_and_location(): @@ -377,14 +382,14 @@ def test_subgraph_input_reference_targets_instance_for_update_and_location(): refs = analysis.analyze_workflow_models(workflow, available_models=[]) assert len(refs) == 1 - assert refs[0]["node_id"] == 1 - assert refs[0]["widget_index"] == 0 - assert refs[0]["subgraph_id"] == "subgraph-1" - assert refs[0]["is_top_level"] is True - assert refs[0]["subgraph_path"] is None - assert refs[0]["locate_node_id"] == 1 - assert refs[0]["locate_subgraph_id"] == "" - assert refs[0]["promoted_inner_node_id"] == 10 + assert refs[0].node_id == 1 + assert refs[0].widget_index == 0 + assert refs[0].subgraph_id == "subgraph-1" + assert refs[0].is_top_level is True + assert refs[0].extra_value("subgraph_path") is None + assert refs[0].extra_value("locate_node_id") == 1 + assert refs[0].extra_value("locate_subgraph_id") == "" + assert refs[0].extra_value("promoted_inner_node_id") == 10 no_proxy_workflow = { **workflow, @@ -398,23 +403,21 @@ def test_subgraph_input_reference_targets_instance_for_update_and_location(): no_proxy_refs = analysis.analyze_workflow_models( no_proxy_workflow, available_models=[] ) - assert no_proxy_refs[0]["node_id"] == 1 - assert no_proxy_refs[0]["promoted_inner_node_id"] == 10 + assert no_proxy_refs[0].node_id == 1 + assert no_proxy_refs[0].extra_value("promoted_inner_node_id") == 10 - update_workflow_nodes( - workflow, - [ - { - "node_id": refs[0]["node_id"], - "widget_index": refs[0]["widget_index"], - "resolved_path": "replacement.safetensors", - "category": "checkpoints", - "subgraph_id": refs[0]["subgraph_id"], - "is_top_level": refs[0]["is_top_level"], - "promoted_widget_name": refs[0]["promoted_widget_name"], - } - ], + resolution = Resolution.from_mapping( + { + "node_id": refs[0].node_id, + "widget_index": refs[0].widget_index, + "resolved_path": "replacement.safetensors", + "category": "checkpoints", + "subgraph_id": refs[0].subgraph_id, + "is_top_level": refs[0].is_top_level, + "promoted_widget_name": refs[0].extra_value("promoted_widget_name"), + } ) + update_workflow_nodes(workflow, [resolution]) assert workflow["nodes"][0]["widgets_values"] == ["replacement.safetensors"] assert workflow["definitions"]["subgraphs"][0]["nodes"][0]["widgets_values"] == [ @@ -479,9 +482,9 @@ def test_reused_subgraph_instance_refreshes_promoted_model_context(): ) assert len(refs) == 1 - assert refs[0]["node_id"] == 1 - assert refs[0]["original_path"] == "new.safetensors" - assert refs[0]["promoted_inner_node_id"] == 10 + assert refs[0].node_id == 1 + assert refs[0].original_path == "new.safetensors" + assert refs[0].extra_value("promoted_inner_node_id") == 10 def test_reference_matching_handles_placeholders_paths_and_scanner_records(tmp_path): diff --git a/tests/test_workflow_service.py b/tests/test_workflow_service.py index e71411a..1e36cc3 100644 --- a/tests/test_workflow_service.py +++ b/tests/test_workflow_service.py @@ -7,6 +7,14 @@ from unittest.mock import AsyncMock, MagicMock import pytest from aiohttp import web +from core.contracts import ( + MissingModel, + ModelCatalogEntry, + ModelMatch, + Resolution, + SearchResult, + WorkflowAnalysisResult, +) from core.path_utils import get_filename_from_path from core.routes.context import RouteContext from core.routes.workflow_analysis import register_workflow_analysis_routes @@ -40,6 +48,21 @@ def _request(payload): return SimpleNamespace(json=AsyncMock(return_value=payload)) +def _analysis_result(missing_models=(), **counts): + typed_missing = tuple( + item + if isinstance(item, MissingModel) + else MissingModel.from_mapping(item) + for item in missing_models + ) + return WorkflowAnalysisResult( + missing_models=typed_missing, + total_resolved=counts.get("total_resolved", 0), + total_missing=counts.get("total_missing", len(typed_missing)), + total_models_analyzed=counts.get("total_models_analyzed", len(typed_missing)), + ) + + def _build_routes(overrides=None): routes = _Routes() extension = SimpleNamespace( @@ -49,7 +72,7 @@ def _build_routes(overrides=None): ) values = { "analyze_and_find_matches": MagicMock( - return_value={"missing_models": [], "total_missing": 0} + return_value=WorkflowAnalysisResult() ), "apply_resolution": MagicMock(side_effect=lambda workflow, _resolutions: workflow), "asyncio": asyncio, @@ -94,6 +117,20 @@ async def test_analyze_route_validates_workflow_payload(payload, message): assert json.loads(response.text) == {"error": message} +@pytest.mark.asyncio +async def test_analyze_route_rejects_non_text_analysis_id(): + handlers, values = _build_routes() + response = await handlers[("POST", "/model_resolver/analyze")]( + _request({"workflow": {"nodes": []}, "analysis_id": 123}) + ) + + assert response.status == 400 + assert "Workflow analysis request analysis_id must be a string" in ( + json.loads(response.text)["error"] + ) + values["analyze_and_find_matches"].assert_not_called() + + @pytest.mark.asyncio async def test_analyze_route_filters_references_and_uses_local_download_sources(): missing_models = [ @@ -114,7 +151,16 @@ async def test_analyze_route_filters_references_and_uses_local_download_sources( }, { "original_path": "already-local.safetensors", - "matches": [{"confidence": 100}], + "matches": [ + { + "model": { + "path": "models/already-local.safetensors", + "filename": "already-local.safetensors", + "category": "checkpoints", + }, + "confidence": 100, + } + ], }, { "original_path": "installed-custom-node.safetensors", @@ -124,33 +170,39 @@ async def test_analyze_route_filters_references_and_uses_local_download_sources( def analyze(_workflow, _threshold, _limit, progress_callback, **_kwargs): progress_callback({"stage": "matching", "current": 1, "total": 5}) - return {"missing_models": missing_models} + return _analysis_result(missing_models) def popular_url(filename): if filename == "popular.safetensors": - return { - "url": "https://example.com/popular.safetensors", - "type": "checkpoint", - "directory": "checkpoints", - "size": 100, - } + return ModelCatalogEntry( + filename=filename, + url="https://example.com/popular.safetensors", + model_type="checkpoint", + directory="checkpoints", + size=100, + ) return None def model_list(filename, exact_only=False): assert exact_only is True if filename == "popular.safetensors": - return {"size": 200} + return SearchResult( + source="model_list", + filename=filename, + size=200, + ) if filename == "listed.safetensors": - return { - "url": "https://example.com/listed.safetensors", - "filename": "listed.safetensors", - "name": "Listed model", - "type": "checkpoint", - "directory": "checkpoints", - "size": 300, - "match_type": "exact", - "confidence": 100, - } + return SearchResult( + source="model_list", + url="https://example.com/listed.safetensors", + filename="listed.safetensors", + name="Listed model", + model_type="checkpoint", + size=300, + match_type="exact", + confidence=100, + extra={"directory": "checkpoints"}, + ) return None handlers, values = _build_routes( @@ -160,7 +212,7 @@ async def test_analyze_route_filters_references_and_uses_local_download_sources( "get_popular_model_url": MagicMock(side_effect=popular_url), "search_model_list": MagicMock(side_effect=model_list), "should_skip_existing_custom_node_reference": MagicMock( - side_effect=lambda item: item["original_path"].startswith( + side_effect=lambda item: item.original_path.startswith( "installed-custom" ) ), @@ -244,7 +296,14 @@ async def test_workflow_resolution_route_returns_updated_workflow(): _request( { "workflow": workflow, - "resolutions": [{"node_id": 1, "widget_index": 0}], + "resolutions": [ + { + "node_id": 1, + "widget_index": 0, + "resolved_path": "resolved.safetensors", + "category": "checkpoints", + } + ], } ) ) @@ -252,13 +311,68 @@ async def test_workflow_resolution_route_returns_updated_workflow(): assert response.status == 200 assert json.loads(response.text) == {"workflow": updated, "success": True} values["apply_resolution"].assert_called_once_with( - workflow, [{"node_id": 1, "widget_index": 0}] + workflow, + [ + Resolution.from_mapping( + { + "node_id": 1, + "widget_index": 0, + "resolved_path": "resolved.safetensors", + "category": "checkpoints", + } + ) + ], ) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "resolution", + [ + { + "node_id": 1, + "widget_index": "0", + "resolved_path": "resolved.safetensors", + }, + { + "node_id": 1, + "widget_index": -1, + "resolved_path": "resolved.safetensors", + }, + { + "node_id": 1, + "widget_index": 0, + "resolved_path": 123, + }, + ], +) +async def test_workflow_resolution_route_rejects_invalid_resolution_types(resolution): + handlers, values = _build_routes() + response = await handlers[("POST", "/model_resolver/resolve")]( + _request( + { + "workflow": {"nodes": []}, + "resolutions": [resolution], + } + ) + ) + + assert response.status == 400 + assert json.loads(response.text)["error"].startswith("Invalid resolution:") + values["apply_resolution"].assert_not_called() + + @pytest.mark.asyncio async def test_local_matches_route_rescans_and_passes_search_options(): - matches = [{"path": os.path.join("models", "model.safetensors")}] + matches = [ + ModelMatch.from_mapping( + { + "model": { + "path": os.path.join("models", "model.safetensors") + } + } + ) + ] handlers, values = _build_routes( {"search_local_matches": MagicMock(return_value=matches)} ) @@ -273,7 +387,9 @@ async def test_local_matches_route_rescans_and_passes_search_options(): ) assert response.status == 200 - assert json.loads(response.text) == {"matches": matches} + assert json.loads(response.text) == { + "matches": [match.to_dict() for match in matches] + } values["invalidate_local_hash_match_cache"].assert_called_once_with() values["search_local_matches"].assert_called_once_with( "model.safetensors", @@ -284,6 +400,20 @@ async def test_local_matches_route_rescans_and_passes_search_options(): ) +@pytest.mark.asyncio +async def test_local_matches_route_rejects_non_text_fields(): + handlers, values = _build_routes() + response = await handlers[("POST", "/model_resolver/local-matches")]( + _request({"filename": 123, "category": "checkpoints"}) + ) + + assert response.status == 400 + assert json.loads(response.text) == { + "error": "Local matches request filename must be a string" + } + values["search_local_matches"].assert_not_called() + + @pytest.mark.asyncio async def test_workflow_progress_route_delegates_analysis_tracker(): handlers, values = _build_routes()