""" Embedding service (FastAPI). API (simple list-in, list-out; aligned by index): - POST /embed/text body: ["text1", "text2", ...] -> [[...], ...] - POST /embed/image body: ["url_or_path1", ...] -> [[...], ...] """ import logging import os import threading from typing import Any, Dict, List, Optional import numpy as np from fastapi import FastAPI, HTTPException from embeddings.config import CONFIG from embeddings.protocols import ImageEncoderProtocol from config.services_config import get_embedding_backend_config logger = logging.getLogger(__name__) app = FastAPI(title="saas-search Embedding Service", version="1.0.0") # Models are loaded at startup, not lazily _text_model: Optional[Any] = None _image_model: Optional[ImageEncoderProtocol] = None _text_backend_name: str = "" open_text_model = True open_image_model = True # Enable image embedding when using clip-as-service _text_encode_lock = threading.Lock() _image_encode_lock = threading.Lock() @app.on_event("startup") def load_models(): """Load models at service startup to avoid first-request latency.""" global _text_model, _image_model, _text_backend_name logger.info("Loading embedding models at startup...") # Load text model if open_text_model: try: backend_name, backend_cfg = get_embedding_backend_config() _text_backend_name = backend_name if backend_name == "tei": from embeddings.tei_model import TEITextModel base_url = ( os.getenv("TEI_BASE_URL") or backend_cfg.get("base_url") or CONFIG.TEI_BASE_URL ) timeout_sec = int( os.getenv("TEI_TIMEOUT_SEC") or backend_cfg.get("timeout_sec") or CONFIG.TEI_TIMEOUT_SEC ) logger.info("Loading text backend: tei (base_url=%s)", base_url) _text_model = TEITextModel( base_url=str(base_url), timeout_sec=timeout_sec, ) elif backend_name == "local_st": from embeddings.qwen3_model import Qwen3TextModel model_id = ( os.getenv("TEXT_MODEL_ID") or backend_cfg.get("model_id") or CONFIG.TEXT_MODEL_ID ) logger.info("Loading text backend: local_st (model=%s)", model_id) _text_model = Qwen3TextModel(model_id=str(model_id)) else: raise ValueError( f"Unsupported embedding backend: {backend_name}. " "Supported: tei, local_st" ) logger.info("Text backend loaded successfully: %s", _text_backend_name) except Exception as e: logger.error(f"Failed to load text model: {e}", exc_info=True) raise # Load image model: clip-as-service (recommended) or local CN-CLIP if open_image_model: try: if CONFIG.USE_CLIP_AS_SERVICE: from embeddings.clip_as_service_encoder import ClipAsServiceImageEncoder logger.info(f"Loading image encoder via clip-as-service: {CONFIG.CLIP_AS_SERVICE_SERVER}") _image_model = ClipAsServiceImageEncoder( server=CONFIG.CLIP_AS_SERVICE_SERVER, batch_size=CONFIG.IMAGE_BATCH_SIZE, ) logger.info("Image model (clip-as-service) loaded successfully") else: from embeddings.clip_model import ClipImageModel logger.info(f"Loading local image model: {CONFIG.IMAGE_MODEL_NAME} (device: {CONFIG.IMAGE_DEVICE})") _image_model = ClipImageModel( model_name=CONFIG.IMAGE_MODEL_NAME, device=CONFIG.IMAGE_DEVICE, ) logger.info("Image model (local CN-CLIP) loaded successfully") except Exception as e: logger.error("Failed to load image model: %s", e, exc_info=True) raise logger.info("All embedding models loaded successfully, service ready") def _as_list(embedding: Optional[np.ndarray]) -> Optional[List[float]]: if embedding is None: return None if not isinstance(embedding, np.ndarray): embedding = np.array(embedding, dtype=np.float32) if embedding.ndim != 1: embedding = embedding.reshape(-1) return embedding.astype(np.float32).tolist() @app.get("/health") def health() -> Dict[str, Any]: """Health check endpoint. Returns status and model loading state.""" return { "status": "ok", "text_model_loaded": _text_model is not None, "text_backend": _text_backend_name, "image_model_loaded": _image_model is not None, } @app.post("/embed/text") def embed_text(texts: List[str]) -> List[Optional[List[float]]]: if _text_model is None: raise RuntimeError("Text model not loaded") normalized: List[str] = [] for i, t in enumerate(texts): if not isinstance(t, str): raise HTTPException(status_code=400, detail=f"Invalid text at index {i}: must be string") s = t.strip() if not s: raise HTTPException(status_code=400, detail=f"Invalid text at index {i}: empty string") normalized.append(s) try: with _text_encode_lock: embs = _text_model.encode_batch( normalized, batch_size=int(CONFIG.TEXT_BATCH_SIZE), device=CONFIG.TEXT_DEVICE, normalize_embeddings=bool(CONFIG.TEXT_NORMALIZE_EMBEDDINGS), ) except Exception as e: logger.error("Text embedding backend failure: %s", e, exc_info=True) raise HTTPException( status_code=502, detail=f"Text embedding backend failure: {e}", ) from e if embs is None or len(embs) != len(normalized): raise RuntimeError( f"Text model response length mismatch: expected {len(normalized)}, " f"got {0 if embs is None else len(embs)}" ) out: List[Optional[List[float]]] = [] for i, emb in enumerate(embs): vec = _as_list(emb) if vec is None: raise RuntimeError(f"Text model returned empty embedding for index {i}") out.append(vec) return out @app.post("/embed/image") def embed_image(images: List[str]) -> List[Optional[List[float]]]: if _image_model is None: raise RuntimeError("Image model not loaded") urls: List[str] = [] for i, url_or_path in enumerate(images): if not isinstance(url_or_path, str): raise HTTPException(status_code=400, detail=f"Invalid image at index {i}: must be string URL/path") s = url_or_path.strip() if not s: raise HTTPException(status_code=400, detail=f"Invalid image at index {i}: empty URL/path") urls.append(s) with _image_encode_lock: vectors = _image_model.encode_image_urls(urls, batch_size=CONFIG.IMAGE_BATCH_SIZE) if vectors is None or len(vectors) != len(urls): raise RuntimeError( f"Image model response length mismatch: expected {len(urls)}, " f"got {0 if vectors is None else len(vectors)}" ) out: List[Optional[List[float]]] = [] for i, vec in enumerate(vectors): out_vec = _as_list(vec) if out_vec is None: raise RuntimeError(f"Image model returned empty embedding for index {i}") out.append(out_vec) return out