src/app.py

"""Assistente RAG do setor de energia.

Dois modelos selecionáveis (só um carregado por vez):
- gemma-search: agêntico (search_document -> document_query -> resposta).
- gemma-naive-rag: RAG simples (consultar_base_conhecimento -> resposta).
"""

import logging

import streamlit as st

import config
from backend import ensure_llm, load_retriever, run_agent, run_naive_agent

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")

st.set_page_config(layout="centered", page_title="Assistente RAG - Energia", page_icon="⚡")
st.title("Case Mobile")

# ---------------------------------------------------------------------------
# Seletor de modelo
# ---------------------------------------------------------------------------
model_key = st.selectbox(
    "Modelo",
    list(config.MODELS.keys()),
    index=list(config.MODELS).index(config.DEFAULT_MODEL),
)

# Trocar de modelo limpa o histórico (contexto/fluxo diferente)
if st.session_state.get("ui_model_key") != model_key:
    st.session_state.history = []
    st.session_state.ui_model_key = model_key

# ---------------------------------------------------------------------------
# Carregamento (retriever cacheado; LLM trocado conforme o seletor)
# ---------------------------------------------------------------------------
if "retriever" not in st.session_state:
    with st.spinner("Carregando banco vetorial (1º start baixa o banco; pode demorar)..."):
        st.session_state.retriever = load_retriever()

if st.session_state.get("llm_key") != model_key:
    with st.spinner(f"Carregando modelo '{model_key}' (descarrega o anterior; pode demorar)..."):
        ensure_llm(model_key)

model = st.session_state.llm_model
tokenizer = st.session_state.llm_tokenizer

if "history" not in st.session_state:
    st.session_state.history = []  # lista de (pergunta, resposta)

# Histórico
for pergunta, resposta in st.session_state.history:
    with st.chat_message("user"):
        st.markdown(pergunta)
    with st.chat_message("assistant"):
        st.markdown(resposta)


def _format_event(name: str, args: dict) -> str:
    if name == "search_document":
        return f"🔎 **search_document** — `{args.get('query', '')}`"
    if name == "document_query":
        return (
            f"📄 **document_query** — `{args.get('document_id', '')}` "
            f"· _{args.get('query', '')}_"
        )
    if name == "consultar_base_conhecimento":
        return f"🔎 **consultar_base_conhecimento** — `{args.get('termo_busca', '')}`"
    return f"🛠️ **{name}** — `{args}`"


if pergunta := st.chat_input("Pergunte sobre as regulamentações do setor de energia..."):
    with st.chat_message("user"):
        st.markdown(pergunta)

    runner = run_naive_agent if config.MODELS[model_key]["flow"] == "naive" else run_agent

    with st.chat_message("assistant"):
        with st.status("Consultando a base de conhecimento...", expanded=True) as status:
            def on_event(name, args):
                status.write(_format_event(name, args))

            try:
                resposta = runner(
                    model,
                    tokenizer,
                    st.session_state.retriever,
                    pergunta,
                    on_event=on_event,
                )
                status.update(label="Concluído", state="complete", expanded=False)
            except Exception as exc:
                logging.exception("Erro no agente")
                resposta = f"Ocorreu um erro ao processar a pergunta: {exc}"
                status.update(label="Erro", state="error")

        st.markdown(resposta)

    st.session_state.history.append((pergunta, resposta))