app.py

import os
from langchain_huggingface import HuggingFaceEndpoint
import streamlit as st
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from transformers import pipeline
from huggingface_hub import login  # For authentication

# --- Nova Importação ---
import re

# Authenticate with Hugging Face (required for private models)
login(token=os.getenv("HF_TOKEN"))

# Model IDs
chat_model_id = "mistralai/Mistral-7B-Instruct-v0.3"
translation_model_id = "Helsinki-NLP/opus-mt-tc-big-en-pt"  # Alternative model

# Initialize translation pipeline
translation_pipeline = pipeline(
    "translation_en_to_pt",
    model=translation_model_id,
    token=os.getenv("HF_TOKEN")  # Required for private models
)

def get_llm_hf_inference(model_id=chat_model_id, max_new_tokens=128, temperature=0.1):
    """
    Returns a language model for HuggingFace inference.
    """
    llm = HuggingFaceEndpoint(
        repo_id=model_id,
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        token=os.getenv("HF_TOKEN")
    )
    return llm

# --- Nova Função para Carregar a Base de Conhecimento ---
def carregar_base_de_conhecimento(caminho):
    """
    Carrega e organiza a base de conhecimento de um arquivo .txt.
    """
    base_de_conhecimento = {}
    with open(caminho, 'r', encoding='utf-8') as arquivo:
        conteudo = arquivo.read()
        # Dividir cada entrada por linhas vazias
        entradas = re.split(r'\n\s*\n', conteudo)
        for entrada in entradas:
            linhas = entrada.strip().split('\n')
            sintoma = None
            diagnostico = None
            for linha in linhas:
                if linha.startswith("Sintoma:"):
                    sintoma = linha.split(":", 1)[1].strip().lower()
                # elif linha.startswith("Diagnóstico:"):
                #     diagnostico = linha.split(":", 1)[1].strip()
            if sintoma and diagnostico:
                base_de_conhecimento[sintoma] = diagnostico
    return base_de_conhecimento

# Carregar a base de conhecimento
base_de_conhecimento = carregar_base_de_conhecimento("base.txt")

# --- Função para Buscar na Base de Conhecimento ---
def buscar_na_base(sintoma, base):
    """
    Busca um diagnóstico na base de conhecimento com base no sintoma.
    """
    sintoma = sintoma.lower()
    for chave in base.keys():
        if sintoma in chave or chave in sintoma:
            return base[chave]
    return None

# Configure the Streamlit app
st.session_state.sidebar_state = "collapsed"
st.set_page_config(page_title="HuggingFace ChatBot", page_icon="🤗")
st.title("POC ChatBot")
st.markdown(f"*This is a simple chatbot with {chat_model_id}.*")

# Initialize session state for avatars
if "avatars" not in st.session_state:
    st.session_state.avatars = {'user': None, 'assistant': None}

# Initialize session state for user text input
if 'user_text' not in st.session_state:
    st.session_state.user_text = ""

# Initialize session state for model parameters
if "max_response_length" not in st.session_state:
    st.session_state.max_response_length = 256

if "system_message" not in st.session_state:
    st.session_state.system_message = "You are a doctor who will help, based on the symptoms, and will give a diagnosis in Brazilian Portuguese. Your answer should be direct, simple and short, you can even ask a question to provide a more accurate answer. You should ask only about health."

if "starter_message" not in st.session_state:
    st.session_state.starter_message = "Olá, como posso ajudar você?"

# Initialize session state for translation
if "translate_to_pt" not in st.session_state:
    st.session_state.translate_to_pt = False  # Default: translation disabled

