modules/ai_analysis/faiss_utils.py

"""
Модуль утиліт для роботи з FAISS векторними індексами.
Цей файл повинен бути розміщений у modules/ai_analysis/faiss_utils.py
"""

import logging
import os
from pathlib import Path
import json
import hashlib
from datetime import datetime
import shutil
import tempfile
import sys

logger = logging.getLogger(__name__)

# Перевірка наявності змінних середовища Hugging Face
IS_HUGGINGFACE = os.environ.get("SPACE_ID") is not None
if IS_HUGGINGFACE:
    logger.info("Виявлено середовище Hugging Face Spaces")

try:
    import faiss
    import numpy as np
    from llama_index.vector_stores.faiss import FaissVectorStore
    from llama_index.core import load_index_from_storage
    from llama_index.core import StorageContext
    
    FAISS_AVAILABLE = True
    logger.info("FAISS успішно імпортовано")
except ImportError as e:
    logger.warning(f"FAISS або llama-index-vector-stores-faiss не встановлено: {e}. Використання FAISS буде вимкнено.")
    FAISS_AVAILABLE = False

def check_faiss_available():
    """
    Перевірка доступності FAISS.
    
    Returns:
        bool: True, якщо FAISS доступний, False інакше
    """
    return FAISS_AVAILABLE

def generate_file_hash(file_path):
    """
    Генерує хеш для файлу на основі його вмісту.
    
    Args:
        file_path (str): Шлях до файлу
        
    Returns:
        str: Хеш файлу або None у випадку помилки
    """
    try:
        if not os.path.exists(file_path):
            logger.error(f"Файл не знайдено: {file_path}")
            return None
            
        # Отримуємо базову інформацію про файл для додавання в хеш
        file_stat = os.stat(file_path)
        file_size = file_stat.st_size
        file_mtime = file_stat.st_mtime
        
        # Створюємо хеш на основі вмісту файлу
        sha256 = hashlib.sha256()
        
        # Додаємо базову інформацію про файл
        sha256.update(f"{file_size}_{file_mtime}".encode())
        
        # Додаємо вміст файлу
        with open(file_path, "rb") as f:
            for byte_block in iter(lambda: f.read(4096), b""):
                sha256.update(byte_block)
        
        return sha256.hexdigest()
        
    except Exception as e:
        logger.error(f"Помилка при генерації хешу файлу: {e}")
        return None

def save_indices_metadata(directory_path, metadata):
    """
    Зберігає метадані індексів у JSON файл.
    
    Args:
        directory_path (str): Шлях до директорії з індексами
        metadata (dict): Метадані для збереження
        
    Returns:
        bool: True, якщо збереження успішне, False інакше
    """
    try:
        # Перевіряємо наявність директорії
        if not os.path.exists(directory_path):
            logger.warning(f"Директорія {directory_path} не існує. Створюємо...")
            os.makedirs(directory_path, exist_ok=True)
            
        metadata_path = Path(directory_path) / "metadata.json"
        
        # Додаємо додаткову інформацію про оточення
        metadata["environment"] = {
            "is_huggingface": IS_HUGGINGFACE,
            "python_version": sys.version,
            "platform": sys.platform
        }
        
        # Додаємо логування для діагностики
        logger.info(f"Збереження метаданих у {metadata_path}")
        logger.info(f"Розмір метаданих: {len(str(metadata))} символів")
        
        with open(metadata_path, "w", encoding="utf-8") as f:
            json.dump(metadata, f, ensure_ascii=False, indent=2)
        
        # Перевіряємо, що файл було створено
        if os.path.exists(metadata_path):
            logger.info(f"Метадані успішно збережено у {metadata_path}")
            return True
        else:
            logger.error(f"Файл {metadata_path} не було створено")
            return False
        
    except Exception as e:
        logger.error(f"Помилка при збереженні метаданих: {e}")
        return False

def load_indices_metadata(directory_path):
    """
    Завантажує метадані індексів з JSON файлу.
    
    Args:
        directory_path (str): Шлях до директорії з індексами
        
    Returns:
        dict: Метадані або пустий словник у випадку помилки
    """
    try:
        metadata_path = Path(directory_path) / "metadata.json"
        
        if not metadata_path.exists():
            logger.warning(f"Файл метаданих не знайдено: {metadata_path}")
            return {}
        
        with open(metadata_path, "r", encoding="utf-8") as f:
            metadata = json.load(f)
        
        logger.info(f"Метадані успішно завантажено з {metadata_path}")
        return metadata
        
    except Exception as e:
        logger.error(f"Помилка при завантаженні метаданих: {e}")
        return {}

