"""Text embedding client for the local embedding HTTP service.""" import logging import os import pickle from datetime import timedelta from typing import Any, List, Optional, Union import numpy as np import redis import requests logger = logging.getLogger(__name__) from config.services_config import get_embedding_base_url # Try to import REDIS_CONFIG, but allow import to fail from config.env_config import REDIS_CONFIG class TextEmbeddingEncoder: """ Text embedding encoder using network service. """ def __init__(self, service_url: Optional[str] = None): resolved_url = service_url or os.getenv("EMBEDDING_SERVICE_URL") or get_embedding_base_url() self.service_url = str(resolved_url).rstrip("/") self.endpoint = f"{self.service_url}/embed/text" self.expire_time = timedelta(days=REDIS_CONFIG.get("cache_expire_days", 180)) self.cache_prefix = str(REDIS_CONFIG.get("embedding_cache_prefix", "embedding")).strip() or "embedding" logger.info("Creating TextEmbeddingEncoder instance with service URL: %s", self.service_url) try: self.redis_client = redis.Redis( host=REDIS_CONFIG.get("host", "localhost"), port=REDIS_CONFIG.get("port", 6479), password=REDIS_CONFIG.get("password"), decode_responses=False, socket_timeout=REDIS_CONFIG.get("socket_timeout", 1), socket_connect_timeout=REDIS_CONFIG.get("socket_connect_timeout", 1), retry_on_timeout=REDIS_CONFIG.get("retry_on_timeout", False), health_check_interval=10, ) self.redis_client.ping() logger.info("Redis cache initialized for embeddings") except Exception as e: logger.warning("Failed to initialize Redis cache for embeddings: %s, continuing without cache", e) self.redis_client = None def _call_service(self, request_data: List[str], normalize_embeddings: bool = True) -> List[Any]: """ Call the embedding service API. Args: request_data: List of texts Returns: List of embeddings (list[float]) or nulls (None), aligned to input order """ try: response = requests.post( self.endpoint, params={"normalize": "true" if normalize_embeddings else "false"}, json=request_data, timeout=60 ) response.raise_for_status() return response.json() except requests.exceptions.RequestException as e: logger.error(f"TextEmbeddingEncoder service request failed: {e}", exc_info=True) raise def encode( self, sentences: Union[str, List[str]], normalize_embeddings: bool = True, device: str = 'cpu', batch_size: int = 32 ) -> np.ndarray: """ Encode text into embeddings via network service with Redis caching. Args: sentences: Single string or list of strings to encode normalize_embeddings: Whether to request normalized embeddings from service device: Device parameter ignored for service compatibility batch_size: Batch size for processing (used for service requests) Returns: numpy array of dtype=object，元素均为有效 np.ndarray 向量。 若任一输入无法生成向量，将直接抛出异常。 """ # Convert single string to list if isinstance(sentences, str): sentences = [sentences] # Check cache first uncached_indices: List[int] = [] uncached_texts: List[str] = [] embeddings: List[Optional[np.ndarray]] = [None] * len(sentences) for i, text in enumerate(sentences): cached = self._get_cached_embedding(text) if cached is not None: embeddings[i] = cached else: uncached_indices.append(i) uncached_texts.append(text) # Prepare request data for uncached texts (after cache check) request_data = list(uncached_texts) # If there are uncached texts, call service if uncached_texts: response_data = self._call_service(request_data, normalize_embeddings=normalize_embeddings) # Process response for i, text in enumerate(uncached_texts): original_idx = uncached_indices[i] if response_data and i < len(response_data): embedding = response_data[i] else: embedding = None if embedding is not None: embedding_array = np.array(embedding, dtype=np.float32) if self._is_valid_embedding(embedding_array): embeddings[original_idx] = embedding_array self._set_cached_embedding(text, embedding_array, normalize_embeddings) else: raise ValueError( f"Invalid embedding returned from service for text index {original_idx}" ) else: raise ValueError(f"No embedding found for text index {original_idx}: {text[:50]}...") # 返回 numpy 数组（dtype=object），元素均为有效 np.ndarray 向量 return np.array(embeddings, dtype=object) def _is_valid_embedding(self, embedding: np.ndarray) -> bool: """ Check if embedding is valid (not None, correct shape, no NaN/Inf). Args: embedding: Embedding array to validate Returns: True if valid, False otherwise """ if embedding is None: return False if not isinstance(embedding, np.ndarray): return False if embedding.size == 0: return False # Check for NaN or Inf values if not np.isfinite(embedding).all(): return False return True def _get_cached_embedding( self, query: str ) -> Optional[np.ndarray]: """Get embedding from cache if exists (with sliding expiration)""" if not self.redis_client: return None try: cache_key = f"{self.cache_prefix}:{query}" cached_data = self.redis_client.get(cache_key) if cached_data: embedding = pickle.loads(cached_data) # Validate cached embedding - if invalid, ignore cache and return None if self._is_valid_embedding(embedding): logger.debug(f"Cache hit for embedding: {query}") # Update expiration time on access (sliding expiration) self.redis_client.expire(cache_key, self.expire_time) return embedding else: logger.warning( f"Invalid embedding found in cache (contains NaN/Inf or invalid shape), " f"ignoring cache for query: {query[:50]}..." ) # Delete invalid cache entry try: self.redis_client.delete(cache_key) except Exception as e: logger.debug(f"Failed to delete invalid cache entry: {e}") return None return None except Exception as e: logger.error(f"Error retrieving embedding from cache: {e}") return None def _set_cached_embedding( self, query: str, embedding: np.ndarray, normalize_embeddings: bool = True, ) -> bool: """Store embedding in cache""" if not self.redis_client: return False try: cache_key = f"{self.cache_prefix}:{query}" serialized_data = pickle.dumps(embedding) self.redis_client.setex( cache_key, self.expire_time, serialized_data ) logger.debug(f"Successfully cached embedding for query: {query}") return True except (redis.exceptions.BusyLoadingError, redis.exceptions.ConnectionError, redis.exceptions.TimeoutError, redis.exceptions.RedisError) as e: logger.warning(f"Redis error storing embedding in cache: {e}") return False except Exception as e: logger.error(f"Error storing embedding in cache: {e}") return False