feat(cache): index versions by model id

This commit is contained in:
pixelpaws
2025-10-28 18:39:37 +08:00
parent ce3adaf831
commit 5264e49f2a
4 changed files with 164 additions and 44 deletions

View File

@@ -4,38 +4,60 @@ from py.services.model_cache import ModelCache
@pytest.mark.asyncio
async def test_name_sort_respects_file_name_display():
items = [
{"model_name": "Bravo", "file_name": "zulu", "folder": "", "size": 1, "modified": 1},
{"model_name": "Alpha", "file_name": "alpha", "folder": "", "size": 1, "modified": 1},
{"model_name": "Charlie", "file_name": "echo", "folder": "", "size": 1, "modified": 1},
async def test_model_cache_tracks_versions_by_model_id():
item_one = {
'file_path': '/models/a.safetensors',
'file_name': 'model-a-v1',
'folder': '',
'civitai': {'id': 101, 'modelId': 1, 'name': 'Alpha'},
}
item_two = {
'file_path': '/models/a_v2.safetensors',
'file_name': 'model-a-v2',
'folder': '',
'civitai': {'id': 102, 'modelId': 1, 'name': 'Beta'},
}
item_three = {
'file_path': '/models/b.safetensors',
'file_name': 'model-b',
'folder': '',
'civitai': {'id': 201, 'modelId': 2, 'name': 'Gamma'},
}
cache = ModelCache(
raw_data=[item_one, item_two, item_three],
folders=[],
name_display_mode='model_name',
)
versions = cache.get_versions_by_model_id(1)
assert versions == [
{'versionId': 101, 'name': 'Alpha', 'fileName': 'model-a-v1'},
{'versionId': 102, 'name': 'Beta', 'fileName': 'model-a-v2'},
]
cache = ModelCache(raw_data=items, folders=[], name_display_mode="file_name")
# Returned descriptors should not allow external mutation of the cache index
versions[0]['name'] = 'mutated'
assert cache.model_id_index[1][0]['name'] == 'Alpha'
sorted_items = await cache.get_sorted_data("name", "asc")
assert [item["file_name"] for item in sorted_items] == [
"alpha",
"echo",
"zulu",
# Removing entries updates both indexes
cache.remove_from_version_index(item_one)
assert cache.get_versions_by_model_id(1) == [
{'versionId': 102, 'name': 'Beta', 'fileName': 'model-a-v2'},
]
cache.remove_from_version_index(item_two)
assert cache.get_versions_by_model_id(1) == []
assert 1 not in cache.model_id_index
@pytest.mark.asyncio
async def test_update_name_display_mode_resorts_cached_name_order():
items = [
{"model_name": "Zulu", "file_name": "alpha", "folder": "", "size": 1, "modified": 1},
{"model_name": "Alpha", "file_name": "zulu", "folder": "", "size": 1, "modified": 1},
# Re-adding should not introduce duplicates
cache.add_to_version_index(item_two)
cache.add_to_version_index(item_two)
assert cache.get_versions_by_model_id('1') == [
{'versionId': 102, 'name': 'Beta', 'fileName': 'model-a-v2'},
]
cache = ModelCache(raw_data=items, folders=[], name_display_mode="model_name")
initial = await cache.get_sorted_data("name", "asc")
assert [item["model_name"] for item in initial] == ["Alpha", "Zulu"]
await cache.update_name_display_mode("file_name")
# The cached name sort should refresh immediately based on the new mode
updated = await cache.get_sorted_data("name", "asc")
assert [item["file_name"] for item in updated] == ["alpha", "zulu"]
# Other model IDs remain accessible
assert cache.get_versions_by_model_id(2) == [
{'versionId': 201, 'name': 'Gamma', 'fileName': 'model-b'},
]