interface/chatbot.py

import os
import sys
from dotenv import load_dotenv
import edgedb
from openai import OpenAI
from qdrant_client import QdrantClient
import streamlit as st

# Add the parent directory to the Python path
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from pipelines.message import send_message # noqa

# st.set_page_config(page_title="Carlos AI Agent")

load_dotenv()

@st.cache_resource
def connect_to_services():
    oa_client = OpenAI(
        api_key=os.environ.get("OPENAI_API_KEY")
    )

    qdrant_client = QdrantClient(
        url=os.environ.get("QDRANT_URL"),
        api_key=os.environ.get("QDRANT_KEY")
    )

    edgedb_client = edgedb.create_client()

    return oa_client, qdrant_client, edgedb_client


class Chatbot:
    def __init__(self):
        # Conteúdo dos botões do sidebar
        if "topics" not in st.session_state:
            st.session_state.topics = [
                "Níveis da conta govbr.",
                "Dúvidas no reconhecimento facial.",
                "Como recuperar minha conta gov.br",
                "Dúvidas para aumentar o nível com a CIN."
            ]

        if "services" not in st.session_state:
            oa_client, qdrant_client, edgedb_client =  connect_to_services()
            st.session_state.services = {
                "oa_client": oa_client, 
                "qdrant_client": qdrant_client, 
                "edgedb_client": edgedb_client
            }
        
        if "chat" not in st.session_state and st.session_state.get('username'):
            chat = st.session_state.services["edgedb_client"].query('''
                Select Chat {
                    id
                } filter .user.username = <str>$username and .prototipo.id = <uuid>$prototipo_id
                ''',
                username=st.session_state.username,
                prototipo_id=st.session_state.prototipo_id
            )

            if not chat:
                chat = self._create_default_chat()
            else:
                chat = chat[0]
                self.update_messages(chat)
            
            st.session_state.chat = chat

        

        if "is_feedback_active" not in st.session_state:
            st.session_state.is_feedback_active = False

        
        self.user_button_input = ""     # Pergunta do usuário no chat
        # self.is_feedback_active = False # Flag para ativar o feedback do o
    
    def _create_default_chat(self):
        chat = st.session_state.services["edgedb_client"].query('''
                    INSERT Chat {
                        user := (SELECT User FILTER .username = <str>$username),
                        prototipo := (SELECT Prototipo FILTER .id = <uuid>$prototipo_id)
                    }
                    ''',
                    username=st.session_state.username,
                    prototipo_id=st.session_state.prototipo_id
                )
        message = st.session_state.services["edgedb_client"].query('''
            SELECT (
                INSERT Message {
                    content := "Como eu posso ajudar?",
                    role := "assistant",
                    chat := (SELECT Chat FILTER .id = <uuid>$chat_id)
                }
            ) {
                content,
                role                                             
            }
        ''', chat_id=chat[0].id)
        st.session_state.chat_history = message
        return chat[0]

    def _save_msg_to_db(self, msg, role):
        st.session_state.services["edgedb_client"].query('''
            INSERT Message {
                content := <str>$content,
                chat := (SELECT Chat FILTER .id = <uuid>$chat_id),
                role := <str>$role
            }
        ''', content=msg, chat_id=st.session_state.chat.id, role=role)
    

    def mount_chatbot(self):
        # Exibição do título e subtítulo
        st.title("Bem-vindo à ajuda do gov.br")
        st.caption("💬 Qual a sua dificuldade hoje? Estou aqui para ajudar!")

        # Exibição do espaço para mandar mensagem
        if user_query := st.chat_input(placeholder="Digite sua mensagem"):
            st.session_state.is_feedback_active = False # Desativando o feedback
            return user_query
    

    def create_sidebar(self):
        st.image('https://www.gov.br/++theme++padrao_govbr/img/govbr-logo-large.png', width=200)
        st.header("Tópicos frequentes")
        
        for topic in st.session_state.topics:
            if st.button(topic, key=topic):
                self.user_button_input = topic
                st.session_state.is_feedback_active = False # Desativando o feedback

        # Espaços em branco para organização
        for _ in range(5):
            st.write("")
        
        # Botão centralizado
        col1, col2, col3 = st.columns([1, 2, 1])
        with col2:
            if st.button("LIMPAR HISTÓRICO"):
                st.session_state.services["edgedb_client"].query('''
                    DELETE Chat
                    FILTER .id = <uuid>$chat_id;
                ''', chat_id=st.session_state.chat.id)
                st.session_state.chat = self._create_default_chat()
                st.session_state.is_feedback_active = False # Desativando o feedback
    

    def send_message_for_ai(self, prompt):
        with st.spinner("Obtendo conteúdo da página..."):
            embedding = st.session_state.services["oa_client"].embeddings.create(
                input=[prompt],
                model=os.environ.get("OPENAI_MODEL_EMBEDDING")
            ).data[0].embedding
        
            child_texts = st.session_state.services["qdrant_client"].search(
                collection_name=os.environ.get("COLLECTION_NAME"), 
                query_vector=embedding,
                limit=3
            )
        
            contexts = []
        
            for child_text in child_texts:
                parent_text = st.session_state.services["edgedb_client"].query('''
                    SELECT Pattern {
                        content,
                        url,
                        parent_id
                    }
                    FILTER .id = <uuid>$parent_id
                    ''', parent_id=child_text.payload["parent_id"])[0]
                context = {
                    "content": parent_text.content,
                    "url": parent_text.url,
                    "parent_id": parent_text.parent_id
                }
                contexts.append(context)
            
            formatted_messages = [{"content": msg.content, "role": msg.role} for msg in st.session_state.chat_history]

            self._save_msg_to_db(prompt, "user")
         
            stream_response = send_message(
                st.session_state.services["oa_client"],
                contexts,
                prompt,
                formatted_messages
            )

            return stream_response


    def generate_answer(self, prompt):
        with st.chat_message("assistant"):
            stream_response = self.send_message_for_ai(prompt)
            response = st.write_stream(stream_response)
        self._save_msg_to_db(response, "assistant")
        self.update_messages(st.session_state.chat)
    

    def display_chat(self):
        for message in st.session_state.chat_history:
            st.chat_message(message.role).write(message.content)
    
    def update_messages(self, chat):
        messages = st.session_state.services["edgedb_client"].query('''
            SELECT Message {
                id,
                content,
                role
            } 
            FILTER .chat.id = <uuid>$chat_id
            ORDER BY .created_at ASC;
        ''', chat_id=chat.id)
        st.session_state.chat_history = messages

    def display_feedback(self):
        user_input = st.session_state.chat_history[-2].content
        bot_output = st.session_state.chat_history[-1].content

        with st.expander("Avaliação do atendimento"):
            st.write(f'O que achou da resposta para a pergunta "{user_input}"?')
            
            rate = st.feedback("stars")
            # rate = st.feedback("faces")

            text_feedback = st.text_input("Comentários extras:")

            # Botão para confirmar a avaliação
            if rate is not None:
                if st.button("Enviar Avaliação"):
                    try:
                        feedback_rate = rate + 1
                        st.session_state.services["edgedb_client"].query('''
                            INSERT Feedback {
                                rating := <int16>$rating,
                                content := <str>$content,
                                message := (SELECT Message FILTER .id = <uuid>$message_id),
                                prototipo := (SELECT Prototipo FILTER .id = <uuid>$prototipo_id)
                            }
                        ''',
                        message_id=st.session_state.chat_history[-1].id,
                        rating=feedback_rate,
                        content=text_feedback,
                        prototipo_id=st.session_state.prototipo_id
                        )
                        st.session_state.chat_history
                        st.success(f"Avaliação enviada!")
                    except Exception as e:
                        print(e)
                        st.error("Erro ao enviar avaliação!")

                    st.session_state.is_feedback_active = False # Desativando o feedback

                    # TODO Colocar nessa parte a estrutura para adicionar o feedback_data ao banco de dados