modules/ai_analysis/google_embeddings_utils.py

"""
Утиліти для роботи з Google Embeddings API.
Цей файл повинен бути розміщений у modules/ai_analysis/google_embeddings_utils.py
"""

import os
import logging
from typing import List, Dict, Any, Optional
import time

logger = logging.getLogger(__name__)

# Перевірка наявності необхідних бібліотек
try:
    from google import genai
    GOOGLE_GENAI_AVAILABLE = True
    logger.info("Google GenAI SDK успішно імпортовано")
except ImportError as e:
    logger.warning(f"Google GenAI SDK не встановлено: {e}. Використання Google Embeddings буде вимкнено.")
    GOOGLE_GENAI_AVAILABLE = False

class GoogleEmbeddingsManager:
    """
    Менеджер для роботи з Google Embeddings API.
    """
    
    def __init__(self, api_key=None, model_name="text-embedding-004", task_type="retrieval_query"):
        """
        Ініціалізація менеджера Google Embeddings.
        
        Args:
            api_key (str, optional): API ключ для Google API. Якщо не вказано, спробує використати GEMINI_API_KEY з середовища.
            model_name (str): Назва моделі ембедингів.
            task_type (str): Тип задачі для ембедингів. Може бути "retrieval_query" або "retrieval_document".
        """
        self.api_key = api_key or os.getenv("GEMINI_API_KEY")
        self.model_name = model_name
        self.task_type = task_type
        self.client = None
        self.initialized = False
        
        # Спроба ініціалізації клієнта
        self._initialize_client()
    
    def _initialize_client(self):
        """
        Ініціалізує клієнт Google GenAI API.
        
        Returns:
            bool: True, якщо ініціалізація успішна, False в іншому випадку.
        """
        if not GOOGLE_GENAI_AVAILABLE:
            logger.error("Google GenAI SDK не встановлено. Встановіть пакет: pip install google-genai")
            return False
        
        if not self.api_key:
            logger.error("API ключ для Google API не вказано. Встановіть змінну GEMINI_API_KEY.")
            return False
        
        try:
            # Ініціалізація клієнта
            genai.configure(api_key=self.api_key)
            self.client = genai.Client()
            self.initialized = True
            logger.info(f"Клієнт Google GenAI успішно ініціалізовано для моделі {self.model_name}")
            return True
        except Exception as e:
            logger.error(f"Помилка при ініціалізації клієнта Google GenAI: {e}")
            return False
    
    def get_embeddings(self, texts, batch_size=8, retry_attempts=3, retry_delay=1):
        """
        Отримує ембединги для списку текстів.
        
        Args:
            texts (list): Список текстів для ембедингу.
            batch_size (int): Розмір батча для обробки.
            retry_attempts (int): Кількість спроб у випадку помилки.
            retry_delay (int): Затримка між спробами в секундах.
            
        Returns:
            list: Список ембедингів для кожного тексту або None у випадку помилки.
        """
        if not self.initialized:
            if not self._initialize_client():
                return None
        
        if not texts:
            logger.warning("Порожній список текстів для ембедингу")
            return []
        
        # Переконуємося, що input завжди список
        if not isinstance(texts, list):
            texts = [texts]
        
        try:
            all_embeddings = []
            
            # Обробка по батчам для ефективності
            for i in range(0, len(texts), batch_size):
                batch = texts[i:i + batch_size]
                
                # Спроби з повторами у випадку помилки
                for attempt in range(retry_attempts):
                    try:
                        result = self.client.models.embed_content(
                            model=self.model_name,
                            contents=batch,
                            config={"task_type": self.task_type}
                        )
                        
                        # Вилучення ембедингів
                        batch_embeddings = [embedding.values for embedding in result.embeddings]
                        all_embeddings.extend(batch_embeddings)
                        break
                    except Exception as e:
                        if attempt == retry_attempts - 1:
                            logger.error(f"Не вдалося отримати ембединги після {retry_attempts} спроб: {e}")
                            return None
                        logger.warning(f"Спроба {attempt+1} невдала: {e}. Повторна спроба через {retry_delay} сек.")
                        time.sleep(retry_delay)
            
            logger.info(f"Успішно отримано {len(all_embeddings)} ембедингів від Google API")
            return all_embeddings
            
        except Exception as e:
            logger.error(f"Помилка при отриманні ембедингів від Google API: {e}")
            return None
    
    def get_embedding_dimension(self):
        """
        Отримує розмірність ембедингів.
        
        Returns:
            int: Розмірність ембедингів або 0 у випадку помилки.
        """
        if not self.initialized:
            if not self._initialize_client():
                return 0
        
        try:
            # Отримуємо ембединг для тестового тексту
            result = self.client.models.embed_content(
                model=self.model_name,
                contents=["Test"],
                config={"task_type": self.task_type}
            )
            
            # Отримуємо розмірність
            [embedding] = result.embeddings
            return len(embedding.values)
            
        except Exception as e:
            logger.error(f"Помилка при отриманні розмірності ембедингів: {e}")
            return 0

# Приклад використання
if __name__ == "__main__":
    # Налаштування логування
    logging.basicConfig(level=logging.INFO)
    
    # Ініціалізація менеджера
    manager = GoogleEmbeddingsManager()
    
    # Отримання розмірності ембедингів
    dimension = manager.get_embedding_dimension()
    print(f"Розмірність ембедингів: {dimension}")
    
    # Отримання ембедингів для текстів
    texts = ["Це тестовий текст", "Це ще один тестовий текст"]
    embeddings = manager.get_embeddings(texts)
    
    if embeddings:
        print(f"Отримано {len(embeddings)} ембедингів")
        print(f"Розмірність першого ембедингу: {len(embeddings[0])}")