modules/data_management/index_utils.py

import logging
import os
import json
import traceback
from pathlib import Path
import pandas as pd
import tiktoken
from typing import List, Dict, Any, Optional, Tuple

from llama_index.core import (
    Document,
)


# Встановлюємо змінну середовища, щоб примусово використовувати CPU
os.environ["CUDA_VISIBLE_DEVICES"] = ""
os.environ["TORCH_DEVICE"] = "cpu"

from modules.config.ai_settings import (
    get_metadata_csv,
    CHUNK_SIZE,
    CHUNK_OVERLAP,
    EXCLUDED_EMBED_METADATA_KEYS,
    EXCLUDED_LLM_METADATA_KEYS,
    GOOGLE_EMBEDDING_MODEL
)

# Налаштування логування
logger = logging.getLogger(__name__)

def initialize_embedding_model():
    """
    Ініціалізує модель ембедингів згідно налаштувань.
    Використовує офіційний пакет GeminiEmbedding для Google Embeddings.
    
    Returns:
        object: Модель ембедингів
    """
    try:
        # ПЕРША СПРОБА: Google Embeddings через офіційний пакет
        google_api_key = os.getenv("GEMINI_API_KEY")
        
        # Перевіряємо наявність API ключа
        if google_api_key:
            try:
                logger.info("Спроба ініціалізації Google Embeddings API через GeminiEmbedding...")
                
                from llama_index.embeddings.gemini import GeminiEmbedding
                
                # Використовуємо модель Gemini для ембедингів
                model_name = "models/embedding-004"  # або "models/text-embedding-001"
                
                # Створюємо модель ембедингів Gemini
                embed_model = GeminiEmbedding(
                    model_name=model_name,
                    api_key=google_api_key,
                    task_type="retrieval_query"  # або "retrieval_document"
                )
                
                # Тестуємо модель
                logger.info("Виконуємо тестовий запит до Gemini Embeddings API...")
                test_embedding = embed_model.get_text_embedding("Тестовий запит до Gemini Embeddings API")
                
                if test_embedding:
                    logger.info(f"Тестовий запит успішний, отримано ембединг розмірністю {len(test_embedding)}")
                    logger.info(f"Успішно ініціалізовано модель ембедингів Google Gemini: {model_name}")
                    return embed_model
                else:
                    raise Exception("Тестове підключення до Google API не вдалося - отримано порожній результат")
                
            except ImportError as imp_err:
                logger.error(f"Помилка імпорту модуля GeminiEmbedding: {imp_err}")
                logger.error("Можливо, потрібно встановити пакет: pip install llama-index-embeddings-gemini")
                logger.warning("Спробуємо альтернативні методи...")
                
                # Спробуємо альтернативний імпорт для Google AI SDK
                try:
                    # Через Google GenAI SDK безпосередньо
                    from google import genai
                    
                    logger.info("Спроба ініціалізації через Google GenAI API безпосередньо...")
                    
                    # Ініціалізуємо клієнт Google GenAI
                    genai.configure(api_key=google_api_key)
                    client = genai.Client()
                    
                    # Функція для отримання ембедингів від Google API
                    def get_google_embeddings(texts):
                        if not isinstance(texts, list):
                            texts = [texts]
                        
                        try:
                            # Використовуємо Google Embeddings API
                            result = client.models.embed_content(
                                model=GOOGLE_EMBEDDING_MODEL,
                                contents=texts,
                                config={"task_type": "retrieval_query"}
                            )
                            
                            # Виймаємо ембединги
                            embeddings = [embedding.values for embedding in result.embeddings]
                            
                            # Повертаємо в правильному форматі для LlamaIndex
                            return embeddings[0] if len(embeddings) == 1 else embeddings
                        except Exception as e:
                            logger.error(f"Помилка при отриманні ембедингів від Google API: {e}")
                            logger.error(traceback.format_exc())
                            raise
                    
                    # Тестуємо
                    test_result = get_google_embeddings(["Тестовий запит до Google GenAI API"])
                    if test_result:
                        # Створюємо кастомну модель ембедингів
                        embed_model = CustomEmbedding(
                            embed_func=get_google_embeddings,
                            embed_batch_size=8
                        )
                        
                        logger.info(f"Успішно ініціалізовано кастомну модель ембедингів Google через GenAI SDK")
                        return embed_model
                    else:
                        raise Exception("Тестове підключення до Google API не вдалося")
                except ImportError:
                    logger.error("Не вдалося імпортувати ні llama-index-embeddings-gemini, ні google.genai")
                except Exception as e:
                    logger.error(f"Помилка при альтернативній ініціалізації Google Embeddings: {e}")
                    logger.error(traceback.format_exc())
                    
            except Exception as e:
                logger.error(f"Не вдалося ініціалізувати Google Embeddings API: {e}")
                logger.error(traceback.format_exc())
                
        else:
            logger.warning("API ключ Google не знайдено (змінна GOOGLE_API_KEY не встановлена)")
            logger.warning("Будь ласка, додайте GOOGLE_API_KEY у файл .env або змінні середовища")
        
        # ДРУГА СПРОБА: HuggingFace ембединги
        logger.info("Використання локальних HuggingFace ембедингів...")
        from llama_index.embeddings.huggingface import HuggingFaceEmbedding
        from modules.config.ai_settings import DEFAULT_EMBEDDING_MODEL, FALLBACK_EMBEDDING_MODEL
        
        try:
            # Явно вказуємо використання CPU
            embed_model = HuggingFaceEmbedding(
                model_name=DEFAULT_EMBEDDING_MODEL,
                device="cpu"  # Явно вказуємо CPU
            )
            logger.info(f"Успішно ініціалізовано модель ембедингів HuggingFace на CPU: {DEFAULT_EMBEDDING_MODEL}")
            return embed_model
        except Exception as e:
            logger.warning(f"Не вдалося ініціалізувати основну модель HuggingFace ембедингів: {e}")
            
            # Спробуємо резервну модель
            try:
                embed_model = HuggingFaceEmbedding(
                    model_name=FALLBACK_EMBEDDING_MODEL,
                    device="cpu"  # Явно вказуємо CPU
                )
                logger.info(f"Успішно ініціалізовано резервну модель HuggingFace ембедингів на CPU: {FALLBACK_EMBEDDING_MODEL}")
                return embed_model
            except Exception as fallback_error:
                logger.error(f"Не вдалося ініціалізувати резервну модель HuggingFace: {fallback_error}")
                
                # Створення найпростішого фальшивого ембедера для аварійної ситуації
                try:
                    from llama_index.embeddings.custom import CustomEmbedding
                except ImportError:
                    # Для сумісності зі старими версіями бібліотеки
                    from llama_index.core.embeddings.custom import CustomEmbedding
                    
                import numpy as np
                
                def fallback_embedding_func(texts):
                    if not isinstance(texts, list):
                        texts = [texts]
                    
                    # Генеруємо фіктивні ембедінги (розмірність 768 - типова)
                    embeddings = [np.random.rand(768).tolist() for _ in texts]
                    return embeddings[0] if len(embeddings) == 1 else embeddings
                
                logger.warning("Використовуємо аварійний фальшивий ембедер")
                return CustomEmbedding(embed_func=fallback_embedding_func)
            
    except Exception as e:
        logger.error(f"Критична помилка при ініціалізації моделей ембедингів: {e}")
        logger.error(traceback.format_exc())
        
        # Аварійний фальшивий ембедер
        try:
            from llama_index.embeddings.custom import CustomEmbedding
        except ImportError:
            # Для сумісності зі старими версіями бібліотеки
            from llama_index.core.embeddings.custom import CustomEmbedding
            
        import numpy as np
        
        def emergency_embedding_func(texts):
            if not isinstance(texts, list):
                texts = [texts]
            return [np.random.rand(768).tolist() for _ in texts]
        
        logger.warning("Використовуємо аварійний фальшивий ембедер через критичну помилку")
        return CustomEmbedding(embed_func=emergency_embedding_func)


