app.py

# ---------- app.py ----------
# Dependências:
# pip install gradio faiss-cpu sentence-transformers openai

import os
from pathlib import Path
import pickle

import gradio as gr
import faiss
from sentence_transformers import SentenceTransformer
from openai import OpenAI

# ========= NVIDIA API =========
# Em local: defina NV_API_KEY ou NVIDIA_API_KEY no ambiente.
# Em Hugging Face Spaces: crie um "Repository secret" chamado NVIDIA_API_KEY.
NV_API_KEY = os.environ.get("NVIDIA_API_KEY") or os.environ.get("NV_API_KEY")
if not NV_API_KEY:
    raise RuntimeError(
        "A chave da NVIDIA não foi encontrada.\n"
        "Defina um secret chamado 'NVIDIA_API_KEY' (ou NV_API_KEY) com a tua chave da NVIDIA.\n"
        "• Localmente: export NVIDIA_API_KEY='SUA_CHAVE'\n"
        "• Hugging Face Spaces: Settings -> Repository secrets -> Add secret."
    )

client = OpenAI(
    base_url="https://integrate.api.nvidia.com/v1",
    api_key=NV_API_KEY,
)
CHAT_MODEL = "meta/llama3-8b-instruct"

# ========= Configuração do App =========
APP_TITLE = "EcoLexIA – Assistente Inteligente de Leis Ambientais de Portugal"

INTRO = (
    "👋 Bem-vindo ao **EcoLexIA**, o teu assistente jurídico especializado em **direito do ambiente em Portugal**.\n\n"
    "Este sistema utiliza **RAG (Retrieval-Augmented Generation)** para consultar automaticamente os documentos legais "
    "carregados (leis, decretos, regulamentos, pareceres, etc.) e responder às tuas perguntas com base nesses textos."
)

SUGGESTION_QUESTIONS = [
    "Resuma os principais princípios da Lei de Bases do Ambiente.",
    "Quais são as obrigações do Estado em matéria de proteção ambiental?",
    "Explique como funciona a Avaliação de Impacte Ambiental em Portugal.",
    "Que legislação regula a gestão de resíduos urbanos?",
    "Existe enquadramento legal para participação pública em decisões ambientais?",
    "Quais são as regras sobre emissões poluentes na indústria?",
]

SUGGESTIONS_THEMES = {
    "Lei de Bases do Ambiente": [
        "Quais são os princípios fundamentais da Lei de Bases do Ambiente?",
        "Como a Lei de Bases do Ambiente enquadra o desenvolvimento sustentável?",
    ],
    "Avaliação de Impacte Ambiental (AIA)": [
        "Explique o que é Avaliação de Impacte Ambiental e quando é obrigatória.",
        "Que entidades estão envolvidas no processo de AIA?",
    ],
    "Resíduos & Poluição": [
        "Que legislação trata da gestão de resíduos em Portugal?",
        "Que obrigações têm as empresas relativamente ao controlo de emissões poluentes?",
    ],
    "Ordenamento do Território & Conservação": [
        "Como o ordenamento do território se articula com a proteção ambiental?",
        "Que diplomas legais regulam áreas protegidas e conservação da natureza?",
    ],
}

# ========= Caminhos do índice =========
INDEX_FILE = "faiss_index.faiss"
EMBEDDINGS_FILE = "embeddings.pkl"

if not Path(INDEX_FILE).exists() or not Path(EMBEDDINGS_FILE).exists():
    raise FileNotFoundError(
        "❌ Índice não encontrado.\n"
        "Certifique-se de que 'faiss_index.faiss' e 'embeddings.pkl' "
        "foram gerados pelo build_index.py na mesma pasta deste app."
    )

index = faiss.read_index(INDEX_FILE)
with open(EMBEDDINGS_FILE, "rb") as f:
    emb_data = pickle.load(f)

texts = emb_data.get("texts", [])
metadatas = emb_data.get("metadatas", [])

# Mesmo modelo de embeddings usado no build_index.py
embedding_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

# Histórico do chat: lista de (user, assistant)
dialog_history = []


