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):
        # Configuração inicial do histórico de chat
        if "chat_history" not in st.session_state:
            st.session_state.chat_history = [{
                "role": "assistant",
                "content": "Como eu posso ajudar?"
            }]
        
        # Conteúdo dos botões do sidebar
        if "topics" not in st.session_state:
            st.session_state["topics"] = [
                "Niveis da conta govbr.",
                "Dúvidas no reconhecimento facial.",
                "Como recuperar minha conta gov.br",
                "Dúvidas para aumentar o nível com a cin."
            ]
        
        # Pergunta do usuário no chatbot
        self.response = ""

        self.oa_client, self.qdrant_client, self.edgedb_client =  connect_to_services()

    
    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"):
            self.add_to_history(user_query, role="user")

            # Simulação de resposta do assistente
            self.send_message_for_ai(user_query)
            # response = self.send_message_for_ai(user_query)
            # self.add_to_history(response, role="assistant")
    
            
    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.response = topic

        
        # 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.chat_history = [{
                    "role": "assistant",
                    "content": "Como eu posso ajudar?",
                }]

    def send_message_for_ai(self, prompt):
        embedding = self.oa_client.embeddings.create(
            input=[prompt],
            model=os.environ.get("OPENAI_MODEL_EMBEDDING")
        ).data[0].embedding
    
        child_texts = self.qdrant_client.search(
            collection_name=os.environ.get("COLLECTION_NAME"), 
            query_vector=embedding,
            limit=3
        )
    
        contexts = []
    
        for child_text in child_texts:
            parent_text = self.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)
    
        
        stream_response = send_message(
            self.oa_client,
            contexts,
            prompt,
            []
        )
        
        # Se a resposta for uma lista ou um gerador, você precisa processá-la
        response_content = []
        for chunk in stream_response:
            content = chunk.choices[0].delta.content if chunk.choices[0].delta.content else ""
            response_content.append(content)
    
        # Junte todas as partes da resposta em uma única string
        full_response = ''.join(response_content)


        self.add_to_history(full_response, role="assistant")



    def add_to_history(self, message, role="user"):
        st.session_state.chat_history.append({
            "role": role,
            "content": message
            })