Spaces:
Running
Running
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)}" |