def find_latest_indices(base_dir="temp/indices"):
    """
    Знаходить найновіші збережені індекси.
    
    Args:
        base_dir (str): Базова директорія з індексами
        
    Returns:
        tuple: (bool, str) - (наявність індексів, шлях до найновіших індексів)
    """
    try:
        # Перевіряємо наявність базової директорії
        indices_dir = Path(base_dir)
        
        if not indices_dir.exists():
            logger.info(f"Директорія {base_dir} не існує")
            return False, None
            
        if not os.path.isdir(indices_dir):
            logger.warning(f"{base_dir} існує, але не є директорією")
            return False, None
            
        if not any(indices_dir.iterdir()):
            logger.info(f"Директорія {base_dir} порожня")
            return False, None
        
        # Отримання списку піддиректорій з індексами
        try:
            subdirs = [d for d in indices_dir.iterdir() if d.is_dir()]
        except Exception as iter_err:
            logger.error(f"Помилка при перегляді директорії {base_dir}: {iter_err}")
            return False, None
            
        if not subdirs:
            logger.info("Індекси не знайдено")
            return False, None
        
        # Знаходимо найновішу директорію
        try:
            latest_dir = max(subdirs, key=lambda x: x.stat().st_mtime)
        except Exception as sort_err:
            logger.error(f"Помилка при сортуванні директорій: {sort_err}")
            return False, None
        
        logger.info(f"Знайдено індекси у директорії {latest_dir}")
        return True, str(latest_dir)
        
    except Exception as e:
        logger.error(f"Помилка при пошуку індексів: {e}")
        return False, None

def find_indices_by_hash(csv_hash, base_dir="temp/indices"):
    """
    Знаходить індекси, що відповідають вказаному хешу CSV файлу.
    
    Args:
        csv_hash (str): Хеш CSV файлу
        base_dir (str): Базова директорія з індексами
        
    Returns:
        tuple: (bool, str) - (наявність індексів, шлях до відповідних індексів)
    """
    try:
        if not csv_hash:
            logger.warning("Не вказано хеш CSV файлу")
            return False, None
            
        # Перевіряємо наявність базової директорії
        indices_dir = Path(base_dir)
        
        if not indices_dir.exists():
            logger.info(f"Директорія {base_dir} не існує")
            return False, None
            
        if not any(indices_dir.iterdir()):
            logger.info(f"Директорія {base_dir} порожня")
            return False, None
        
        # Отримання списку піддиректорій з індексами
        subdirs = [d for d in indices_dir.iterdir() if d.is_dir()]
        if not subdirs:
            logger.info("Індекси не знайдено")
            return False, None
        
        # Перевіряємо кожну директорію на відповідність хешу
        for directory in subdirs:
            metadata_path = directory / "metadata.json"
            if metadata_path.exists():
                try:
                    with open(metadata_path, "r", encoding="utf-8") as f:
                        metadata = json.load(f)
                    
                    if "csv_hash" in metadata and metadata["csv_hash"] == csv_hash:
                        # Додатково перевіряємо наявність файлів індексів
                        if (directory / "docstore.json").exists():
                            logger.info(f"Знайдено індекси для CSV з хешем {csv_hash} у {directory}")
                            return True, str(directory)
                        else:
                            logger.warning(f"Знайдено метадані для CSV з хешем {csv_hash}, але файли індексів відсутні у {directory}")
                except Exception as md_err:
                    logger.warning(f"Помилка при читанні метаданих {metadata_path}: {md_err}")
        
        # Якщо відповідних індексів не знайдено, повертаємо найновіші
        logger.info(f"Не знайдено індексів для CSV з хешем {csv_hash}, спроба знайти найновіші")
        return find_latest_indices(base_dir)
        
    except Exception as e:
        logger.error(f"Помилка при пошуку індексів за хешем: {e}")
        return False, None

