app.py

import gradio as gr
import os
from llama_index.core import Settings
from documents_prep import load_json_documents, load_table_data, load_image_data, load_csv_chunks
from utils import get_llm_model, get_embedding_model, get_reranker_model, log_message, answer_question
from index_retriever import create_vector_index, create_query_engine
import sys
from config import (
    HF_REPO_ID, HF_TOKEN, DOWNLOAD_DIR, CHUNKS_FILENAME,
    JSON_FILES_DIR, TABLE_DATA_DIR, IMAGE_DATA_DIR, DEFAULT_MODEL, AVAILABLE_MODELS
)


def initialize_system(repo_id, hf_token, download_dir, chunks_filename=None, 
                     json_files_dir=None, table_data_dir=None, image_data_dir=None, 
                     use_json_instead_csv=False):
    try:
        log_message("Инициализация системы")
        os.makedirs(download_dir, exist_ok=True)
        
        embed_model = get_embedding_model()
        llm = get_llm_model(DEFAULT_MODEL)
        reranker = get_reranker_model()
        
        Settings.embed_model = embed_model
        Settings.llm = llm
        
        all_documents = []
        chunks_df = None
        
        if use_json_instead_csv and json_files_dir:
            log_message("Используем JSON файлы вместо CSV")
            json_documents = load_json_documents(repo_id, hf_token, json_files_dir, download_dir)
            all_documents.extend(json_documents)
        else:
            if chunks_filename:
                log_message("Загружаем данные из CSV")
                csv_documents, chunks_df = load_csv_chunks(repo_id, hf_token, chunks_filename, download_dir)
                all_documents.extend(csv_documents)
        
        if table_data_dir:
            log_message("Добавляю табличные данные")
            table_documents = load_table_data(repo_id, hf_token, table_data_dir)
            all_documents.extend(table_documents)
        
        if image_data_dir:
            log_message("Добавляю данные изображений")
            image_documents = load_image_data(repo_id, hf_token, image_data_dir)
            all_documents.extend(image_documents)
        
        log_message(f"Всего документов: {len(all_documents)}")
        
        vector_index = create_vector_index(all_documents)
        query_engine = create_query_engine(vector_index)
        
        log_message(f"Система успешно инициализирована")
        return query_engine, chunks_df, reranker, vector_index
        
    except Exception as e:
        log_message(f"Ошибка инициализации: {str(e)}")
        return None, None, None, None

def switch_model(model_name, vector_index):
    from llama_index.core import Settings
    from index_retriever import create_query_engine
    
    try:
        log_message(f"Переключение на модель: {model_name}")
        
        new_llm = get_llm_model(model_name)
        Settings.llm = new_llm
        
        if vector_index is not None:
            new_query_engine = create_query_engine(vector_index)
            log_message(f"Модель успешно переключена на: {model_name}")
            return new_query_engine, f"✅ Модель переключена на: {model_name}"
        else:
            return None, "❌ Ошибка: система не инициализирована"
            
    except Exception as e:
        error_msg = f"Ошибка переключения модели: {str(e)}"
        log_message(error_msg)
        return None, f"❌ {error_msg}"

def create_demo_interface(answer_question_func, switch_model_func, current_model):
    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):
                    model_dropdown = gr.Dropdown(
                        choices=list(AVAILABLE_MODELS.keys()),
                        value=current_model,
                        label="🤖 Выберите языковую модель",
                        info="Выберите модель для генерации ответов"
                    )
                with gr.Column(scale=1):
                    switch_btn = gr.Button("🔄 Переключить модель", variant="secondary")
                    model_status = gr.Textbox(
                        value=f"Текущая модель: {current_model}",
                        label="Статус модели",
                        interactive=False
                    )
            
            with gr.Row():
                with gr.Column(scale=3):
                    question_input = gr.Textbox(
                        label="Ваш вопрос к базе знаний",
                        placeholder="Введите вопрос по нормативным документам...",
                        lines=3
                    )
                    ask_btn = gr.Button("🔍 Найти ответ", variant="primary", size="lg")
                    
                    gr.Examples(
                        examples=[
                            "О чем этот рисунок: ГОСТ Р 50.04.07-2022 Приложение Л. Л.1.5 Рисунок Л.2",
                            "Л.9 Формула в ГОСТ Р 50.04.07 - 2022 что и о чем там?", 
                            "Какой стандарт устанавливает порядок признания протоколов испытаний продукции в области использования атомной энергии?",
                            "Кто несет ответственность за организацию и проведение признания протоколов испытаний продукции?",
                            "В каких случаях могут быть признаны протоколы испытаний, проведенные лабораториями?",
                        ],
                        inputs=question_input
                    )
            
            with gr.Row():
                with gr.Column(scale=2):
                    answer_output = gr.HTML(
                        label="",
                        value=f"<div style='background-color: #2d3748; color: white; padding: 20px; border-radius: 10px; text-align: center;'>Здесь появится ответ на ваш вопрос...<br><small>Текущая модель: {current_model}</small></div>",
                    )
                
                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>",
                    )
            
            switch_btn.click(
                fn=switch_model_func,
                inputs=[model_dropdown],
                outputs=[model_status]
            )
            
            ask_btn.click(
                fn=answer_question_func,
                inputs=[question_input],
                outputs=[answer_output, sources_output]
            )
            
            question_input.submit(
                fn=answer_question_func,
                inputs=[question_input],
                outputs=[answer_output, sources_output]
            )
    
    return demo

query_engine = None
chunks_df = None
reranker = None
vector_index = None
current_model = DEFAULT_MODEL

def main_answer_question(question):
    global query_engine, reranker, current_model, chunks_df
    return answer_question(question, query_engine, reranker, current_model, chunks_df)

def main_switch_model(model_name):
    global query_engine, vector_index, current_model
    
    new_query_engine, status_message = switch_model(model_name, vector_index)
    if new_query_engine:
        query_engine = new_query_engine
        current_model = model_name
    
    return status_message

def main():
    global query_engine, chunks_df, reranker, vector_index, current_model
    
    log_message("Запуск AIEXP - AI Expert для нормативной документации")
    
    query_engine, chunks_df, reranker, vector_index = initialize_system(
        repo_id=HF_REPO_ID,
        hf_token=HF_TOKEN,
        download_dir=DOWNLOAD_DIR,
        json_files_dir=JSON_FILES_DIR,
        table_data_dir=TABLE_DATA_DIR,
        image_data_dir=IMAGE_DATA_DIR,
    )
    
    if query_engine:
        log_message("Запуск веб-интерфейса")
        demo = create_demo_interface(
            answer_question_func=main_answer_question,
            switch_model_func=main_switch_model,
            current_model=current_model
        )
        demo.launch(
            server_name="0.0.0.0",
            server_port=7860,
            share=True,
            debug=False
        )
    else:
        log_message("Невозможно запустить приложение из-за ошибки инициализации")
        sys.exit(1)

if __name__ == "__main__":
    main()