Update app.py

app.py

@@ -13,10 +13,10 @@ import validators
 import numpy as np
 import pandas as pd
 import streamlit as st
-from typing import List
 from pathlib import Path
 from numpy import ndarray
 from outlines import models
+from typing import List, Dict
 from qdrant_client import QdrantClient
 from optimum_encoder import OptimumEncoder
 from huggingface_hub import snapshot_download
@@ -442,22 +442,29 @@ if __name__ == '__main__':
     try:
         with open(conversations_path, 'rb') as fp:
             packed_bytes = fp.read()
-        conversations: dict = msgpack.unpackb(packed_bytes, raw=False)
+        conversations: Dict[str, list] = msgpack.unpackb(packed_bytes, raw=False)
     except:
         conversations = {}
 
-    if st.session_state.menu_id == 'ChatBot':
+    def options_list(conversations: Dict[str, list]):
         if 'chat_id' not in st.session_state:
-            st.session_state.chat_id = 'New Chat'
-            
+            return 'New Chat'
+        elif st.session_state.chat_id is None:
+            return 'New Chat'
+        elif st.session_state.chat_id not in list(conversations.keys()):
+            return [st.session_state.chat_id] + list(conversations.keys())
+        else:
+            return list(conversations.keys())
+
+    if st.session_state.menu_id == 'ChatBot':            
         with st.sidebar:
-            st.session_state.chat_id = st.selectbox(
+            st.selectbox(
                 label='Choose a conversation',
-                options=[st.session_state.chat_id] + list(conversations.keys()) if st.session_state.chat_id not in conversations.keys() else list(conversations.keys()),
+                options=options_list(conversations),
                 index=0,
                 placeholder='_',
+                key='chat_id'
             )
-            print(f'KEY SELECTBOX: {st.session_state.chat_id}')
             
             col1, col2 = st.columns([1,1])
 
@@ -504,11 +511,30 @@ if __name__ == '__main__':
             for message in st.session_state.messages:
                 with st.chat_message(message["role"]):
                     st.markdown(message["content"])
+
+        def generate_conv_title(llm: vllm.LLM, template):
+            if st.session_state.chat_id == 'New Chat':
+                outputs = llm.generate(
+                    prompts=template.format(system='You are a helpful assistant.', user=f"""Understand the question of the user. 
+                    Resume in one single sentence what is the subject of the conversation and what is the user talking about.
+                    
+                    Question : {st.session_state.user_input}"""), 
+                    sampling_params=vllm.SamplingParams(
+                        temperature=0.3, 
+                        max_tokens=30
+                    )
+                )
+                st.session_state.chat_id = outputs[0].outputs[0].text
                 
-        if prompt := st.chat_input("Message Video Game Assistant"):
+        if prompt := st.chat_input(
+            "Message Video Game Assistant", 
+            on_submit=generate_conv_title, 
+            key='user_input',
+            args=(llm, template)
+        ):
             st.chat_message("user").markdown(prompt)
             st.session_state.messages.append({"role": "user", "content": prompt})
-            print(f'PROMPT: {prompt}')
+            
             ai_response = main(prompt, client, collection_name, template, llm, dense_model, sparse_model)
             with st.chat_message("assistant"):
                 message_placeholder = st.empty()
@@ -520,19 +546,6 @@ if __name__ == '__main__':
                 message_placeholder.write(re.sub('▌', '', full_response))
                                           
             st.session_state.messages.append({"role": "assistant", "content": full_response})
-
-            if st.session_state.chat_id == 'New Chat':
-                outputs = llm.generate(
-                    prompts=template.format(system='You are a helpful assistant.', user=f"""Understand the question of the user. 
-                    Resume in one single sentence what is the subject of the conversation and what is the user talking about.
-                    
-                    Question : {prompt}"""), 
-                    sampling_params=vllm.SamplingParams(
-                        temperature=0.3, 
-                        max_tokens=30
-                    )
-                )
-                st.session_state.chat_id = outputs[0].outputs[0].text
                 
             conversations.update({st.session_state.chat_id: st.session_state.messages})
             with open(conversations_path, 'wb') as fp: