refactor(search): extract orchestrator, providers, and cache from SearchService
Replace the 1095-line monolithic SearchService with a thin facade that delegates to SearchOrchestrator. Split the implementation into three focused modules: - search_orchestrator.py: request parsing, provider coordination, and response construction - search_providers.py: per-source execution, progress callbacks, and async task creation - search_cache.py: timestamping and stable cache stamping for results Add unit tests for the new components in tests/test_search_service_components.py. Preserve the existing HTTP route behavior by keeping SearchService.search_sources as the stable route-facing API.
This commit is contained in:
Dariusz L
@@ -0,0 +1,104 @@
|
||||
"""Timestamping and cache helpers for source search results."""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
SEARCH_RESULT_KEYS = (
|
||||
"popular",
|
||||
"model_list",
|
||||
"huggingface",
|
||||
"civitai",
|
||||
"civarchive",
|
||||
"lora_manager_archive",
|
||||
)
|
||||
|
||||
|
||||
class SearchResultCache:
|
||||
"""Keep stable timestamps for equivalent search results."""
|
||||
|
||||
def __init__(self, timestamps: dict[str, str]):
|
||||
self._timestamps = timestamps
|
||||
|
||||
@staticmethod
|
||||
def _current_timestamp() -> str:
|
||||
return (
|
||||
datetime.now(timezone.utc)
|
||||
.replace(microsecond=0)
|
||||
.isoformat()
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _result_signature(source_key: str, result: Any) -> str:
|
||||
if isinstance(result, list):
|
||||
return "|".join(
|
||||
SearchResultCache._result_signature(source_key, item)
|
||||
for item in result
|
||||
)
|
||||
if not isinstance(result, dict):
|
||||
return ""
|
||||
|
||||
parts = [
|
||||
source_key,
|
||||
result.get("download_url")
|
||||
or result.get("url")
|
||||
or result.get("model_url")
|
||||
or "",
|
||||
result.get("filename") or result.get("path") or "",
|
||||
result.get("repo_id") or result.get("repo") or "",
|
||||
result.get("model_id") or "",
|
||||
result.get("version_id") or "",
|
||||
result.get("name") or "",
|
||||
]
|
||||
return "::".join(str(part).strip() for part in parts)
|
||||
|
||||
def stamp_result(
|
||||
self,
|
||||
source_key: str,
|
||||
result: Any,
|
||||
*,
|
||||
force_search: bool,
|
||||
) -> Any:
|
||||
if isinstance(result, list):
|
||||
return [
|
||||
self.stamp_result(
|
||||
source_key,
|
||||
item,
|
||||
force_search=force_search,
|
||||
)
|
||||
for item in result
|
||||
]
|
||||
if not isinstance(result, dict):
|
||||
return result
|
||||
|
||||
signature = self._result_signature(source_key, result)
|
||||
if not signature:
|
||||
return result
|
||||
|
||||
timestamp = (
|
||||
result.get("searched_at")
|
||||
or result.get("searchedAt")
|
||||
or (
|
||||
None
|
||||
if force_search
|
||||
else self._timestamps.get(signature)
|
||||
)
|
||||
)
|
||||
if not timestamp:
|
||||
timestamp = self._current_timestamp()
|
||||
if force_search:
|
||||
self._timestamps[signature] = timestamp
|
||||
else:
|
||||
self._timestamps.setdefault(signature, timestamp)
|
||||
result["searched_at"] = timestamp
|
||||
return result
|
||||
|
||||
def stamp_results(self, payload: dict[str, Any], *, force_search: bool):
|
||||
"""Stamp all provider result collections in a search response."""
|
||||
for source_key in SEARCH_RESULT_KEYS:
|
||||
if payload.get(source_key):
|
||||
payload[source_key] = self.stamp_result(
|
||||
source_key,
|
||||
payload[source_key],
|
||||
force_search=force_search,
|
||||
)
|
||||
return payload
|
||||
@@ -0,0 +1,511 @@
|
||||
"""Request orchestration for model source searches."""
|
||||
|
||||
from ..routes.context import RouteContext
|
||||
from .search_cache import SearchResultCache
|
||||
from .search_providers import (
|
||||
SearchCancelled,
|
||||
SearchProviderRunner,
|
||||
SearchRequest,
|
||||
)
|
||||
|
||||
|
||||
class SearchOrchestrator:
|
||||
"""Parse search requests, coordinate providers, and build responses."""
|
||||
|
||||
def __init__(self, context: RouteContext):
|
||||
extension = context.get("self")
|
||||
self.search_tracker = extension.search_tracker
|
||||
self.search_result_timestamps = extension.search_result_timestamps
|
||||
self.logger = extension.logger
|
||||
self.CivArchiveSearchError = context.get("CivArchiveSearchError")
|
||||
self.asyncio = context.get("asyncio")
|
||||
self.build_search_result = context.get("build_search_result")
|
||||
self.cancel_progress_response = context.get("cancel_progress_response")
|
||||
self.clear_civarchive_search_cache = context.get(
|
||||
"clear_civarchive_search_cache"
|
||||
)
|
||||
self.clear_civitai_search_cache = context.get(
|
||||
"clear_civitai_search_cache"
|
||||
)
|
||||
self.clear_huggingface_search_cache = context.get(
|
||||
"clear_huggingface_search_cache"
|
||||
)
|
||||
self.clear_lora_manager_archive_search_cache = context.get(
|
||||
"clear_lora_manager_archive_search_cache"
|
||||
)
|
||||
self.extract_sha256_from_metadata = context.get(
|
||||
"extract_sha256_from_metadata"
|
||||
)
|
||||
self.format_size_bytes = context.get("format_size_bytes")
|
||||
self.get_civitai_download_url = context.get("get_civitai_download_url")
|
||||
self.get_popular_model_url = context.get("get_popular_model_url")
|
||||
self.get_progress_response = context.get("get_progress_response")
|
||||
self.json_api_endpoint = context.get("json_api_endpoint")
|
||||
self.reload_model_list = context.get("reload_model_list")
|
||||
self.reload_popular_databases = context.get("reload_popular_databases")
|
||||
self.resolve_civarchive_model_version = context.get(
|
||||
"resolve_civarchive_model_version"
|
||||
)
|
||||
self.resolve_urn = context.get("resolve_urn")
|
||||
self.routes = context.get("routes")
|
||||
self.search_civarchive_for_file = context.get(
|
||||
"search_civarchive_for_file"
|
||||
)
|
||||
self.search_civitai = context.get("search_civitai")
|
||||
self.search_civitai_for_file = context.get("search_civitai_for_file")
|
||||
self.search_huggingface_for_file = context.get(
|
||||
"search_huggingface_for_file"
|
||||
)
|
||||
self.search_local_matches_by_hash = context.get(
|
||||
"search_local_matches_by_hash"
|
||||
)
|
||||
self.search_lora_manager_archive_for_file = context.get(
|
||||
"search_lora_manager_archive_for_file"
|
||||
)
|
||||
self.search_model_list = context.get("search_model_list")
|
||||
self.to_bool = context.get("to_bool")
|
||||
self.to_int = context.get("to_int")
|
||||
self.web = context.get("web")
|
||||
self.search_cache = SearchResultCache(self.search_result_timestamps)
|
||||
self.provider_runner = SearchProviderRunner(self)
|
||||
|
||||
def format_log_value(self, value):
|
||||
if value is None or value == "":
|
||||
return None
|
||||
if isinstance(value, bool):
|
||||
return "yes" if value else "no"
|
||||
if isinstance(value, (list, tuple, set)):
|
||||
return ",".join(str(item) for item in value)
|
||||
|
||||
text = str(value)
|
||||
if any(char.isspace() for char in text):
|
||||
return '"' + text.replace("\\", "\\\\").replace('"', '\\"') + '"'
|
||||
return text
|
||||
|
||||
def format_log_size(self, value):
|
||||
return self.format_size_bytes(value, include_space=False)
|
||||
|
||||
def format_log_fields(self, **fields):
|
||||
parts = []
|
||||
for key, value in fields.items():
|
||||
formatted = self.format_log_value(value)
|
||||
if formatted is not None:
|
||||
parts.append(f"{key}={formatted}")
|
||||
return " ".join(parts)
|
||||
|
||||
def normalize_result_extra(self, extra):
|
||||
if not extra:
|
||||
return {}
|
||||
normalized = dict(extra)
|
||||
model_id = normalized.pop("model_id", None)
|
||||
version_id = normalized.pop("version_id", None)
|
||||
if model_id or version_id:
|
||||
normalized["ids"] = (
|
||||
f"{model_id}@{version_id}"
|
||||
if model_id and version_id
|
||||
else model_id or version_id
|
||||
)
|
||||
if "files_count" in normalized:
|
||||
normalized["files"] = normalized.pop("files_count")
|
||||
return normalized
|
||||
|
||||
def format_result_details(self, result, extra=None):
|
||||
if isinstance(result, list):
|
||||
return self.format_log_fields(count=len(result))
|
||||
if not isinstance(result, dict):
|
||||
fields = {"result": "none" if result is None else result}
|
||||
fields.update(self.normalize_result_extra(extra))
|
||||
return self.format_log_fields(**fields)
|
||||
|
||||
model_id = result.get("model_id")
|
||||
version_id = result.get("version_id")
|
||||
ids = (
|
||||
f"{model_id}@{version_id}"
|
||||
if model_id and version_id
|
||||
else model_id or version_id
|
||||
)
|
||||
fields = {
|
||||
"name": result.get("name"),
|
||||
"file": result.get("filename") or result.get("path"),
|
||||
"match": result.get("match_type"),
|
||||
"repo": result.get("repo_id") or result.get("repo"),
|
||||
"ids": ids,
|
||||
"size": self.format_log_size(result.get("size")),
|
||||
"files": result.get("files_count"),
|
||||
}
|
||||
fields.update(self.normalize_result_extra(extra))
|
||||
return self.format_log_fields(**fields)
|
||||
|
||||
def log_search_result(self, source_name, result, extra=None):
|
||||
details = self.format_result_details(result, extra)
|
||||
if result and (not isinstance(result, list) or len(result) > 0):
|
||||
self.logger.info(f"Search [{source_name}] found {details}")
|
||||
else:
|
||||
self.logger.info(f"Search [{source_name}] miss {details}")
|
||||
|
||||
async def search_sources(self, request):
|
||||
"""Execute all requested sources and return a JSON response."""
|
||||
try:
|
||||
data = await request.json()
|
||||
filename = data.get("filename", "")
|
||||
category = data.get("category", "")
|
||||
base_model_context = data.get("base_model_context", "")
|
||||
progress_id = str(data.get("progress_id") or "").strip()
|
||||
progress_source = str(data.get("progress_source") or "").strip()
|
||||
civitai_candidate_limit = self.to_int(
|
||||
data.get("civitai_candidate_limit"),
|
||||
5,
|
||||
)
|
||||
civitai_candidate_limit = max(1, min(civitai_candidate_limit, 20))
|
||||
civarchive_candidate_limit = self.to_int(
|
||||
data.get("civarchive_candidate_limit"),
|
||||
10,
|
||||
)
|
||||
civarchive_candidate_limit = max(
|
||||
1,
|
||||
min(civarchive_candidate_limit, 30),
|
||||
)
|
||||
|
||||
is_urn = self.to_bool(data.get("is_urn", False), False)
|
||||
civitai_key = data.get("civitai_key", "")
|
||||
civitai_session_token = data.get("civitai_session_token", "")
|
||||
hf_token = data.get("hf_token", "")
|
||||
brave_search_api_key = data.get("brave_search_api_key", "")
|
||||
civitai_use_trpc_search = self.to_bool(
|
||||
data.get("civitai_use_trpc_search", True),
|
||||
True,
|
||||
)
|
||||
civitai_use_api_search = self.to_bool(
|
||||
data.get("civitai_use_api_search", True),
|
||||
True,
|
||||
)
|
||||
civitai_use_html_fallback = self.to_bool(
|
||||
data.get("civitai_use_html_fallback", True),
|
||||
True,
|
||||
)
|
||||
hf_use_api_search = self.to_bool(
|
||||
data.get("hf_use_api_search", True),
|
||||
True,
|
||||
)
|
||||
hf_use_comfy_org_fallback = self.to_bool(
|
||||
data.get("hf_use_comfy_org_fallback", True),
|
||||
True,
|
||||
)
|
||||
hf_use_brave_fallback = self.to_bool(
|
||||
data.get("hf_use_brave_fallback", True),
|
||||
True,
|
||||
)
|
||||
|
||||
model_id = data.get("model_id")
|
||||
version_id = data.get("version_id")
|
||||
if not filename and not (is_urn and model_id and version_id):
|
||||
return self.web.json_response(
|
||||
{
|
||||
"error": (
|
||||
"Filename is required for non-URN, or "
|
||||
"model_id+version_id for URN"
|
||||
)
|
||||
},
|
||||
status=400,
|
||||
)
|
||||
|
||||
raw_sources = data.get("sources", ["all"])
|
||||
if isinstance(raw_sources, str):
|
||||
raw_sources = [raw_sources]
|
||||
elif not isinstance(raw_sources, list):
|
||||
raw_sources = ["all"]
|
||||
|
||||
normalized_sources = frozenset(
|
||||
str(source).strip().lower()
|
||||
for source in raw_sources
|
||||
if str(source).strip()
|
||||
)
|
||||
if not normalized_sources:
|
||||
normalized_sources = frozenset({"all"})
|
||||
if "all" in normalized_sources:
|
||||
normalized_sources = frozenset(
|
||||
{
|
||||
"local",
|
||||
"huggingface",
|
||||
"civitai",
|
||||
"civarchive",
|
||||
"lora_manager_archive",
|
||||
}
|
||||
)
|
||||
|
||||
if not progress_source:
|
||||
progress_source = (
|
||||
next(iter(normalized_sources))
|
||||
if len(normalized_sources) == 1
|
||||
else "all"
|
||||
)
|
||||
force_search = self.to_bool(data.get("force_search"), False)
|
||||
search_request = SearchRequest(
|
||||
data=data,
|
||||
filename=filename,
|
||||
category=category,
|
||||
base_model_context=base_model_context,
|
||||
progress_id=progress_id,
|
||||
progress_source=progress_source,
|
||||
civitai_candidate_limit=civitai_candidate_limit,
|
||||
civarchive_candidate_limit=civarchive_candidate_limit,
|
||||
is_urn=is_urn,
|
||||
civitai_key=civitai_key,
|
||||
civitai_session_token=civitai_session_token,
|
||||
hf_token=hf_token,
|
||||
brave_search_api_key=brave_search_api_key,
|
||||
civitai_use_trpc_search=civitai_use_trpc_search,
|
||||
civitai_use_api_search=civitai_use_api_search,
|
||||
civitai_use_html_fallback=civitai_use_html_fallback,
|
||||
hf_use_api_search=hf_use_api_search,
|
||||
hf_use_comfy_org_fallback=hf_use_comfy_org_fallback,
|
||||
hf_use_brave_fallback=hf_use_brave_fallback,
|
||||
force_search=force_search,
|
||||
normalized_sources=normalized_sources,
|
||||
)
|
||||
raise_if_search_cancelled = (
|
||||
self.provider_runner.raise_if_search_cancelled
|
||||
)
|
||||
|
||||
self.search_tracker.update(
|
||||
progress_id,
|
||||
progress_source,
|
||||
"starting",
|
||||
"Preparing search",
|
||||
8,
|
||||
)
|
||||
raise_if_search_cancelled(search_request, progress_source)
|
||||
|
||||
if force_search:
|
||||
self.search_tracker.update(
|
||||
progress_id,
|
||||
progress_source,
|
||||
"cache",
|
||||
"Refreshing search caches",
|
||||
12,
|
||||
)
|
||||
if search_request.search_local:
|
||||
self.reload_popular_databases()
|
||||
self.reload_model_list()
|
||||
if search_request.search_huggingface_source:
|
||||
self.clear_huggingface_search_cache()
|
||||
if search_request.search_civitai_source:
|
||||
self.clear_civitai_search_cache()
|
||||
if search_request.search_civarchive_source:
|
||||
self.clear_civarchive_search_cache()
|
||||
if search_request.search_lora_manager_archive_source:
|
||||
self.clear_lora_manager_archive_search_cache()
|
||||
self.logger.debug(
|
||||
"Force search enabled: cleared cache "
|
||||
+ self.format_log_fields(
|
||||
sources=sorted(normalized_sources),
|
||||
)
|
||||
)
|
||||
raise_if_search_cancelled(search_request, progress_source)
|
||||
|
||||
self.logger.info(
|
||||
f"Search [{','.join(sorted(normalized_sources))}] request "
|
||||
+ self.format_log_fields(
|
||||
file=filename,
|
||||
cat=category,
|
||||
urn=is_urn,
|
||||
ids=(
|
||||
f"{model_id}@{version_id}"
|
||||
if model_id and version_id
|
||||
else model_id or version_id
|
||||
),
|
||||
base=base_model_context,
|
||||
force=force_search,
|
||||
)
|
||||
)
|
||||
|
||||
results = {
|
||||
"popular": None,
|
||||
"model_list": None,
|
||||
"huggingface": None,
|
||||
"civitai": None,
|
||||
"civarchive": None,
|
||||
"lora_manager_archive": None,
|
||||
"local_hash_matches": [],
|
||||
"found": False,
|
||||
"searched_sources": sorted(normalized_sources),
|
||||
"source_errors": {},
|
||||
}
|
||||
|
||||
def iter_result_items(result):
|
||||
if isinstance(result, list):
|
||||
for item in result:
|
||||
if isinstance(item, dict):
|
||||
yield item
|
||||
elif isinstance(result, dict):
|
||||
yield result
|
||||
|
||||
def collect_local_hash_matches(payload):
|
||||
matches = []
|
||||
seen_match_paths = set()
|
||||
seen_hashes = set()
|
||||
for source_key in (
|
||||
"huggingface",
|
||||
"civitai",
|
||||
"civarchive",
|
||||
"lora_manager_archive",
|
||||
"popular",
|
||||
"model_list",
|
||||
):
|
||||
for source_result in iter_result_items(payload.get(source_key)):
|
||||
raise_if_search_cancelled(search_request, source_key)
|
||||
sha256 = self.extract_sha256_from_metadata(source_result)
|
||||
if not sha256 or sha256 in seen_hashes:
|
||||
continue
|
||||
seen_hashes.add(sha256)
|
||||
|
||||
self.search_tracker.update(
|
||||
progress_id,
|
||||
progress_source,
|
||||
"local_hash",
|
||||
"Checking local metadata hashes",
|
||||
94,
|
||||
)
|
||||
try:
|
||||
hash_matches = self.search_local_matches_by_hash(
|
||||
sha256,
|
||||
category=category or None,
|
||||
max_matches=20,
|
||||
force_rescan=force_search,
|
||||
)
|
||||
except Exception as hash_error:
|
||||
self.logger.warning(
|
||||
"Local metadata hash lookup failed "
|
||||
f"for {source_key}:{sha256}: {hash_error}"
|
||||
)
|
||||
continue
|
||||
|
||||
for match in hash_matches:
|
||||
model_path = (
|
||||
match.get("model", {}).get("path")
|
||||
or match.get("path")
|
||||
or ""
|
||||
)
|
||||
path_key = model_path.lower()
|
||||
if path_key and path_key in seen_match_paths:
|
||||
continue
|
||||
if path_key:
|
||||
seen_match_paths.add(path_key)
|
||||
matches.append(
|
||||
{
|
||||
**match,
|
||||
"hash_lookup_source": source_key,
|
||||
"hash_lookup_filename": (
|
||||
source_result.get("filename")
|
||||
or source_result.get("path")
|
||||
or filename
|
||||
),
|
||||
"hash_lookup_sha256": sha256,
|
||||
}
|
||||
)
|
||||
|
||||
if matches:
|
||||
self.logger.info(
|
||||
"Search local hash matches "
|
||||
+ self.format_log_fields(count=len(matches))
|
||||
)
|
||||
return matches
|
||||
|
||||
search_tasks = self.provider_runner.create_search_tasks(
|
||||
search_request
|
||||
)
|
||||
if len(search_tasks) > 1:
|
||||
self.logger.debug(
|
||||
"Search sources async "
|
||||
+ self.format_log_fields(count=len(search_tasks))
|
||||
)
|
||||
self.search_tracker.update(
|
||||
progress_id,
|
||||
progress_source,
|
||||
"running",
|
||||
"Waiting for search sources",
|
||||
18,
|
||||
)
|
||||
raise_if_search_cancelled(search_request, progress_source)
|
||||
|
||||
for source_results, source_found in await self.asyncio.gather(
|
||||
*search_tasks
|
||||
):
|
||||
raise_if_search_cancelled(search_request, progress_source)
|
||||
for source_key, source_result in source_results.items():
|
||||
if source_key == "source_errors":
|
||||
results["source_errors"].update(source_result or {})
|
||||
elif source_result:
|
||||
results[source_key] = source_result
|
||||
if source_found:
|
||||
results["found"] = True
|
||||
|
||||
raise_if_search_cancelled(search_request, progress_source)
|
||||
results["local_hash_matches"] = collect_local_hash_matches(results)
|
||||
raise_if_search_cancelled(search_request, progress_source)
|
||||
|
||||
self.logger.info(
|
||||
f"Search [{','.join(results['searched_sources'])}] done "
|
||||
+ self.format_log_fields(found=results["found"])
|
||||
)
|
||||
self.search_tracker.update(
|
||||
progress_id,
|
||||
progress_source,
|
||||
"completed",
|
||||
"Search complete",
|
||||
100,
|
||||
status="completed",
|
||||
)
|
||||
self.search_cache.stamp_results(
|
||||
results,
|
||||
force_search=force_search,
|
||||
)
|
||||
return self.web.json_response(results)
|
||||
|
||||
except SearchCancelled:
|
||||
self.search_tracker.update(
|
||||
progress_id if "progress_id" in locals() else "",
|
||||
progress_source if "progress_source" in locals() else "",
|
||||
"cancelled",
|
||||
"Cancelled",
|
||||
100,
|
||||
status="cancelled",
|
||||
cancelled=True,
|
||||
)
|
||||
self.logger.info(
|
||||
"Search cancelled "
|
||||
+ self.format_log_fields(
|
||||
source=(
|
||||
progress_source
|
||||
if "progress_source" in locals()
|
||||
else ""
|
||||
),
|
||||
progress_id=(
|
||||
progress_id if "progress_id" in locals() else ""
|
||||
),
|
||||
)
|
||||
)
|
||||
return self.web.json_response(
|
||||
{
|
||||
"cancelled": True,
|
||||
"found": False,
|
||||
"searched_sources": (
|
||||
sorted(normalized_sources)
|
||||
if "normalized_sources" in locals()
|
||||
else []
|
||||
),
|
||||
"source_errors": {},
|
||||
}
|
||||
)
|
||||
|
||||
except Exception as exc:
|
||||
self.search_tracker.update(
|
||||
progress_id if "progress_id" in locals() else "",
|
||||
progress_source if "progress_source" in locals() else "",
|
||||
"error",
|
||||
str(exc),
|
||||
100,
|
||||
status="error",
|
||||
)
|
||||
self.logger.exception(f"Model Resolver search error: {exc}")
|
||||
return self.web.json_response({"error": str(exc)}, status=500)
|
||||
@@ -0,0 +1,619 @@
|
||||
"""Provider-specific search execution and progress handling."""
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
|
||||
class SearchCancelled(BaseException):
|
||||
"""Signal cancellation through provider worker boundaries."""
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SearchRequest:
|
||||
"""Normalized values shared by the provider search tasks."""
|
||||
|
||||
data: dict[str, Any]
|
||||
filename: str
|
||||
category: str
|
||||
base_model_context: str
|
||||
progress_id: str
|
||||
progress_source: str
|
||||
civitai_candidate_limit: int
|
||||
civarchive_candidate_limit: int
|
||||
is_urn: bool
|
||||
civitai_key: str
|
||||
civitai_session_token: str
|
||||
hf_token: str
|
||||
brave_search_api_key: str
|
||||
civitai_use_trpc_search: bool
|
||||
civitai_use_api_search: bool
|
||||
civitai_use_html_fallback: bool
|
||||
hf_use_api_search: bool
|
||||
hf_use_comfy_org_fallback: bool
|
||||
hf_use_brave_fallback: bool
|
||||
force_search: bool
|
||||
normalized_sources: frozenset[str]
|
||||
|
||||
@property
|
||||
def search_local(self):
|
||||
return "local" in self.normalized_sources
|
||||
|
||||
@property
|
||||
def search_huggingface_source(self):
|
||||
return "huggingface" in self.normalized_sources
|
||||
|
||||
@property
|
||||
def search_civitai_source(self):
|
||||
return "civitai" in self.normalized_sources
|
||||
|
||||
@property
|
||||
def search_civarchive_source(self):
|
||||
return "civarchive" in self.normalized_sources
|
||||
|
||||
@property
|
||||
def search_lora_manager_archive_source(self):
|
||||
return "lora_manager_archive" in self.normalized_sources
|
||||
|
||||
|
||||
class SearchProviderRunner:
|
||||
"""Run source searches while keeping provider concerns out of orchestration."""
|
||||
|
||||
def __init__(self, owner):
|
||||
self.owner = owner
|
||||
|
||||
def raise_if_search_cancelled(self, request: SearchRequest, source=""):
|
||||
if self.owner.search_tracker.is_cancelled(request.progress_id):
|
||||
raise SearchCancelled("Search cancelled")
|
||||
|
||||
def make_source_progress_callback(
|
||||
self,
|
||||
request: SearchRequest,
|
||||
source_key,
|
||||
percent_min=None,
|
||||
percent_max=None,
|
||||
):
|
||||
def source_progress_callback(payload):
|
||||
self.raise_if_search_cancelled(request, source_key)
|
||||
if not isinstance(payload, dict):
|
||||
return
|
||||
|
||||
progress_payload = dict(payload)
|
||||
stage = progress_payload.pop("stage", "running")
|
||||
message = progress_payload.pop("message", "Searching...")
|
||||
percent = progress_payload.pop("percent", None)
|
||||
status = progress_payload.pop("status", "running")
|
||||
progress_payload.pop("source", None)
|
||||
|
||||
if (
|
||||
percent is not None
|
||||
and percent_min is not None
|
||||
and percent_max is not None
|
||||
):
|
||||
try:
|
||||
normalized_percent = max(0.0, min(100.0, float(percent)))
|
||||
percent = percent_min + (
|
||||
normalized_percent / 100.0
|
||||
) * (percent_max - percent_min)
|
||||
except (TypeError, ValueError):
|
||||
percent = None
|
||||
|
||||
self.owner.search_tracker.update(
|
||||
request.progress_id,
|
||||
source_key,
|
||||
stage,
|
||||
message,
|
||||
percent,
|
||||
status=status,
|
||||
**progress_payload,
|
||||
)
|
||||
|
||||
return source_progress_callback
|
||||
|
||||
def run_source_search(
|
||||
self,
|
||||
request: SearchRequest,
|
||||
source_key,
|
||||
search_task_fn,
|
||||
initial_stage="query",
|
||||
initial_message=None,
|
||||
initial_percent=30,
|
||||
log_start_fields=None,
|
||||
error_handlers=None,
|
||||
):
|
||||
if initial_message is None:
|
||||
initial_message = f"Querying {source_key.capitalize()}"
|
||||
self.raise_if_search_cancelled(request, source_key)
|
||||
self.owner.search_tracker.update(
|
||||
request.progress_id,
|
||||
source_key,
|
||||
initial_stage,
|
||||
initial_message,
|
||||
initial_percent,
|
||||
)
|
||||
start_fields = log_start_fields or {"file": request.filename}
|
||||
self.owner.logger.info(
|
||||
f"Search [{source_key}] start "
|
||||
+ self.owner.format_log_fields(**start_fields)
|
||||
)
|
||||
try:
|
||||
self.raise_if_search_cancelled(request, source_key)
|
||||
source_results, source_found = search_task_fn()
|
||||
self.raise_if_search_cancelled(request, source_key)
|
||||
done_messages = {
|
||||
"local": "Local database checked",
|
||||
"huggingface": "HuggingFace checked",
|
||||
"civitai": "CivitAI checked",
|
||||
"civarchive": "CivArchive checked",
|
||||
"lora_manager_archive": "LoRA Manager archive checked",
|
||||
}
|
||||
done_msg = done_messages.get(
|
||||
source_key,
|
||||
f"{source_key} checked",
|
||||
)
|
||||
self.owner.search_tracker.update(
|
||||
request.progress_id,
|
||||
source_key,
|
||||
"done",
|
||||
done_msg,
|
||||
92,
|
||||
)
|
||||
return source_results, source_found
|
||||
except Exception as exc:
|
||||
if error_handlers:
|
||||
for exc_type, handler_fn in error_handlers.items():
|
||||
if isinstance(exc, exc_type):
|
||||
return handler_fn(exc)
|
||||
raise
|
||||
|
||||
def mark_any_model_fallback(self, request: SearchRequest, result):
|
||||
if isinstance(result, list):
|
||||
return [
|
||||
self.mark_any_model_fallback(request, item)
|
||||
for item in result
|
||||
]
|
||||
if not isinstance(result, dict):
|
||||
return result
|
||||
|
||||
marked = dict(result)
|
||||
marked["any_model_match"] = True
|
||||
marked["base_model_fallback"] = True
|
||||
marked["requested_base_model"] = request.base_model_context
|
||||
return marked
|
||||
|
||||
def execute_search_with_fallback(
|
||||
self,
|
||||
request: SearchRequest,
|
||||
source_key,
|
||||
search_fn,
|
||||
any_model_label,
|
||||
):
|
||||
self.raise_if_search_cancelled(request, source_key)
|
||||
result = search_fn(
|
||||
request.base_model_context or None,
|
||||
self.make_source_progress_callback(request, source_key),
|
||||
)
|
||||
self.raise_if_search_cancelled(request, source_key)
|
||||
self.owner.log_search_result(source_key, result)
|
||||
|
||||
if not result and request.base_model_context:
|
||||
self.raise_if_search_cancelled(request, source_key)
|
||||
self.owner.search_tracker.update(
|
||||
request.progress_id,
|
||||
source_key,
|
||||
"any_model",
|
||||
f"Retrying {any_model_label} any model",
|
||||
72,
|
||||
)
|
||||
self.owner.logger.info(
|
||||
f"Search [{source_key}] retry any model "
|
||||
+ self.owner.format_log_fields(
|
||||
file=request.filename,
|
||||
cat=request.category,
|
||||
base=request.base_model_context,
|
||||
)
|
||||
)
|
||||
result = search_fn(
|
||||
None,
|
||||
self.make_source_progress_callback(request, source_key, 72, 92),
|
||||
)
|
||||
self.raise_if_search_cancelled(request, source_key)
|
||||
self.owner.log_search_result(f"{source_key}/any_model", result)
|
||||
if result:
|
||||
result = self.mark_any_model_fallback(request, result)
|
||||
return result
|
||||
|
||||
def search_local_sources(self, request: SearchRequest):
|
||||
def task():
|
||||
source_results = {"popular": None, "model_list": None}
|
||||
source_found = False
|
||||
|
||||
popular_info = self.owner.get_popular_model_url(request.filename)
|
||||
self.owner.log_search_result("popular", popular_info)
|
||||
self.owner.search_tracker.update(
|
||||
request.progress_id,
|
||||
"local",
|
||||
"model_list",
|
||||
"Checking local model database",
|
||||
58,
|
||||
)
|
||||
model_list_result = self.owner.search_model_list(request.filename)
|
||||
self.owner.log_search_result(
|
||||
"model_list",
|
||||
model_list_result,
|
||||
{
|
||||
"confidence": (
|
||||
model_list_result.get("confidence")
|
||||
if model_list_result
|
||||
else None
|
||||
)
|
||||
},
|
||||
)
|
||||
if popular_info:
|
||||
popular_result = {
|
||||
"source": "popular",
|
||||
"filename": request.filename,
|
||||
**popular_info,
|
||||
}
|
||||
if (
|
||||
model_list_result
|
||||
and model_list_result.get("filename", "").lower()
|
||||
== request.filename.lower()
|
||||
and model_list_result.get("size")
|
||||
):
|
||||
popular_result["size"] = model_list_result.get("size")
|
||||
source_results["popular"] = popular_result
|
||||
source_found = True
|
||||
|
||||
if model_list_result:
|
||||
confidence = model_list_result.get("confidence", 0)
|
||||
if (request.is_urn and confidence >= 70) or not request.is_urn:
|
||||
source_results["model_list"] = model_list_result
|
||||
source_found = True
|
||||
|
||||
return source_results, source_found
|
||||
|
||||
return self.run_source_search(
|
||||
request,
|
||||
"local",
|
||||
task,
|
||||
initial_stage="popular",
|
||||
initial_message="Checking popular models",
|
||||
initial_percent=28,
|
||||
log_start_fields={
|
||||
"file": request.filename,
|
||||
"cat": request.category,
|
||||
},
|
||||
)
|
||||
|
||||
def search_huggingface_source_task(self, request: SearchRequest):
|
||||
def task():
|
||||
hf_result = self.owner.search_huggingface_for_file(
|
||||
request.filename,
|
||||
token=request.hf_token or None,
|
||||
brave_api_key=request.brave_search_api_key or None,
|
||||
use_api_search=request.hf_use_api_search,
|
||||
use_comfy_org_fallback=request.hf_use_comfy_org_fallback,
|
||||
use_brave_fallback=request.hf_use_brave_fallback,
|
||||
force_refresh=request.force_search,
|
||||
progress_callback=self.make_source_progress_callback(
|
||||
request,
|
||||
"huggingface",
|
||||
),
|
||||
)
|
||||
self.owner.log_search_result("huggingface", hf_result)
|
||||
return {"huggingface": hf_result}, bool(hf_result)
|
||||
|
||||
return self.run_source_search(
|
||||
request,
|
||||
"huggingface",
|
||||
task,
|
||||
initial_stage="query",
|
||||
initial_message="Querying HuggingFace",
|
||||
initial_percent=32,
|
||||
log_start_fields={"file": request.filename},
|
||||
)
|
||||
|
||||
def search_civitai_source_task(self, request: SearchRequest):
|
||||
def task():
|
||||
source_results = {"civitai": None}
|
||||
source_found = False
|
||||
|
||||
if request.is_urn:
|
||||
model_id = request.data.get("model_id")
|
||||
version_id = request.data.get("version_id")
|
||||
|
||||
if model_id and version_id:
|
||||
self.owner.search_tracker.update(
|
||||
request.progress_id,
|
||||
"civitai",
|
||||
"urn",
|
||||
"Resolving CivitAI URN",
|
||||
46,
|
||||
)
|
||||
model_info = self.owner.resolve_urn(model_id, version_id)
|
||||
if model_info:
|
||||
self.owner.search_tracker.update(
|
||||
request.progress_id,
|
||||
"civitai",
|
||||
"file",
|
||||
"Selecting CivitAI file",
|
||||
76,
|
||||
)
|
||||
primary_file = None
|
||||
for file_info in model_info.get("files", []):
|
||||
if file_info.get("name") == model_info.get(
|
||||
"expected_filename"
|
||||
):
|
||||
primary_file = file_info
|
||||
break
|
||||
if primary_file is None:
|
||||
primary_file = (model_info.get("files") or [{}])[0]
|
||||
|
||||
download_url = self.owner.get_civitai_download_url(
|
||||
version_id
|
||||
)
|
||||
source_results["civitai"] = self.owner.build_search_result(
|
||||
"civitai",
|
||||
model_id=model_id,
|
||||
version_id=version_id,
|
||||
name=model_info.get("model_name"),
|
||||
version_name=model_info.get("version_name"),
|
||||
filename=model_info.get("expected_filename"),
|
||||
type=request.category,
|
||||
download_url=download_url,
|
||||
url=(
|
||||
"https://civitai.com/models/"
|
||||
f"{model_id}?modelVersionId={version_id}"
|
||||
),
|
||||
size=primary_file.get("size"),
|
||||
base_model=model_info.get("base_model"),
|
||||
tags=model_info.get("tags", []),
|
||||
sha256=(
|
||||
primary_file.get("sha256")
|
||||
or (primary_file.get("hashes") or {}).get(
|
||||
"SHA256"
|
||||
)
|
||||
or (primary_file.get("hashes") or {}).get(
|
||||
"sha256"
|
||||
)
|
||||
),
|
||||
hashes=primary_file.get("hashes") or {},
|
||||
match_type="exact",
|
||||
confidence=100.0,
|
||||
)
|
||||
self.owner.log_search_result(
|
||||
"civitai/urn",
|
||||
source_results["civitai"],
|
||||
{"files_count": len(model_info.get("files", []))},
|
||||
)
|
||||
source_found = True
|
||||
else:
|
||||
self.owner.log_search_result(
|
||||
"civitai/urn",
|
||||
None,
|
||||
{"model_id": model_id, "version_id": version_id},
|
||||
)
|
||||
elif request.category:
|
||||
self.owner.search_tracker.update(
|
||||
request.progress_id,
|
||||
"civitai",
|
||||
"fallback",
|
||||
"Searching CivitAI fallback",
|
||||
58,
|
||||
)
|
||||
self.owner.logger.info(
|
||||
"Search [civitai] URN ids missing; falling back"
|
||||
)
|
||||
civitai_results = self.owner.search_civitai(
|
||||
request.filename,
|
||||
model_type=request.category,
|
||||
)
|
||||
self.owner.log_search_result(
|
||||
"civitai/fallback",
|
||||
civitai_results[0] if civitai_results else None,
|
||||
{"results_count": len(civitai_results)},
|
||||
)
|
||||
if civitai_results:
|
||||
first_result = civitai_results[0]
|
||||
source_results["civitai"] = self.owner.build_search_result(
|
||||
"civitai",
|
||||
model_id=first_result.get("model_id"),
|
||||
version_id=first_result.get("version_id"),
|
||||
name=first_result.get("name"),
|
||||
filename=first_result.get("filename"),
|
||||
type=first_result.get("type"),
|
||||
download_url=first_result.get("download_url"),
|
||||
url=first_result.get("url"),
|
||||
size=first_result.get("size"),
|
||||
base_model=first_result.get("base_model"),
|
||||
tags=first_result.get("tags", []),
|
||||
)
|
||||
source_found = True
|
||||
else:
|
||||
civitai_result = self.execute_search_with_fallback(
|
||||
request,
|
||||
"civitai",
|
||||
lambda base_ctx, callback: self.owner.search_civitai_for_file(
|
||||
request.filename,
|
||||
api_key=request.civitai_key or None,
|
||||
model_type=request.category,
|
||||
base_model_context=base_ctx,
|
||||
session_token=request.civitai_session_token or None,
|
||||
candidate_limit=request.civitai_candidate_limit,
|
||||
use_trpc_search=request.civitai_use_trpc_search,
|
||||
use_api_search=request.civitai_use_api_search,
|
||||
use_html_fallback=request.civitai_use_html_fallback,
|
||||
progress_callback=callback,
|
||||
),
|
||||
"CivitAI",
|
||||
)
|
||||
if civitai_result:
|
||||
source_results["civitai"] = civitai_result
|
||||
source_found = True
|
||||
|
||||
return source_results, source_found
|
||||
|
||||
return self.run_source_search(
|
||||
request,
|
||||
"civitai",
|
||||
task,
|
||||
initial_stage="query",
|
||||
initial_message="Querying CivitAI",
|
||||
initial_percent=30,
|
||||
log_start_fields={
|
||||
"file": request.filename,
|
||||
"cat": request.category,
|
||||
"urn": request.is_urn,
|
||||
},
|
||||
)
|
||||
|
||||
def search_civarchive_source_task(self, request: SearchRequest):
|
||||
def task():
|
||||
source_results = {"civarchive": None}
|
||||
source_found = False
|
||||
|
||||
if request.is_urn:
|
||||
model_id = request.data.get("model_id")
|
||||
version_id = request.data.get("version_id")
|
||||
if model_id and version_id:
|
||||
self.owner.search_tracker.update(
|
||||
request.progress_id,
|
||||
"civarchive",
|
||||
"urn",
|
||||
"Resolving CivArchive version",
|
||||
50,
|
||||
)
|
||||
civarchive_result = self.owner.resolve_civarchive_model_version(
|
||||
model_id,
|
||||
version_id,
|
||||
query=request.filename,
|
||||
)
|
||||
self.owner.log_search_result(
|
||||
"civarchive/urn",
|
||||
civarchive_result,
|
||||
{"model_id": model_id, "version_id": version_id},
|
||||
)
|
||||
if civarchive_result:
|
||||
source_results["civarchive"] = civarchive_result
|
||||
source_found = True
|
||||
else:
|
||||
self.owner.log_search_result(
|
||||
"civarchive/urn",
|
||||
None,
|
||||
{"model_id": model_id, "version_id": version_id},
|
||||
)
|
||||
else:
|
||||
civarchive_result = self.execute_search_with_fallback(
|
||||
request,
|
||||
"civarchive",
|
||||
lambda base_ctx, callback: self.owner.search_civarchive_for_file(
|
||||
request.filename,
|
||||
model_type=request.category,
|
||||
base_model_context=base_ctx,
|
||||
limit=request.civarchive_candidate_limit,
|
||||
progress_callback=callback,
|
||||
),
|
||||
"CivArchive",
|
||||
)
|
||||
if civarchive_result:
|
||||
source_results["civarchive"] = civarchive_result
|
||||
source_found = True
|
||||
|
||||
return source_results, source_found
|
||||
|
||||
def handle_civarchive_error(error):
|
||||
error_message = f"CivArchive search failed: {error}"
|
||||
self.owner.logger.warning(error_message)
|
||||
self.owner.search_tracker.update(
|
||||
request.progress_id,
|
||||
"civarchive",
|
||||
"error",
|
||||
error_message,
|
||||
100,
|
||||
status="error",
|
||||
)
|
||||
return (
|
||||
{
|
||||
"civarchive": None,
|
||||
"source_errors": {"civarchive": error_message},
|
||||
},
|
||||
False,
|
||||
)
|
||||
|
||||
return self.run_source_search(
|
||||
request,
|
||||
"civarchive",
|
||||
task,
|
||||
initial_stage="query",
|
||||
initial_message="Querying CivArchive",
|
||||
initial_percent=30,
|
||||
log_start_fields={
|
||||
"file": request.filename,
|
||||
"cat": request.category,
|
||||
"urn": request.is_urn,
|
||||
},
|
||||
error_handlers={
|
||||
self.owner.CivArchiveSearchError: handle_civarchive_error,
|
||||
},
|
||||
)
|
||||
|
||||
def search_lora_manager_archive_source_task(self, request: SearchRequest):
|
||||
def task():
|
||||
result = self.execute_search_with_fallback(
|
||||
request,
|
||||
"lora_manager_archive",
|
||||
lambda base_ctx, callback: self.owner.search_lora_manager_archive_for_file(
|
||||
request.filename,
|
||||
model_type=request.category,
|
||||
base_model_context=base_ctx,
|
||||
progress_callback=callback,
|
||||
),
|
||||
"LoRA archive",
|
||||
)
|
||||
return {"lora_manager_archive": result}, bool(result)
|
||||
|
||||
return self.run_source_search(
|
||||
request,
|
||||
"lora_manager_archive",
|
||||
task,
|
||||
initial_stage="query",
|
||||
initial_message="Searching LoRA Manager archive",
|
||||
initial_percent=36,
|
||||
log_start_fields={"file": request.filename, "cat": request.category},
|
||||
)
|
||||
|
||||
def create_search_tasks(self, request: SearchRequest):
|
||||
tasks = []
|
||||
if request.search_local:
|
||||
tasks.append(self.owner.asyncio.to_thread(self.search_local_sources, request))
|
||||
if request.search_huggingface_source:
|
||||
tasks.append(
|
||||
self.owner.asyncio.to_thread(
|
||||
self.search_huggingface_source_task,
|
||||
request,
|
||||
)
|
||||
)
|
||||
if request.search_civitai_source:
|
||||
tasks.append(
|
||||
self.owner.asyncio.to_thread(
|
||||
self.search_civitai_source_task,
|
||||
request,
|
||||
)
|
||||
)
|
||||
if request.search_civarchive_source and (
|
||||
request.filename or (request.is_urn and request.data.get("model_id") and request.data.get("version_id"))
|
||||
):
|
||||
tasks.append(
|
||||
self.owner.asyncio.to_thread(
|
||||
self.search_civarchive_source_task,
|
||||
request,
|
||||
)
|
||||
)
|
||||
if request.search_lora_manager_archive_source and request.filename:
|
||||
tasks.append(
|
||||
self.owner.asyncio.to_thread(
|
||||
self.search_lora_manager_archive_source_task,
|
||||
request,
|
||||
)
|
||||
)
|
||||
return tasks
|
||||
+6
-1086
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,100 @@
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from core.services.search_cache import SearchResultCache
|
||||
from core.services.search_providers import SearchCancelled, SearchProviderRunner
|
||||
|
||||
|
||||
def _request(**overrides):
|
||||
values = {
|
||||
"filename": "model.safetensors",
|
||||
"category": "checkpoints",
|
||||
"base_model_context": "SDXL",
|
||||
"progress_id": "search-1",
|
||||
"progress_source": "civitai",
|
||||
"force_search": False,
|
||||
}
|
||||
values.update(overrides)
|
||||
return SimpleNamespace(**values)
|
||||
|
||||
|
||||
def test_search_result_cache_reuses_timestamp_for_same_result():
|
||||
timestamps = {}
|
||||
cache = SearchResultCache(timestamps)
|
||||
|
||||
first = cache.stamp_results(
|
||||
{"civitai": {"name": "Model", "download_url": "https://example.test/model"}},
|
||||
force_search=False,
|
||||
)
|
||||
second = cache.stamp_results(
|
||||
{"civitai": {"name": "Model", "download_url": "https://example.test/model"}},
|
||||
force_search=False,
|
||||
)
|
||||
|
||||
assert first["civitai"]["searched_at"] == second["civitai"]["searched_at"]
|
||||
assert len(timestamps) == 1
|
||||
|
||||
|
||||
def test_search_result_cache_stamps_nested_results_and_preserves_empty_values():
|
||||
cache = SearchResultCache({})
|
||||
payload = {
|
||||
"popular": [
|
||||
{"filename": "one.safetensors", "url": "https://example.test/one"},
|
||||
{"filename": "two.safetensors", "url": "https://example.test/two"},
|
||||
],
|
||||
"civitai": None,
|
||||
"source_errors": {"civitai": "unavailable"},
|
||||
}
|
||||
|
||||
stamped = cache.stamp_results(payload, force_search=False)
|
||||
|
||||
assert all(item["searched_at"] for item in stamped["popular"])
|
||||
assert stamped["civitai"] is None
|
||||
assert stamped["source_errors"] == {"civitai": "unavailable"}
|
||||
|
||||
|
||||
def test_search_provider_runner_reports_progress_and_result():
|
||||
owner = SimpleNamespace(
|
||||
logger=MagicMock(),
|
||||
search_tracker=MagicMock(),
|
||||
format_log_fields=MagicMock(return_value="file=model.safetensors"),
|
||||
)
|
||||
owner.search_tracker.is_cancelled.return_value = False
|
||||
runner = SearchProviderRunner(owner)
|
||||
|
||||
result = runner.run_source_search(
|
||||
_request(),
|
||||
"civitai",
|
||||
lambda: ({"civitai": {"name": "Model"}}, True),
|
||||
initial_stage="query",
|
||||
initial_message="Querying CivitAI",
|
||||
initial_percent=30,
|
||||
)
|
||||
|
||||
assert result == ({"civitai": {"name": "Model"}}, True)
|
||||
progress_calls = owner.search_tracker.update.call_args_list
|
||||
assert progress_calls[0].args[:5] == (
|
||||
"search-1",
|
||||
"civitai",
|
||||
"query",
|
||||
"Querying CivitAI",
|
||||
30,
|
||||
)
|
||||
assert progress_calls[-1].args[:5] == (
|
||||
"search-1",
|
||||
"civitai",
|
||||
"done",
|
||||
"CivitAI checked",
|
||||
92,
|
||||
)
|
||||
|
||||
|
||||
def test_search_provider_runner_raises_cancellation_before_provider_call():
|
||||
owner = SimpleNamespace(search_tracker=MagicMock())
|
||||
owner.search_tracker.is_cancelled.return_value = True
|
||||
runner = SearchProviderRunner(owner)
|
||||
|
||||
with pytest.raises(SearchCancelled):
|
||||
runner.raise_if_search_cancelled(_request(), "civitai")
|
||||
Reference in New Issue
Block a user