# ========= Recuperação de contexto =========
def retrieve_context(query: str, k: int = 4) -> str:
    """Busca k trechos mais relevantes no índice FAISS para a pergunta."""
    if not query or not query.strip():
        return ""

    # Proteções básicas
    if index.ntotal == 0 or not texts:
        return ""

    q_emb = embedding_model.encode([query], convert_to_numpy=True)
    _, indices = index.search(q_emb, k)

    parts = []
    for idx in indices[0]:
        if idx < 0 or idx >= len(texts):
            continue
        chunk = texts[idx]
        meta = metadatas[idx] if idx < len(metadatas) else {}
        src = meta.get("source", "documento desconhecido")
        parts.append(f"[Documento: {src}]\n{chunk}")

    return "\n\n---\n\n".join(parts)


# ========= Streaming da NVIDIA (ROBUSTO) =========
def nv_stream(messages, temperature: float, top_p: float, max_tokens: int):
    """
    Stream da resposta da NVIDIA (LLaMA 3) com proteções:
    - chunks podem vir sem `choices` (ou choices vazia)
    - `delta` pode não existir
    - `delta.content` pode ser None
    """
    reply = ""

    stream = client.chat.completions.create(
        model=CHAT_MODEL,
        messages=messages,
        temperature=temperature,
        top_p=top_p,
        max_tokens=max_tokens,
        stream=True,
    )

    for chunk in stream:
        choices = getattr(chunk, "choices", None)
        if not choices:
            continue

        choice0 = choices[0]
        delta = getattr(choice0, "delta", None)
        if not delta:
            continue

        content = getattr(delta, "content", None)
        if content:
            reply += content
            yield reply


# ========= Lógica do chat =========
def chatbot(user_input: str, temperature: float, top_p: float, max_tokens: int):
    global dialog_history

    if not user_input or not user_input.strip():
        return dialog_history, ""

    context = retrieve_context(user_input, k=6)

    # ✅ Se não houver contexto, não deixa o modelo inventar
    if not context.strip():
        reply_full = (
            "Não encontrei base suficiente nos documentos carregados para responder com segurança.\n\n"
            "Sugestões:\n"
            "• Verifica se os PDFs/leis foram realmente indexados.\n"
            "• Faz uma pergunta mais específica (ex.: diploma/ano/artigo).\n"
            "• Consulta a legislação oficial (Diário da República Eletrónico).\n"
        )
        dialog_history.append((user_input, reply_full))
        return dialog_history, ""

    system_msg = {
        "role": "system",
        "content": (
            "És um assistente jurídico especializado em direito do ambiente em Portugal. "
            "Responde SEMPRE em português europeu, de forma clara e estruturada.\n\n"
            "Regras:\n"
            "1. Usa apenas o contexto abaixo para responder.\n"
            "2. Se não houver informação suficiente, diz que não encontras base nos documentos e "
            "sugere consultar a legislação oficial.\n"
            "3. Indica o nome do documento (PDF) sempre que fizer sentido.\n\n"
            f"=== CONTEXTO RECUPERADO ===\n{context}\n\n"
        ),
    }

    messages = [system_msg]

    # Reconstroi histórico (limita para evitar prompt gigante)
    MAX_TURNS = 8  # últimas 8 interações
    for u, a in dialog_history[-MAX_TURNS:]:
        messages.append({"role": "user", "content": u})
        messages.append({"role": "assistant", "content": a})

    messages.append({"role": "user", "content": user_input})

    reply_full = ""
    try:
        for partial in nv_stream(messages, temperature, top_p, max_tokens):
            reply_full = partial

        # ✅ Se o stream não devolveu conteúdo, evita choices[0] e afins
        if not reply_full.strip():
            reply_full = (
                "⚠️ A API devolveu uma resposta vazia (sem conteúdo). "
                "Isto pode acontecer por limite de contexto, rate limit, ou erro de input.\n"
                "Tenta reduzir o tamanho dos trechos recuperados (k) ou o histórico."
            )

        dialog_history.append((user_input, reply_full))

    except Exception as e:
        reply_full = f"⚠️ Erro na API NVIDIA: {type(e).__name__}: {e}"
        dialog_history.append((user_input, reply_full))

    return dialog_history, ""


def clear_history():
    global dialog_history
    dialog_history = []
    return [], ""


