diff --git a/core/sources/civarchive.py b/core/sources/civarchive.py index c42dcd4..2dd4cc3 100644 --- a/core/sources/civarchive.py +++ b/core/sources/civarchive.py @@ -48,7 +48,11 @@ from ..type_utils import ( select_primary_model_file, to_int, ) -from .common import build_unified_search_result, is_remote_link_marked_dead +from .common import ( + build_custom_url_result, + build_unified_search_result, + is_remote_link_marked_dead, +) log = create_module_logger(__name__) @@ -2419,34 +2423,30 @@ def build_civarchive_custom_result( else details.get("url") ) ) - return { - "source": "civarchive", - "details_source": "civarchive", - "model_id": model_id, - "version_id": version_id, - "name": details.get("name") or filename, - "version_name": selected_version.get("name") or "", - "type": details.get("type") or file_info.get("type") or "", - "filename": filename, - "url": version_url or details.get("url"), - "version_url": version_url or details.get("url"), - "platform_url": details.get("platform_url") - or selected_version.get("platform_url"), - "download_url": download_url, - "download_urls": [url for url in download_urls if url], - "size": file_info.get("size"), - "base_model": selected_version.get("base_model"), - "tags": details.get("tags") or [], - "trained_words": selected_version.get("trained_words") or [], - "images": details.get("images") or selected_version.get("images") or [], - "description": selected_version.get("description") + return build_custom_url_result( + source="civarchive", + model_id=model_id, + version_id=version_id, + name=details.get("name") or filename, + version_name=selected_version.get("name") or "", + type=details.get("type") or file_info.get("type") or "", + filename=filename, + url=version_url or details.get("url"), + download_url=download_url, + size=file_info.get("size"), + base_model=selected_version.get("base_model"), + tags=details.get("tags") or [], + trained_words=selected_version.get("trained_words") or [], + images=details.get("images") or selected_version.get("images") or [], + description=selected_version.get("description") or details.get("description") or "", - "sha256": sha256, - "hash": sha256, - "hashes": hashes, - "platform": details.get("platform"), - "match_type": "custom_url", - "custom_url": True, - } + sha256=sha256, + hashes=hashes, + platform_url=details.get("platform_url") + or selected_version.get("platform_url"), + download_urls=[url for url in download_urls if url], + hash=sha256, + platform=details.get("platform"), + ) diff --git a/core/sources/civitai.py b/core/sources/civitai.py index 6b729dc..358b68b 100644 --- a/core/sources/civitai.py +++ b/core/sources/civitai.py @@ -142,7 +142,7 @@ def check_civitai_api_key(api_key: Optional[str]) -> Dict[str, Any]: ) -from .common import build_unified_search_result +from .common import build_custom_url_result, build_unified_search_result def _build_civitai_result_from_version( @@ -2272,32 +2272,28 @@ def build_civitai_custom_result( else details.get("url") ) ) - return { - "source": "civitai", - "details_source": "civitai", - "model_id": model_id, - "version_id": version_id, - "name": details.get("name") or filename, - "version_name": selected_version.get("name") or "", - "type": details.get("type") or file_info.get("type") or "", - "filename": filename, - "url": version_url or details.get("url"), - "version_url": version_url or details.get("url"), - "download_url": download_url, - "size": file_info.get("size"), - "base_model": selected_version.get("base_model"), - "tags": details.get("tags") or [], - "trained_words": selected_version.get("trained_words") or [], - "images": details.get("images") or selected_version.get("images") or [], - "description": selected_version.get("description") + return build_custom_url_result( + source="civitai", + model_id=model_id, + version_id=version_id, + name=details.get("name") or filename, + version_name=selected_version.get("name") or "", + type=details.get("type") or file_info.get("type") or "", + filename=filename, + url=version_url or details.get("url"), + download_url=download_url, + size=file_info.get("size"), + base_model=selected_version.get("base_model"), + tags=details.get("tags") or [], + trained_words=selected_version.get("trained_words") or [], + images=details.get("images") or selected_version.get("images") or [], + description=selected_version.get("description") or details.get("description") or "", - "model_description": details.get("description") or "", - "sha256": sha256, - "hashes": hashes, - "match_type": "custom_url", - "custom_url": True, - } + sha256=sha256, + hashes=hashes, + model_description=details.get("description") or "", + ) def resolve_civitai_version_custom_result( diff --git a/core/sources/common.py b/core/sources/common.py index 84197e9..9aa529a 100644 --- a/core/sources/common.py +++ b/core/sources/common.py @@ -103,3 +103,52 @@ def build_unified_search_result( images=images, **extra, ) + + +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/tests/test_source_internals.py b/tests/test_source_internals.py index c361f4d..9c5ef96 100644 --- a/tests/test_source_internals.py +++ b/tests/test_source_internals.py @@ -27,6 +27,7 @@ from core.sources.civarchive import ( _normalize_archive_version, _normalize_download_url, _resolve_file_size_bytes, + build_civarchive_custom_result, parse_civarchive_url, ) @@ -44,6 +45,7 @@ from core.sources.civitai import ( _search_civitai_public_api_candidates, _search_civitai_red_candidates, _search_civitai_trpc_candidates, + build_civitai_custom_result, build_civitai_session_cookie, clear_search_cache, get_civitai_model_details, @@ -99,6 +101,124 @@ class CivitaiResultBuilderTests(unittest.TestCase): self.assertEqual(10 * 1024, result["size"]) +class CustomUrlResultBuilderTests(unittest.TestCase): + def test_civitai_custom_result_preserves_all_result_fields(self): + result = build_civitai_custom_result({ + "model_id": 123, + "version_id": 456, + "name": "Example model", + "type": "LORA", + "tags": ["style"], + "images": [{"url": "https://example.test/image"}], + "description": "Model description", + "version_url": "https://civitai.com/models/123?modelVersionId=456", + "selected_version": { + "id": 456, + "name": "v1", + "base_model": "SD 1.5", + "trained_words": ["example"], + "description": "Version notes", + "files": [{ + "name": "example.safetensors", + "type": "Model", + "download_url": "https://civitai.com/api/download/models/456", + "size": 123, + "sha256": "abc123", + "hashes": {"SHA256": "abc123"}, + }], + }, + }) + + self.assertEqual({ + "source": "civitai", + "details_source": "civitai", + "model_id": 123, + "version_id": 456, + "name": "Example model", + "version_name": "v1", + "type": "LORA", + "filename": "example.safetensors", + "url": "https://civitai.com/models/123?modelVersionId=456", + "version_url": "https://civitai.com/models/123?modelVersionId=456", + "download_url": "https://civitai.com/api/download/models/456", + "size": 123, + "base_model": "SD 1.5", + "tags": ["style"], + "trained_words": ["example"], + "images": [{"url": "https://example.test/image"}], + "description": "Version notes", + "model_description": "Model description", + "sha256": "abc123", + "hashes": {"SHA256": "abc123"}, + "match_type": "custom_url", + "custom_url": True, + }, result) + + def test_civarchive_custom_result_preserves_provider_fields(self): + result = build_civarchive_custom_result({ + "model_id": 789, + "version_id": 987, + "name": "Archived model", + "type": "Checkpoint", + "tags": ["archive"], + "platform": "civitai", + "platform_url": "https://civitai.com/models/789", + "description": "Archive description", + "version_url": "https://civarchive.com/models/789?modelVersionId=987", + "selected_version": { + "id": 987, + "name": "archive-v1", + "base_model": "SDXL", + "trained_words": ["archived"], + "description": "Archive notes", + "platform_url": "https://civitai.com/models/789?modelVersionId=987", + "files": [{ + "name": "archived.safetensors", + "type": "Model", + "download_url": "https://civarchive.com/download/987", + "download_urls": [ + "https://civarchive.com/download/987", + "https://mirror.example.test/987", + ], + "size": 456, + "sha256": "def456", + "hashes": {"SHA256": "def456"}, + }], + }, + }) + + self.assertEqual({ + "source": "civarchive", + "details_source": "civarchive", + "model_id": 789, + "version_id": 987, + "name": "Archived model", + "version_name": "archive-v1", + "type": "Checkpoint", + "filename": "archived.safetensors", + "url": "https://civarchive.com/models/789?modelVersionId=987", + "version_url": "https://civarchive.com/models/789?modelVersionId=987", + "platform_url": "https://civitai.com/models/789", + "download_url": "https://civarchive.com/download/987", + "download_urls": [ + "https://civarchive.com/download/987", + "https://mirror.example.test/987", + ], + "size": 456, + "base_model": "SDXL", + "tags": ["archive"], + "trained_words": ["archived"], + "images": [], + "description": "Archive notes", + "sha256": "def456", + "hash": "def456", + "hashes": {"SHA256": "def456"}, + "platform": "civitai", + "match_type": "custom_url", + "custom_url": True, + }, result) + + class CivitaiModelDetailsTests(unittest.TestCase): def test_preserves_full_selected_version_for_metadata_sidecar(self):