nodes/Azornes-Comfyui-Model-Resolver
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: unify model result builders

Browse Source
This commit is contained in:
Dariusz L
2026-08-06 16:20:31 +02:00
parent 7afcd52f3f
commit 348f7ed8e0
15 changed files with 172 additions and 116 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/routes/registry.py
+1 -1
View File
@@ -105,7 +105,7 @@ def register_routes(self):
save_settings as save_resolver_settings,
)
from ..type_utils import (
build_search_result,
build_model_result,
extract_sha256_from_metadata,
fetch_remote_file_size_cached,
first_non_empty,
core/services/custom_url_service.py
+26 -23
View File
@@ -14,6 +14,7 @@ class CustomUrlService(ModelServiceDependencies):
"""Resolve a user-provided provider URL into a normal search result."""
UnsafeUrlError = self.UnsafeUrlError
asyncio = self.asyncio
build_model_result = self.build_model_result
build_civarchive_custom_result = self.build_civarchive_custom_result
build_civitai_custom_result = self.build_civitai_custom_result
build_huggingface_custom_result = self.build_huggingface_custom_result
@@ -64,17 +65,18 @@ class CustomUrlService(ModelServiceDependencies):
"civitai": "CivitAI",
"huggingface": "HuggingFace",
}.get(source, "Custom URL")
return {
"source": source,
"details_source": source,
"name": source_label,
"filename": filename,
"url": url,
"version_url": url,
"download_url": url,
"match_type": "custom_url",
"custom_url": True,
}
return build_model_result(
source,
name=source_label,
filename=filename,
url=url,
version_url=url,
download_url=url,
match_type="custom_url",
details_source=source,
custom_url=True,
result_mode="compact_custom_url",
)
def _collect_custom_url_local_hash_matches(result, category):
source_key = str(result.get("source") or "custom").strip().lower()
@@ -171,21 +173,22 @@ class CustomUrlService(ModelServiceDependencies):
civitai_key or None,
)
if not result and version_id:
result = {
"source": "civitai",
"details_source": "civitai",
"version_id": version_id,
"name": expected_filename or f"CivitAI version {version_id}",
"filename": expected_filename or f"civitai-{version_id}",
"url": normalized_url,
"version_url": normalized_url,
"download_url": get_civitai_download_url(
result = build_model_result(
"civitai",
version_id=version_id,
name=expected_filename or f"CivitAI version {version_id}",
filename=expected_filename or f"civitai-{version_id}",
url=normalized_url,
version_url=normalized_url,
download_url=get_civitai_download_url(
version_id,
civitai_key or None,
),
"match_type": "custom_url",
"custom_url": True,
}
match_type="custom_url",
details_source="civitai",
custom_url=True,
result_mode="compact_custom_url",
)
elif civarchive_parsed:
source = "civarchive"
if civarchive_parsed.get("sha256"):
core/services/model_utils.py
+6
View File
@@ -46,6 +46,10 @@ class NormalizeSha256Protocol(Protocol):
def __call__(self, value: Any) -> str: ...
class ModelResultBuilderProtocol(Protocol):
def __call__(self, source: str, **fields: Any) -> Dict[str, Any]: ...
class ReadJsonProtocol(Protocol):
def __call__(self, file_path: str, default: Any = None) -> Any: ...
@@ -276,6 +280,7 @@ class CustomUrlDependencies:
logger: Any
UnsafeUrlError: type[Exception]
asyncio: Any
build_model_result: ModelResultBuilderProtocol
build_civarchive_custom_result: Optional[BuildCivarchiveCustomResultProtocol]
build_civitai_custom_result: Optional[BuildCivitaiCustomResultProtocol]
build_huggingface_custom_result: Optional[BuildHuggingFaceCustomResultProtocol]
@@ -305,6 +310,7 @@ class CustomUrlDependencies:
logger=extension.logger,
UnsafeUrlError=context.require("UnsafeUrlError"),
asyncio=context.require("asyncio"),
build_model_result=context.require("build_model_result"),
build_civarchive_custom_result=context.get(
"build_civarchive_custom_result"
),
core/services/search_dependencies.py
+3 -3
View File
@@ -30,7 +30,7 @@ class IntParserProtocol(Protocol):
def __call__(self, value: Any, default: int = 0) -> int: ...
class SearchResultBuilderProtocol(Protocol):
class ModelResultBuilderProtocol(Protocol):
def __call__(self, source: str, **fields: Any) -> dict[str, Any]: ...
@@ -100,7 +100,7 @@ class SearchDependencies:
logger: LoggerProtocol
civarchive_search_error: type[BaseException]
asyncio: Any
build_search_result: SearchResultBuilderProtocol
build_model_result: ModelResultBuilderProtocol
clear_civarchive_search_cache: ClearCacheProtocol
clear_civitai_search_cache: ClearCacheProtocol
clear_huggingface_search_cache: ClearCacheProtocol
@@ -144,7 +144,7 @@ class SearchDependencies:
logger=cls._extension_value(extension, "logger"),
asyncio=context.require("asyncio"),
civarchive_search_error=context.require("CivArchiveSearchError"),
build_search_result=context.require("build_search_result"),
build_model_result=context.require("build_model_result"),
clear_civarchive_search_cache=context.require(
"clear_civarchive_search_cache"
),
core/services/search_orchestrator.py
+1 -1
View File
@@ -22,7 +22,7 @@ class SearchOrchestrator:
self.logger = dependencies.logger
self.CivArchiveSearchError = dependencies.civarchive_search_error
self.asyncio = dependencies.asyncio
self.build_search_result = dependencies.build_search_result
self.build_model_result = dependencies.build_model_result
self.clear_civarchive_search_cache = (
dependencies.clear_civarchive_search_cache
)
core/services/search_providers.py
+2 -2
View File
@@ -381,7 +381,7 @@ class SearchProviderRunner:
version_name=model_info.get("version_name"),
confidence=100.0,
)
source_results["civitai"] = self.owner.build_search_result(
source_results["civitai"] = self.owner.build_model_result(
"civitai", **result_payload
)
self.owner.log_search_result(
@@ -418,7 +418,7 @@ class SearchProviderRunner:
)
if civitai_results:
first_result = civitai_results[0]
source_results["civitai"] = self.owner.build_search_result(
source_results["civitai"] = self.owner.build_model_result(
"civitai",
model_id=first_result.get("model_id"),
version_id=first_result.get("version_id"),
core/sources/civarchive.py
+9 -6
View File
@@ -34,7 +34,7 @@ from ..path_utils import get_filename_from_path
from ..progress import get_progress_reporter
from ..type_utils import (
DEFAULT_BROWSER_USER_AGENT,
build_search_result,
build_model_result,
clear_remote_size_cache,
extract_file_size,
extract_trained_words,
@@ -50,7 +50,6 @@ from ..type_utils import (
to_int,
)
from .common import (
build_custom_url_result,
collect_download_urls,
is_remote_link_marked_dead,
)
@@ -1204,7 +1203,7 @@ def _build_result_from_normalized_version(
if size is None:
size = _resolve_file_size_bytes(file_info, download_urls)
return build_search_result(
return build_model_result(
source="civarchive",
model_id=model_id,
version_id=version_id,
@@ -1655,7 +1654,7 @@ def _build_result_from_payload(
else ""
)
return build_search_result(
return build_model_result(
source="civarchive",
model_id=model_id,
version_id=version_id,
@@ -1813,7 +1812,7 @@ def _build_result_from_search_candidate(
open_url = urljoin(CIVARCHIVE_BASE_URL, str(candidate.get("url") or ""))
sha256 = (parsed or {}).get("sha256")
return build_search_result(
return build_model_result(
source="civarchive",
model_id=candidate.get("model_id") or candidate.get("modelId"),
version_id=candidate.get("version_id") or candidate.get("modelVersionId"),
@@ -2383,7 +2382,7 @@ def build_civarchive_custom_result(
else details.get("url")
)
)
return build_custom_url_result(
return build_model_result(
source="civarchive",
model_id=model_id,
version_id=version_id,
@@ -2403,6 +2402,10 @@ def build_civarchive_custom_result(
or "",
sha256=sha256,
hashes=hashes,
details_source="civarchive",
version_url=version_url or details.get("url"),
custom_url=True,
result_mode="custom_url",
platform_url=details.get("platform_url")
or selected_version.get("platform_url"),
download_urls=[url for url in download_urls if url],
core/sources/civitai.py
+8 -5
View File
@@ -44,7 +44,7 @@ from ..type_utils import (
DEFAULT_BROWSER_USER_AGENT,
as_dict,
as_list,
build_search_result,
build_model_result,
check_credential_http,
extract_file_size,
extract_trained_words,
@@ -145,7 +145,6 @@ def check_civitai_api_key(api_key: Optional[str]) -> Dict[str, Any]:
)
from .common import build_custom_url_result
def build_civitai_result_payload(
@@ -202,7 +201,7 @@ def _build_civitai_result_from_version(
version_id = version.get("id")
size = extract_file_size(file_info)
download_url = file_info.get("downloadUrl") or get_civitai_download_url(version_id)
return build_search_result(
return build_model_result(
source="civitai",
**build_civitai_result_payload(
model_id=model_id,
@@ -971,7 +970,7 @@ def search_civitai_for_file(
api_key=api_key,
)
if hash_result:
result = build_search_result(
result = build_model_result(
"civitai",
model_id=hash_result.get("model_id"),
version_id=hash_result.get("version_id"),
@@ -2385,7 +2384,7 @@ def build_civitai_custom_result(
else details.get("url")
)
)
return build_custom_url_result(
return build_model_result(
source="civitai",
model_id=model_id,
version_id=version_id,
@@ -2405,6 +2404,10 @@ def build_civitai_custom_result(
or "",
sha256=sha256,
hashes=hashes,
details_source="civitai",
version_url=version_url or details.get("url"),
custom_url=True,
result_mode="custom_url",
model_description=details.get("description") or "",
)
core/sources/common.py
+1 -50
View File
@@ -1,5 +1,5 @@
"""
Common utilities and result builder functions for external model metadata sources.
Common utilities for external model metadata sources.
"""
from typing import Any, Callable, Dict, List, Optional, Tuple
@@ -91,52 +91,3 @@ def collect_download_urls(
urls.append(url)
return urls
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
core/sources/huggingface.py
+3 -3
View File
@@ -22,7 +22,7 @@ from ..matcher import build_filename_search_queries
from ..network_utils import execute_provider_json_request, host_matches_domain, request_source_response
from ..path_utils import METADATA_DIR, get_filename_from_path, read_json_safe, write_json_atomic
from ..type_utils import (
build_search_result,
build_model_result,
check_credential_http,
clear_remote_size_cache,
extract_file_size,
@@ -724,7 +724,7 @@ def _build_huggingface_result(
size = extract_file_size(file_info)
if not size:
size = _fetch_remote_file_size_bytes(download_url, headers=headers)
return build_search_result(
return build_model_result(
"huggingface",
model_id=repo_id,
version_id=None,
@@ -1567,7 +1567,7 @@ def build_huggingface_custom_result(
page_url = (
f"https://huggingface.co/{repo_id}/blob/{branch}/{quote_url_path(file_path)}"
)
return build_search_result(
return build_model_result(
"huggingface",
model_id=repo_id,
version_id=None,
core/sources/lora_manager_archive.py
+2 -2
View File
@@ -29,7 +29,7 @@ from ..matcher import (
from ..path_utils import get_filename_from_path
from ..progress import get_progress_reporter
from ..type_utils import (
build_search_result,
build_model_result,
extract_file_size,
get_generic_filename_tokens,
normalize_lora_manager_type,
@@ -489,7 +489,7 @@ def _build_result_from_row(
)
match_type = "exact" if confidence == 100.0 else "similar"
return build_search_result(
return build_model_result(
source="lora_manager_archive",
model_id=model_id,
version_id=version_id,
core/type_utils.py
+49 -4
View File
@@ -1272,11 +1272,11 @@ def normalize_model_file_info(
}
def build_search_result(
def build_model_result(
source: str,
*,
model_id: Any,
version_id: Any,
model_id: Any = None,
version_id: Any = None,
name: str = "",
version_name: str = "",
type: str = "",
@@ -1293,14 +1293,40 @@ def build_search_result(
normalize_hashes: bool = False,
trained_words: Optional[List[str]] = None,
images: Optional[List[Dict[str, Any]]] = None,
details_source: Optional[str] = None,
version_url: Optional[str] = None,
custom_url: bool = False,
result_mode: str = "search",
**extra: Any,
) -> Dict[str, Any]:
"""Helper to unify search result dictionary creation across different sources."""
"""Build the shared model-result contract for search and custom URL results."""
if normalize_hashes:
hashes = normalize_hashes_dict(hashes)
if not sha256 and "sha256" in hashes:
sha256 = hashes["sha256"]
if result_mode not in {"search", "custom_url", "compact_custom_url"}:
raise ValueError(f"Unsupported model result mode: {result_mode}")
if result_mode == "compact_custom_url":
result = {
"source": source,
"details_source": details_source or source,
"name": name,
"filename": filename,
"url": url,
"version_url": version_url or url,
"download_url": download_url,
"match_type": match_type,
"custom_url": custom_url,
}
if model_id is not None:
result["model_id"] = model_id
if version_id is not None:
result["version_id"] = version_id
result.update(extra)
return result
result = {
"source": source,
"model_id": model_id,
@@ -1321,6 +1347,25 @@ def build_search_result(
"trained_words": trained_words or [],
"images": images or [],
}
if result_mode == "custom_url":
result.pop("confidence", None)
result.update(
{
"details_source": details_source or source,
"version_url": version_url or url,
"match_type": "custom_url",
"custom_url": custom_url,
}
)
else:
if details_source is not None:
result["details_source"] = details_source
if version_url is not None:
result["version_url"] = version_url
if custom_url:
result["custom_url"] = True
result.update(extra)
return result
tests/test_refactoring_targets.py
+1
View File
@@ -1226,6 +1226,7 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase):
"logger",
"UnsafeUrlError",
"asyncio",
"build_model_result",
"build_civarchive_custom_result",
"build_civitai_custom_result",
"build_huggingface_custom_result",
tests/test_refactoring_unification.py
+56 -12
View File
@@ -1,6 +1,6 @@
import ast
import sys
import os
import sys
import unittest
from pathlib import Path
from unittest.mock import MagicMock, patch
@@ -10,11 +10,10 @@ if parent_dir not in sys.path:
sys.path.insert(0, parent_dir)
from core.matcher import MODEL_TITLE_MATCH_THRESHOLD
from core.sources.lora_manager_archive import _normalize_model_type
from core.type_utils import (
build_model_result,
normalize_lora_manager_type,
normalize_model_file_info,
build_search_result,
)
@@ -50,8 +49,8 @@ class TestRefactoringUnification(unittest.TestCase):
self.assertEqual(res["model_id"], 10)
self.assertEqual(res["version_id"], 20)
def test_build_search_result(self):
res = build_search_result(
def test_build_model_result(self):
res = build_model_result(
source="civitai",
model_id=100,
version_id=200,
@@ -66,6 +65,53 @@ class TestRefactoringUnification(unittest.TestCase):
self.assertEqual(res["filename"], "test.safetensors")
self.assertEqual(res["extra_field"], "hello_world")
def test_build_model_result_custom_url_contract(self):
res = build_model_result(
"civitai",
model_id=100,
version_id=200,
name="Test Model",
filename="test.safetensors",
url="https://civitai.com/models/100?modelVersionId=200",
download_url="https://civitai.com/api/download/models/200",
match_type="custom_url",
details_source="civitai",
version_url="https://civitai.com/models/100?modelVersionId=200",
custom_url=True,
result_mode="custom_url",
)
self.assertEqual(res["details_source"], "civitai")
self.assertEqual(res["version_url"], res["url"])
self.assertTrue(res["custom_url"])
self.assertNotIn("confidence", res)
def test_build_model_result_compact_custom_url_contract(self):
res = build_model_result(
"civitai",
name="CivitAI",
filename="test.safetensors",
url="https://civitai.com/api/download/models/200",
download_url="https://civitai.com/api/download/models/200",
details_source="civitai",
version_url="https://civitai.com/api/download/models/200",
match_type="custom_url",
custom_url=True,
result_mode="compact_custom_url",
)
self.assertEqual({
"source": "civitai",
"details_source": "civitai",
"name": "CivitAI",
"filename": "test.safetensors",
"url": "https://civitai.com/api/download/models/200",
"version_url": "https://civitai.com/api/download/models/200",
"download_url": "https://civitai.com/api/download/models/200",
"match_type": "custom_url",
"custom_url": True,
}, res)
@patch("requests.get")
def test_request_page_text_success(self, mock_get):
from core.sources.civarchive import _request_page_text
@@ -195,8 +241,8 @@ class TestRefactoringUnification(unittest.TestCase):
def test_alphanumeric_normalizers(self):
from core.matcher import normalize_base_model
from core.path_utils import _normalize_base_model_token
from core.path_templates import _normalize_token
from core.path_utils import _normalize_base_model_token
from core.settings import _normalize_tag
from core.type_utils import normalize_alphanumeric_key
@@ -254,15 +300,15 @@ class TestRefactoringUnification(unittest.TestCase):
self.assertNotIn(".", res)
def test_listification(self):
from core.type_utils import as_list
from core.settings import _listify_tags
from core.type_utils import as_list
self.assertEqual(as_list("a,b, c"), ["a", "b", "c"])
self.assertEqual(list(_listify_tags("a,b; c")), ["a", "b", "c"])
self.assertEqual(as_list(["a", "", None, "b"]), ["a", "b"])
self.assertEqual(list(_listify_tags(["a", "", None, "b"])), ["a", "b"])
def test_build_search_result_normalizes_hashes(self):
def test_build_model_result_normalizes_hashes(self):
from core.type_utils import normalize_hashes_dict
valid_sha = "a" * 64
@@ -270,7 +316,7 @@ class TestRefactoringUnification(unittest.TestCase):
self.assertEqual(norm_hashes["sha256"], valid_sha.lower())
self.assertEqual(norm_hashes["autoV2"], "1234567890")
res = build_search_result(
res = build_model_result(
"civitai",
model_id=123,
version_id=456,
@@ -279,7 +325,7 @@ class TestRefactoringUnification(unittest.TestCase):
download_url="https://civitai.com/api/download/123",
normalize_hashes=True,
)
expected = build_search_result(
expected = build_model_result(
"civitai",
model_id=123,
version_id=456,
@@ -363,5 +409,3 @@ if __name__ == "__main__":
unittest.main()
tests/test_search_service_components.py
+4 -4
View File
@@ -55,7 +55,7 @@ def _build_search_orchestrator():
"self": extension,
"asyncio": asyncio,
"CivArchiveSearchError": Exception,
"build_search_result": lambda source, **fields: {
"build_model_result": lambda source, **fields: {
"source": source,
**fields,
},
@@ -358,7 +358,7 @@ def test_search_provider_runner_resolves_civitai_urn():
}
),
get_civitai_download_url=MagicMock(return_value="https://example.test/download"),
build_search_result=MagicMock(side_effect=lambda source, **fields: {"source": source, **fields}),
build_model_result=MagicMock(side_effect=lambda source, **fields: {"source": source, **fields}),
)
owner.search_tracker.is_cancelled.return_value = False
owner.log_search_result = MagicMock()
@@ -413,7 +413,7 @@ def test_search_provider_runner_keeps_first_file_when_expected_name_is_missing()
}
),
get_civitai_download_url=MagicMock(return_value="https://example.test/download"),
build_search_result=MagicMock(side_effect=lambda source, **fields: {"source": source, **fields}),
build_model_result=MagicMock(side_effect=lambda source, **fields: {"source": source, **fields}),
)
owner.search_tracker.is_cancelled.return_value = False
owner.log_search_result = MagicMock()
@@ -452,7 +452,7 @@ def test_search_provider_runner_falls_back_for_civitai_urn_without_ids():
}
]
),
build_search_result=MagicMock(
build_model_result=MagicMock(
side_effect=lambda source, **fields: {"source": source, **fields}
),
log_search_result=MagicMock(),