app.py

import gradio as gr
import os
import shutil
import pandas as pd
from datetime import datetime
from scripts.document_processor import process_multiple_documents, save_processed_chunks, load_processed_chunks
from scripts.rag_engine import build_rag_system, query_documents, format_response_with_sources, add_new_document_to_system, load_rag_system
import json
import tempfile
from scripts.config import *


if not os.path.exists(UPLOAD_FOLDER):
    os.makedirs(UPLOAD_FOLDER)

if not os.path.exists("processed_data"):
    os.makedirs("processed_data")

if not os.path.exists(RAG_FILES_DIR):
    os.makedirs(RAG_FILES_DIR)



def initialize_system():
    global query_engine
    query_engine = None
    
    # IMPORTANT: Setup LLM settings at the very beginning
    from scripts.config import setup_llm_settings, download_pretrained_files
    setup_llm_settings()
    
    # Check if local RAG system exists
    local_rag_exists = os.path.exists(os.path.join(RAG_FILES_DIR, 'faiss_index.index'))
    local_csv_exists = os.path.exists(PROCESSED_DATA_FILE)
    
    # If no local system exists, try to download from HuggingFace
    if not local_rag_exists and not local_csv_exists:
        print("No local RAG system found. Attempting to download from HuggingFace...")
        download_success = download_pretrained_files()
        
        if download_success:
            print("✅ Downloaded pre-trained files from HuggingFace Hub")
            # Update existence flags after download
            local_rag_exists = os.path.exists(os.path.join(RAG_FILES_DIR, 'faiss_index.index'))
            local_csv_exists = os.path.exists(PROCESSED_DATA_FILE)
        else:
            print("⚠️ Failed to download pre-trained files. System will start empty.")
    
    # Try to load existing RAG system
    if local_rag_exists:
        try:
            print("Found existing RAG system files, loading...")
            query_engine = load_rag_system()  
            if query_engine is not None:
                chunk_count = 0
                if os.path.exists(PROCESSED_DATA_FILE):
                    processed_chunks = load_processed_chunks(PROCESSED_DATA_FILE)
                    chunk_count = len(processed_chunks)
                else:
                    try:
                        import pickle
                        with open(os.path.join(RAG_FILES_DIR, 'documents.pkl'), 'rb') as f:
                            documents = pickle.load(f)
                            chunk_count = len(documents)
                    except Exception as e:
                        print(f"Could not count documents: {e}")
                        chunk_count = "неизвестно"
                
                return f"✅ AIEXP система инициализирована с {chunk_count} фрагментами нормативных документов (загружена из индекса)"
        except Exception as e:
            print(f"Не удалось загрузить сохраненную систему: {str(e)}")
    
    # If no RAG system but CSV exists, build from CSV
    if local_csv_exists and query_engine is None:
        try:
            print("Building RAG system from CSV file...")
            processed_chunks_df = load_processed_chunks(PROCESSED_DATA_FILE)
            
            # Check for required columns
            required_columns = {'document_id', 'file_link', 'chunk_text', 'chunk_id'}
            missing_columns = required_columns - set(processed_chunks_df.columns)
            if missing_columns:
                return f"❌ Ошибка при инициализации из CSV: отсутствуют необходимые столбцы: {missing_columns}"
            
            # Fill missing optional columns
            if 'txt_file_id' not in processed_chunks_df.columns:
                processed_chunks_df['txt_file_id'] = processed_chunks_df['document_id']
            if 'section' not in processed_chunks_df.columns:
                processed_chunks_df['section'] = ''
            if 'subsection' not in processed_chunks_df.columns:
                processed_chunks_df['subsection'] = ''
            if 'chunk_length' not in processed_chunks_df.columns:
                processed_chunks_df['chunk_length'] = processed_chunks_df['chunk_text'].str.len()
            
            processed_chunks = processed_chunks_df.to_dict('records')
            if processed_chunks:
                print(f"Building RAG system with {len(processed_chunks)} chunks...")
                query_engine = build_rag_system(processed_chunks)
                return f"✅ AIEXP система инициализирована с {len(processed_chunks)} фрагментами нормативных документов (построена из CSV)"
        except Exception as e:
            return f"❌ Ошибка при инициализации из CSV: {str(e)}"
    
    return "🔄 AIEXP система готова к работе. Загрузите нормативные документы для создания базы знаний."

