From 348f7ed8e01c3c370ffa5164d1708140dce2c141 Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Thu, 6 Aug 2026 16:20:31 +0200 Subject: [PATCH] refactor: unify model result builders --- core/routes/registry.py | 2 +- core/services/custom_url_service.py | 49 +++++++++--------- core/services/model_utils.py | 6 +++ core/services/search_dependencies.py | 6 +-- core/services/search_orchestrator.py | 2 +- core/services/search_providers.py | 4 +- core/sources/civarchive.py | 15 +++--- core/sources/civitai.py | 13 +++-- core/sources/common.py | 51 +------------------ core/sources/huggingface.py | 6 +-- core/sources/lora_manager_archive.py | 4 +- core/type_utils.py | 53 +++++++++++++++++-- tests/test_refactoring_targets.py | 1 + tests/test_refactoring_unification.py | 68 ++++++++++++++++++++----- tests/test_search_service_components.py | 8 +-- 15 files changed, 172 insertions(+), 116 deletions(-) diff --git a/core/routes/registry.py b/core/routes/registry.py index 111ad94..4bf289d 100644 --- a/core/routes/registry.py +++ b/core/routes/registry.py @@ -105,7 +105,7 @@ def register_routes(self): save_settings as save_resolver_settings, ) from ..type_utils import ( - build_search_result, + build_model_result, extract_sha256_from_metadata, fetch_remote_file_size_cached, first_non_empty, diff --git a/core/services/custom_url_service.py b/core/services/custom_url_service.py index 57ddc59..e776539 100644 --- a/core/services/custom_url_service.py +++ b/core/services/custom_url_service.py @@ -14,6 +14,7 @@ class CustomUrlService(ModelServiceDependencies): """Resolve a user-provided provider URL into a normal search result.""" UnsafeUrlError = self.UnsafeUrlError asyncio = self.asyncio + build_model_result = self.build_model_result build_civarchive_custom_result = self.build_civarchive_custom_result build_civitai_custom_result = self.build_civitai_custom_result build_huggingface_custom_result = self.build_huggingface_custom_result @@ -64,17 +65,18 @@ class CustomUrlService(ModelServiceDependencies): "civitai": "CivitAI", "huggingface": "HuggingFace", }.get(source, "Custom URL") - return { - "source": source, - "details_source": source, - "name": source_label, - "filename": filename, - "url": url, - "version_url": url, - "download_url": url, - "match_type": "custom_url", - "custom_url": True, - } + return build_model_result( + source, + name=source_label, + filename=filename, + url=url, + version_url=url, + download_url=url, + match_type="custom_url", + details_source=source, + custom_url=True, + result_mode="compact_custom_url", + ) def _collect_custom_url_local_hash_matches(result, category): source_key = str(result.get("source") or "custom").strip().lower() @@ -171,21 +173,22 @@ class CustomUrlService(ModelServiceDependencies): civitai_key or None, ) if not result and version_id: - result = { - "source": "civitai", - "details_source": "civitai", - "version_id": version_id, - "name": expected_filename or f"CivitAI version {version_id}", - "filename": expected_filename or f"civitai-{version_id}", - "url": normalized_url, - "version_url": normalized_url, - "download_url": get_civitai_download_url( + result = build_model_result( + "civitai", + version_id=version_id, + name=expected_filename or f"CivitAI version {version_id}", + filename=expected_filename or f"civitai-{version_id}", + url=normalized_url, + version_url=normalized_url, + download_url=get_civitai_download_url( version_id, civitai_key or None, ), - "match_type": "custom_url", - "custom_url": True, - } + match_type="custom_url", + details_source="civitai", + custom_url=True, + result_mode="compact_custom_url", + ) elif civarchive_parsed: source = "civarchive" if civarchive_parsed.get("sha256"): diff --git a/core/services/model_utils.py b/core/services/model_utils.py index 91b5621..7ea743d 100644 --- a/core/services/model_utils.py +++ b/core/services/model_utils.py @@ -46,6 +46,10 @@ class NormalizeSha256Protocol(Protocol): def __call__(self, value: Any) -> str: ... +class ModelResultBuilderProtocol(Protocol): + def __call__(self, source: str, **fields: Any) -> Dict[str, Any]: ... + + class ReadJsonProtocol(Protocol): def __call__(self, file_path: str, default: Any = None) -> Any: ... @@ -276,6 +280,7 @@ class CustomUrlDependencies: logger: Any UnsafeUrlError: type[Exception] asyncio: Any + build_model_result: ModelResultBuilderProtocol build_civarchive_custom_result: Optional[BuildCivarchiveCustomResultProtocol] build_civitai_custom_result: Optional[BuildCivitaiCustomResultProtocol] build_huggingface_custom_result: Optional[BuildHuggingFaceCustomResultProtocol] @@ -305,6 +310,7 @@ class CustomUrlDependencies: logger=extension.logger, UnsafeUrlError=context.require("UnsafeUrlError"), asyncio=context.require("asyncio"), + build_model_result=context.require("build_model_result"), build_civarchive_custom_result=context.get( "build_civarchive_custom_result" ), diff --git a/core/services/search_dependencies.py b/core/services/search_dependencies.py index 74b2fff..3ada55e 100644 --- a/core/services/search_dependencies.py +++ b/core/services/search_dependencies.py @@ -30,7 +30,7 @@ class IntParserProtocol(Protocol): def __call__(self, value: Any, default: int = 0) -> int: ... -class SearchResultBuilderProtocol(Protocol): +class ModelResultBuilderProtocol(Protocol): def __call__(self, source: str, **fields: Any) -> dict[str, Any]: ... @@ -100,7 +100,7 @@ class SearchDependencies: logger: LoggerProtocol civarchive_search_error: type[BaseException] asyncio: Any - build_search_result: SearchResultBuilderProtocol + build_model_result: ModelResultBuilderProtocol clear_civarchive_search_cache: ClearCacheProtocol clear_civitai_search_cache: ClearCacheProtocol clear_huggingface_search_cache: ClearCacheProtocol @@ -144,7 +144,7 @@ class SearchDependencies: logger=cls._extension_value(extension, "logger"), asyncio=context.require("asyncio"), civarchive_search_error=context.require("CivArchiveSearchError"), - build_search_result=context.require("build_search_result"), + build_model_result=context.require("build_model_result"), clear_civarchive_search_cache=context.require( "clear_civarchive_search_cache" ), diff --git a/core/services/search_orchestrator.py b/core/services/search_orchestrator.py index a1d8884..a9384c2 100644 --- a/core/services/search_orchestrator.py +++ b/core/services/search_orchestrator.py @@ -22,7 +22,7 @@ class SearchOrchestrator: self.logger = dependencies.logger self.CivArchiveSearchError = dependencies.civarchive_search_error self.asyncio = dependencies.asyncio - self.build_search_result = dependencies.build_search_result + self.build_model_result = dependencies.build_model_result self.clear_civarchive_search_cache = ( dependencies.clear_civarchive_search_cache ) diff --git a/core/services/search_providers.py b/core/services/search_providers.py index a701efd..19b8366 100644 --- a/core/services/search_providers.py +++ b/core/services/search_providers.py @@ -381,7 +381,7 @@ class SearchProviderRunner: version_name=model_info.get("version_name"), confidence=100.0, ) - source_results["civitai"] = self.owner.build_search_result( + source_results["civitai"] = self.owner.build_model_result( "civitai", **result_payload ) self.owner.log_search_result( @@ -418,7 +418,7 @@ class SearchProviderRunner: ) if civitai_results: first_result = civitai_results[0] - source_results["civitai"] = self.owner.build_search_result( + source_results["civitai"] = self.owner.build_model_result( "civitai", model_id=first_result.get("model_id"), version_id=first_result.get("version_id"), diff --git a/core/sources/civarchive.py b/core/sources/civarchive.py index 3689d29..dba792d 100644 --- a/core/sources/civarchive.py +++ b/core/sources/civarchive.py @@ -34,7 +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, + build_model_result, clear_remote_size_cache, extract_file_size, extract_trained_words, @@ -50,7 +50,6 @@ from ..type_utils import ( to_int, ) from .common import ( - build_custom_url_result, collect_download_urls, is_remote_link_marked_dead, ) @@ -1204,7 +1203,7 @@ def _build_result_from_normalized_version( if size is None: size = _resolve_file_size_bytes(file_info, download_urls) - return build_search_result( + return build_model_result( source="civarchive", model_id=model_id, version_id=version_id, @@ -1655,7 +1654,7 @@ def _build_result_from_payload( else "" ) - return build_search_result( + return build_model_result( source="civarchive", model_id=model_id, version_id=version_id, @@ -1813,7 +1812,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_search_result( + return build_model_result( source="civarchive", model_id=candidate.get("model_id") or candidate.get("modelId"), version_id=candidate.get("version_id") or candidate.get("modelVersionId"), @@ -2383,7 +2382,7 @@ def build_civarchive_custom_result( else details.get("url") ) ) - return build_custom_url_result( + return build_model_result( source="civarchive", model_id=model_id, version_id=version_id, @@ -2403,6 +2402,10 @@ def build_civarchive_custom_result( or "", sha256=sha256, hashes=hashes, + details_source="civarchive", + version_url=version_url or details.get("url"), + custom_url=True, + result_mode="custom_url", platform_url=details.get("platform_url") or selected_version.get("platform_url"), download_urls=[url for url in download_urls if url], diff --git a/core/sources/civitai.py b/core/sources/civitai.py index 0e71a34..5efb537 100644 --- a/core/sources/civitai.py +++ b/core/sources/civitai.py @@ -44,7 +44,7 @@ from ..type_utils import ( DEFAULT_BROWSER_USER_AGENT, as_dict, as_list, - build_search_result, + build_model_result, check_credential_http, extract_file_size, extract_trained_words, @@ -145,7 +145,6 @@ def check_civitai_api_key(api_key: Optional[str]) -> Dict[str, Any]: ) -from .common import build_custom_url_result def build_civitai_result_payload( @@ -202,7 +201,7 @@ def _build_civitai_result_from_version( version_id = version.get("id") size = extract_file_size(file_info) download_url = file_info.get("downloadUrl") or get_civitai_download_url(version_id) - return build_search_result( + return build_model_result( source="civitai", **build_civitai_result_payload( model_id=model_id, @@ -971,7 +970,7 @@ def search_civitai_for_file( api_key=api_key, ) if hash_result: - result = build_search_result( + result = build_model_result( "civitai", model_id=hash_result.get("model_id"), version_id=hash_result.get("version_id"), @@ -2385,7 +2384,7 @@ def build_civitai_custom_result( else details.get("url") ) ) - return build_custom_url_result( + return build_model_result( source="civitai", model_id=model_id, version_id=version_id, @@ -2405,6 +2404,10 @@ def build_civitai_custom_result( or "", sha256=sha256, hashes=hashes, + details_source="civitai", + version_url=version_url or details.get("url"), + custom_url=True, + result_mode="custom_url", model_description=details.get("description") or "", ) diff --git a/core/sources/common.py b/core/sources/common.py index 2e87be7..4a7161a 100644 --- a/core/sources/common.py +++ b/core/sources/common.py @@ -1,5 +1,5 @@ """ -Common utilities and result builder functions for external model metadata sources. +Common utilities for external model metadata sources. """ from typing import Any, Callable, Dict, List, Optional, Tuple @@ -91,52 +91,3 @@ def collect_download_urls( urls.append(url) return urls - - -def build_custom_url_result( - source: str, - *, - model_id: Any, - version_id: Any, - name: str, - version_name: str, - type: str, - filename: str, - url: str, - download_url: str, - size: Optional[int], - base_model: Optional[str], - tags: List[Any], - trained_words: List[Any], - images: List[Dict[str, Any]], - description: str, - sha256: Optional[str], - hashes: Dict[str, Any], - **extra: Any, -) -> Dict[str, Any]: - """Build the shared, unnormalized result contract for custom URL resolvers.""" - result = { - "source": source, - "details_source": source, - "model_id": model_id, - "version_id": version_id, - "name": name, - "version_name": version_name, - "type": type, - "filename": filename, - "url": url, - "version_url": url, - "download_url": download_url, - "size": size, - "base_model": base_model, - "tags": tags, - "trained_words": trained_words, - "images": images, - "description": description, - "sha256": sha256, - "hashes": hashes, - "match_type": "custom_url", - "custom_url": True, - } - result.update(extra) - return result diff --git a/core/sources/huggingface.py b/core/sources/huggingface.py index f792510..2f9ab80 100644 --- a/core/sources/huggingface.py +++ b/core/sources/huggingface.py @@ -22,7 +22,7 @@ from ..matcher import build_filename_search_queries from ..network_utils import execute_provider_json_request, host_matches_domain, request_source_response from ..path_utils import METADATA_DIR, get_filename_from_path, read_json_safe, write_json_atomic from ..type_utils import ( - build_search_result, + build_model_result, check_credential_http, clear_remote_size_cache, extract_file_size, @@ -724,7 +724,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_search_result( + return build_model_result( "huggingface", model_id=repo_id, version_id=None, @@ -1567,7 +1567,7 @@ def build_huggingface_custom_result( page_url = ( f"https://huggingface.co/{repo_id}/blob/{branch}/{quote_url_path(file_path)}" ) - return build_search_result( + return build_model_result( "huggingface", model_id=repo_id, version_id=None, diff --git a/core/sources/lora_manager_archive.py b/core/sources/lora_manager_archive.py index 0a12b76..e91c8f9 100644 --- a/core/sources/lora_manager_archive.py +++ b/core/sources/lora_manager_archive.py @@ -29,7 +29,7 @@ from ..matcher import ( from ..path_utils import get_filename_from_path from ..progress import get_progress_reporter from ..type_utils import ( - build_search_result, + build_model_result, extract_file_size, get_generic_filename_tokens, normalize_lora_manager_type, @@ -489,7 +489,7 @@ def _build_result_from_row( ) match_type = "exact" if confidence == 100.0 else "similar" - return build_search_result( + return build_model_result( source="lora_manager_archive", model_id=model_id, version_id=version_id, diff --git a/core/type_utils.py b/core/type_utils.py index 7bebdf8..e8444e0 100644 --- a/core/type_utils.py +++ b/core/type_utils.py @@ -1272,11 +1272,11 @@ def normalize_model_file_info( } -def build_search_result( +def build_model_result( source: str, *, - model_id: Any, - version_id: Any, + model_id: Any = None, + version_id: Any = None, name: str = "", version_name: str = "", type: str = "", @@ -1293,14 +1293,40 @@ def build_search_result( normalize_hashes: bool = False, trained_words: Optional[List[str]] = None, images: Optional[List[Dict[str, Any]]] = None, + details_source: Optional[str] = None, + version_url: Optional[str] = None, + custom_url: bool = False, + result_mode: str = "search", **extra: Any, ) -> Dict[str, Any]: - """Helper to unify search result dictionary creation across different sources.""" + """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 = { + "source": source, + "details_source": details_source or source, + "name": name, + "filename": filename, + "url": url, + "version_url": version_url or url, + "download_url": download_url, + "match_type": match_type, + "custom_url": custom_url, + } + 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, @@ -1321,6 +1347,25 @@ def build_search_result( "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 diff --git a/tests/test_refactoring_targets.py b/tests/test_refactoring_targets.py index 87425a8..c1a4df4 100644 --- a/tests/test_refactoring_targets.py +++ b/tests/test_refactoring_targets.py @@ -1226,6 +1226,7 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase): "logger", "UnsafeUrlError", "asyncio", + "build_model_result", "build_civarchive_custom_result", "build_civitai_custom_result", "build_huggingface_custom_result", diff --git a/tests/test_refactoring_unification.py b/tests/test_refactoring_unification.py index 8245fb9..8e77b29 100644 --- a/tests/test_refactoring_unification.py +++ b/tests/test_refactoring_unification.py @@ -1,6 +1,6 @@ import ast -import sys import os +import sys import unittest from pathlib import Path from unittest.mock import MagicMock, patch @@ -10,11 +10,10 @@ if parent_dir not in sys.path: sys.path.insert(0, parent_dir) from core.matcher import MODEL_TITLE_MATCH_THRESHOLD -from core.sources.lora_manager_archive import _normalize_model_type from core.type_utils import ( + build_model_result, normalize_lora_manager_type, normalize_model_file_info, - build_search_result, ) @@ -50,8 +49,8 @@ class TestRefactoringUnification(unittest.TestCase): self.assertEqual(res["model_id"], 10) self.assertEqual(res["version_id"], 20) - def test_build_search_result(self): - res = build_search_result( + def test_build_model_result(self): + res = build_model_result( source="civitai", model_id=100, version_id=200, @@ -66,6 +65,53 @@ class TestRefactoringUnification(unittest.TestCase): self.assertEqual(res["filename"], "test.safetensors") self.assertEqual(res["extra_field"], "hello_world") + def test_build_model_result_custom_url_contract(self): + res = build_model_result( + "civitai", + model_id=100, + version_id=200, + name="Test Model", + filename="test.safetensors", + url="https://civitai.com/models/100?modelVersionId=200", + download_url="https://civitai.com/api/download/models/200", + match_type="custom_url", + details_source="civitai", + version_url="https://civitai.com/models/100?modelVersionId=200", + custom_url=True, + 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) + + def test_build_model_result_compact_custom_url_contract(self): + res = build_model_result( + "civitai", + name="CivitAI", + filename="test.safetensors", + url="https://civitai.com/api/download/models/200", + download_url="https://civitai.com/api/download/models/200", + details_source="civitai", + version_url="https://civitai.com/api/download/models/200", + match_type="custom_url", + custom_url=True, + result_mode="compact_custom_url", + ) + + self.assertEqual({ + "source": "civitai", + "details_source": "civitai", + "name": "CivitAI", + "filename": "test.safetensors", + "url": "https://civitai.com/api/download/models/200", + "version_url": "https://civitai.com/api/download/models/200", + "download_url": "https://civitai.com/api/download/models/200", + "match_type": "custom_url", + "custom_url": True, + }, res) + @patch("requests.get") def test_request_page_text_success(self, mock_get): from core.sources.civarchive import _request_page_text @@ -195,8 +241,8 @@ class TestRefactoringUnification(unittest.TestCase): def test_alphanumeric_normalizers(self): from core.matcher import normalize_base_model - from core.path_utils import _normalize_base_model_token from core.path_templates import _normalize_token + from core.path_utils import _normalize_base_model_token from core.settings import _normalize_tag from core.type_utils import normalize_alphanumeric_key @@ -254,15 +300,15 @@ class TestRefactoringUnification(unittest.TestCase): self.assertNotIn(".", res) def test_listification(self): - from core.type_utils import as_list from core.settings import _listify_tags + from core.type_utils import as_list self.assertEqual(as_list("a,b, c"), ["a", "b", "c"]) self.assertEqual(list(_listify_tags("a,b; c")), ["a", "b", "c"]) self.assertEqual(as_list(["a", "", None, "b"]), ["a", "b"]) self.assertEqual(list(_listify_tags(["a", "", None, "b"])), ["a", "b"]) - def test_build_search_result_normalizes_hashes(self): + def test_build_model_result_normalizes_hashes(self): from core.type_utils import normalize_hashes_dict valid_sha = "a" * 64 @@ -270,7 +316,7 @@ class TestRefactoringUnification(unittest.TestCase): self.assertEqual(norm_hashes["sha256"], valid_sha.lower()) self.assertEqual(norm_hashes["autoV2"], "1234567890") - res = build_search_result( + res = build_model_result( "civitai", model_id=123, version_id=456, @@ -279,7 +325,7 @@ class TestRefactoringUnification(unittest.TestCase): download_url="https://civitai.com/api/download/123", normalize_hashes=True, ) - expected = build_search_result( + expected = build_model_result( "civitai", model_id=123, version_id=456, @@ -363,5 +409,3 @@ if __name__ == "__main__": unittest.main() - - diff --git a/tests/test_search_service_components.py b/tests/test_search_service_components.py index de24fda..2e1b9ed 100644 --- a/tests/test_search_service_components.py +++ b/tests/test_search_service_components.py @@ -55,7 +55,7 @@ def _build_search_orchestrator(): "self": extension, "asyncio": asyncio, "CivArchiveSearchError": Exception, - "build_search_result": lambda source, **fields: { + "build_model_result": lambda source, **fields: { "source": source, **fields, }, @@ -358,7 +358,7 @@ def test_search_provider_runner_resolves_civitai_urn(): } ), get_civitai_download_url=MagicMock(return_value="https://example.test/download"), - build_search_result=MagicMock(side_effect=lambda source, **fields: {"source": source, **fields}), + build_model_result=MagicMock(side_effect=lambda source, **fields: {"source": source, **fields}), ) owner.search_tracker.is_cancelled.return_value = False owner.log_search_result = MagicMock() @@ -413,7 +413,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_search_result=MagicMock(side_effect=lambda source, **fields: {"source": source, **fields}), + build_model_result=MagicMock(side_effect=lambda source, **fields: {"source": source, **fields}), ) owner.search_tracker.is_cancelled.return_value = False owner.log_search_result = MagicMock() @@ -452,7 +452,7 @@ def test_search_provider_runner_falls_back_for_civitai_urn_without_ids(): } ] ), - build_search_result=MagicMock( + build_model_result=MagicMock( side_effect=lambda source, **fields: {"source": source, **fields} ), log_search_result=MagicMock(),