# -*- coding: utf-8 -*- """ Сервис генерации эмбеддингов через GigaChat API с кэшированием в Redis. Поддерживает одиночные запросы с кэшированием (Redis + локальный кэш) и пакетную обработку. """ import asyncio import logging import json import hashlib import time from functools import lru_cache from typing import List, Optional, Any from gigachat import GigaChat from gigachat.exceptions import AuthenticationError, RateLimitError import redis.asyncio as aioredis logger = logging.getLogger(__name__) try: from fastembed import SparseTextEmbedding HAS_SPARSE = True except ImportError: HAS_SPARSE = False logging.warning("fastembed не установлен, sparse-эмбеддинги отключены") class EmbeddingService: def __init__( self, api_key: str, model: str = "GigaChat-Embeddings", timeout: int = 30, cache_size: int = 4096, verify_ssl: bool = False, redis_url: Optional[str] = None ): self.api_key = api_key self.model = model self.timeout = timeout self.cache_size = cache_size self.verify_ssl = verify_ssl self._giga = None # Локальный кэш (fallback, если Redis недоступен) self._local_cache = {} self._local_cache_max = cache_size # Redis клиент self.redis = None if redis_url: try: self.redis = aioredis.from_url(redis_url, decode_responses=True) logger.info("Redis подключён для кэширования эмбеддингов") except Exception as e: logger.warning(f"Не удалось подключиться к Redis: {e}, используем только локальный кэш") # Синхронный кэш для fallback self.embed_sync = lru_cache(maxsize=cache_size)(self._embed_sync_impl) self._sparse_model = None def _get_client(self) -> GigaChat: if self._giga is None: self._giga = GigaChat( credentials=self.api_key, model=self.model, verify_ssl_certs=self.verify_ssl, timeout=self.timeout ) logger.debug(f"Клиент GigaChat создан (модель {self.model}, timeout {self.timeout}s)") return self._giga def _get_sparse_model(self): if not HAS_SPARSE: return None if self._sparse_model is None: self._sparse_model = SparseTextEmbedding(model_name="Qdrant/bm25") return self._sparse_model # ---------- Синхронное получение эмбеддинга (с обрезкой) ---------- def _embed_sync_impl(self, text: str) -> Optional[List[float]]: if not text or not text.strip(): logger.warning("Попытка получить эмбеддинг пустого текста") return None MAX_CHARS = 800 if len(text) > MAX_CHARS: text = text[:MAX_CHARS] logger.debug(f"Текст обрезан до {MAX_CHARS} символов для эмбеддинга") client = self._get_client() try: try: response = client.embeddings(input=[text]) except TypeError: response = client.embeddings([text]) if response.data and len(response.data) > 0: logger.debug(f"Эмбеддинг получен, размерность {len(response.data[0].embedding)}") return response.data[0].embedding else: logger.error(f"Пустой ответ эмбеддингов для текста: {text[:100]}") return None except (AuthenticationError, RateLimitError) as e: logger.error(f"Ошибка авторизации/лимитов GigaChat: {e}") raise except Exception as e: logger.error(f"Ошибка GigaChat при эмбеддинге: {e}") raise async def embed(self, text: str, retries: int = 3, delay: float = 1.0) -> Optional[List[float]]: if not isinstance(text, str): logger.error(f"embed получил не строку, а {type(text)}: {text!r}") return None # 1. Проверка в Redis if self.redis: try: key = f"emb:{hashlib.sha256(text.encode()).hexdigest()}" cached = await self.redis.get(key) if cached: return json.loads(cached) except Exception as e: logger.warning(f"Redis недоступен, используем локальный кэш: {e}") # 2. Локальный кэш if text in self._local_cache: return self._local_cache[text] # 3. Вызов API for attempt in range(retries): try: loop = asyncio.get_running_loop() result = await loop.run_in_executor(None, self.embed_sync, text) if result is not None: # Сохраняем в Redis if self.redis: try: key = f"emb:{hashlib.sha256(text.encode()).hexdigest()}" await self.redis.setex(key, 3600, json.dumps(result)) except: pass # Сохраняем в локальный кэш if len(self._local_cache) >= self._local_cache_max: # Удаляем первый элемент (не самый эффективный, но простой) self._local_cache.pop(next(iter(self._local_cache))) self._local_cache[text] = result return result else: raise ValueError("Эмбеддинг вернул None") except Exception as e: wait = delay * (2 ** attempt) logger.warning( f"Попытка {attempt+1}/{retries} получить эмбеддинг не удалась: {e}. " f"Повтор через {wait:.1f} сек." ) await asyncio.sleep(wait) logger.error(f"Не удалось получить эмбеддинг после {retries} попыток") return None # ---------- Пакетная обработка (для индексации) ---------- def _embed_batch_sync(self, texts: List[str]) -> Optional[List[List[float]]]: if not texts: return [] MAX_CHARS = 800 truncated_texts = [t[:MAX_CHARS] if t and len(t) > MAX_CHARS else (t or "") for t in texts] client = self._get_client() for attempt in range(5): try: try: response = client.embeddings(input=truncated_texts) except TypeError: response = client.embeddings(truncated_texts) if response.data and len(response.data) == len(truncated_texts): sorted_embeddings = [None] * len(truncated_texts) for item in response.data: idx = item.index if idx < len(sorted_embeddings): sorted_embeddings[idx] = item.embedding return sorted_embeddings else: logger.error(f"Неполный ответ эмбеддингов: ожидалось {len(truncated_texts)}, получено {len(response.data) if response.data else 0}") return None except RateLimitError as e: if attempt < 4: wait = 2 ** attempt logger.warning(f"Rate limit, повтор через {wait} секунд") time.sleep(wait) else: logger.error("Превышено количество попыток при 429") return None except Exception as e: logger.error(f"Ошибка пакетного эмбеддинга: {e}") return None return None async def embed_batch(self, texts: List[str], retries: int = 3, delay: float = 1.0) -> Optional[List[List[float]]]: if not texts: return [] BATCH_SIZE = 50 results = [] for i in range(0, len(texts), BATCH_SIZE): batch = texts[i:i+BATCH_SIZE] loop = asyncio.get_running_loop() vectors = await loop.run_in_executor(None, self._embed_batch_sync, batch) if vectors is None: for attempt in range(retries): wait = delay * (2 ** attempt) await asyncio.sleep(wait) vectors = await loop.run_in_executor(None, self._embed_batch_sync, batch) if vectors is not None: break else: return None results.extend(vectors) await asyncio.sleep(0.5) return results async def embed_sparse_batch(self, texts: List[str]) -> List[Optional[Any]]: model = self._get_sparse_model() if model is None: return [None] * len(texts) loop = asyncio.get_running_loop() embeddings = await loop.run_in_executor(None, list, model.embed(texts)) return embeddings async def embed_sparse(self, text: str) -> Optional[Any]: model = self._get_sparse_model() if model is None: return None loop = asyncio.get_running_loop() embeddings = await loop.run_in_executor(None, lambda: list(model.embed([text]))) if embeddings and len(embeddings) > 0: return embeddings[0] return None