app.py

"""
Веб-интерфейс RAG-системы — Gradio
Запуск:  python3 app.py
Публичная ссылка создаётся автоматически (--share).
"""

import os
import shutil
import threading
from pathlib import Path

import gradio as gr
from rag_system import RAGSystem, CHROMA_PERSIST_DIR

# ─── Глобальный экземпляр системы (thread-safe) ───────────────────────────
_rag: RAGSystem | None = None
_rag_lock   = threading.Lock()   # инициализация RAGSystem
_index_lock = threading.Lock()   # индексация PDF


def _get_rag(provider: str, openai_key: str) -> RAGSystem:
    """
    Возвращает единственный экземпляр RAGSystem (singleton).
    Пересоздаёт только при смене провайдера.
    Lock предотвращает двойную инициализацию при параллельных запросах Gradio.
    """
    global _rag
    key = openai_key.strip() if openai_key else None
    if key:
        if provider == "groq":
            os.environ["GROQ_API_KEY"] = key
        elif provider == "gemini":
            os.environ["GOOGLE_API_KEY"] = key
        else:
            os.environ["OPENAI_API_KEY"] = key

    with _rag_lock:
        if _rag is None:
            _rag = RAGSystem(llm_provider=provider)
        elif _rag.llm_provider != provider:
            # Меняем только LLM — vectorstore и индексированные данные сохраняются
            _rag.switch_llm(provider)
    return _rag


# ─── Обработчики ──────────────────────────────────────────────────────────

def index_pdf(files, provider: str, openai_key: str) -> str:
    """Принимает список загруженных PDF, индексирует их."""
    if not files:
        return "⚠️ Файлы не выбраны."
    with _index_lock:
        try:
            rag = _get_rag(provider, openai_key)
            log_lines = []
            for file in files:
                # Gradio 5 возвращает строку-путь, Gradio 4 — объект с .name
                tmp_path = Path(file if isinstance(file, str) else file.name)
                added = rag.add_documents(str(tmp_path))
                if added == 0:
                    log_lines.append(f"⏭️ {tmp_path.name} — уже в базе, пропущен")
                else:
                    log_lines.append(f"✅ {tmp_path.name} — добавлено {added} чанков")
            stats = rag.get_stats()
            log_lines.append(f"\nВсего в базе: {stats['total_chunks']} чанков")
            return "\n".join(log_lines)
        except Exception as e:
            return f"❌ Ошибка индексации: {e}"


def ask(question: str, provider: str, openai_key: str):
    """Обрабатывает вопрос, возвращает (ответ, источники)."""
    if not question.strip():
        return "Введите вопрос.", ""
    try:
        rag = _get_rag(provider, openai_key)
        result = rag.ask_question(question)

        # Форматируем источники отдельно
        sources_lines = []
        seen: set = set()
        for src in result.sources:
            key = (src["file"], src["page"])
            if key not in seen:
                seen.add(key)
                score = f"{src['score']:.0%}" if src.get("score") is not None else "—"
                sources_lines.append(f"• **{src['file']}**, стр. {src['page']}  `[{score}]`")

        if result.conflicts:
            sources_lines.append("")
            sources_lines.append("⚠️ **Обнаружены противоречия между источниками:**")
            for c in result.conflicts:
                sources_lines.append(f"  — {c}")

        sources_md = "\n".join(sources_lines) if sources_lines else "_Источники не найдены_"
        return result.answer, sources_md

    except Exception as e:
        return f"❌ Ошибка: {e}", ""


def clear_db(provider: str, openai_key: str) -> str:
    """Очищает векторную базу данных."""
    global _rag
    db_path = Path(CHROMA_PERSIST_DIR)
    if db_path.exists():
        shutil.rmtree(db_path)
    _rag = None
    return "🗑️ База данных очищена. Загрузите PDF заново."


def toggle_key_visibility(provider: str):
    """Показывает/скрывает поле API-ключа в зависимости от провайдера."""
    return gr.update(visible=(provider in ("openai", "groq", "gemini")))


