Unify custom URL result builders
This commit is contained in:
Dariusz L
+29
-29
@@ -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"),
|
||||
)
|
||||
|
||||
|
||||
+21
-25
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user