app.py

import os
import glob
import nltk
from huggingface_hub import login
import gradio as gr

# ШАГ 1: Гарантированная загрузка необходимого компонента NLTK
# Эта команда выполняется в самом начале и решает ошибку 'LookupError'.
try:
    print("Проверяем наличие NLTK компонента 'punkt'...")
    nltk.data.find('tokenizers/punkt')
    print("'punkt' уже на месте.")
except LookupError:
    print("Компонент 'punkt' не найден. Запускаю принудительную загрузку...")
    nltk.download('punkt', quiet=False)
    print("Загрузка 'punkt' завершена.")


from langchain_community.document_loaders import UnstructuredWordDocumentLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain.chains import RetrievalQA
from langchain_huggingface import HuggingFaceEndpoint

# --- Константы и настройки ---
DOCX_FILE_PATH = "" # Путь будет определен автоматически
FAISS_INDEX_PATH = "faiss_index"
MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
REPO_ID = "mistralai/Mistral-7B-Instruct-v0.2"

# --- ШАГ 2: Автоматическое создание базы знаний ---

def create_vector_db_if_not_exists():
    """Создает векторную базу данных, если она еще не создана."""
    if os.path.exists(FAISS_INDEX_PATH):
        print(f"База знаний найдена в '{FAISS_INDEX_PATH}'. Пропускаю создание.")
        return

    print("База знаний не найдена. Запускаю процесс создания...")
    
    # Автоматически находим первый .docx файл в репозитории
    docx_files = glob.glob("*.docx")
    if not docx_files:
        raise FileNotFoundError("Ошибка: Не найден .docx файл в репозитории. Пожалуйста, загрузите ваш документ.")
    
    global DOCX_FILE_PATH
    DOCX_FILE_PATH = docx_files[0]
    print(f"Найден документ для обработки: {DOCX_FILE_PATH}")

    # Загрузка и обработка документа
    loader = UnstructuredWordDocumentLoader(DOCX_FILE_PATH)
    documents = loader.load()
    
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=150)
    docs = text_splitter.split_documents(documents)
    
    print(f"Документ разделен на {len(docs)} частей.")

    # Создание эмбеддингов и сохранение базы
    embeddings = HuggingFaceEmbeddings(model_name=MODEL_NAME)
    db = FAISS.from_documents(docs, embeddings)
    db.save_local(FAISS_INDEX_PATH)
    
    print(f"База знаний успешно создана и сохранена в '{FAISS_INDEX_PATH}'.")

# --- ШАГ 3: Загрузка AI и настройка логики чат-бота ---

def initialize_qa_chain():
    """Инициализирует и возвращает готовую к работе цепочку QA."""
    # Проверка наличия токена
    HF_TOKEN = os.environ.get("HUGGINGFACEHUB_API_TOKEN")
    if not HF_TOKEN:
        raise ValueError("Секрет 'HUGGINGFACEHUB_API_TOKEN' не найден. Пожалуйста, добавьте его в настройках Space.")
    
    login(token=HF_TOKEN)
    print("Успешная аутентификация по токену.")

    # Загрузка векторной базы
    embeddings = HuggingFaceEmbeddings(model_name=MODEL_NAME)
    db = FAISS.load_local(FAISS_INDEX_PATH, embeddings, allow_dangerous_deserialization=True)

    # Настройка языковой модели (LLM) через новый класс HuggingFaceEndpoint
    llm = HuggingFaceEndpoint(
        repo_id=REPO_ID,
        temperature=0.3,
        max_new_tokens=512,
        repetition_penalty=1.1,
        huggingfacehub_api_token=HF_TOKEN
    )

    # Создание цепочки Вопрос-Ответ
    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=db.as_retriever(search_kwargs={"k": 3}),
        return_source_documents=True
    )
    print("Цепочка QA успешно инициализирована.")
    return qa_chain

# --- Основной код ---
create_vector_db_if_not_exists()
qa_chain = initialize_qa_chain()

def chatbot_response(message, history):
    """Основная функция для генерации ответа чат-бота."""
    response = qa_chain.invoke({"query": message})
    return response["result"]

# --- ШАГ 4: Создание веб-интерфейса ---

with gr.Blocks(theme='gradio/soft', title="AI-Ассистент по ВКР") as demo:
    gr.Markdown("# 🤖 AI-Ассистент по вопросам ВКР")
    gr.Markdown(f"Этот бот отвечает на вопросы на основе документа: **{os.path.basename(DOCX_FILE_PATH)}**.")
    
    chatbot = gr.ChatInterface(
        fn=chatbot_response,
        chatbot=gr.Chatbot(height=400),
        textbox=gr.Textbox(placeholder="Задайте свой вопрос о правилах оформления, сроках или структуре ВКР...", container=False, scale=7),
        title=None,
        examples=[
            "Какие требования к объему магистерской диссертации?",
            "Как правильно оформить список литературы?",
            "Какие сроки сдачи и защиты ВКР в этом году?",
            "Что должно быть во введении?",
            "Какой процент оригинальности требуется?"
        ],
        clear_btn="🗑️ Очистить диалог",
        retry_btn=None,
        undo_btn=None,
    )

# Запуск приложения
demo.launch()