# ─── Интерфейс ────────────────────────────────────────────────────────────

CSS = """
#title { text-align: center; }
#answer-box textarea { font-size: 14px; min-height: 260px; }
#sources-box { font-size: 13px; }
.gradio-container { max-width: 900px !important; margin: auto; }
"""

with gr.Blocks(title="RAG-чатбот для документов") as demo:

    gr.Markdown("# 📄 RAG-чатбот для документов", elem_id="title")
    gr.Markdown(
        "Загрузите PDF или DOCX документ → задайте вопрос → получите ответ со ссылкой на источник.",
        elem_id="title",
    )

    with gr.Row():
        provider = gr.Radio(
            choices=["gemini", "groq", "openai", "ollama"],
            value="gemini",
            label="LLM-провайдер",
            scale=1,
        )
        openai_key = gr.Textbox(
            label="API Key (Groq: gsk_... / OpenAI: sk-...)",
            placeholder="gsk_... или sk-...",
            type="password",
            visible=True,
            scale=2,
        )

    provider.change(toggle_key_visibility, inputs=provider, outputs=openai_key)

    with gr.Tab("📄 Загрузка документов"):
        pdf_input = gr.File(
            label="PDF / DOCX файлы",
            file_types=["pdf", ".pdf", ".docx", ".doc"],
            file_count="multiple",
        )
        with gr.Row():
            index_btn = gr.Button("📥 Индексировать", variant="primary")
            clear_btn = gr.Button("🗑️ Очистить базу", variant="stop")
        index_log = gr.Textbox(label="Лог индексации", lines=6, interactive=False)

        index_btn.click(index_pdf, inputs=[pdf_input, provider, openai_key], outputs=index_log)
        clear_btn.click(clear_db, inputs=[provider, openai_key], outputs=index_log)

    with gr.Tab("💬 Вопрос-ответ"):
        question = gr.Textbox(
            label="Вопрос",
            placeholder="Например: Какие услуги в разделе 2.1 предоставляются бесплатно?",
            lines=2,
        )
        ask_btn = gr.Button("🔍 Спросить", variant="primary")

        answer_box = gr.Textbox(
            label="Ответ",
            lines=12,
            interactive=False,
            elem_id="answer-box",
        )
        sources_box = gr.Markdown(
            label="Источники",
            elem_id="sources-box",
        )

        ask_btn.click(
            ask,
            inputs=[question, provider, openai_key],
            outputs=[answer_box, sources_box],
        )
        question.submit(
            ask,
            inputs=[question, provider, openai_key],
            outputs=[answer_box, sources_box],
        )

    with gr.Tab("ℹ️ Инструкция"):
        gr.Markdown("""
### Как пользоваться

1. **Выберите провайдер** — `gemini` (по умолчанию) или `groq`, `openai`, `ollama`.
2. **Введите API-ключ** — Gemini / Groq: `gsk_...` / OpenAI: `sk-...`.
3. **Вкладка «Загрузка»** — загрузите PDF или DOCX и нажмите **Индексировать**.
4. **Вкладка «Вопрос-ответ»** — введите вопрос и нажмите **Спросить** (или Enter).

### Возможности системы

| Функция | Описание |
|---|---|
| 📄 PDF + DOCX | Загружайте документы любого объёма |
| 🔍 Гибридный поиск | BM25 + векторный — находит точные коды и семантику |
| 🧠 Chain of Thought | Рассуждает пошагово, анализирует данные из документов |
| ⚠️ Конфликты | Обнаруживает противоречия между документами |
| 🛡️ Защита | Не выдумывает — если нет в документе, так и скажет |

### Провайдеры LLM
- **Gemini** (по умолчанию): ключ на [aistudio.google.com](https://aistudio.google.com)
- **Groq**: бесплатный ключ на [console.groq.com](https://console.groq.com)
- **OpenAI**: ключ вида `sk-...`
- **Ollama**: локальный запуск (ключ не нужен)
        """)


if __name__ == "__main__":
    demo.queue()
    demo.launch(css=CSS)