232 lines
9.9 KiB
Python
232 lines
9.9 KiB
Python
# -*- 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 |