From 016138bf07740d5e929bc7effc109941c7ea19a6 Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Thu, 6 Aug 2026 20:40:14 +0200 Subject: [PATCH] refactor: share CivArchive result builder --- core/sources/civarchive.py | 85 +++++++++++++++++++++++++++++----- tests/test_source_internals.py | 85 ++++++++++++++++++++++++++++++++++ 2 files changed, 159 insertions(+), 11 deletions(-) diff --git a/core/sources/civarchive.py b/core/sources/civarchive.py index 2623033..eec5ec8 100644 --- a/core/sources/civarchive.py +++ b/core/sources/civarchive.py @@ -1166,6 +1166,72 @@ def _select_primary_model_file(files: List[Dict[str, Any]]) -> Optional[Dict[str return select_primary_model_file(valid_files) +_RESULT_HASH_UNSET = object() + + +def _build_civarchive_result( + *, + model_id: Any, + version_id: Any, + name: str, + version_name: str, + model_type: str, + filename: str, + url: str, + platform_url: Optional[str], + civitai_model_id: Any, + civitai_model_version_id: Any, + download_urls: List[str], + size: Optional[int], + base_model: Optional[str], + tags: List[str], + trained_words: List[Any], + images: List[Any], + creator: Dict[str, Any], + platform: Optional[str], + is_deleted: bool, + match_type: str, + confidence: float = 0.0, + sha256: Optional[str] = None, + hashes: Optional[Dict[str, Any]] = None, + hash_value: Any = _RESULT_HASH_UNSET, +) -> Dict[str, Any]: + extra_fields = { + "civitai_model_id": civitai_model_id, + "civitai_model_version_id": civitai_model_version_id, + "creator": creator, + "platform": platform, + "is_deleted": is_deleted, + } + if hash_value is not _RESULT_HASH_UNSET: + extra_fields["hash"] = hash_value + + return build_model_result( + source="civarchive", + model_id=model_id, + version_id=version_id, + name=name, + version_name=version_name, + type=model_type, + filename=filename, + url=url, + platform_url=platform_url, + download_url=download_urls[0], + download_urls=download_urls, + size=size, + base_model=base_model, + tags=tags, + trained_words=trained_words, + images=images, + match_type=match_type, + confidence=confidence, + sha256=sha256, + hashes=hashes, + normalize_hashes=True, + **extra_fields, + ) + + def _build_result_from_normalized_version( model_details: Dict[str, Any], version: Dict[str, Any], @@ -1196,19 +1262,17 @@ def _build_result_from_normalized_version( if size is None: size = _resolve_file_size_bytes(file_info, download_urls) - return build_model_result( - source="civarchive", + return _build_civarchive_result( model_id=model_id, version_id=version_id, name=model_details.get("name") or "", version_name=version.get("name") or "", - type=model_details.get("type") or file_info.get("type"), + model_type=model_details.get("type") or file_info.get("type"), filename=filename, url=url, platform_url=version.get("platform_url") or version.get("platformUrl"), civitai_model_id=version.get("civitai_model_id") or version.get("civitaiModelId"), civitai_model_version_id=version.get("civitai_model_version_id") or version.get("civitaiModelVersionId"), - download_url=download_urls[0], download_urls=download_urls, size=size, base_model=version.get("base_model") or version.get("baseModel") or version.get("baseModelType"), @@ -1219,7 +1283,9 @@ def _build_result_from_normalized_version( platform=model_details.get("platform"), is_deleted=False, match_type=match_type, - normalize_hashes=True, + confidence=0.0, + sha256=None, + hashes=None, ) @@ -1631,19 +1697,17 @@ def _build_result_from_payload( else "" ) - return build_model_result( - source="civarchive", + return _build_civarchive_result( model_id=model_id, version_id=version_id, name=model_name, version_name=version_name, - type=context.get("type") or selected_file.get("type"), + model_type=context.get("type") or selected_file.get("type"), filename=filename, url=civarchive_url, platform_url=version.get("platform_url") or version.get("platformUrl"), civitai_model_id=version.get("civitai_model_id") or version.get("civitaiModelId"), civitai_model_version_id=version.get("civitai_model_version_id") or version.get("civitaiModelVersionId"), - download_url=download_urls[0], download_urls=download_urls, size=_resolve_file_size_bytes(selected_file, download_urls), base_model=version.get("baseModel") or version.get("base_model") or version.get("baseModelType"), @@ -1660,9 +1724,8 @@ def _build_result_from_payload( match_type=match_type, confidence=best_confidence, sha256=sha256, - hash=sha256, hashes=hashes, - normalize_hashes=True, + hash_value=sha256, ) diff --git a/tests/test_source_internals.py b/tests/test_source_internals.py index 76d5ab2..77efc9c 100644 --- a/tests/test_source_internals.py +++ b/tests/test_source_internals.py @@ -22,6 +22,8 @@ import requests # civarchive internals # --------------------------------------------------------------------------- from core.sources.civarchive import ( + _build_result_from_normalized_version, + _build_result_from_payload, _collect_archive_download_urls, _fetch_remote_file_size_bytes, _normalize_archive_version, @@ -130,6 +132,89 @@ class CivitaiResultBuilderTests(unittest.TestCase): self.assertEqual(10 * 1024, result["size"]) +class CivarchiveResultBuilderTests(unittest.TestCase): + + def test_normalized_version_builder_preserves_search_result_contract(self): + result = _build_result_from_normalized_version( + model_details={ + "model_id": 123, + "name": "Archive model", + "type": "LORA", + "tags": ["style"], + "creator": {"username": "archive-user"}, + "platform": "civarchive", + }, + version={ + "id": 456, + "name": "v1", + "base_model": "SDXL", + "trained_words": ["trigger"], + "images": [], + "url": "https://civarchive.com/models/123?modelVersionId=456", + }, + file_info={ + "name": "archive.safetensors", + "download_url": "https://civarchive.com/api/download/models/456.safetensors", + "size": 2048, + }, + match_type="similar", + ) + + self.assertEqual("civarchive", result["source"]) + self.assertEqual(123, result["model_id"]) + self.assertEqual(456, result["version_id"]) + self.assertEqual("archive.safetensors", result["filename"]) + self.assertEqual(2048, result["size"]) + self.assertEqual("SDXL", result["base_model"]) + self.assertEqual("archive-user", result["creator"]["username"]) + self.assertFalse(result["is_deleted"]) + self.assertEqual(0.0, result["confidence"]) + self.assertNotIn("hash", result) + + def test_payload_builder_preserves_confidence_hash_and_deleted_fields(self): + sha256 = "a" * 64 + result = _build_result_from_payload( + { + "id": 123, + "name": "Archive model", + "type": "LORA", + "creator_username": "archive-user", + "creator_name": "Archive User", + "creator_url": "https://civarchive.com/users/archive-user", + "deletedAt": "2025-01-01T00:00:00Z", + "tags": ["style"], + "version": { + "id": 456, + "name": "v1", + "baseModel": "SDXL", + "trainedWords": ["trigger"], + "files": [ + { + "name": "archive.safetensors", + "downloadUrl": "https://civarchive.com/api/download/models/456.safetensors", + "sizeBytes": 4096, + "hashes": {"SHA256": sha256}, + } + ], + }, + }, + query="archive.safetensors", + preferred_filename="archive.safetensors", + ) + + self.assertEqual("civarchive", result["source"]) + self.assertEqual(123, result["model_id"]) + self.assertEqual(456, result["version_id"]) + self.assertEqual(4096, result["size"]) + self.assertEqual(100.0, result["confidence"]) + self.assertEqual("exact", result["match_type"]) + self.assertEqual(sha256, result["sha256"]) + self.assertEqual(sha256, result["hash"]) + self.assertEqual(sha256, result["hashes"]["sha256"]) + self.assertTrue(result["is_deleted"]) + self.assertEqual("archive-user", result["creator"]["username"]) + + class CivitaiCredentialCheckTests(unittest.TestCase): def test_empty_credentials_keep_provider_specific_precheck_messages(self):