diff --git a/core/metadata_utils.py b/core/metadata_utils.py index e8b82c9..56debf0 100644 --- a/core/metadata_utils.py +++ b/core/metadata_utils.py @@ -11,6 +11,43 @@ from .path_utils import ( ) +def build_metadata_source_payload( + *, + source: str, + details_source: str, + filename: str, + category: str, + model_name: str, + model_type: str, + sha256: str, + size: Any, + result: Dict[str, Any], + description: str, + version_description: str, + extra_fields: Optional[Dict[str, Any]] = None, +) -> Dict[str, Any]: + """Build the normalized source payload passed to the sidecar writer.""" + payload = { + "source": source, + "details_source": details_source, + "filename": filename, + "category": category, + "model_name": model_name, + "name": model_name, + "model_type": model_type, + "sha256": sha256, + "size": size, + "base_model": result.get("base_model"), + "base_model_source": result.get("base_model_source"), + "base_model_inferred": bool(result.get("base_model_inferred")), + "description": description, + "model_description": description, + "version_description": version_description, + } + payload.update(extra_fields or {}) + return payload + + def build_sidecar_file_identity( model_path: str, *, diff --git a/core/services/civitai_search_service.py b/core/services/civitai_search_service.py index 973b7f4..d4b7286 100644 --- a/core/services/civitai_search_service.py +++ b/core/services/civitai_search_service.py @@ -1,48 +1,13 @@ """CivitAI model search service.""" from .. import path_utils +from ..metadata_utils import build_metadata_source_payload from ..request_utils import extract_request_sha256 from ..routes.context import RouteContext from ..sources.common import collect_download_urls from .model_utils import CivitAISearchDependencies, ModelServiceDependencies -def _build_metadata_sidecar_payload( - *, - source, - details_source, - filename, - category, - model_name, - model_type, - sha256, - size, - result, - description, - version_description, - extra_fields=None, -): - payload = { - "source": source, - "details_source": details_source, - "filename": filename, - "category": category, - "model_name": model_name, - "name": model_name, - "model_type": model_type, - "sha256": sha256, - "size": size, - "base_model": result.get("base_model"), - "base_model_source": result.get("base_model_source"), - "base_model_inferred": bool(result.get("base_model_inferred")), - "description": description, - "model_description": description, - "version_description": version_description, - } - payload.update(extra_fields or {}) - return payload - - class CivitAISearchService(ModelServiceDependencies): """Execute exact-match CivitAI and model metadata searches.""" @@ -489,7 +454,7 @@ class CivitAISearchService(ModelServiceDependencies): model_type = result.get("model_type", "") or result.get( "type", "" ) - metadata_payload = _build_metadata_sidecar_payload( + metadata_payload = build_metadata_source_payload( source=source_name, details_source=result.get("details_source") or source_name, filename=filename, @@ -591,7 +556,7 @@ class CivitAISearchService(ModelServiceDependencies): split_result_descriptions(result) ) model_type = result.get("model_type", "") - metadata_payload = _build_metadata_sidecar_payload( + metadata_payload = build_metadata_source_payload( source="metadata_import", details_source=result.get("source") or "metadata", filename=filename, @@ -858,7 +823,7 @@ class CivitAISearchService(ModelServiceDependencies): model_description = response.get("model_description") or response.get( "description" ) or "" - metadata_payload = _build_metadata_sidecar_payload( + metadata_payload = build_metadata_source_payload( source="local", details_source=response.get("details_source") or "", filename=filename, diff --git a/tests/test_metadata_builder.py b/tests/test_metadata_builder.py index 37db30a..2aa3b7a 100644 --- a/tests/test_metadata_builder.py +++ b/tests/test_metadata_builder.py @@ -13,11 +13,57 @@ from core.metadata_builder import ( build_missing_local_metadata, normalize_metadata_build_mode, ) -from core.metadata_utils import build_sidecar_file_identity, merge_counted_payload +from core.metadata_utils import ( + build_metadata_source_payload, + build_sidecar_file_identity, + merge_counted_payload, +) from core.path_utils import get_model_resolver_sidecar_path class MetadataBuilderTests(unittest.TestCase): + def test_build_metadata_source_payload_preserves_common_source_fields(self): + payload = build_metadata_source_payload( + source="civitai", + details_source="civitai", + filename="model.safetensors", + category="checkpoints", + model_name="Example Model", + model_type="Checkpoint", + sha256="a" * 64, + size=1234, + result={ + "base_model": "SDXL", + "base_model_source": "civitai", + "base_model_inferred": 1, + }, + description="Model description", + version_description="Version description", + extra_fields={"author": "Author", "size": 5678}, + ) + + self.assertEqual( + { + "source": "civitai", + "details_source": "civitai", + "filename": "model.safetensors", + "category": "checkpoints", + "model_name": "Example Model", + "name": "Example Model", + "model_type": "Checkpoint", + "sha256": "a" * 64, + "size": 5678, + "base_model": "SDXL", + "base_model_source": "civitai", + "base_model_inferred": True, + "description": "Model description", + "model_description": "Model description", + "version_description": "Version description", + "author": "Author", + }, + payload, + ) + def test_build_sidecar_file_identity_returns_common_schema_fields(self): identity = build_sidecar_file_identity( r"models\example.safetensors",