feat : add stream, copy response, model display

model/selector.py

@@ -38,5 +38,5 @@ def ModelSelector():
 
     if(st.session_state["assistant"]):
         splitter = model_mapping[selected_model_option].split(".")
-        st.session_state["assistant"].setModel(ModelManager().get_model(splitter[0], splitter[1]))
+        st.session_state["assistant"].setModel(ModelManager().get_model(splitter[0], splitter[1]), splitter[1])
     

pages/chatbot.py

@@ -1,27 +1,39 @@
 import streamlit as st
-from streamlit_chat import message
+from langchain_core.messages import AIMessage, HumanMessage
 from model import selector
 from util import getYamlConfig
-
+from st_copy_to_clipboard import st_copy_to_clipboard
 
 def display_messages():
-    for i, (msg, is_user) in enumerate(st.session_state["messages"]):
-        message(msg, is_user=is_user, key=str(i))
-    st.session_state["thinking_spinner"] = st.empty()
 
+    for i, message in enumerate(st.session_state.chat_history):
+        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_{i}")
+        
+        elif isinstance(message, HumanMessage):
+            with st.chat_message("Moi"):
+                st.write(message.content)
 
-def process_input():
-    if "user_input" in st.session_state and st.session_state["user_input"] and len(st.session_state["user_input"].strip()) > 0:
-        user_text = st.session_state["user_input"].strip()
 
-        prompt_sys = st.session_state.prompt_system if 'prompt_system' in st.session_state and st.session_state.prompt_system != '' else ""
-    
-        with st.session_state["thinking_spinner"], st.spinner(f"Je réfléchis"):
-            agent_text = st.session_state["assistant"].ask(user_text, prompt_system=prompt_sys, messages=st.session_state["messages"] if "messages" in st.session_state else [], variables=st.session_state["data_dict"])
-            
-        st.session_state["messages"].append((user_text, True))
-        st.session_state["messages"].append((agent_text, False))
-        st.session_state["user_input"] = ""
+def launchQuery(query: str = None):
+
+    # Initialize the assistant's response
+    full_response = st.write_stream(
+        st.session_state["assistant"].ask(
+            query,
+            prompt_system=st.session_state.prompt_system,
+            messages=st.session_state["chat_history"] if "chat_history" in st.session_state else [],
+            variables=st.session_state["data_dict"]
+        ))
+
+    # Temporary placeholder AI message in chat history
+    st.session_state["chat_history"].append(AIMessage(content=full_response, kwargs={"model": st.session_state["assistant"].getReadableModel()}))
+    st.rerun()
 
 
 def show_prompts():
@@ -34,30 +46,37 @@ def show_prompts():
 
         for item in yaml_data[categroy]:
             if expander.button(item, key=f"button_{item}"):
-                st.session_state["user_input"] = item
-                process_input()
+                launchQuery(item)
 
 
 def page():
     st.subheader("Posez vos questions")
 
-    if "user_input" in st.session_state:
-        process_input()
-
     if "assistant" not in st.session_state:
         st.text("Assistant non initialisé")
 
+    if "chat_history" not in st.session_state:
+        st.session_state["chat_history"] = []
+
+    st.markdown("<style>iframe{height:50px;}</style>", unsafe_allow_html=True)
+
     # Collpase for default prompts
     show_prompts()
 
     # Models selector
     selector.ModelSelector()
-    
+
     # Displaying messages
     display_messages()
 
-    # Input user query
-    st.text_input("Message", key="user_input", on_change=process_input)
+
+    user_query = st.chat_input("")
+    if user_query is not None and user_query != "":
+
+        st.session_state["chat_history"].append(HumanMessage(content=user_query))
+        
+        # Stream and display response
+        launchQuery(user_query)
 
 
 page()

rag.py

@@ -23,6 +23,7 @@ class Rag:
     document_vector_store = None
     retriever = None
     chain = None
+    readableModelName = ""
 
     def __init__(self, vectore_store=None):
         
@@ -36,9 +37,13 @@ class Rag:
 
         self.vector_store = vectore_store
 
-    def setModel(self, model):
+    def setModel(self, model, readableModelName = ""):
         self.model = model
+        self.readableModelName = readableModelName
 
+    def getReadableModel(self):
+        return self.readableModelName
+    
     def ingestToDb(self, file_path: str, filename: str):
 
         docs = PyPDFLoader(file_path=file_path).load()
@@ -105,7 +110,7 @@ class Rag:
             chain_input.update(extra_vars)
         
 
-        return self.chain.invoke(chain_input)
+        return self.chain.stream(chain_input)
 
     def clear(self):
         self.document_vector_store = None