app.py

import gradio as gr
import os
import sys
from pathlib import Path
from dotenv import load_dotenv

# Add project root to sys.path
sys.path.append(str(Path(__file__).parent))

import config 
from modules.rag_system import rag_system

load_dotenv()


def format_sources(sources):
    if not sources:
        return ""
    
    html = "<div style='margin-top: 3rem;'><h3 style='color: #3E8BF7; border-bottom: 1px solid rgba(62,139,247,0.3); padding-bottom: 0.5rem;'>📚 Використані джерела:</h3>"
    for i, res in enumerate(sources):
        chunk = res.get('chunk', {})
        meta = chunk.get('metadata', {})
        source_name = meta.get('source_name', 'Законодавство')
        article = meta.get('article_title', '')
        content = chunk.get('text', '')
        
        html += f"""
        <div class='source-card'>
            <div class='source-meta'>[{i+1}] {source_name} {f'— {article}' if article else ''}</div>
            <div style='font-size: 0.95rem; line-height: 1.6; color: rgba(255,255,255,0.85);'>{content}</div>
        </div>
        """
    html += "</div>"
    return html

def run_chat(query, api_key, search_method, use_reranker, legal_area, top_k, temperature):
    if not query.strip():
        return "Будь ласка, введіть запитання.", ""
    
    # Mapping display names to internal values
    method_map = {
        "🔄 Гібридний (Рекомендовано)": "hybrid",
        "🔤 BM25 (Ключові слова)": "bm25",
        "🧠 Семантичний": "semantic"
    }
    
    # Mapping Ukrainian legal area names to internal keys
    area_map = {
        "Всі": "Всі",
        "Сімейне право": "сімейне_право",
        "Трудове право": "трудове_право",
        "Земельне право": "земельне_право",
        "Цивільне право": "цивільне_право",
        "Податкове право": "податкове_право",
        "Кримінальне право": "кримінальне_право",
        "Конституційне право": "конституційне_право",
        "Адміністративне судочинство": "адміністративне_судочинство"
    }
    
    internal_method = method_map.get(search_method, "hybrid")
    internal_area = area_map.get(legal_area, "Всі")
    
    try:
        answer, sources = rag_system.process_query(
            query=query,
            api_key=api_key,
            use_reranker=use_reranker,
            top_k_rerank=int(top_k),
            temperature=float(temperature),
            search_method=internal_method,
            legal_area=internal_area
        )
        
        sources_html = format_sources(sources)
        return answer, sources_html
    except Exception as e:
        return f"Помилка при обробці запиту: {str(e)}", ""

# --- Gradio UI Construction ---
css_path = Path("assets/style.css")
with open(css_path, "r", encoding="utf-8") as f:
    custom_css = f.read()

with gr.Blocks(title="Асистент із Законодавства", css=custom_css, theme=gr.themes.Soft()) as demo:
    # Header
    with gr.Row(elem_classes="header-container"):
        with gr.Column(scale=0, min_width=80):
            gr.HTML("<div style='font-size: 40px;'>🇺🇦</div>")
        
        with gr.Column(scale=1):
            gr.HTML("""
                <div class='header-text'>
                    <h1>Асистент із Законодавства України</h1>
                    <p>Інтелектуальна система пошуку та аналізу юридичних документів на базі ШІ</p>
                </div>
            """)

    with gr.Row():
        # --- Left Sidebar ---
        with gr.Column(scale=1, elem_classes="sidebar"):
            gr.Markdown("### ⚙️ Конфігурація")
            
            with gr.Group():
                api_key_btn = gr.Button("🔑 Встановити API ключ", variant="secondary")
                api_key_input = gr.Textbox(
                    type="password", 
                    label="Groq API Key", 
                    placeholder="Введіть ваш ключ...", 
                    value=os.getenv("GROQ_API_KEY", "")
                )
            
            search_method = gr.Radio(
                label="Метод пошуку",
                choices=["🔄 Гібридний (Рекомендовано)", "🔤 BM25 (Ключові слова)", "🧠 Семантичний"],
                value="🔄 Гібридний (Рекомендовано)"
            )
            
            use_reranker = gr.Checkbox(label="Використовувати Reranker", value=True)

            legal_area = gr.Dropdown(
                label="Галузь права (Фільтр)",
                choices=[
                    "Всі", 
                    "Сімейне право", 
                    "Трудове право", 
                    "Земельне право", 
                    "Цивільне право", 
                    "Податкове право", 
                    "Кримінальне право", 
                    "Конституційне право", 
                    "Адміністративне судочинство"
                ],
                value="Всі"
            )

            with gr.Accordion("🛠️ Розширені параметри", open=False):
                top_k = gr.Slider(label="Кількість джерел", minimum=1, maximum=20, step=1, value=config.DEFAULT_TOP_K_RERANK)
                temperature = gr.Slider(label="Температура генерації", minimum=0.0, maximum=1.0, step=0.1, value=0.5)

        # --- Main Content Area ---
        with gr.Column(scale=3):
            with gr.Group():
                query_input = gr.Textbox(
                    label="💬 Ваше запитання", 
                    placeholder="Наприклад: Які підстави для звільнення працівника за прогул?",
                    lines=3
                )
                submit_btn = gr.Button("🔍 Знайти відповідь", elem_classes="primary-btn")
            
            with gr.Column(visible=True):
                output_answer = gr.Markdown(label="📝 Відповідь асистента")
                output_sources = gr.HTML(label="📚 Джерела")

            gr.Markdown("### 🔍 Швидкі приклади")
            example_questions = [
                "Що загрожує за найманство згідно з Кримінальним кодексом України?",
                "Які підстави для звільнення працівника за прогул?",
                "Які права має затримана особа під час затримання?",
                "Що таке розбій і яка стаття КК України його регулює?",
                "Які строки притягнення до адміністративної відповідальності?"
            ]
            
            for q in example_questions:
                btn = gr.Button(q, elem_classes="example-btn")
                btn.click(lambda x=q: x, outputs=[query_input]).then(
                    fn=run_chat,
                    inputs=[query_input, api_key_input, search_method, use_reranker, legal_area, top_k, temperature],
                    outputs=[output_answer, output_sources]
                )
            
    # --- Interactions ---
    submit_btn.click(
        fn=run_chat,
        inputs=[query_input, api_key_input, search_method, use_reranker, legal_area, top_k, temperature],
        outputs=[output_answer, output_sources]
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)