From 151c246cde25fd1aecb6e96b658e7e506198dc86 Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Wed, 5 Aug 2026 00:52:08 +0200 Subject: [PATCH] refactor model list result construction --- core/sources/model_list.py | 44 ++++++++++++-------------------- tests/test_matcher.py | 51 +++++++++++++++++++++++++++++++++++++- 2 files changed, 66 insertions(+), 29 deletions(-) diff --git a/core/sources/model_list.py b/core/sources/model_list.py index 7cef5fb..6334e29 100644 --- a/core/sources/model_list.py +++ b/core/sources/model_list.py @@ -177,6 +177,18 @@ 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 search_model_list( filename: str, exact_only: bool = False ) -> Optional[Dict[str, Any]]: @@ -207,13 +219,7 @@ def search_model_list( url = model.get("url", "") if url: return { - "source": "model_list", - "filename": model_filename, - "url": url, - "name": model.get("name", ""), - "type": model.get("type", ""), - "directory": model.get("save_path", "checkpoints"), - "size": model.get("size", ""), + **_build_model_list_result(model), "match_type": "exact", } @@ -237,13 +243,7 @@ def search_model_list( if url: score = calculate_similarity(filename_norm, normalize_filename(model_filename)) return { - "source": "model_list", - "filename": model_filename, - "url": url, - "name": model.get("name", ""), - "type": model.get("type", ""), - "directory": model.get("save_path", "checkpoints"), - "size": model.get("size", ""), + **_build_model_list_result(model), "match_type": "fuzzy", "confidence": round(score * 100, 1), } @@ -265,13 +265,7 @@ def search_model_list( if url: best_score = score best_match = { - "source": "model_list", - "filename": model_filename, - "url": url, - "name": model.get("name", ""), - "type": model.get("type", ""), - "directory": model.get("save_path", "checkpoints"), - "size": model.get("size", ""), + **_build_model_list_result(model), "match_type": "similar", "confidence": round(score * 100, 1), } @@ -310,13 +304,7 @@ def search_model_list_multiple(filename: str, limit: int = 5) -> List[Dict[str, if url: results.append( { - "source": "model_list", - "filename": model_filename, - "url": url, - "name": model.get("name", ""), - "type": model.get("type", ""), - "directory": model.get("save_path", "checkpoints"), - "size": model.get("size", ""), + **_build_model_list_result(model), "confidence": round(score * 100, 1), } ) diff --git a/tests/test_matcher.py b/tests/test_matcher.py index 51f64f8..bbfdb33 100644 --- a/tests/test_matcher.py +++ b/tests/test_matcher.py @@ -1,4 +1,6 @@ import unittest +from unittest.mock import patch + from core.matcher import ( build_filename_search_queries, calculate_archived_model_confidence, @@ -8,8 +10,8 @@ from core.matcher import ( normalize_filename, normalize_model_family_filename, ) -from core.sources.model_list import search_model_list, search_model_list_multiple from core.progress import report_progress +from core.sources.model_list import search_model_list, search_model_list_multiple from core.type_utils import ( MODEL_CONTAINER_SUFFIXES, MODEL_VARIANT_SUFFIXES, @@ -21,6 +23,8 @@ from core.type_utils import ( as_list, format_size_bytes, ) + + class MatcherTests(unittest.TestCase): def test_normalize_filename_removes_extension_and_converts_lowercase(self): self.assertEqual("test model", normalize_filename("Test_Model.safetensors")) @@ -171,6 +175,51 @@ class MatcherTests(unittest.TestCase): except Exception as e: self.fail(f"search_model_list failed with exception: {e}") + def test_model_list_result_fields_preserve_match_mode_contracts(self): + models = [ + { + "filename": "exact.safetensors", + "url": "https://example.test/exact", + "name": "Exact model", + "type": "checkpoints", + "save_path": "checkpoints", + "size": "12 MB", + }, + { + "filename": "target-model.safetensors", + "url": "https://example.test/target", + "name": "Target model", + "type": "loras", + "save_path": "loras", + "size": 24, + }, + ] + + with patch("core.sources.model_list._load_model_list", return_value=models): + exact = search_model_list("exact.safetensors", exact_only=True) + 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.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]) + def test_build_filename_search_queries(self): # Verify that common precisions are extracted correctly queries = build_filename_search_queries("sdxl_ep10_fp16.safetensors")