def count_tokens(text, model="gpt-3.5-turbo"):
    """
    Підраховує приблизну кількість токенів для тексту.
    
    Args:
        text (str): Текст для підрахунку токенів
        model (str): Назва моделі для вибору енкодера
        
    Returns:
        int: Кількість токенів
    """
    try:
        encoding = tiktoken.encoding_for_model(model)
        tokens = encoding.encode(text)
        return len(tokens)
    except Exception as e:
        logger.warning(f"Не вдалося підрахувати токени через tiktoken: {e}")
        # Якщо не можемо використати tiktoken, робимо просту оцінку
        return len(text) // 3  # Приблизна оцінка

def convert_dataframe_to_documents(df: pd.DataFrame) -> List[Document]:
    """
    Перетворює DataFrame з даними Jira в документи для індексування.
    
    Args:
        df (pd.DataFrame): DataFrame з даними Jira
        
    Returns:
        List[Document]: Список документів для індексування
    """
    logger.info("Перетворення даних DataFrame в документи для LlamaIndex...")
    
    jira_documents = []
    total_tokens = 0
    
    for idx, row in df.iterrows():
        # Основний текст - опис тікета
        text = ""
        if 'Description' in row and pd.notnull(row['Description']):
            text = str(row['Description'])
            
        # Додавання коментарів, якщо вони є
        for col in df.columns:
            if col.startswith('Comment') and pd.notnull(row[col]):
                text += f"\n\nКоментар: {str(row[col])}"
        
        # Метадані для документа
        metadata = metadata = get_metadata_csv(row, idx)
        
        # Додатково перевіряємо поле зв'язків, якщо воно є
        if 'Outward issue link (Relates)' in row and pd.notnull(row['Outward issue link (Relates)']):
            metadata["related_issues"] = row['Outward issue link (Relates)']
        
        # Додатково перевіряємо інші можливі поля зв'язків
        for col in df.columns:
            if col.startswith('Outward issue link') and col != 'Outward issue link (Relates)' and pd.notnull(row[col]):
                link_type = col.replace('Outward issue link (', '').replace(')', '')
                if "links" not in metadata:
                    metadata["links"] = {}
                metadata["links"][link_type] = str(row[col])
        
        # Створюємо документ з вказаними виключеннями
        doc = Document(
            text=text,
            metadata=metadata,
            excluded_embed_metadata_keys=EXCLUDED_EMBED_METADATA_KEYS,
            excluded_llm_metadata_keys=EXCLUDED_LLM_METADATA_KEYS
        )
        
        # Підраховуємо токени
        token_count = count_tokens(text)
        total_tokens += token_count
        
        # Додаємо документ до списку
        jira_documents.append(doc)
    
    logger.info(f"Створено {len(jira_documents)} документів з {total_tokens} токенами")
    return jira_documents

