From f199c9b591c0629dffa5cfa1b533b3a45799c5b9 Mon Sep 17 00:00:00 2001 From: Will Miao <13051207myq@gmail.com> Date: Tue, 6 Jan 2026 08:46:55 +0800 Subject: [PATCH] feat: propagate version info to model update record creation. Fixes #756 - Pass `version_info` parameter through download manager to model update service - Enhance `_create_record` to use version info when creating records for missing versions - Add `_extract_single_version` helper method for consistent version extraction - Improve handling of version metadata during library synchronization --- py/services/download_manager.py | 1 + py/services/model_update_service.py | 93 +++++++++++++-------- tests/services/test_model_update_service.py | 32 ++++++- 3 files changed, 87 insertions(+), 39 deletions(-) diff --git a/py/services/download_manager.py b/py/services/download_manager.py index f4b10b89..1715937b 100644 --- a/py/services/download_manager.py +++ b/py/services/download_manager.py @@ -487,6 +487,7 @@ class DownloadManager: model_type, resolved_model_id, sorted(version_ids), + version_info=version_info, ) except Exception as exc: logger.debug( diff --git a/py/services/model_update_service.py b/py/services/model_update_service.py index c8053635..4e686d02 100644 --- a/py/services/model_update_service.py +++ b/py/services/model_update_service.py @@ -589,6 +589,8 @@ class ModelUpdateService: model_type: str, model_id: int, version_ids: Sequence[int], + *, + version_info: Optional[Mapping] = None, ) -> ModelUpdateRecord: """Persist a new set of in-library version identifiers.""" @@ -600,6 +602,7 @@ class ModelUpdateService: normalized_versions, model_type=model_type, model_id=model_id, + version_info=version_info, ) self._upsert_record(record) return record @@ -944,6 +947,7 @@ class ModelUpdateService: model_type: Optional[str] = None, model_id: Optional[int] = None, last_checked_at: Optional[float] = None, + version_info: Optional[Mapping] = None, ) -> ModelUpdateRecord: local_set = set(normalized_local) versions: List[ModelVersionRecord] = [] @@ -965,19 +969,26 @@ class ModelUpdateService: seen_ids = {version.version_id for version in versions} for missing_id in sorted(local_set - seen_ids): - versions.append( - ModelVersionRecord( - version_id=missing_id, - name=None, - base_model=None, - released_at=None, - size_bytes=None, - preview_url=None, - is_in_library=True, - should_ignore=ignore_map.get(missing_id, False), - sort_index=len(versions), + new_version: Optional[ModelVersionRecord] = None + if version_info and _normalize_int(version_info.get("id")) == missing_id: + new_version = self._extract_single_version(version_info, index=len(versions)) + + if new_version: + versions.append(replace(new_version, is_in_library=True)) + else: + versions.append( + ModelVersionRecord( + version_id=missing_id, + name=None, + base_model=None, + released_at=None, + size_bytes=None, + preview_url=None, + is_in_library=True, + should_ignore=ignore_map.get(missing_id, False), + sort_index=len(versions), + ) ) - ) return ModelUpdateRecord( model_type=model_type, @@ -1083,33 +1094,45 @@ class ModelUpdateService: return [] if not isinstance(versions, Iterable): return None + extracted: List[ModelVersionRecord] = [] for index, entry in enumerate(versions): - if not isinstance(entry, Mapping): - continue - version_id = _normalize_int(entry.get("id")) - if version_id is None: - continue - name = _normalize_string(entry.get("name")) - base_model = _normalize_string(entry.get("baseModel")) - released_at = _normalize_string(entry.get("publishedAt") or entry.get("createdAt")) - size_bytes = self._extract_size_bytes(entry.get("files")) - preview_url = self._extract_preview_url(entry.get("images")) - extracted.append( - ModelVersionRecord( - version_id=version_id, - name=name, - base_model=base_model, - released_at=released_at, - size_bytes=size_bytes, - preview_url=preview_url, - is_in_library=False, - should_ignore=False, - sort_index=index, - ) - ) + version_record = self._extract_single_version(entry, index) + if version_record: + extracted.append(version_record) + return extracted + def _extract_single_version( + self, entry: Any, index: int = 0 + ) -> Optional[ModelVersionRecord]: + """Convert a raw metadata entry into a structured record.""" + + if not isinstance(entry, Mapping): + return None + + version_id = _normalize_int(entry.get("id")) + if version_id is None: + return None + + name = _normalize_string(entry.get("name")) + base_model = _normalize_string(entry.get("baseModel")) + released_at = _normalize_string(entry.get("publishedAt") or entry.get("createdAt")) + size_bytes = self._extract_size_bytes(entry.get("files")) + preview_url = self._extract_preview_url(entry.get("images")) + + return ModelVersionRecord( + version_id=version_id, + name=name, + base_model=base_model, + released_at=released_at, + size_bytes=size_bytes, + preview_url=preview_url, + is_in_library=False, + should_ignore=False, + sort_index=index, + ) + def _extract_size_bytes(self, files) -> Optional[int]: if not isinstance(files, Iterable): return None diff --git a/tests/services/test_model_update_service.py b/tests/services/test_model_update_service.py index d378c5f1..6538ac17 100644 --- a/tests/services/test_model_update_service.py +++ b/tests/services/test_model_update_service.py @@ -515,7 +515,31 @@ async def test_refresh_rewrites_remote_preview_urls(tmp_path): assert record is not None assert record.versions preview_url = record.versions[0].preview_url - assert ( - preview_url - == "https://image.civitai.com/safe/width=450,optimized=true/preview.png" - ) + +@pytest.mark.asyncio +async def test_update_in_library_versions_populates_metadata(tmp_path): + db_path = tmp_path / "updates.sqlite" + service = ModelUpdateService(str(db_path)) + + version_info = { + "id": 123, + "name": "v1.0", + "baseModel": "SD 1.5", + "publishedAt": "2024-03-01T00:00:00Z", + "files": [{"sizeKB": 1024, "type": "Model", "primary": True}], + "images": [{"url": "https://example.com/preview.png"}], + } + + await service.update_in_library_versions("lora", 1, [123], version_info=version_info) + record = await service.get_record("lora", 1) + + assert record is not None + assert len(record.versions) == 1 + version = record.versions[0] + assert version.version_id == 123 + assert version.name == "v1.0" + assert version.base_model == "SD 1.5" + assert version.size_bytes == 1024 * 1024 + assert version.preview_url == "https://example.com/preview.png" + assert version.is_in_library is True +