def get_uploaded_files_info():
    if not os.path.exists(UPLOAD_FOLDER):
        return "Нет загруженных файлов в базе знаний"
    
    files = os.listdir(UPLOAD_FOLDER)
    if not files:
        return "Нет загруженных файлов в базе знаний"
    
    file_info = []
    file_count = len(files)
    for file in files:
        file_path = os.path.join(UPLOAD_FOLDER, file)
        size = os.path.getsize(file_path)
        modified = datetime.fromtimestamp(os.path.getmtime(file_path)).strftime("%Y-%m-%d %H:%M")
        file_info.append(f"📄 {file} ({size} байт, добавлен: {modified})")
    
    return f"Всего нормативных документов в базе: {file_count}\n\n" + "\n".join(file_info)

def upload_files(files):
    global query_engine
    
    if not files:
        return "Файлы не выбраны", get_uploaded_files_info()
    
    uploaded_count = 0
    errors = []
    
    for file in files:
        try:
            filename = os.path.basename(file.name)
            destination = os.path.join(UPLOAD_FOLDER, filename)
            shutil.copy2(file.name, destination)
            uploaded_count += 1
            
            if query_engine is not None:
                try:
                    query_engine = add_new_document_to_system(destination, query_engine)
                except Exception as e:
                    errors.append(f"Ошибка добавления {filename} в систему: {str(e)}")
                    
        except Exception as e:
            errors.append(f"Ошибка загрузки {file.name}: {str(e)}")
    
    result_message = f"Загружено нормативных документов: {uploaded_count}"
    if errors:
        result_message += f"\nОшибки:\n" + "\n".join(errors)
    else:
        result_message += f"\nДокументы автоматически добавлены в базу знаний"
    
    return result_message, get_uploaded_files_info()

def process_all_documents():
    global query_engine
    
    if not os.path.exists(UPLOAD_FOLDER):
        return "Папка с нормативными документами не найдена"
    
    files = os.listdir(UPLOAD_FOLDER)
    if not files:
        return "Нет нормативных документов для обработки"
    
    file_paths = [os.path.join(UPLOAD_FOLDER, f) for f in files]
    
    try:
        processed_chunks = process_multiple_documents(file_paths)
        
        if not processed_chunks:
            return "Не удалось создать фрагменты нормативных документов"
        
        save_processed_chunks(processed_chunks, PROCESSED_DATA_FILE)
        
        query_engine = build_rag_system(processed_chunks)
        
        with open(INDEX_STATE_FILE, 'w', encoding='utf-8') as f:
            json.dump({
                'processed_files': files,
                'chunks_count': len(processed_chunks),
                'last_update': datetime.now().isoformat()
            }, f, ensure_ascii=False, indent=2)
        
        return f"Обработка базы знаний завершена успешно!\nОбработано нормативных документов: {len(files)}\nСоздано фрагментов: {len(processed_chunks)}\nAIEXP система готова для работы с нормативной документацией."
        
    except Exception as e:
        return f"Ошибка при обработке нормативных документов: {str(e)}"

def get_system_status():
    status_info = []
    
    files_count = len(os.listdir(UPLOAD_FOLDER)) if os.path.exists(UPLOAD_FOLDER) else 0
    
    if os.path.exists(INDEX_STATE_FILE):
        with open(INDEX_STATE_FILE, 'r', encoding='utf-8') as f:
            state = json.load(f)
        status_info.append(f"🟢 AIEXP система активна")
        status_info.append(f"📊 Нормативных документов в базе: {files_count}")
        status_info.append(f"📝 Фрагментов в индексе: {state.get('chunks_count', 0)}")
        status_info.append(f"🕒 Последнее обновление: {state.get('last_update', 'Неизвестно')}")
        
        if state.get('processed_files'):
            status_info.append(f"📋 Обработанные документы:")
            for file in state['processed_files'][:10]:
                status_info.append(f"   • {file}")
            if len(state['processed_files']) > 10:
                status_info.append(f"   ... и еще {len(state['processed_files']) - 10} документов")
    else:
        status_info.append("🔴 AIEXP система не инициализирована")
        status_info.append(f"📊 Нормативных документов загружено: {files_count}")
        status_info.append("Обработайте документы для создания базы знаний")
    
    return "\n".join(status_info)