def check_index_integrity(indices_path: str) -> Tuple[bool, str]:
    """
    Перевіряє цілісність індексів.
    
    Args:
        indices_path (str): Шлях до директорії з індексами
        
    Returns:
        Tuple[bool, str]: (True, '') якщо індекси валідні, (False, 'повідомлення про помилку') в іншому випадку
    """
    try:
        indices_path = Path(indices_path)
        
        # Перевірка наявності директорії
        if not indices_path.exists() or not indices_path.is_dir():
            return False, f"Директорія з індексами не існує: {indices_path}"
        
        # Перевірка наявності маркера валідності
        valid_marker = indices_path / "indices.valid"
        if not valid_marker.exists():
            return False, f"Маркер валідності індексів не знайдено в {indices_path}"
        
        # Перевірка наявності файлів індексів
        required_files = ["docstore.json"]
        for file in required_files:
            if not (indices_path / file).exists():
                return False, f"Файл {file} не знайдено в {indices_path}"
        
        # Перевірка наявності BM25 індексу
        bm25_path = indices_path / "bm25"
        if not bm25_path.exists() or not bm25_path.is_dir():
            return False, f"Директорія з BM25 індексом не знайдено в {indices_path}"
        
        return True, ""
        
    except Exception as e:
        return False, f"Помилка при перевірці цілісності індексів: {str(e)}"

