refactor: share CivArchive result builder

This commit is contained in:
Dariusz L
2026-08-06 20:40:14 +02:00
parent 233731acce
commit 016138bf07
2 changed files with 159 additions and 11 deletions
+74 -11
View File
@@ -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,
)
+85
View File
@@ -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):