modules/ai_analysis/ai_assistant.py

import logging
import traceback
import os
import json
from pathlib import Path
import time
from typing import Dict, List, Any, Optional, Tuple

# Імпорт необхідних модулів для роботи з індексами
from llama_index.core import (
    StorageContext,
    load_index_from_storage
)
from llama_index.retrievers.bm25 import BM25Retriever
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.retrievers import QueryFusionRetriever
from llama_index.core.llms import ChatMessage

# Імпорт утиліт для роботи з індексами
from modules.data_management.index_utils import (
    check_indexing_availability, 
    check_index_integrity,
    count_tokens
)

# Імпорт налаштувань
from modules.config.ai_settings import (
    SIMILARITY_TOP_K,
    HYBRID_SEARCH_MODE
)
from prompts import system_prompt_hybrid_chat

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


class JiraAIAssistant:
    """
    Клас для роботи з AI асистентом для аналізу даних Jira.
    """
    
    def __init__(self, indices_dir=None, model_name="gpt-3.5-turbo", temperature=0.7):
        """
        Ініціалізація асистента.
        
        Args:
            indices_dir (str): Шлях до директорії з індексами
            model_name (str): Назва моделі для використання
            temperature (float): Температура для генерації
        """
        self.indices_dir = indices_dir
        self.model_name = model_name
        self.temperature = temperature
        self.index = None
        self.bm25_retriever = None
        
        # Завантажуємо індекси, якщо вказано шлях
        if indices_dir:
            self.load_indices(indices_dir)
    
    def load_indices(self, indices_path):
        """
        Завантаження індексів з директорії.
        
        Args:
            indices_path (str): Шлях до директорії з індексами
            
        Returns:
            bool: True, якщо індекси успішно завантажено
        """
        try:
            logger.info(f"Завантаження індексів з {indices_path}")
            
            # Перевіряємо наявність директорії
            if not os.path.exists(indices_path):
                logger.error(f"Директорія з індексами не існує: {indices_path}")
                return False
            
            # Перевіряємо наявність файлу-маркера
            marker_path = os.path.join(indices_path, "indices.valid")
            if not os.path.exists(marker_path):
                logger.error(f"Файл-маркер індексів не знайдено: {marker_path}")
                return False
            
            # Імпортуємо необхідні модулі
            from llama_index.core import VectorStoreIndex, StorageContext
            from llama_index.retrievers.bm25 import BM25Retriever
            
            try:
                # Завантажуємо індекс
                storage_context = StorageContext.from_defaults(persist_dir=str(indices_path))
                self.index = VectorStoreIndex.from_storage_context(storage_context)
                
                # Завантажуємо BM25 retriever
                docstore = storage_context.docstore
                
                # Завантажуємо параметри BM25
                bm25_dir = os.path.join(indices_path, "bm25")
                bm25_params_path = os.path.join(bm25_dir, "params.json")
                
                if os.path.exists(bm25_params_path):
                    with open(bm25_params_path, "r", encoding="utf-8") as f:
                        bm25_params = json.load(f)
                    
                    similarity_top_k = bm25_params.get("similarity_top_k", 10)
                else:
                    similarity_top_k = 10
                
                self.bm25_retriever = BM25Retriever.from_defaults(
                    docstore=docstore,
                    similarity_top_k=similarity_top_k
                )
                
                logger.info(f"Індекси успішно завантажено з {indices_path}")
                return True
                
            except Exception as e:
                logger.error(f"Помилка при завантаженні індексів: {e}")
                logger.error(traceback.format_exc())
                return False
                
        except Exception as e:
            logger.error(f"Помилка при завантаженні індексів: {e}")
            logger.error(traceback.format_exc())
            return False
    
    def chat(self, query, history=None):
        """
        Відповідь на запит користувача з використанням індексів.
        
        Args:
            query (str): Запит користувача
            history (list, optional): Історія чату
            
        Returns:
            str: Відповідь асистента
        """
        try:
            if not self.index or not self.bm25_retriever:
                return "Індекси не завантажено. Будь ласка, завантажте дані."
            
            # Отримуємо відповідні документи
            bm25_results = self.bm25_retriever.retrieve(query)
            vector_results = self.index.as_retriever().retrieve(query)
            
            # Об'єднуємо результати
            all_results = list(bm25_results) + list(vector_results)
            
            # Видаляємо дублікати
            unique_results = []
            seen_ids = set()
            for result in all_results:
                if result.node_id not in seen_ids:
                    unique_results.append(result)
                    seen_ids.add(result.node_id)
            
            # Обмежуємо кількість результатів
            unique_results = unique_results[:10]
            
            # Формуємо контекст
            context = "\n\n".join([result.get_content() for result in unique_results])
            
            
            # Формуємо промпт
            prompt = f"""Використовуй надану інформацію для відповіді на запитання.

Контекст:
{context}

Запитання: {query}

Дай детальну відповідь на запитання, використовуючи тільки інформацію з контексту. Якщо інформації недостатньо, скажи про це.
"""
            
            # Отримуємо відповідь від моделі
            from llama_index.llms.openai import OpenAI
            
            llm = OpenAI(model=self.model_name, temperature=self.temperature)
            response = llm.complete(prompt)
            
            return response.text
            
        except Exception as e:
            logger.error(f"Помилка при обробці запиту: {e}")
            logger.error(traceback.format_exc())
            return f"Виникла помилка при обробці запиту: {str(e)}"