# Sidebar for settings
with st.sidebar:
    st.header("Configurações")

  
    # AI Settings
    st.session_state.system_message = st.text_area(
        "System Message", value="You are a doctor who will help, based on the symptoms, and will give a diagnosis in Brazilian Portuguese. Your answer should be direct, simple and short, you can even ask a question to provide a more accurate answer. You should ask only about health."
    )
    st.session_state.starter_message = st.text_area(
        'Primeira mensagem', value="Olá, como posso ajudar você?"
    )

    # Model Settings
    st.session_state.max_response_length = st.number_input(
        "Tamanho máximo da resposta", value=256
    )

    # Translation Toggle
    # st.session_state.translate_to_pt = st.checkbox(
    #     "Translate response to Brazilian Portuguese", value=False
    # )

    # Avatar Selection
    st.markdown("*Avatars:*")
    col1, col2 = st.columns(2)
    with col1:
        st.session_state.avatars['assistant'] = st.selectbox(
            "AI Avatar", options=["😷", "💬", "🤖"], index=0
        )
    with col2:
        st.session_state.avatars['user'] = st.selectbox(
            "User Avatar", options=["👤", "👱‍♂️", "👨🏾", "👩", "👧🏾"], index=0
        )

    # Reset Chat History
    reset_history = st.button("Reset Chat History")

# Initialize or reset chat history
if "chat_history" not in st.session_state or reset_history:
    st.session_state.chat_history = [{"role": "assistant", "content": st.session_state.starter_message}]

def translate_to_portuguese(text):
    """
    Translates the given text to Brazilian Portuguese using the translation model.
    """
    translation = translation_pipeline(text)
    return translation[0]['translation_text']

# --- Modificação na Função get_response ---
def get_response(system_message, chat_history, user_text, max_new_tokens=256):
    """
    Gera uma resposta do chatbot, consultando a base de conhecimento antes.
    """
    # Consultar a base de conhecimento
    diagnostico = buscar_na_base(user_text, base_de_conhecimento)
    if diagnostico:
        # Retornar o diagnóstico da base
        chat_history.append({'role': 'user', 'content': user_text})
        chat_history.append({'role': 'assistant', 'content': diagnostico})
        return diagnostico, chat_history

    # Caso o diagnóstico não seja encontrado, usar o modelo
    hf = get_llm_hf_inference(max_new_tokens=max_new_tokens, temperature=0.1)
    prompt = PromptTemplate.from_template(
        (
            "[INST] {system_message}"
            "\nCurrent Conversation:\n{chat_history}\n\n"
            "\nUser: {user_text}.\n [/INST]"
            "\nAI:"
        )
    )
    chat = prompt | hf.bind(skip_prompt=True) | StrOutputParser(output_key='content')
    response = chat.invoke(input=dict(system_message=system_message, user_text=user_text, chat_history=chat_history))
    response = response.split("AI:")[-1].strip()

    # Traduzir a resposta se necessário
    if st.session_state.translate_to_pt:
        response = translate_to_portuguese(response)

    # Atualizar o histórico de conversas
    chat_history.append({'role': 'user', 'content': user_text})
    chat_history.append({'role': 'assistant', 'content': response})
    return response, chat_history

# Chat interface
chat_interface = st.container()
with chat_interface:
    output_container = st.container()
    st.session_state.user_text = st.chat_input(placeholder="Digite seus sintomas aqui.")

# Display chat messages
with output_container:
    # Para cada mensagem no histórico
    for message in st.session_state.chat_history:
        # Pular a mensagem do sistema
        if message['role'] == 'system':
            continue

        # Exibir a mensagem do chat usando o avatar correto
        with st.chat_message(message['role'], avatar=st.session_state['avatars'][message['role']]):
            st.markdown(message['content'])

    # Quando o usuário insere um novo texto:
    if st.session_state.user_text:
        # Exibir a nova mensagem do usuário imediatamente
        with st.chat_message("user", avatar=st.session_state.avatars['user']):
            st.markdown(st.session_state.user_text)

        # Exibir um spinner enquanto espera pela resposta
        with st.chat_message("assistant", avatar=st.session_state.avatars['assistant']):
            with st.spinner("Pensando..."):
                # Chamar a função de resposta
                response, st.session_state.chat_history = get_response(
                    system_message=st.session_state.system_message,
                    user_text=st.session_state.user_text,
                    chat_history=st.session_state.chat_history,
                    max_new_tokens=st.session_state.max_response_length,
                )
                st.markdown(response)