sollai/chat.py

import streamlit as st
import os
from pathlib import Path
from typing import List, Dict
import tempfile
from pinecone import Pinecone
from datetime import datetime

def save_uploaded_file(file) -> str:
    """Salva um arquivo temporariamente e retorna o caminho"""
    with tempfile.NamedTemporaryFile(delete=False, suffix=Path(file.name).suffix) as tmp:
        tmp.write(file.getvalue())
        return tmp.name

def format_citation(citation: dict, index: int) -> str:
    """Formata uma citação com estilo de referência acadêmica"""
    try:
        references = citation.get("references", [])
        position = citation.get("position", 0)
        citation_text = []
        
        for ref in references:
            if not ref:
                continue
                
            file_data = ref.get("file", {}) or {}
            metadata = file_data.get("metadata", {}) or {}
            
            # Usar nome original do arquivo se disponível
            file_name = (metadata.get("original_name") 
                        or file_data.get("name", "Documento"))
            pages = ref.get("pages", [])
            
            # Criar link de referência
            citation_text.append(f"[{index}] {file_name}")
            if pages:
                citation_text.append(f"Page{'s' if len(pages) > 1 else ''} {', '.join(map(str, pages))}")
        
        return " | ".join(citation_text)
    except Exception as e:
        st.error(f"Erro ao formatar citação: {str(e)}")
        return f"[{index}] Referência"

def show_references(citations: List[dict]):
    """Mostra a seção de referências no estilo acadêmico com trechos relevantes"""
    if not citations:
        return
    
    st.markdown("---")
    st.markdown("### References:")
    
    for i, citation in enumerate(citations, 1):
        try:
            references = citation.get("references", [])
            position = citation.get("position", 0)
            
            for ref in references:
                if not ref:
                    continue
                    
                file_data = ref.get("file", {}) or {}
                metadata = file_data.get("metadata", {}) or {}
                
                # Usar nome original do arquivo
                file_name = (metadata.get("original_name") 
                            or file_data.get("name", "Documento"))
                pages = ref.get("pages", [])
                
                with st.container():
                    # Cabeçalho da referência
                    col1, col2 = st.columns([1, 4])
                    with col1:
                        st.markdown(f"**[{i}]**")
                    with col2:
                        st.markdown(f"**{file_name}**")
                        if pages:
                            st.caption(f"Page {', '.join(map(str, pages))}")
                    
                    # Trecho relevante
                    if position is not None and citation.get("text"):
                        with st.expander("Ver trecho citado"):
                            st.caption("Trecho relevante:")
                            st.markdown("> " + citation["text"].strip())
                    
                    st.divider()
        except Exception as e:
            st.error(f"Erro ao mostrar referência {i}: {str(e)}")

def format_snippet(snippet: dict) -> str:
    """Formata um snippet de contexto"""
    try:
        content = snippet.get("content", "").strip()
        score = snippet.get("score", 0)
        reference = snippet.get("reference", {}) or {}
        file_data = reference.get("file", {}) or {}
        metadata = file_data.get("metadata", {}) or {}
        
        # Pegar nome original do arquivo se disponível
        file_name = (metadata.get("original_name") 
                    or file_data.get("name", "Documento"))
        
        # Pegar páginas
        pages = reference.get("pages", [])
        page_info = f"Page {', '.join(map(str, pages))}" if pages else ""
        
        # Formatar texto
        return f"""
        📄 **Fonte:** {file_name} {page_info}
        📊 **Relevância:** {score:.2%}
        
        > {content}
        """
    except Exception as e:
        return f"❌ Erro ao formatar snippet: {str(e)}"

def show_context(snippets: List[dict]):
    """Mostra snippets de contexto relevantes"""
    if not snippets:
        return
        
    with st.expander("🔍 **Contexto Relevante**"):
        try:
            st.info("ℹ️ Trechos relevantes encontrados nos documentos:")
            
            for snippet in snippets:
                if snippet:  # Verificar se o snippet é válido
                    st.markdown(format_snippet(snippet))
                    st.divider()
        except Exception as e:
            st.error(f"❌ Erro ao mostrar contexto: {str(e)}")