def answer_question(question):
    global query_engine
    
    if not question.strip():
        return "Пожалуйста, введите вопрос по нормативной документации", ""
    
    if query_engine is None:
        return "База знаний не готова. Сначала загрузите и обработайте нормативные документы.", ""
    
    try:
        response = query_documents(query_engine, question)
        formatted_response = format_response_with_sources(response)
        
        answer = formatted_response['answer']
        
        sources_info = []
        sources_info.append("📚 Источники из нормативной документации:")
        for i, source in enumerate(formatted_response['sources'][:5], 1):
            sources_info.append(f"\n{i}. Документ: {source['document_id']}")
            if source['section']:
                sources_info.append(f"   Раздел: {source['section']}")
            if source['subsection']:
                sources_info.append(f"   Подраздел: {source['subsection']}")
            sources_info.append(f"   Фрагмент: ...{source['text_preview'][:150]}...")
        
        return answer, "\n".join(sources_info)
        
    except Exception as e:
        return f"Ошибка при обработке вопроса: {str(e)}", ""

def clear_all_data():
    global query_engine
    
    try:
        if os.path.exists(UPLOAD_FOLDER):
            shutil.rmtree(UPLOAD_FOLDER)
            os.makedirs(UPLOAD_FOLDER)
        
        if os.path.exists("processed_data"):
            shutil.rmtree("processed_data")
            os.makedirs("processed_data")
        
        if os.path.exists(RAG_FILES_DIR):
            shutil.rmtree(RAG_FILES_DIR)
            os.makedirs(RAG_FILES_DIR)
        
        query_engine = None
        
        return "Вся база знаний успешно очищена", get_uploaded_files_info(), get_system_status()
        
    except Exception as e:
        return f"Ошибка при очистке базы знаний: {str(e)}", get_uploaded_files_info(), get_system_status()

