diff --git a/core/services/embedding_service.py b/core/services/embedding_service.py new file mode 100644 index 0000000..121e3ce --- /dev/null +++ b/core/services/embedding_service.py @@ -0,0 +1,195 @@ +# -*- coding: utf-8 -*- +""" +Сервис генерации эмбеддингов через GigaChat API. +Поддерживает одиночные запросы с кэшированием (lru_cache) и пакетную обработку. +Все параметры (модель, таймаут, размер кэша, verify_ssl) вынесены в конфиг. +""" + +import asyncio +import logging +from functools import lru_cache +from typing import List, Optional, Any + +from gigachat import GigaChat +from gigachat.exceptions import AuthenticationError, RateLimitError + +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 + ): + self.api_key = api_key + self.model = model + self.timeout = timeout + self.cache_size = cache_size + self.verify_ssl = verify_ssl + self._giga = None + # Кэширование одиночных эмбеддингов (синхронная обёртка с lru_cache) + self.embed_sync = lru_cache(maxsize=cache_size)(self._embed_sync_impl) + self._sparse_model = None + + def _get_client(self) -> GigaChat: + """Ленивое создание клиента 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]]: + """Синхронный вызов GigaChat API для одного текста. Результат кэшируется.""" + if not text or not text.strip(): + logger.warning("Попытка получить эмбеддинг пустого текста") + return None + + # Обрезка до 800 символов (безопасно для лимита 514 токенов GigaChat) + MAX_CHARS = 800 + if len(text) > MAX_CHARS: + text = text[:MAX_CHARS] + logger.debug(f"Текст обрезан до {MAX_CHARS} символов для эмбеддинга") + + client = self._get_client() + try: + # Совместимость с разными версиями SDK: пробуем input=[text] или просто [text] + 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 + 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: + 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 # 1,2,4,8 секунд + logger.warning(f"Rate limit, повтор через {wait} секунд") + time.sleep(wait) # синхронный sleep, так как мы в синхронной функции + 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): + await asyncio.sleep(delay * (2 ** attempt)) + 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) # 500 мс + return results + + async def embed_sparse_batch(self, texts: List[str]) -> List[Optional[Any]]: + """Возвращает список объектов SparseEmbedding для каждого текста""" + model = self._get_sparse_model() + if model is None: + # Возвращаем список из 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 \ No newline at end of file