from __future__ import annotations import logging from pathlib import Path from typing import Any import gradio as gr from PIL import Image from src.config import ALL_CATEGORIES_LABEL, ConfigurationError, METADATA_CSV from src.dataset import DatasetError, get_categories, load_metadata from src.favorites import add_result_to_favorites, prepare_favorites_export from src.runtime import AppRuntime from src.ui.constants import IMAGE_QUERY_EMPTY_PREVIEW, TEXT_QUERY_EMPTY_PREVIEW from src.ui.formatters import ( error_overview, favorite_gallery, image_query_preview, initial_overview, result_gallery, result_to_state, search_overview, text_query_preview, ) from src.ui.types import FavoriteState, GalleryItem, ResultState logger = logging.getLogger(__name__) def category_to_backend(category: str | None) -> str: return category or ALL_CATEGORIES_LABEL def current_categories() -> list[str]: try: categories = get_categories(METADATA_CSV) except Exception as exc: logger.warning("Could not load categories: %s", exc) return [ALL_CATEGORIES_LABEL] return categories or [ALL_CATEGORIES_LABEL] def metadata_status() -> tuple[bool, str]: if not METADATA_CSV.exists(): return ( False, "metadata.csv was not found. Put images under data/images// and run `python scripts/build_metadata.py`.", ) try: metadata = load_metadata(METADATA_CSV) except Exception as exc: return False, f"metadata.csv could not be loaded: {exc}" if metadata.empty: return ( False, "metadata.csv does not contain any images. Add images and run `python scripts/build_metadata.py`.", ) return True, "" def configuration_message(exc: Exception) -> str: text = str(exc) if "UPSTASH_VECTOR" in text or "Upstash Vector" in text: return "Upstash Vector configuration was not detected. Check UPSTASH_VECTOR_REST_URL and UPSTASH_VECTOR_REST_TOKEN in .env." return f"Configuration could not be loaded: {text}" def dataset_message(exc: Exception) -> str: text = str(exc) if "metadata.csv" in text: return f"Dataset configuration error: {text}" return f"Dataset could not be read: {text}" class AppCallbacks: def __init__(self, runtime: AppRuntime) -> None: self.runtime = runtime @staticmethod def preview_text_query(text: str) -> str: return text_query_preview(text) @staticmethod def preview_image_query(image: Image.Image | None) -> str: return image_query_preview(image is not None) @staticmethod def select_text_mode() -> str: return "text" @staticmethod def select_image_mode() -> str: return "image" def run_search( self, query_mode: str, text_query: str, image_query: Image.Image | None, top_k: int, category_filter: str, min_similarity: float, ) -> tuple[str, list[GalleryItem], str | None, list[ResultState]]: mode = query_mode if query_mode in {"text", "image"} else "text" text = (text_query or "").strip() category = category_to_backend(category_filter) threshold = float(min_similarity or 0) if mode == "text": if not text: return ( error_overview( "Search cannot be completed", "Enter an image description before running text search.", ), [], None, [], ) current_query = text elif image_query is not None: current_query = "Uploaded image" else: return ( error_overview( "Search cannot be completed", "Upload a query image before running image search.", ), [], None, [], ) ready, message = metadata_status() if not ready: return error_overview("Dataset error", message), [], None, [] try: service = self.runtime.get_search_service() if mode == "text": results = service.search_by_text(text, int(top_k), category) else: assert image_query is not None results = service.search_by_image(image_query, int(top_k), category) except ConfigurationError as exc: self.runtime.reset_search_service() return error_overview("Configuration error", configuration_message(exc)), [], None, [] except (DatasetError, FileNotFoundError, ValueError) as exc: self.runtime.reset_search_service() return error_overview("Search cannot be completed", dataset_message(exc)), [], None, [] except Exception: self.runtime.reset_search_service() logger.exception("Search failed") preload_note = ( " The CLIP model also failed to preload earlier." if self.runtime.preload_error else "" ) return ( error_overview( "Search cannot be completed", f"An unexpected error occurred during search. Check the application logs for details.{preload_note}", ), [], None, [], ) filtered_results = [ result_to_state(result) for result in results if float(result.score) >= threshold ] empty_message = None if not filtered_results: empty_message = ( "No images matched the current filters. Try changing the query, category, or minimum similarity." ) return ( search_overview( current_query=current_query, returned_results=len(filtered_results), message=empty_message, ), result_gallery(filtered_results), None, filtered_results, ) @staticmethod def select_result_from_gallery( current_results: list[ResultState] | None, evt: gr.SelectData, ) -> str | None: results = list(current_results or []) index: Any = evt.index if isinstance(index, (list, tuple)): index = index[0] if index else None try: selected_index = int(index) except (TypeError, ValueError): return None if 0 <= selected_index < len(results): return str(results[selected_index].get("id") or "") return None @staticmethod def add_selected_to_favorites( selected_result_id: str | None, current_results: list[ResultState] | None, favorites: list[FavoriteState] | None, ) -> tuple[list[FavoriteState], list[GalleryItem], str, Path | None]: updated_favorites, message = add_result_to_favorites( selected_result_id, current_results, favorites, ) return ( updated_favorites, favorite_gallery(updated_favorites), message, prepare_favorites_export(updated_favorites), ) @staticmethod def clear_search( top_k: int, category_filter: str, min_similarity: float, ) -> tuple[ str, None, str, str, str, list[GalleryItem], None, list[ResultState], ]: return ( "", None, TEXT_QUERY_EMPTY_PREVIEW, IMAGE_QUERY_EMPTY_PREVIEW, initial_overview(top_k, category_filter, min_similarity), [], None, [], )