Редактировать embedding_service.py

rag/services/embedding_service.py

@@ -1,18 +1,20 @@
 # -*- coding: utf-8 -*-
 """
-Сервис генерации эмбеддингов через GigaChat API.
-Поддерживает одиночные запросы с кэшированием (lru_cache) и пакетную обработку.
-Все параметры (модель, таймаут, размер кэша, verify_ssl) вынесены в конфиг.
+Сервис генерации эмбеддингов через GigaChat API с кэшированием в Redis.
+Поддерживает одиночные запросы с кэшированием (Redis + локальный кэш) и пакетную обработку.
 """
 
 import asyncio
 import logging
-import time  # используется только в синхронной функции _embed_batch_sync
+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__)
 
@@ -23,6 +25,7 @@ except ImportError:
     HAS_SPARSE = False
     logging.warning("fastembed не установлен, sparse-эмбеддинги отключены")
 
+
 class EmbeddingService:
     def __init__(
         self,
@@ -30,7 +33,8 @@ class EmbeddingService:
         model: str = "GigaChat-Embeddings",
         timeout: int = 30,
         cache_size: int = 4096,
-        verify_ssl: bool = False
+        verify_ssl: bool = False,
+        redis_url: Optional[str] = None
     ):
         self.api_key = api_key
         self.model = model
@@ -38,6 +42,18 @@ class EmbeddingService:
         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
 
@@ -93,11 +109,39 @@ class EmbeddingService:
         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")
@@ -107,7 +151,7 @@ class EmbeddingService:
                     f"Попытка {attempt+1}/{retries} получить эмбеддинг не удалась: {e}. "
                     f"Повтор через {wait:.1f} сек."
                 )
-                await asyncio.sleep(wait)  # <-- исправлено: await asyncio.sleep
+                await asyncio.sleep(wait)
         logger.error(f"Не удалось получить эмбеддинг после {retries} попыток")
         return None
 
@@ -136,9 +180,9 @@ class EmbeddingService:
                     return None
             except RateLimitError as e:
                 if attempt < 4:
-                    wait = 2 ** attempt  # 1,2,4,8 секунд
+                    wait = 2 ** attempt
                     logger.warning(f"Rate limit, повтор через {wait} секунд")
-                    time.sleep(wait)   # оставляем, так как функция синхронная и выполняется в отдельном потоке
+                    time.sleep(wait)
                 else:
                     logger.error("Превышено количество попыток при 429")
                     return None
@@ -159,14 +203,14 @@ class EmbeddingService:
             if vectors is None:
                 for attempt in range(retries):
                     wait = delay * (2 ** attempt)
-                    await asyncio.sleep(wait)  # <-- исправлено: await asyncio.sleep
+                    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)   # <-- исправлено: await asyncio.sleep
+            await asyncio.sleep(0.5)
         return results
 
     async def embed_sparse_batch(self, texts: List[str]) -> List[Optional[Any]]: