modules/ai_analysis/jira_qa_assistant.py

import os
import logging
import tiktoken
from datetime import datetime
from pathlib import Path
from typing import List, Dict, Any, Optional

from modules.config.ai_settings import DEFAULT_EMBEDDING_MODEL, FALLBACK_EMBEDDING_MODEL
from prompts import system_prompt_qa_assistant

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

# Імпорти з LlamaIndex
try:
    from llama_index.core import Document
    from llama_index.core.llms import ChatMessage
    LLAMA_INDEX_AVAILABLE = True
except ImportError:
    logger.warning("Не вдалося імпортувати LlamaIndex. Встановіть необхідні залежності для використання Q/A асистента.")
    LLAMA_INDEX_AVAILABLE = False

class JiraQAAssistant:
    """
    Клас асистента для режиму Q/A з повним контекстом для даних Jira.
    Дозволяє задавати питання по всім документам Jira без використання пошукових індексів.
    """
    def __init__(self, api_key_openai=None, api_key_gemini=None, model_type="gemini", temperature=0.2):
        """
        Ініціалізація Q/A асистента.
        
        Args:
            api_key_openai (str): API ключ для OpenAI
            api_key_gemini (str): API ключ для Google Gemini
            model_type (str): Тип моделі ("openai" або "gemini")
            temperature (float): Параметр температури для генерації відповідей
        """
        self.model_type = model_type.lower()
        self.temperature = temperature
        self.api_key_openai = api_key_openai or os.getenv("OPENAI_API_KEY", "")
        self.api_key_gemini = api_key_gemini or os.getenv("GEMINI_API_KEY", "")
        
        # Перевірка наявності LlamaIndex
        if not LLAMA_INDEX_AVAILABLE:
            logger.error("LlamaIndex не доступний. Встановіть пакети: pip install llama-index-llms-gemini llama-index")
            raise ImportError("LlamaIndex не встановлено. Необхідний для роботи Q/A асистента.")
        
        # Ініціалізація моделі LLM
        self.llm = None
        
        # Дані Jira
        self.df = None
        self.jira_documents = []
        
        # Ініціалізуємо модель LLM
        self._initialize_llm()
    
    def _initialize_llm(self):
        """Ініціалізує модель LLM відповідно до налаштувань."""
        try:
            # Ініціалізація LLM моделі
            if self.model_type == "gemini" and self.api_key_gemini:
                os.environ["GEMINI_API_KEY"] = self.api_key_gemini
                
                from llama_index.llms.gemini import Gemini
                self.llm = Gemini(
                    model="models/gemini-2.0-flash",
                    temperature=self.temperature,
                    max_tokens=4096,
                )
                logger.info("Успішно ініціалізовано Gemini 2.0 Flash модель")
            elif self.model_type == "openai" and self.api_key_openai:
                os.environ["OPENAI_API_KEY"] = self.api_key_openai
                
                from llama_index.llms.openai import OpenAI
                self.llm = OpenAI(
                    model="gpt-4o-mini",
                    temperature=self.temperature,
                    max_tokens=4096
                )
                logger.info("Успішно ініціалізовано OpenAI GPT-4o-mini модель")
            else:
                error_msg = f"Не вдалося ініціалізувати LLM модель типу {self.model_type}. Перевірте API ключі."
                logger.error(error_msg)
                raise ValueError(error_msg)
                
        except Exception as e:
            logger.error(f"Помилка ініціалізації моделі LLM: {e}")
            raise
    
    def load_documents_from_dataframe(self, df):
        """
        Завантаження документів прямо з DataFrame без створення індексів.
        
        Args:
            df (pandas.DataFrame): DataFrame з даними Jira
            
        Returns:
            bool: True якщо дані успішно завантажено
        """
        try:
            logger.info("Завантаження даних з DataFrame для Q/A")
            
            # Зберігаємо оригінальний DataFrame
            self.df = df.copy()
            
            # Конвертуємо дані в документи
            self._convert_dataframe_to_documents()
            
            return True
            
        except Exception as e:
            logger.error(f"Помилка при завантаженні даних з DataFrame: {e}")
            return False
    
    def _convert_dataframe_to_documents(self):
        """
        Перетворює дані DataFrame в об'єкти Document для роботи з моделлю LLM.
        """
        import pandas as pd
        
        if self.df is None:
            logger.error("Не вдалося створити документи: відсутні дані DataFrame")
            return
            
        logger.info("Перетворення даних DataFrame в документи для Q/A...")
        
        self.jira_documents = []
        
        for idx, row in self.df.iterrows():
            # Основний текст - опис тікета
            text = ""
            if 'Description' in row and pd.notna(row['Description']):
                text = str(row['Description'])
                
            # Додавання коментарів, якщо вони є
            for col in self.df.columns:
                if col.startswith('Comment') and pd.notna(row[col]):
                    text += f"\n\nКоментар: {str(row[col])}"
            
            # Метадані для документа
            metadata = {
                "issue_key": row['Issue key'] if 'Issue key' in row and pd.notna(row['Issue key']) else "",
                "issue_type": row['Issue Type'] if 'Issue Type' in row and pd.notna(row['Issue Type']) else "",
                "status": row['Status'] if 'Status' in row and pd.notna(row['Status']) else "",
                "priority": row['Priority'] if 'Priority' in row and pd.notna(row['Priority']) else "",
                "assignee": row['Assignee'] if 'Assignee' in row and pd.notna(row['Assignee']) else "",
                "reporter": row['Reporter'] if 'Reporter' in row and pd.notna(row['Reporter']) else "",
                "created": str(row['Created']) if 'Created' in row and pd.notna(row['Created']) else "",
                "updated": str(row['Updated']) if 'Updated' in row and pd.notna(row['Updated']) else "",
                "summary": row['Summary'] if 'Summary' in row and pd.notna(row['Summary']) else "",
                "project": row['Project name'] if 'Project name' in row and pd.notna(row['Project name']) else ""
            }
            
            # Додатково перевіряємо поле зв'язків, якщо воно є
            if 'Outward issue link (Relates)' in row and pd.notna(row['Outward issue link (Relates)']):
                metadata["related_issues"] = row['Outward issue link (Relates)']
            
            # Додатково перевіряємо інші можливі поля зв'язків
            for col in self.df.columns:
                if col.startswith('Outward issue link') and col != 'Outward issue link (Relates)' and pd.notna(row[col]):
                    link_type = col.replace('Outward issue link ', '').strip('()')
                    if "links" not in metadata:
                        metadata["links"] = {}
                    metadata["links"][link_type] = str(row[col])
            
            # Створення документа
            doc = Document(
                text=text,
                metadata=metadata
            )
            
            self.jira_documents.append(doc)
        
        logger.info(f"Створено {len(self.jira_documents)} документів для Q/A")
        
    def _count_tokens(self, text: str, model: str = "gpt-3.5-turbo") -> int:
        """
        Підраховує приблизну кількість токенів для тексту.
        
        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, робимо просту оцінку
            # В середньому 1 токен ≈ 3 символи для змішаного тексту
            return len(text) // 3  # Приблизна оцінка

    def run_qa(self, question: str) -> Dict[str, Any]:
        """
        Запускає режим Q/A з повним контекстом.
        
        Args:
            question (str): Питання користувача
            
        Returns:
            Dict[str, Any]: Словник з результатами, включаючи відповідь та метадані
        """
        if not self.jira_documents or not self.llm:
            error_msg = "Не вдалося виконати запит: відсутні документи або LLM"
            logger.error(error_msg)
            return {"error": error_msg}
        
        try:
            logger.info(f"Запуск режиму Q/A з повним контекстом для питання: {question}")
            
            # Підготовка повного контексту з усіх документів
            full_context = "ПОВНИЙ КОНТЕКСТ JIRA ТІКЕТІВ:\n\n"
            
            # Додаємо статистику по тікетах
            status_counts = {}
            type_counts = {}
            priority_counts = {}
            assignee_counts = {}
            
            for doc in self.jira_documents:
                status = doc.metadata.get("status", "")
                issue_type = doc.metadata.get("issue_type", "")
                priority = doc.metadata.get("priority", "")
                assignee = doc.metadata.get("assignee", "")
                
                if status:
                    status_counts[status] = status_counts.get(status, 0) + 1
                if issue_type:
                    type_counts[issue_type] = type_counts.get(issue_type, 0) + 1
                if priority:
                    priority_counts[priority] = priority_counts.get(priority, 0) + 1
                if assignee:
                    assignee_counts[assignee] = assignee_counts.get(assignee, 0) + 1
            
            # Додаємо статистику до контексту
            full_context += f"Всього тікетів: {len(self.jira_documents)}\n\n"
            
            full_context += "Статуси:\n"
            for status, count in sorted(status_counts.items(), key=lambda x: x[1], reverse=True):
                full_context += f"- {status}: {count}\n"
            
            full_context += "\nТипи тікетів:\n"
            for issue_type, count in sorted(type_counts.items(), key=lambda x: x[1], reverse=True):
                full_context += f"- {issue_type}: {count}\n"
            
            full_context += "\nПріоритети:\n"
            for priority, count in sorted(priority_counts.items(), key=lambda x: x[1], reverse=True):
                full_context += f"- {priority}: {count}\n"
            
            # Додаємо топ-5 виконавців
            if assignee_counts:
                full_context += "\nТоп виконавців:\n"
                for assignee, count in sorted(assignee_counts.items(), key=lambda x: x[1], reverse=True)[:5]:
                    full_context += f"- {assignee}: {count} тікетів\n"
            
            # Додаємо всі тікети з метаданими та текстом
            full_context += "\nДЕТАЛЬНА ІНФОРМАЦІЯ ПРО ТІКЕТИ:\n\n"
            
            for i, doc in enumerate(self.jira_documents):
                # Використовуємо ключ тікета, якщо доступний, інакше номер
                ticket_id = doc.metadata.get("issue_key", f"TICKET-{i+1}")
                summary = doc.metadata.get("summary", "Без опису")
                
                full_context += f"ТІКЕТ {ticket_id}: {summary}\n"
                
                # Додаємо всі метадані
                for key, value in doc.metadata.items():
                    # Пропускаємо виключені поля та вже виведені поля
                    if key == "project" or key == "summary" or key == "issue_key":
                        continue
                    
                    if isinstance(value, dict):
                        # Обробка вкладених словників (наприклад, links)
                        full_context += f"{key}:\n"
                        for sub_key, sub_value in value.items():
                            if sub_value:
                                full_context += f"  - {sub_key}: {sub_value}\n"
                    elif value:  # Додаємо тільки непорожні значення
                        full_context += f"{key}: {value}\n"
                
                # Додаємо текст документа
                if doc.text:
                    # Якщо текст дуже довгий, обмежуємо його для економії токенів
                    if len(doc.text) > 1000:
                        truncated_text = doc.text[:1000] + "... [текст скорочено]"
                        full_context += f"Опис: {truncated_text}\n"
                    else:
                        full_context += f"Опис: {doc.text}\n"
                
                full_context += "\n" + "-"*40 + "\n\n"
            
            # Підрахуємо токени для повного контексту
            full_context_tokens = self._count_tokens(full_context)
            logger.info(f"Приблизна кількість токенів у повному контексті: {full_context_tokens}")
            
            # Перевірка на перевищення ліміту токенів (для Gemini 2.0 Flash - 1,048,576 вхідних токенів)
            max_input_tokens = 1048576
            if full_context_tokens > max_input_tokens:
                logger.warning(f"Контекст перевищує ліміт вхідних токенів моделі ({full_context_tokens} > {max_input_tokens}).")
                logger.info("Виконується скорочення контексту...")
                
                # Обчислюємо, скільки можна включити тікетів
                tokens_per_ticket = full_context_tokens / len(self.jira_documents)
                safe_ticket_count = int(max_input_tokens * 0.8 / tokens_per_ticket)  # 80% від ліміту для безпеки
                
                # Обчислюємо новий контекст з меншою кількістю тікетів
                full_context = full_context.split("ДЕТАЛЬНА ІНФОРМАЦІЯ ПРО ТІКЕТИ:")[0]
                full_context += "\nДЕТАЛЬНА ІНФОРМАЦІЯ ПРО ТІКЕТИ (скорочено):\n\n"
                
                for i, doc in enumerate(self.jira_documents[:safe_ticket_count]):
                    ticket_id = doc.metadata.get("issue_key", f"TICKET-{i+1}")
                    summary = doc.metadata.get("summary", "Без опису")
                    
                    full_context += f"ТІКЕТ {ticket_id}: {summary}\n"
                    
                    # Додаємо найважливіші метадані
                    important_fields = ["status", "priority", "assignee", "created", "updated"]
                    for key in important_fields:
                        value = doc.metadata.get(key, "")
                        if value:
                            full_context += f"{key}: {value}\n"
                    
                    # Додаємо скорочений опис
                    if doc.text:
                        short_text = doc.text[:300] + "..." if len(doc.text) > 300 else doc.text
                        full_context += f"Опис: {short_text}\n"
                    
                    full_context += "\n" + "-"*30 + "\n\n"
                
                full_context += f"\n[Показано {safe_ticket_count} з {len(self.jira_documents)} тікетів через обмеження контексту]\n"
                
                # Перераховуємо токени для скороченого контексту
                full_context_tokens = self._count_tokens(full_context)
                logger.info(f"Скорочений контекст: {full_context_tokens} токенів")
            
            # Системний промпт для режиму Q/A
            system_prompt = system_prompt_qa_assistant
            
            # Підрахунок токенів для питання
            question_tokens = self._count_tokens(question)
            
            # Формуємо повідомлення для чату
            messages = [
                ChatMessage(role="system", content=system_prompt),
                ChatMessage(role="system", content=full_context),
                ChatMessage(role="user", content=question)
            ]
            
            # Отримуємо відповідь від LLM
            logger.info("Генерація відповіді...")
            response = self.llm.chat(messages)
            
            # Підрахунок токенів для відповіді
            response_text = str(response)
            response_tokens = self._count_tokens(response_text)
            
            logger.info(f"Відповідь успішно згенеровано, токенів: {response_tokens}")
            
            return {
                "answer": response_text,
                "metadata": {
                    "question_tokens": question_tokens,
                    "context_tokens": full_context_tokens,
                    "response_tokens": response_tokens,
                    "total_tokens": question_tokens + full_context_tokens + response_tokens,
                    "documents_used": len(self.jira_documents)
                }
            }
        
        except Exception as e:
            error_msg = f"Помилка при виконанні Q/A з повним контекстом: {e}"
            logger.error(error_msg)
            return {"error": error_msg}
        
    def get_statistics(self) -> Dict[str, Any]:
        """
        Повертає загальну статистику за документами.
        
        Returns:
            Dict[str, Any]: Словник зі статистикою
        """
        if not self.jira_documents:
            return {"error": "Немає завантажених документів"}
        
        # Статистика по тікетах
        status_counts = {}
        type_counts = {}
        priority_counts = {}
        assignee_counts = {}
        
        for doc in self.jira_documents:
            status = doc.metadata.get("status", "")
            issue_type = doc.metadata.get("issue_type", "")
            priority = doc.metadata.get("priority", "")
            assignee = doc.metadata.get("assignee", "")
            
            if status:
                status_counts[status] = status_counts.get(status, 0) + 1
            if issue_type:
                type_counts[issue_type] = type_counts.get(issue_type, 0) + 1
            if priority:
                priority_counts[priority] = priority_counts.get(priority, 0) + 1
            if assignee:
                assignee_counts[assignee] = assignee_counts.get(assignee, 0) + 1
        
        # Формуємо результат
        return {
            "document_count": len(self.jira_documents),
            "status_counts": status_counts,
            "type_counts": type_counts,
            "priority_counts": priority_counts,
            "top_assignees": dict(sorted(assignee_counts.items(), key=lambda x: x[1], reverse=True)[:5])
        }