def main():
    st.title("💬 Chat")
    
    if 'assistant' not in st.session_state:
        st.error("❌ Nenhum assistente configurado!")
        return
    
    # Área de documentos
    with st.sidebar:
        st.subheader("📁 Documentos")
        
        # Mostrar contagem de documentos
        try:
            files = st.session_state['assistant'].get_files()
            total_files = len(files)
            
            # Mostrar status dos documentos
            if total_files > 0:
                st.info(f"📊 {total_files} documento{'s' if total_files > 1 else ''} carregado{'s' if total_files > 1 else ''}")
                
                # Mostrar detalhes em um expander
                with st.expander("Ver detalhes dos documentos"):
                    for file in files:
                        with st.container():
                            st.text(f"📄 {file['name'] if 'name' in file else 'Documento'}")
                            if 'size' in file:
                                st.caption(f"📦 {file['size']/1024:.1f} KB")
                            if 'created_on' in file:
                                st.caption(f"📅 {file['created_on']}")
                            st.divider()
            else:
                st.warning("⚠️ Nenhum documento carregado")
            
        except Exception as e:
            st.error(f"❌ Erro ao carregar documentos: {str(e)}")
            st.warning("⚠️ Nenhum documento carregado")
        
        st.info("ℹ️ Carregue documentos para permitir que o assistente os utilize nas respostas.")
        
        # Upload de arquivos
        uploaded_files = st.file_uploader(
            "Carregar documentos",
            accept_multiple_files=True,
            type=['pdf', 'txt', 'doc', 'docx']
        )
        
        if uploaded_files:
            if st.button("📤 Enviar Documentos"):
                with st.spinner("Carregando documentos..."):
                    try:
                        for uploaded_file in uploaded_files:
                            with st.status(f"Processando {uploaded_file.name}..."):
                                temp_path = save_uploaded_file(uploaded_file)
                                try:
                                    response = st.session_state['assistant'].upload_file(temp_path)
                                    st.success("✅ Arquivo processado!")
                                finally:
                                    try:
                                        os.unlink(temp_path)
                                    except:
                                        pass
                                        
                        st.success("✅ Todos os documentos foram carregados!")
                        st.rerun()
                        
                    except Exception as e:
                        st.error(f"❌ Erro: {str(e)}")

    # Interface do chat
    if "messages" not in st.session_state:
        st.session_state.messages = []

    # Mostrar mensagens anteriores
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])
            
            # Mostrar citações se existirem
            if message.get("citations"):
                with st.expander("📚 **Fontes Consultadas**"):
                    st.info("ℹ️ O assistente baseou sua resposta nos seguintes documentos:")
                    for i, citation in enumerate(message["citations"], 1):
                        st.markdown(format_citation(citation, i))

    # Input do usuário
    if prompt := st.chat_input("Digite sua mensagem..."):
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.markdown(prompt)
            
        with st.chat_message("assistant"):
            try:
                response = st.session_state['assistant'].fazer_pergunta(prompt)
                
                # Mostrar contexto relevante primeiro
                if response.get("snippets"):
                    show_context(response["snippets"])
                
                # Mostrar resposta com citações inline
                content = response["content"]
                citations = response.get("citations", [])
                
                # Processar citações e extrair trechos
                processed_citations = []
                for citation in citations:
                    # Extrair o trecho do texto original
                    position = citation.get("position", 0)
                    text_before = content[max(0, position-100):position]
                    text_after = content[position:position+100]
                    citation["text"] = f"...{text_before}**{text_after}**..."
                    processed_citations.append(citation)
                
                # Adicionar números de referência ao conteúdo
                for i, citation in enumerate(processed_citations, 1):
                    position = citation.get("position", 0)
                    content = content[:position] + f" [{i}]" + content[position:]
                
                # Mostrar conteúdo
                st.markdown(content)
                
                # Mostrar seção de referências
                if processed_citations:
                    show_references(processed_citations)
                
                # Salvar mensagem
                st.session_state.messages.append({
                    "role": "assistant",
                    "content": content,
                    "citations": processed_citations
                })
                
            except Exception as e:
                st.error(f"❌ Erro: {str(e)}")

if __name__ == "__main__":
    main()