# Fix 1: Update file_types parameter format
def create_demo_interface():
    with gr.Blocks(title="AIEXP - AI Expert для нормативной документации", theme=gr.themes.Soft()) as demo:
        
        gr.Markdown("""
        # 🤖 AIEXP - Artificial Intelligence Expert
        
        ## Инструмент для работы с нормативной документацией
        
        **Возможности системы:**
        - 🔍 Поиск информации по запросу с указанием источников среди нормативной документации
        - 📋 Цитирование пунктов нормативной документации из базы знаний  
        - 📝 Краткий пересказ содержания разделов или целых нормативных документов
        - 🔎 Семантический анализ соответствия информации требованиям НД
        - 📋 Формирование пошаговых планов действий на основании требований НД
        
        **Поддерживаемые форматы:** PDF, DOCX, TXT, CSV, XLSX, JSON
        """)
        
        with gr.Tab("🏠 Поиск по нормативным документам"):
            gr.Markdown("### Задайте вопрос по нормативной документации")
            
            with gr.Row():
                with gr.Column(scale=3):
                    question_input = gr.Textbox(
                        label="Ваш вопрос к базе знаний",
                        placeholder="Введите вопрос по нормативным документам...",
                        lines=3
                    )
                    ask_btn = gr.Button("🔍 Найти ответ", variant="primary", size="lg")
                    
                    # Fix 2: Simplify Examples component
                    gr.Examples(
                        examples=[
                            ["Какой стандарт устанавливает порядок признания протоколов испытаний продукции в области использования атомной энергии?"],
                            ["Кто несет ответственность за организацию и проведение признания протоколов испытаний продукции?"],
                            ["В каких случаях могут быть признаны протоколы испытаний, проведенные лабораториями, не включенными в перечисления?"],
                            ["Какие критерии используются органом по сертификации для анализа документации на втором этапе признания протоколов испытаний?"]
                        ],
                        inputs=question_input
                    )
                
                with gr.Column(scale=4):
                    answer_output = gr.Textbox(
                        label="Ответ на основе нормативных документов",
                        lines=8,
                        interactive=False
                    )
                    
                    sources_output = gr.Textbox(
                        label="Источники из нормативной документации",
                        lines=10,
                        interactive=False
                    )
        
        with gr.Tab("📤 Управление базой знаний (Администратор)"):
            gr.Markdown("### Загрузка и обработка нормативных документов")
            
            with gr.Row():
                with gr.Column(scale=2):
                    # Fix 3: Update File component parameters
                    file_upload = gr.File(
                        label="Выберите нормативные документы для загрузки",
                        file_count="multiple",
                        file_types=[".pdf", ".docx", ".txt", ".csv", ".xlsx", ".json"],
                        type="filepath"  # Add explicit type
                    )
                    
                    with gr.Row():
                        upload_btn = gr.Button("📤 Загрузить в базу знаний", variant="primary")
                        process_btn = gr.Button("⚙️ Переобработать всю базу", variant="secondary")
                        clear_btn = gr.Button("🗑️ Очистить базу знаний", variant="stop")
                
                with gr.Column(scale=2):
                    upload_status = gr.Textbox(
                        label="Статус загрузки",
                        lines=3,
                        interactive=False
                    )
                    
                    processing_status = gr.Textbox(
                        label="Статус обработки базы знаний",
                        lines=5,
                        interactive=False
                    )
            
            gr.Markdown("### База нормативных документов")
            files_info = gr.Textbox(
                label="Документы в базе знаний",
                lines=8,
                interactive=False,
                value=get_uploaded_files_info()
            )
        
        with gr.Tab("📊 Статус AIEXP системы"):
            gr.Markdown("### Информация о состоянии базы знаний")
            
            system_status = gr.Textbox(
                label="Статус AIEXP системы",
                lines=10,
                interactive=False,
                value=get_system_status()
            )
            
            refresh_status_btn = gr.Button("🔄 Обновить статус системы")
        
        # Fix 4: Event handlers with proper error handling
        def safe_upload_files(files):
            try:
                if files is None:
                    return "Файлы не выбраны", get_uploaded_files_info()
                return upload_files(files)
            except Exception as e:
                return f"Ошибка при загрузке файлов: {str(e)}", get_uploaded_files_info()
        
        def safe_answer_question(question):
            try:
                if not question or not question.strip():
                    return "Пожалуйста, введите вопрос по нормативной документации", ""
                return answer_question(question)
            except Exception as e:
                return f"Ошибка при обработке вопроса: {str(e)}", ""
        
        upload_btn.click(
            fn=safe_upload_files,
            inputs=[file_upload],
            outputs=[upload_status, files_info]
        )
        
        process_btn.click(
            fn=process_all_documents,
            outputs=[processing_status]
        )
        
        ask_btn.click(
            fn=safe_answer_question,
            inputs=[question_input],
            outputs=[answer_output, sources_output]
        )
        
        question_input.submit(
            fn=safe_answer_question,
            inputs=[question_input],
            outputs=[answer_output, sources_output]
        )
        
        clear_btn.click(
            fn=clear_all_data,
            outputs=[processing_status, files_info, system_status]
        )
        
        refresh_status_btn.click(
            fn=get_system_status,
            outputs=[system_status]
        )
    
    return demo

# Fix 5: Update launch parameters for compatibility
if __name__ == "__main__":
    print("Инициализация AIEXP системы...")
    init_message = initialize_system()
    print(init_message)
    
    demo = create_demo_interface()
    demo.launch(
        share=True,
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
        debug=False,  # Add debug flag
        quiet=False   # Add quiet flag for better error visibility
    )