def create_indices_directory(csv_hash=None, base_dir="temp/indices"):
    """
    Створює директорію для зберігання індексів з часовою міткою.
    
    Args:
        csv_hash (str, optional): Хеш CSV файлу для метаданих
        base_dir (str): Базова директорія для індексів
        
    Returns:
        str: Шлях до створеної директорії
    """
    try:
        # Створення базової директорії, якщо вона не існує
        indices_dir = Path(base_dir)
        
        # Очищаємо старі індекси перед створенням нових, якщо ми на Hugging Face
        if IS_HUGGINGFACE and indices_dir.exists():
            logger.info("Очищення старих індексів перед створенням нових на Hugging Face")
            try:
                # Видаляємо тільки старі директорії, якщо їх більше 1
                subdirs = [d for d in indices_dir.iterdir() if d.is_dir()]
                if len(subdirs) > 1:
                    # Сортуємо за часом модифікації (від найстаріших до найновіших)
                    sorted_dirs = sorted(subdirs, key=lambda x: x.stat().st_mtime)
                    
                    # Залишаємо тільки найновішу директорію
                    for directory in sorted_dirs[:-1]:
                        try:
                            shutil.rmtree(directory)
                            logger.info(f"Видалено стару директорію: {directory}")
                        except Exception as del_err:
                            logger.warning(f"Не вдалося видалити директорію {directory}: {del_err}")
            except Exception as clean_err:
                logger.warning(f"Помилка при очищенні старих індексів: {clean_err}")
        
        # Створюємо базову директорію
        indices_dir.mkdir(exist_ok=True, parents=True)
        
        # Створення унікальної директорії з часовою міткою
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        index_dir = indices_dir / timestamp
        
        # Спроба створити директорію
        try:
            index_dir.mkdir(exist_ok=True)
        except Exception as mkdir_err:
            logger.error(f"Не вдалося створити директорію {index_dir}: {mkdir_err}")
            
            # Створюємо тимчасову директорію як запасний варіант
            try:
                temp_dir = tempfile.mkdtemp(prefix="faiss_indices_")
                logger.info(f"Створено тимчасову директорію: {temp_dir}")
                return temp_dir
            except Exception as temp_err:
                logger.error(f"Не вдалося створити тимчасову директорію: {temp_err}")
                return str(indices_dir / "fallback")
        
        # Зберігаємо базові метадані
        metadata = {
            "created_at": timestamp,
            "timestamp": datetime.now().timestamp(),
            "csv_hash": csv_hash
        }
        
        save_indices_metadata(str(index_dir), metadata)
        
        logger.info(f"Створено директорію для індексів: {index_dir}")
        return str(index_dir)
        
    except Exception as e:
        logger.error(f"Помилка при створенні директорії індексів: {e}")
        
        # Створюємо тимчасову директорію як запасний варіант
        try:
            temp_dir = tempfile.mkdtemp(prefix="faiss_indices_")
            logger.info(f"Створено тимчасову директорію для індексів: {temp_dir}")
            return temp_dir
        except Exception:
            # Якщо і це не вдалося, використовуємо директорію temp
            logger.error("Не вдалося створити навіть тимчасову директорію, використовуємо базову temp")
            os.makedirs("temp", exist_ok=True)
            return "temp"

def cleanup_old_indices(max_indices=3, base_dir="temp/indices"):
    """
    Видаляє старі індекси, залишаючи тільки вказану кількість найновіших.
    
    Args:
        max_indices (int): Максимальна кількість індексів для зберігання
        base_dir (str): Базова директорія з індексами
        
    Returns:
        int: Кількість видалених директорій
    """
    try:
        # На Hugging Face Space обмежуємо максимальну кількість індексів до 1
        if IS_HUGGINGFACE:
            max_indices = 1
            logger.info("На Hugging Face Space обмежуємо кількість індексів до 1")
            
        indices_dir = Path(base_dir)
        
        if not indices_dir.exists():
            logger.warning(f"Директорія {base_dir} не існує")
            return 0
        
        # Отримання списку піддиректорій з індексами
        try:
            subdirs = [d for d in indices_dir.iterdir() if d.is_dir()]
        except Exception as iter_err:
            logger.error(f"Помилка при скануванні директорії {base_dir}: {iter_err}")
            return 0
            
        if len(subdirs) <= max_indices:
            logger.info(f"Кількість індексів ({len(subdirs)}) не перевищує ліміт ({max_indices})")
            return 0
        
        # Сортуємо директорії за часом модифікації (від найновіших до найстаріших)
        try:
            sorted_dirs = sorted(subdirs, key=lambda x: x.stat().st_mtime, reverse=True)
        except Exception as sort_err:
            logger.error(f"Помилка при сортуванні директорій: {sort_err}")
            return 0
        
        # Залишаємо тільки max_indices найновіших директорій
        dirs_to_delete = sorted_dirs[max_indices:]
        
        # Видаляємо старі директорії
        deleted_count = 0
        for directory in dirs_to_delete:
            try:
                shutil.rmtree(directory)
                deleted_count += 1
                logger.info(f"Видалено стару директорію індексів: {directory}")
            except Exception as del_err:
                logger.warning(f"Не вдалося видалити директорію {directory}: {del_err}")
        
        return deleted_count
        
    except Exception as e:
        logger.error(f"Помилка при очищенні старих індексів: {e}")
        return 0