refactor: share CivArchive result builder
This commit is contained in:
+74
-11
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user