def check_indexing_availability(indices_path=None):
    """
    Перевіряє доступність функціональності індексування.
    
    Returns:
        bool: True, якщо функціональність доступна, False - інакше
    """
    try:
        # Перевіряємо наявність необхідних модулів
        import importlib
        
        # Список необхідних модулів
        required_modules = [
            "llama_index.core",
            "llama_index.retrievers.bm25",
            "llama_index.vector_stores.faiss",
            "llama_index.embeddings.huggingface"
        ]
        
        # Додаємо Google Embeddings до списку, якщо встановлено змінну середовища
        if os.getenv("GEMINI_API_KEY"):
            required_modules.append("google.genai")
        
        # Перевіряємо кожен модуль
        for module_name in required_modules:
            try:
                importlib.import_module(module_name)
            except ImportError:
                logger.warning(f"Модуль {module_name} не знайдено")
                return False
        
        # Всі модулі доступні
        logger.info("Всі необхідні модулі для індексування доступні")
        return True
        
    except Exception as e:
        logger.error(f"Помилка при перевірці доступності індексування: {e}")
        return False
    
def validate_index_directory(indices_path):
    """
    Перевіряє, чи директорія з індексами існує та містить необхідні файли.
    
    Args:
        indices_path (str): Шлях до директорії з індексами
        
    Returns:
        bool: True, якщо директорія валідна, False - інакше
    """
    try:
        from pathlib import Path
        
        indices_path = Path(indices_path)
        
        # Перевірка наявності директорії
        if not indices_path.exists() or not indices_path.is_dir():
            return False
        
        # Перевірка наявності необхідних файлів
        required_files = ["docstore.json"]
        for file in required_files:
            if not (indices_path / file).exists():
                return False
        
        return True
        
    except Exception as e:
        logger.error(f"Помилка при валідації директорії індексів: {str(e)}")
        return False

def test_google_embeddings():
    """
    Функція для тестування та відлагодження Google Embeddings API.
    Можна запустити як окремий скрипт для перевірки роботи API.
    
    Запуск з командного рядка:
    python -c "from modules.data_management.index_utils import test_google_embeddings; test_google_embeddings()"
    """
    import os
    import logging
    
    # Налаштування логування
    logging.basicConfig(level=logging.INFO)
    logger = logging.getLogger(__name__)
    
    logger.info("Тестування Google Embeddings API...")
    
    # Отримання API ключа
    api_key = os.getenv("GEMINI_API_KEY")
    if not api_key:
        logger.error("GEMINI_API_KEY не знайдений. Перевірте ваш .env файл або змінні середовища.")
        return False
    
    logger.info(f"API ключ Google знайдено: {api_key[:5]}...{api_key[-5:] if len(api_key) > 10 else '***'}")
    
    try:
        from google import genai
        
        # Ініціалізація клієнта
        genai.configure(api_key=api_key)
        client = genai.Client()
        logger.info("Google GenAI клієнт успішно ініціалізовано")
        
        # Спроба отримати ембединги
        text = ["Тестовий текст українською мовою"]
        model = "text-embedding-004"
        
        logger.info(f"Запит до моделі {model} з текстом: {text}")
        
        result = client.models.embed_content(
            model=model,
            contents=text,
            config={"task_type": "retrieval_query"}
        )
        
        # Отримання ембедингів
        [embedding] = result.embeddings
        embedding_values = embedding.values
        
        logger.info(f"Ембединг успішно отримано, розмірність: {len(embedding_values)}")
        logger.info(f"Перші 5 значень: {embedding_values[:5]}")
        
        return True
        
    except ImportError:
        logger.error("Модуль google.genai не знайдено. Будь ласка, встановіть його: pip install google-genai")
        return False
    except Exception as e:
        import traceback
        logger.error(f"Помилка при тестуванні Google Embeddings API: {e}")
        logger.error(traceback.format_exc())
        return False

if __name__ == "__main__":
    test_google_embeddings()