From 7a219bd2d28adb2cd3c3942b81c9788d79953b13 Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Wed, 5 Aug 2026 14:29:02 +0200 Subject: [PATCH] refactor: unify primary model file selection --- core/services/search_providers.py | 15 +++--- core/sources/civarchive.py | 6 +-- core/sources/civitai.py | 38 +++++--------- core/type_utils.py | 14 +++++ tests/test_refactoring_targets.py | 20 ++++++++ tests/test_search_service_components.py | 42 +++++++++++++++ tests/test_source_internals.py | 68 +++++++++++++++++++++++++ 7 files changed, 164 insertions(+), 39 deletions(-) diff --git a/core/services/search_providers.py b/core/services/search_providers.py index 2f89137..a701efd 100644 --- a/core/services/search_providers.py +++ b/core/services/search_providers.py @@ -4,6 +4,7 @@ from dataclasses import dataclass from typing import Any from ..sources.civitai import build_civitai_result_payload +from ..type_utils import select_primary_model_file class SearchCancelled(BaseException): @@ -358,15 +359,11 @@ class SearchProviderRunner: "Selecting CivitAI file", 76, ) - primary_file = None - for file_info in model_info.get("files", []): - if file_info.get("name") == model_info.get( - "expected_filename" - ): - primary_file = file_info - break - if primary_file is None: - primary_file = (model_info.get("files") or [{}])[0] + primary_file = select_primary_model_file( + model_info.get("files") or [], + expected_filename=model_info.get("expected_filename"), + fallback_to_first=True, + ) or {} download_url = self.owner.get_civitai_download_url(version_id) result_payload = build_civitai_result_payload( diff --git a/core/sources/civarchive.py b/core/sources/civarchive.py index b98bf7c..3689d29 100644 --- a/core/sources/civarchive.py +++ b/core/sources/civarchive.py @@ -871,10 +871,6 @@ def _extract_version_files( return _merge_top_file_mirrors(files, top_files or []) -def _select_primary_file(files: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]: - return select_primary_model_file(files) - - def _collect_archive_download_urls( file_info: Dict[str, Any], prioritize_civitai_last: bool = False, @@ -1603,7 +1599,7 @@ def _build_result_from_payload( selected_file = file_info if selected_file is None: - selected_file = _select_primary_file(files) + selected_file = select_primary_model_file(files) best_confidence = calculate_archived_model_confidence( query_value, context.get("name", ""), diff --git a/core/sources/civitai.py b/core/sources/civitai.py index a9311e5..2649dd4 100644 --- a/core/sources/civitai.py +++ b/core/sources/civitai.py @@ -738,13 +738,11 @@ def _find_civitai_file_in_model( resolved: Dict[str, Any], version_id: int ) -> Dict[str, Any]: expected_filename = resolved.get("expected_filename", "") - primary_file = None - for file_info in resolved.get("files", []): - if file_info.get("name") == expected_filename: - primary_file = file_info - break - if primary_file is None: - primary_file = (resolved.get("files") or [{}])[0] + primary_file = select_primary_model_file( + resolved.get("files") or [], + expected_filename=expected_filename, + fallback_to_first=True, + ) or {} return { "source": "civitai", @@ -1248,14 +1246,10 @@ def search_civitai( version_id = latest.get("id") files = latest.get("files", []) - primary_file = None - for f in files: - if f.get("primary", False) or f.get("type") == "Model": - primary_file = f - break - - if not primary_file and files: - primary_file = files[0] + primary_file = select_primary_model_file( + files, + prefer_first_marked=True, + ) result = { "source": "civitai", @@ -1533,16 +1527,10 @@ def resolve_urn( return None files = target_version.get("files", []) - primary_file = None - - # Prefer primary file or type=='Model' - for f in files: - if f.get("primary") or f.get("type") == "Model": - primary_file = f - break - - if not primary_file and files: - primary_file = files[0] # Fallback to first + primary_file = select_primary_model_file( + files, + prefer_first_marked=True, + ) if not primary_file: log.warning(f"No files found for version {version_id}") diff --git a/core/type_utils.py b/core/type_utils.py index ce4a772..4bd1022 100644 --- a/core/type_utils.py +++ b/core/type_utils.py @@ -409,10 +409,16 @@ def select_primary_model_file( files: List[Dict[str, Any]], expected_filename: Optional[str] = None, require_download: bool = False, + fallback_to_first: bool = False, + prefer_first_marked: bool = False, ) -> Optional[Dict[str, Any]]: """ Selects the primary file from a list of model version files, optionally filtering by expected filename or requiring a download URL. + + ``fallback_to_first`` preserves callers that use an exact filename match + followed by the first available file. ``prefer_first_marked`` preserves + callers that select the first file marked primary or of type ``Model``. """ if not isinstance(files, list): files = [files] if files else [] @@ -440,6 +446,14 @@ def select_primary_model_file( for file_info in candidates: if get_filename_from_path(file_name(file_info)).lower() == expected: return file_info + if fallback_to_first: + return candidates[0] + + if prefer_first_marked: + for file_info in candidates: + if file_info.get("primary") or str(file_info.get("type") or "").lower() == "model": + return file_info + return candidates[0] # 1. Look for file marked primary and of type Model for file_info in candidates: diff --git a/tests/test_refactoring_targets.py b/tests/test_refactoring_targets.py index d811b7a..87425a8 100644 --- a/tests/test_refactoring_targets.py +++ b/tests/test_refactoring_targets.py @@ -1082,6 +1082,26 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase): # Require download excludes main.safetensors because it has no download url self.assertEqual(select_primary_model_file(files, require_download=True)["id"], 1) + marked_files = [ + {"id": 4, "name": "first.safetensors", "type": "Model"}, + {"id": 5, "name": "primary.safetensors", "primary": True, "type": "Model"}, + ] + self.assertEqual( + select_primary_model_file( + marked_files, + prefer_first_marked=True, + )["id"], + 4, + ) + self.assertEqual( + select_primary_model_file( + marked_files, + expected_filename="missing.safetensors", + fallback_to_first=True, + )["id"], + 4, + ) + def _build_model_service_for_unit_test( self, service_module="model_service", diff --git a/tests/test_search_service_components.py b/tests/test_search_service_components.py index dac5078..de24fda 100644 --- a/tests/test_search_service_components.py +++ b/tests/test_search_service_components.py @@ -389,6 +389,48 @@ def test_search_provider_runner_resolves_civitai_urn(): assert result["civitai"]["hashes"] == {} +def test_search_provider_runner_keeps_first_file_when_expected_name_is_missing(): + owner = SimpleNamespace( + CivArchiveSearchError=Exception, + logger=MagicMock(), + search_tracker=MagicMock(), + format_log_fields=MagicMock(return_value="ids=10@20"), + resolve_urn=MagicMock( + return_value={ + "model_name": "URN Model", + "version_name": "v1", + "expected_filename": "missing.safetensors", + "base_model": "SDXL", + "files": [ + {"name": "fallback.safetensors", "sha256": "a" * 64}, + { + "name": "primary.safetensors", + "primary": True, + "type": "Model", + "sha256": "b" * 64, + }, + ], + } + ), + get_civitai_download_url=MagicMock(return_value="https://example.test/download"), + build_search_result=MagicMock(side_effect=lambda source, **fields: {"source": source, **fields}), + ) + owner.search_tracker.is_cancelled.return_value = False + owner.log_search_result = MagicMock() + runner = SearchProviderRunner(owner) + + result, found = runner.search_civitai_source_task( + _request( + data={"model_id": 10, "version_id": 20}, + is_urn=True, + base_model_context="", + ) + ) + + assert found is True + assert result["civitai"]["sha256"] == "a" * 64 + + def test_search_provider_runner_falls_back_for_civitai_urn_without_ids(): owner = SimpleNamespace( logger=MagicMock(), diff --git a/tests/test_source_internals.py b/tests/test_source_internals.py index 518b8a0..57b94ad 100644 --- a/tests/test_source_internals.py +++ b/tests/test_source_internals.py @@ -53,6 +53,8 @@ from core.sources.civitai import ( get_model_info_by_hash, get_model_info_for_file, parse_civitai_url, + resolve_urn, + search_civitai, search_civitai_for_file, ) from core.sources.common import is_remote_link_marked_dead @@ -125,6 +127,72 @@ class CivitaiResultBuilderTests(unittest.TestCase): self.assertEqual(10 * 1024, result["size"]) +class CivitaiPrimaryFileSelectionTests(unittest.TestCase): + + @patch("core.sources.civitai.execute_provider_json_request") + def test_resolve_urn_preserves_first_marked_file_order(self, mock_request): + clear_search_cache() + mock_request.return_value = { + "name": "Example model", + "tags": [], + "modelVersions": [ + { + "id": 456, + "name": "v1", + "baseModel": "SDXL", + "files": [ + {"name": "first.safetensors", "type": "Model", "sizeKB": 1}, + { + "name": "primary.safetensors", + "primary": True, + "type": "Model", + "sizeKB": 2, + }, + ], + } + ], + } + + result = resolve_urn(123, 456) + + self.assertEqual("first.safetensors", result["expected_filename"]) + clear_search_cache() + + @patch("core.sources.civitai.execute_provider_json_request") + def test_general_search_preserves_first_marked_file_order(self, mock_request): + clear_search_cache() + mock_request.return_value = { + "items": [ + { + "id": 123, + "name": "Example model", + "type": "Checkpoint", + "stats": {"downloadCount": 1}, + "modelVersions": [ + { + "id": 456, + "baseModel": "SDXL", + "files": [ + {"name": "first.safetensors", "type": "Model", "sizeKB": 1}, + { + "name": "primary.safetensors", + "primary": True, + "type": "Model", + "sizeKB": 2, + }, + ], + } + ], + } + ] + } + + results = search_civitai("example", limit=1) + + self.assertEqual("first.safetensors", results[0]["filename"]) + clear_search_cache() + + class CustomUrlResultBuilderTests(unittest.TestCase): def test_civitai_custom_result_preserves_all_result_fields(self): result = build_civitai_custom_result({