# ========= CSS simples / layout padrão =========
custom_css = r"""
body, .gradio-container {
    background: #ffffff;
    font-family: system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif;
}

/* Header azul */
#header-box {
    max-width: 1100px;
    margin: 1.5rem auto 1rem auto;
}

.header-card {
    background: linear-gradient(135deg, #0b3c91 0%, #1565c0 40%, #1e88e5 100%);
    border-radius: 16px;
    padding: 1.4rem 1.8rem;
    color: #ffffff;
    box-shadow: 0 14px 30px rgba(15, 23, 42, 0.18);
}

.header-title {
    font-size: 1.6rem;
    font-weight: 700;
    margin: 0;
    color: #ffffff !important;
}

.header-subtitle {
    margin: 0.35rem 0 0 0;
    font-size: 0.96rem;
    opacity: 0.95;
    color: #ffffff !important;
}

/* Cartões principais */
.card {
    background: #ffffff;
    border-radius: 16px;
    border: 1px solid #e0e0e0;
    box-shadow: 0 8px 18px rgba(15, 23, 42, 0.06);
    padding: 1rem 1.1rem;
}

/* Chat */
#chat-window {
    height: 60vh;
}

/* Botões */
.gr-button-primary {
    background: #1565c0 !important;
    color: #ffffff !important;
    border: none !important;
}

.gr-button-secondary {
    background: #f5f5f5 !important;
    color: #333333 !important;
    border: 1px solid #e0e0e0 !important;
}

/* Sugestões */
.suggestion-btn {
    width: 100%;
    justify-content: flex-start;
    font-size: 0.88rem;
}

/* Rodapé */
.app-footer {
    margin-top: 1rem;
    font-size: 0.8rem;
    text-align: center;
    color: #555555;
}
"""


# ========= Layout Gradio =========
with gr.Blocks(title=APP_TITLE, css=custom_css, theme=gr.themes.Soft()) as demo:
    # Header
    with gr.Group(elem_id="header-box"):
        gr.HTML(
            f"""
            <div class="header-card">
                <div class="header-title">{APP_TITLE}</div>
                <div class="header-subtitle">
                    Consultor jurídico inteligente com RAG sobre legislação ambiental portuguesa.
                </div>
            </div>
            """
        )

    gr.Markdown(INTRO)

    with gr.Row():
        # Coluna principal (chat)
        with gr.Column(scale=3):
            with gr.Group(elem_classes="card"):
                gr.Markdown("### 💬 Conversa Jurídica")
                chatbot_ui = gr.Chatbot(
                    type="tuples",
                    elem_id="chat-window",
                    label="Chatbot",
                )

                txt = gr.Textbox(
                    placeholder="Escreve aqui a tua pergunta sobre leis do ambiente em Portugal…",
                    lines=3,
                    show_label=False,
                )

                with gr.Row():
                    btn_send = gr.Button("Enviar", variant="primary")
                    btn_clear = gr.Button("Limpar", variant="secondary")

                with gr.Accordion("Parâmetros avançados", open=False):
                    temperature = gr.Slider(0, 1, value=0.5, label="Temperature")
                    top_p = gr.Slider(0, 1, value=0.9, label="Top-p")
                    max_tokens = gr.Slider(64, 2048, value=512, step=64, label="Max Tokens")

            btn_send.click(
                chatbot,
                [txt, temperature, top_p, max_tokens],
                [chatbot_ui, txt],
            )
            txt.submit(
                chatbot,
                [txt, temperature, top_p, max_tokens],
                [chatbot_ui, txt],
            )
            btn_clear.click(
                clear_history,
                [],
                [chatbot_ui, txt],
            )

        # Sidebar
        with gr.Column(scale=2):
            with gr.Group(elem_classes="card"):
                gr.Markdown("### 💡 Sugestões rápidas")
                for q in SUGGESTION_QUESTIONS:
                    gr.Button(q, elem_classes="suggestion-btn").click(
                        lambda s=q: s, outputs=[txt]
                    )

                gr.Markdown("---")
                gr.Markdown("### 📚 Explorar por tema")

                for theme, qs in SUGGESTIONS_THEMES.items():
                    with gr.Accordion(theme, open=False):
                        for q in qs:
                            gr.Button(q, elem_classes="suggestion-btn").click(
                                lambda s=q: s, outputs=[txt]
                            )

    gr.Markdown(
        '<div class="app-footer">EcoLexIA · Sistema RAG para legislação ambiental em Portugal</div>'
    )

# Para Hugging Face Spaces basta que a variável `demo` exista;
# ainda assim manter o launch permite rodar localmente.
if __name__ == "__main__":
    demo.launch()