main.py

import streamlit as st
from vector_db import setup_vector_db, get_local_content
from query_processing import process_query
from web_scrapping import get_web_content
import tempfile
import os

def main():
    # Configuração da página do Streamlit
    st.set_page_config(page_title="Agentic RAG Query System", page_icon="🤖", layout="wide")
    st.title("Agentic RAG Query System")
    st.markdown("""
        Bem-vindo ao sistema de consulta baseado em RAG (Retrieval-Augmented Generation)! 
        Faça uma pergunta e o sistema buscará informações em documentos locais ou na web para fornecer uma resposta.
    """)

    # Upload do PDF pelo usuário
    uploaded_file = st.file_uploader("Carregue seu arquivo PDF", type="pdf")
    
    # Inicialização do banco de dados vetorial
    if uploaded_file is not None:
        if 'vector_db' not in st.session_state:
            with st.spinner("Configurando o banco de dados vetorial..."):
                # Salva o arquivo carregado temporariamente
                with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file:
                    tmp_file.write(uploaded_file.getvalue())
                    tmp_file_path = tmp_file.name

                # Passa o caminho do arquivo temporário para a função setup_vector_db
                st.session_state.vector_db = setup_vector_db(tmp_file_path)
                st.session_state.local_context = get_local_content(st.session_state.vector_db, "")

                # Remove o arquivo temporário após o uso
                os.unlink(tmp_file_path)

            st.success("Banco de dados vetorial configurado com sucesso!")
    else:
        st.warning("Por favor, carregue um arquivo PDF para continuar.")

    # Entrada do usuário
    query = st.text_input("Faça sua pergunta:", placeholder="Ex: O que é Agentic RAG?")

    # Botão para processar a consulta
    if st.button("Buscar Resposta"):
        if query and 'vector_db' in st.session_state:
            with st.spinner("Processando sua consulta..."):
                try:
                    # Processa a consulta
                    result = process_query(query, st.session_state.vector_db, st.session_state.local_context)
                    
                    # Exibe o resultado
                    st.subheader("Resposta:")
                    st.write(result)
                except Exception as e:
                    st.error(f"Ocorreu um erro ao processar sua consulta: {e}")
        else:
            st.warning("Por favor, insira uma pergunta e carregue um PDF.")

if __name__ == "__main__":
    main()