pages/chatbot.py

import streamlit as st
import uuid
from db.db import DatabaseHandler

from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from model import selector
from st_copy_to_clipboard import st_copy_to_clipboard

# Instanciation de la base de données
db = DatabaseHandler()

chapter_num = 0
chapter_session_key = f"chapter_{chapter_num}"

def setChapter(num: int):
    global chapter_num, chapter_session_key
    chapter_num = num
    chapter_session_key = f"chapter_{chapter_num}"


def display_messages():

    for i, message in enumerate(st.session_state[chapter_session_key]["messages"]):
        if isinstance(message, AIMessage):
            with st.chat_message("AI"):
                # Display the model from the kwargs
                model = message.kwargs.get("model", "Unknown Model")  # Get the model, default to "Unknown Model"
                st.write(f"**Model :** {model}")
                st.markdown(message.content)
                st_copy_to_clipboard(message.content,key=f"message_{chapter_num}_{i}_{uuid.uuid4()}")
        
        elif isinstance(message, HumanMessage):
            with st.chat_message("Moi"):
                st.write(message.content)

        elif isinstance(message, SystemMessage):
            with st.chat_message("System"):
                st.write(message.content)



def show_prompts(prompts: list[str]):
    
    expander = st.expander("Prompts pré-définis")

    for i, p in enumerate(prompts):
        button_key = f"button_{chapter_session_key}_{i}"

        # Réinitialisation si la clé existe déjà (quand on change de page)
        if button_key in st.session_state:
            del st.session_state[button_key]
            
        expander.button(p, key=button_key, on_click=promptSubmit, args=(p,))

def promptSubmit(key,):
    inputSubmut(key)


def inputSubmut(query: str = None):
    # Get query or find in session state
    user_query = query if query is not None else st.session_state[f"input_{chapter_session_key}"]

    if user_query is not None and user_query != "":
        st.session_state[chapter_session_key]["messages"].append(HumanMessage(content=user_query))
        
        # Stream and display response
        launchQuery(user_query)

def page():

    chapters = []
    for chapter in st.session_state["chapters"]:
        chapters.append(st.session_state[f"chapter_{chapter['num']}"])

    ###############################################
    # Récupération du chapitre depuis YAML Config #
    ###############################################

    chapter_num = st.session_state.get('current_page', 0)
    if chapter_num == 0:
        return
    
    setChapter(int(chapter_num))
    
    chapter = next((ch for ch in chapters if ch['num'] == chapter_num), None)
    if not chapter:
        st.text("Chapitre non trouvé")
        return


    # Chargement de l'enregistrement correspondant aux filtres
    chapterDB = db.get_prompt_by_filters(num=chapter_num)
    if len(chapterDB) == 0:
        st.text("Chapitre non trouvé")
        return
    
    chapterDB = chapterDB[0]

    #################
    # Some controls #
    #################

    if "assistant" not in st.session_state:
        st.text("Assistant non initialisé")

    if "messages" not in st.session_state[chapter_session_key]:
        st.session_state[chapter_session_key]["messages"] = [ ]

    if len(st.session_state[chapter_session_key]["messages"]) < 2 :
        st.session_state[chapter_session_key]["messages"] = [ SystemMessage(content=chapterDB['prompt_system']) ]


    ############
    #### UI ####
    ############

    st.subheader(chapter['title'])

    st.markdown("<style>iframe{height:50px;}</style>", unsafe_allow_html=True)

    # Collpase for default prompts
    show_prompts(chapter['prompts'])

    # Models selector
    selector.ModelSelector()

    if(len(st.session_state[chapter_session_key]["messages"]) > 1):
        if st.button("Effacer l'historique"):
            st.session_state[chapter_session_key]["messages"] = [ SystemMessage(content=chapterDB['prompt_system']) ]

    # Displaying messages
    display_messages()

    chat_input_key = f"input_{chapter_session_key}"

    # Réinitialisation si la clé existe déjà (quand on change de page)
    if chat_input_key in st.session_state:
        del st.session_state[chat_input_key]


    st.chat_input("", key=chat_input_key, on_submit=inputSubmut)
    

    global response_placeholder
    response_placeholder = st.empty()


    # Check if we need to rerun the app
    if st.session_state.get('rerun', False):
        st.session_state['rerun'] = False
        st.rerun()


def launchQuery(query: str = None):

    # Initialize the assistant's response
    full_response = response_placeholder.write_stream(
            st.session_state["assistant"].ask(
            query,
            prompt_system=st.session_state.prompt_system,
            messages=st.session_state[chapter_session_key]["messages"] if "messages" in st.session_state[chapter_session_key] else [],
        )
    )

    # Temporary placeholder AI message in chat history
    st.session_state[chapter_session_key]["messages"].append(AIMessage(content=full_response, kwargs={"model": st.session_state["assistant"].getReadableModel()}))
    st.session_state['rerun'] = True


page()