RAG_AIEXP_01 / app_1.py
MrSimple07's picture
added text_json files + added the document prep
600d58a
raw
history blame
28.4 kB
import gradio as gr
from huggingface_hub import hf_hub_download
import faiss
import pandas as pd
import os
import json
from llama_index.core import Document, VectorStoreIndex, Settings
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.google_genai import GoogleGenAI
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.core.response_synthesizers import get_response_synthesizer, ResponseMode
from llama_index.core.prompts import PromptTemplate
import time
import sys
from config import *
REPO_ID = "MrSimple01/AIEXP_RAG_FILES"
faiss_index_filename = "cleaned_faiss_index.index"
chunks_filename = "processed_chunks.csv"
download_dir = "rag_files"
table_data_dir = "Табличные данные_JSON"
HF_TOKEN = os.getenv('HF_TOKEN')
GOOGLE_API_KEY = os.getenv('GOOGLE_API_KEY')
CUSTOM_PROMPT_NEW = """
Вы являетесь высокоспециализированным Ассистентом для анализа документов (AIEXP). Ваша цель - предоставлять точные, корректные и контекстно релевантные ответы на основе анализа нормативной документации (НД). Все ваши ответы должны основываться исключительно на предоставленном контексте без использования внешних знаний или предположений.
КРИТИЧЕСКИ ВАЖНО: ВСЕ ОТВЕТЫ ДОЛЖНЫ БЫТЬ ТОЛЬКО НА РУССКОМ ЯЗЫКЕ! НИКОГДА НЕ ОТВЕЧАЙТЕ НА АНГЛИЙСКОМ!
История чата:
{chat_history}
ИНСТРУКЦИИ ПО ОБРАБОТКЕ КОНТЕКСТА:
1. АНАЛИЗ ТАБЛИЧНЫХ ДАННЫХ:
- Если в контексте есть информация начинающаяся с "Таблица", внимательно изучите её содержимое
- Извлекайте данные из строк с заголовками и данными таблицы
- Указывайте номер и название таблицы при ответе
- Структурируйте ответ на основе табличных данных
2. ОПРЕДЕЛЕНИЕ ТИПА ЗАДАЧИ:
Проанализируйте запрос пользователя и определите тип задачи:
1. КРАТКОЕ САММАРИ (ключевые слова: "кратко", "суммировать", "резюме", "основные моменты", "в двух словах"):
- Предоставьте структурированное резюме запрашиваемого раздела/пункта
- Выделите ключевые требования, процедуры или положения
- Используйте нумерованный список для лучшей читаемости
- Сохраняйте терминологию НД
2. ПОИСК ДОКУМЕНТА И ПУНКТА (ключевые слова: "найти", "где", "какой документ", "в каком разделе", "ссылка"):
- Укажите конкретный документ и его структурное расположение
- Предоставьте точные номера разделов/подразделов/пунктов
- Процитируйте релевантные фрагменты
- Если найдено несколько документов, перечислите все с указанием специфики каждого
3. ПРОВЕРКА КОРРЕКТНОСТИ (ключевые слова: "правильно ли", "соответствует ли", "проверить", "корректно", "нарушение"):
- Сопоставьте предоставленную информацию с требованиями НД
- Четко укажите: "СООТВЕТСТВУЕТ" или "НЕ СООТВЕТСТВУЕТ"
- Перечислите конкретные требования НД
- Укажите выявленные расхождения или подтвердите соответствие
- Процитируйте релевантные пункты НД
4. ПЛАН ДЕЙСТВИЙ (ключевые слова: "план", "алгоритм", "последовательность", "как действовать", "пошагово"):
- Создайте пронумерованный пошаговый план
- Каждый шаг должен содержать ссылку на соответствующий пункт НД
- Укажите необходимые документы или формы
- Добавьте временные рамки, если они указаны в НД
- Выделите критические требования или ограничения
5. УТОЧНЯЮЩИЕ ВОПРОСЫ (ключевые слова: "что это значит", "что означает", "объясните", "расскажите подробнее"):
- Используйте историю чата для понимания контекста
- Если вопрос относится к предыдущему обсуждению, опирайтесь на него
- Предоставьте подробное объяснение на основе НД
- Если контекст неясен, попросите уточнения
ПРАВИЛА ФОРМИРОВАНИЯ ОТВЕТОВ:
1. ОБЯЗАТЕЛЬНОЕ УКАЗАНИЕ ИСТОЧНИКОВ:
- Для каждого ответа указывайте: "Согласно [Название документа], раздел [X], пункт [X.X]: [Ваш ответ]"
- В конце ответа добавляйте: "Подробнее об этом можно узнать в документе [Название документа], раздел [X]."
- При отсутствии точного раздела: "Согласно документу [Название]: [Ваш ответ]"
2. СТРОГОЕ СЛЕДОВАНИЕ КОНТЕКСТУ:
- Если информация не найдена: "Информация по вашему запросу не была найдена в нормативной документации."
- НЕ используйте английский язык ни при каких обстоятельствах
- Используйте историю чата для понимания контекста вопросов
3. ИСПОЛЬЗОВАНИЕ ТЕРМИНОЛОГИИ НД:
- Применяйте официальную терминологию из документов
- Сохраняйте оригинальные формулировки ключевых требований
- При необходимости разъясняйте специальные термины на основе НД
4. СТРУКТУРИРОВАНИЕ ОТВЕТОВ:
- Основной ответ на русском языке
- Указание источника
- Дополнительная информация о документе
Контекст: {context_str}
Вопрос: {query_str}
Ответ (ТОЛЬКО НА РУССКОМ ЯЗЫКЕ):
"""
query_engine = None
chunks_df = None
chat_history = []
def log_message(message):
print(message, flush=True)
sys.stdout.flush()
def table_to_document(table_json):
document_id = table_json.get("document_id") or table_json.get("document", "unknown")
metadata = {
"document_id": document_id,
"section": table_json.get("section", ""),
"table_number": table_json.get("table_number", ""),
"table_title": table_json.get("table_title", ""),
}
description = table_json.get("table_description", "")
headers = table_json.get("headers", [])
table_text = f"ТАБЛИЦА: {table_json.get('table_number', '')} - {table_json.get('table_title', '')}\n"
table_text += f"ДОКУМЕНТ: {document_id}\n"
table_text += f"РАЗДЕЛ: {table_json.get('section', '')}\n"
if description:
table_text += f"ОПИСАНИЕ: {description}\n"
if headers:
table_text += f"ЗАГОЛОВКИ ТАБЛИЦЫ: {' | '.join(headers)}\n"
data = table_json.get("data", [])
if data:
table_text += "ДАННЫЕ ТАБЛИЦЫ:\n"
for i, row in enumerate(data):
if isinstance(row, dict):
row_str = " | ".join([f"{k}: {v}" for k,v in row.items()])
table_text += f"Строка {i+1}: {row_str}\n"
return Document(text=table_text, metadata=metadata)
def download_table_data():
log_message("📥 Загрузка табличных данных...")
from huggingface_hub import list_repo_files
table_files = []
try:
files = list_repo_files(repo_id=REPO_ID, repo_type="dataset", token=HF_TOKEN)
for file in files:
if file.startswith(table_data_dir) and file.endswith('.json'):
table_files.append(file)
log_message(f"📊 Найдено {len(table_files)} JSON файлов с таблицами")
table_documents = []
for file_path in table_files:
try:
log_message(f"🔄 Обработка файла: {file_path}")
local_path = hf_hub_download(
repo_id=REPO_ID,
filename=file_path,
local_dir=download_dir,
repo_type="dataset",
token=HF_TOKEN
)
with open(local_path, 'r', encoding='utf-8') as f:
table_data = json.load(f)
log_message(f"📋 Структура JSON: {list(table_data.keys()) if isinstance(table_data, dict) else 'Список'}")
if isinstance(table_data, dict):
if 'sheets' in table_data:
for sheet in table_data['sheets']:
doc = table_to_document(sheet)
table_documents.append(doc)
log_message(f"✅ Создан документ из таблицы: {sheet.get('table_number', 'unknown')}")
else:
doc = table_to_document(table_data)
table_documents.append(doc)
log_message(f"✅ Создан документ из JSON объекта")
elif isinstance(table_data, list):
for table_json in table_data:
doc = table_to_document(table_json)
table_documents.append(doc)
log_message(f"✅ Создан документ из элемента списка")
except Exception as e:
log_message(f"❌ Ошибка обработки файла {file_path}: {str(e)}")
continue
log_message(f"✅ Создано {len(table_documents)} документов из таблиц")
return table_documents
except Exception as e:
log_message(f"❌ Ошибка загрузки табличных данных: {str(e)}")
return []
def improve_query_with_history(question, chat_history_list):
try:
log_message("🔄 Улучшение запроса с учетом истории...")
if not chat_history_list:
log_message("📝 История чата пуста, используем оригинальный запрос")
return question
history_context = ""
for i, (user_msg, bot_msg) in enumerate(chat_history_list[-3:], 1):
history_context += f"Вопрос {i}: {user_msg}\nОтвет {i}: {bot_msg[:200]}...\n\n"
improvement_prompt = f"""Проанализируй историю диалога и улучши текущий запрос пользователя.
ИСТОРИЯ ДИАЛОГА:
{history_context}
ТЕКУЩИЙ ЗАПРОС: {question}
ПРАВИЛА:
1. Если запрос неполный или ссылается на предыдущий контекст (например: "что это", "о чем это", "объясни это"), дополни его информацией из истории
2. Если запрос самодостаточный, верни его без изменений
3. Сохраняй ключевые термины и названия документов из истории
4. Отвечай только улучшенным запросом без дополнительных пояснений
УЛУЧШЕННЫЙ ЗАПРОС:"""
from llama_index.llms.google_genai import GoogleGenAI
llm = GoogleGenAI(model="gemini-2.0-flash", api_key=GOOGLE_API_KEY)
improved_query = llm.complete(improvement_prompt).text.strip()
log_message(f"✨ Улучшенный запрос: {improved_query}")
return improved_query
except Exception as e:
log_message(f"❌ Ошибка улучшения запроса: {str(e)}")
return question
def format_chat_history():
if not chat_history:
return "История чата пуста."
history_text = ""
for i, (user_msg, bot_msg) in enumerate(chat_history[-5:], 1):
history_text += f"Сообщение {i}:\nПользователь: {user_msg}\nАссистент: {bot_msg}\n\n"
return history_text
def answer_question(question, history):
global query_engine, chunks_df, chat_history
if query_engine is None:
return history + [["", "❌ Система не инициализирована"]], ""
try:
start_time = time.time()
log_message(f"🔍 Получен вопрос: {question}")
log_message(f"📜 История чата: {len(chat_history)} сообщений")
# Улучшаем запрос с учетом истории
improved_question = improve_query_with_history(question, chat_history)
log_message(f"🎯 Обработка улучшенного запроса: {improved_question}")
# Форматируем историю чата для промпта
chat_history_text = format_chat_history()
log_message(f"📝 Сформированная история для промпта: {len(chat_history_text)} символов")
log_message("🔎 Поиск релевантных чанков...")
retrieved_nodes = query_engine.retriever.retrieve(improved_question)
log_message(f"📊 Найдено {len(retrieved_nodes)} релевантных чанков")
# Логируем найденные чанки
for i, node in enumerate(retrieved_nodes[:3]):
log_message(f"📄 Чанк {i+1}: {node.text[:100]}...")
log_message(f"🏷️ Метаданные: {node.metadata}")
log_message("🤖 Отправка запроса в LLM...")
# Создаем контекст с историей чата
query_with_context = f"""
История чата:
{chat_history_text}
Текущий вопрос: {question}
"""
response = query_engine.query(query_with_context)
end_time = time.time()
processing_time = end_time - start_time
bot_response = response.response
log_message(f"✅ Получен ответ: {bot_response[:100]}...")
# Проверяем, что ответ на русском языке
if any(english_word in bot_response.lower() for english_word in ['i am sorry', 'i cannot', 'the query', 'this request']):
log_message("⚠️ Обнаружен ответ на английском языке, форсируем русский ответ")
# Принудительно запрашиваем ответ на русском
russian_prompt = f"""
ВАЖНО: Отвечай ТОЛЬКО на русском языке!
Вопрос: {question}
История: {chat_history_text}
Контекст: {retrieved_nodes[0].text if retrieved_nodes else 'Нет контекста'}
Если информации недостаточно для ответа, скажи: "Недостаточно информации для ответа на ваш вопрос в предоставленной документации."
Ответ на русском языке:
"""
from llama_index.llms.google_genai import GoogleGenAI
llm = GoogleGenAI(model="gemini-2.0-flash", api_key=GOOGLE_API_KEY)
bot_response = llm.complete(russian_prompt).text.strip()
log_message(f"🔄 Исправленный ответ на русском: {bot_response[:100]}...")
# Обновляем историю чата
chat_history.append((question, bot_response))
if len(chat_history) > 10:
chat_history = chat_history[-10:]
log_message(f"💾 История чата обновлена. Всего сообщений: {len(chat_history)}")
sources_html = generate_sources_html(retrieved_nodes)
response_with_time = f"{bot_response}\n\n⏱️ Время обработки: {processing_time:.2f} сек"
history.append([question, response_with_time])
return history, sources_html
except Exception as e:
error_msg = f"❌ Ошибка обработки вопроса: {str(e)}"
log_message(f"❌ Ошибка: {str(e)}")
history.append([question, error_msg])
return history, ""
def initialize_models():
global query_engine, chunks_df
try:
log_message("🔄 Инициализация системы...")
os.makedirs(download_dir, exist_ok=True)
log_message("📥 Загрузка основных файлов...")
faiss_index_path = hf_hub_download(
repo_id=REPO_ID,
filename=faiss_index_filename,
local_dir=download_dir,
repo_type="dataset",
token=HF_TOKEN
)
chunks_csv_path = hf_hub_download(
repo_id=REPO_ID,
filename=chunks_filename,
local_dir=download_dir,
repo_type="dataset",
token=HF_TOKEN
)
log_message("📚 Загрузка индекса и данных...")
index_faiss = faiss.read_index(faiss_index_path)
chunks_df = pd.read_csv(chunks_csv_path)
log_message(f"📄 Загружено {len(chunks_df)} основных чанков")
log_message(f"📋 Колонки в chunks_df: {list(chunks_df.columns)}")
table_documents = download_table_data()
log_message("🤖 Настройка моделей...")
embed_model = HuggingFaceEmbedding(model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
llm = GoogleGenAI(model="gemini-2.0-flash", api_key=GOOGLE_API_KEY)
Settings.embed_model = embed_model
Settings.llm = llm
text_column = None
for col in chunks_df.columns:
if 'text' in col.lower() or 'content' in col.lower() or 'chunk' in col.lower():
text_column = col
break
if text_column is None:
text_column = chunks_df.columns[0]
log_message(f"📝 Используется колонка для текста: {text_column}")
documents = []
for i, (_, row) in enumerate(chunks_df.iterrows()):
doc = Document(
text=str(row[text_column]),
metadata={
"chunk_id": row.get('chunk_id', i),
"document_id": row.get('document_id', 'unknown')
}
)
documents.append(doc)
documents.extend(table_documents)
log_message(f"📋 Всего создано {len(documents)} документов ({len(chunks_df)} чанков + {len(table_documents)} таблиц)")
log_message("🔍 Построение векторного индекса...")
vector_index = VectorStoreIndex.from_documents(documents)
retriever = VectorIndexRetriever(
index=vector_index,
similarity_top_k=20,
similarity_cutoff=0.7
)
custom_prompt_template = PromptTemplate(CUSTOM_PROMPT_NEW)
response_synthesizer = get_response_synthesizer(
response_mode=ResponseMode.TREE_SUMMARIZE,
text_qa_template=custom_prompt_template
)
query_engine = RetrieverQueryEngine(
retriever=retriever,
response_synthesizer=response_synthesizer
)
log_message("✅ Система успешно инициализирована!")
return True
except Exception as e:
log_message(f"❌ Ошибка инициализации: {str(e)}")
return False
def generate_sources_html(nodes):
html = "<div style='background-color: #2d3748; color: white; padding: 20px; border-radius: 10px; max-height: 400px; overflow-y: auto;'>"
html += "<h3 style='color: #63b3ed; margin-top: 0;'>📚 Источники:</h3>"
unique_docs = {}
for node in nodes:
metadata = node.metadata if hasattr(node, 'metadata') else {}
doc_id = metadata.get('document_id', 'unknown')
if doc_id not in unique_docs:
unique_docs[doc_id] = []
unique_docs[doc_id].append(node)
for doc_id, doc_nodes in unique_docs.items():
if doc_id == 'unknown' or doc_id == 'Раздел документа':
continue
file_link = None
if chunks_df is not None and 'file_link' in chunks_df.columns:
doc_rows = chunks_df[chunks_df['document_id'] == doc_id]
if not doc_rows.empty:
file_link = doc_rows.iloc[0]['file_link']
html += f"<div style='margin-bottom: 15px; padding: 15px; border: 1px solid #4a5568; border-radius: 8px; background-color: #1a202c;'>"
html += f"<h4 style='margin: 0 0 10px 0; color: #63b3ed;'>📄 {doc_id}</h4>"
if file_link:
html += f"<a href='{file_link}' target='_blank' style='color: #68d391; text-decoration: none; font-size: 14px; display: inline-block; margin-bottom: 10px;'>🔗 Ссылка на документ</a><br>"
table_nodes = [node for node in doc_nodes if 'table_number' in node.metadata]
if table_nodes:
for node in table_nodes[:3]:
metadata = node.metadata
table_num = metadata.get('table_number', '')
table_title = metadata.get('table_title', 'Без названия')
if table_num and table_title != 'Без названия':
html += f"<p style='font-size: 12px; color: #a0aec0; margin: 5px 0;'>📊 {table_num}: {table_title}</p>"
html += "</div>"
html += "</div>"
return html
def clear_chat():
global chat_history
chat_history = []
log_message("🗑️ История чата очищена")
return [], ""
def handle_submit(message, history):
if not message.strip():
return history, ""
updated_history, sources = answer_question(message, history)
return updated_history, sources
def create_demo_interface():
with gr.Blocks(title="AIEXP - AI Expert для нормативной документации", theme=gr.themes.Soft()) as demo:
gr.Markdown("""
# AIEXP - Artificial Intelligence Expert
## Инструмент для работы с нормативной документацией
""")
with gr.Tab("💬 Чат с документами"):
gr.Markdown("### Задайте вопрос по нормативной документации")
with gr.Row():
with gr.Column(scale=2):
chatbot = gr.Chatbot(
label="Диалог с AIEXP",
height=500,
show_copy_button=True
)
with gr.Row():
msg = gr.Textbox(
label="Ваш вопрос",
placeholder="Введите вопрос по нормативным документам...",
lines=2,
scale=4
)
send_btn = gr.Button("📤 Отправить", variant="primary", scale=1)
with gr.Row():
clear_btn = gr.Button("🗑️ Очистить чат", variant="secondary")
gr.Examples(
examples=[
"Какой стандарт устанавливает порядок признания протоколов испытаний продукции в области использования атомной энергии?",
"Кто несет ответственность за организацию и проведение признания протоколов испытаний продукции?",
"В каких случаях могут быть признаны протоколы испытаний, проведенные лабораториями, не включенными в перечисления?",
],
inputs=msg
)
with gr.Column(scale=1):
sources_output = gr.HTML(
label="Источники",
value="<div style='background-color: #2d3748; color: white; padding: 20px; border-radius: 10px; text-align: center;'>Здесь появятся источники...</div>",
)
def submit_message(message, history):
return handle_submit(message, history)
msg.submit(submit_message, [msg, chatbot], [chatbot, sources_output]).then(
lambda: "", None, msg
)
send_btn.click(submit_message, [msg, chatbot], [chatbot, sources_output]).then(
lambda: "", None, msg
)
clear_btn.click(clear_chat, outputs=[chatbot, sources_output])
return demo
if __name__ == "__main__":
log_message("🚀 Запуск AIEXP - AI Expert для нормативной документации")
if initialize_models():
log_message("🌟 Запуск веб-интерфейса...")
demo = create_demo_interface()
demo.launch(
server_name="0.0.0.0",
server_port=7860,
share=True,
debug=False
)
else:
log_message("❌ Невозможно запустить приложение из-за ошибки инициализации")
sys.exit(1)