DocUA's picture
Єдиний коміт - очищення історії
4ad5efa
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)}"