From bb01732b89862976c61c9f58d0f3fa35a2204169 Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Wed, 5 Aug 2026 06:13:18 +0200 Subject: [PATCH] Unify provider search result builders --- core/sources/civarchive.py | 11 ++-- core/sources/civitai.py | 9 ++- core/sources/common.py | 83 --------------------------- core/sources/huggingface.py | 4 +- core/type_utils.py | 33 ++++++++++- tests/test_refactoring_unification.py | 19 ++++-- 6 files changed, 62 insertions(+), 97 deletions(-) diff --git a/core/sources/civarchive.py b/core/sources/civarchive.py index c194271..b98bf7c 100644 --- a/core/sources/civarchive.py +++ b/core/sources/civarchive.py @@ -34,6 +34,7 @@ from ..path_utils import get_filename_from_path from ..progress import get_progress_reporter from ..type_utils import ( DEFAULT_BROWSER_USER_AGENT, + build_search_result, clear_remote_size_cache, extract_file_size, extract_trained_words, @@ -50,7 +51,6 @@ from ..type_utils import ( ) from .common import ( build_custom_url_result, - build_unified_search_result, collect_download_urls, is_remote_link_marked_dead, ) @@ -1208,7 +1208,7 @@ def _build_result_from_normalized_version( if size is None: size = _resolve_file_size_bytes(file_info, download_urls) - return build_unified_search_result( + return build_search_result( source="civarchive", model_id=model_id, version_id=version_id, @@ -1231,6 +1231,7 @@ def _build_result_from_normalized_version( platform=model_details.get("platform"), is_deleted=False, match_type=match_type, + normalize_hashes=True, ) @@ -1658,7 +1659,7 @@ def _build_result_from_payload( else "" ) - return build_unified_search_result( + return build_search_result( source="civarchive", model_id=model_id, version_id=version_id, @@ -1689,6 +1690,7 @@ def _build_result_from_payload( sha256=sha256, hash=sha256, hashes=hashes, + normalize_hashes=True, ) @@ -1815,7 +1817,7 @@ def _build_result_from_search_candidate( open_url = urljoin(CIVARCHIVE_BASE_URL, str(candidate.get("url") or "")) sha256 = (parsed or {}).get("sha256") - return build_unified_search_result( + return build_search_result( source="civarchive", model_id=candidate.get("model_id") or candidate.get("modelId"), version_id=candidate.get("version_id") or candidate.get("modelVersionId"), @@ -1839,6 +1841,7 @@ def _build_result_from_search_candidate( sha256=sha256, hash=sha256, hashes={"SHA256": sha256} if sha256 else {}, + normalize_hashes=True, ) diff --git a/core/sources/civitai.py b/core/sources/civitai.py index 2844149..e844ca9 100644 --- a/core/sources/civitai.py +++ b/core/sources/civitai.py @@ -44,6 +44,7 @@ from ..type_utils import ( DEFAULT_BROWSER_USER_AGENT, as_dict, as_list, + build_search_result, check_credential_http, extract_trained_words, first_non_empty, @@ -143,7 +144,7 @@ def check_civitai_api_key(api_key: Optional[str]) -> Dict[str, Any]: ) -from .common import build_custom_url_result, build_unified_search_result +from .common import build_custom_url_result def _build_civitai_result_from_version( @@ -161,7 +162,7 @@ def _build_civitai_result_from_version( sha256 = file_info.get("sha256") or hashes.get("SHA256") or hashes.get("sha256") size = file_info.get("sizeKB", 0) * 1024 if file_info.get("sizeKB") else file_info.get("size") download_url = file_info.get("downloadUrl") or get_civitai_download_url(version_id) - return build_unified_search_result( + return build_search_result( source="civitai", model_id=model_id, version_id=version_id, @@ -176,6 +177,7 @@ def _build_civitai_result_from_version( match_type=match_type, sha256=sha256, hashes=hashes, + normalize_hashes=True, ) @@ -934,7 +936,7 @@ def search_civitai_for_file( api_key=api_key, ) if hash_result: - result = build_unified_search_result( + result = build_search_result( "civitai", model_id=hash_result.get("model_id"), version_id=hash_result.get("version_id"), @@ -953,6 +955,7 @@ def search_civitai_for_file( confidence=100.0, sha256=requested_sha256, hashes=hash_result.get("hashes") or {"SHA256": requested_sha256}, + normalize_hashes=True, description=hash_result.get("description") or "", model_description=hash_result.get("model_description") or "", ) diff --git a/core/sources/common.py b/core/sources/common.py index 681f516..2e87be7 100644 --- a/core/sources/common.py +++ b/core/sources/common.py @@ -4,8 +4,6 @@ Common utilities and result builder functions for external model metadata source from typing import Any, Callable, Dict, List, Optional, Tuple -from ..type_utils import build_search_result, normalize_sha256 - def is_remote_link_marked_dead(item: Any) -> bool: """Return whether remote metadata marks a link as unavailable.""" @@ -95,87 +93,6 @@ def collect_download_urls( return urls -def normalize_hashes_dict(hashes: Optional[Dict[str, Any]]) -> Dict[str, str]: - """ - Normalize a hashes dictionary (e.g. from CivitAI API or CivArchive) so that - hash algorithm keys are standardized and SHA256 values are uppercase. - """ - if not isinstance(hashes, dict): - return {} - - normalized: Dict[str, str] = {} - for k, v in hashes.items(): - if not k or not v: - continue - key_str = str(k).strip() - val_str = str(v).strip() - if key_str.lower() in ("sha256", "sha-256"): - normalized["sha256"] = normalize_sha256(val_str) - elif key_str.lower() in ("autov2", "auto_v2"): - normalized["autoV2"] = val_str - elif key_str.lower() in ("autov1", "auto_v1"): - normalized["autoV1"] = val_str - elif key_str.lower() == "blake3": - normalized["blake3"] = val_str - else: - normalized[key_str] = val_str - - return normalized - - -def build_unified_search_result( - source: str, - *, - model_id: Any, - version_id: Any, - name: str = "", - version_name: str = "", - type: str = "", - filename: str = "", - url: str = "", - download_url: Optional[str] = None, - size: Optional[int] = None, - base_model: Optional[str] = None, - tags: Optional[List[str]] = None, - match_type: str = "similar", - confidence: float = 0.0, - sha256: Optional[str] = None, - hashes: Optional[Dict[str, Any]] = None, - trained_words: Optional[List[str]] = None, - images: Optional[List[Dict[str, Any]]] = None, - **extra: Any, -) -> Dict[str, Any]: - """ - Unified result builder for all model metadata sources. - Normalizes hash dictionaries and formats a standard model search result. - """ - norm_hashes = normalize_hashes_dict(hashes) - if not sha256 and "sha256" in norm_hashes: - sha256 = norm_hashes["sha256"] - - return build_search_result( - 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, - match_type=match_type, - confidence=confidence, - sha256=sha256, - hashes=norm_hashes, - trained_words=trained_words, - images=images, - **extra, - ) - - def build_custom_url_result( source: str, *, diff --git a/core/sources/huggingface.py b/core/sources/huggingface.py index 913d2cf..f6acf48 100644 --- a/core/sources/huggingface.py +++ b/core/sources/huggingface.py @@ -31,7 +31,6 @@ from ..type_utils import ( normalize_sha256, prepare_remote_size_probe_url, ) -from .common import build_unified_search_result HF_API_URL = "https://huggingface.co/api" HF_AUTHOR_FALLBACKS = ["Comfy-Org", "Kijai"] @@ -717,7 +716,7 @@ def _build_huggingface_result( size = extract_file_size(file_info) if not size: size = _fetch_remote_file_size_bytes(download_url, headers=headers) - return build_unified_search_result( + return build_search_result( "huggingface", model_id=repo_id, version_id=None, @@ -727,6 +726,7 @@ def _build_huggingface_result( size=size, match_type=match_type, sha256=sha256, + normalize_hashes=True, repo_id=repo_id, path=file_path, branch="main", diff --git a/core/type_utils.py b/core/type_utils.py index 5017b64..ce4a772 100644 --- a/core/type_utils.py +++ b/core/type_utils.py @@ -874,6 +874,31 @@ def normalize_sha256(value: Any) -> str: return text.lower() if SHA256_PATTERN.match(text) else "" +def normalize_hashes_dict(hashes: Optional[Dict[str, Any]]) -> Dict[str, str]: + """Normalize common model hash keys and values into a stable mapping.""" + if not isinstance(hashes, dict): + return {} + + normalized: Dict[str, str] = {} + for key, value in hashes.items(): + if not key or not value: + continue + key_str = str(key).strip() + value_str = str(value).strip() + if key_str.lower() in ("sha256", "sha-256"): + normalized["sha256"] = normalize_sha256(value_str) + elif key_str.lower() in ("autov2", "auto_v2"): + normalized["autoV2"] = value_str + elif key_str.lower() in ("autov1", "auto_v1"): + normalized["autoV1"] = value_str + elif key_str.lower() == "blake3": + normalized["blake3"] = value_str + else: + normalized[key_str] = value_str + + return normalized + + def unique_ordered_strings(values: List[Any]) -> List[str]: """Return unique, non-empty strings while preserving their original order.""" seen = set() @@ -1245,12 +1270,18 @@ def build_search_result( match_type: str = "similar", confidence: float = 0.0, sha256: Optional[str] = None, - hashes: Optional[Dict[str, str]] = None, + hashes: Optional[Dict[str, Any]] = None, + normalize_hashes: bool = False, trained_words: Optional[List[str]] = None, images: Optional[List[Dict[str, Any]]] = None, **extra: Any, ) -> Dict[str, Any]: """Helper to unify search result dictionary creation across different sources.""" + if normalize_hashes: + hashes = normalize_hashes_dict(hashes) + if not sha256 and "sha256" in hashes: + sha256 = hashes["sha256"] + result = { "source": source, "model_id": model_id, diff --git a/tests/test_refactoring_unification.py b/tests/test_refactoring_unification.py index 5a24014..8245fb9 100644 --- a/tests/test_refactoring_unification.py +++ b/tests/test_refactoring_unification.py @@ -262,22 +262,33 @@ class TestRefactoringUnification(unittest.TestCase): self.assertEqual(as_list(["a", "", None, "b"]), ["a", "b"]) self.assertEqual(list(_listify_tags(["a", "", None, "b"])), ["a", "b"]) - def test_build_unified_search_result(self): - from core.sources.common import build_unified_search_result, normalize_hashes_dict - + def test_build_search_result_normalizes_hashes(self): + from core.type_utils import normalize_hashes_dict + valid_sha = "a" * 64 norm_hashes = normalize_hashes_dict({"sha256": valid_sha, "autoV2": "1234567890"}) self.assertEqual(norm_hashes["sha256"], valid_sha.lower()) self.assertEqual(norm_hashes["autoV2"], "1234567890") - res = build_unified_search_result( + res = build_search_result( "civitai", model_id=123, version_id=456, name="Test Model", hashes={"sha256": valid_sha}, download_url="https://civitai.com/api/download/123", + normalize_hashes=True, ) + expected = build_search_result( + "civitai", + model_id=123, + version_id=456, + name="Test Model", + download_url="https://civitai.com/api/download/123", + sha256=valid_sha.lower(), + 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)