Update app.py

app.py

@@ -201,7 +201,7 @@ def create_vector_db(final_items):
 def generate_response(db, query_text, previous_context):
     query_results = db.query(
         query_texts=query_text,
-        n_results=5,
+        n_results=7,
     )
     
     if not query_results.get('documents'):
@@ -263,8 +263,7 @@ def streamlit_app():
     st.title("BioModelsRAG")
     
     search_str = st.text_input("Enter search query:")
-
-    # Keep the search input field visible even after submission
+    
     if search_str:
         models = search_models(search_str)
         
@@ -277,7 +276,7 @@ def streamlit_app():
             )
             
             if st.button("Analyze Selected Models"):
-                all_final_items = []
+                final_items = []
                 for model_id in selected_models:
                     model_data = models[model_id]
                     
@@ -289,35 +288,39 @@ def streamlit_app():
                     
                     convert_sbml_to_antimony(model_file_path, antimony_file_path)
                     
-                    final_items = split_biomodels(antimony_file_path)
+                    items = split_biomodels(antimony_file_path)
                     if not final_items:
                         st.write("No content found in the biomodel.")
                         continue
                     
-                    all_final_items.extend(final_items)
+                    final_items.extend(items)
                 
-                global db
-                db = create_vector_db(all_final_items)
+                db = create_vector_db(final_items)
                 
                 if db:
                     st.write("Models have been processed and added to the database.")
                     
-    # Check if the database is created before showing the query input
     if db:
-        user_query = st.text_input("Ask a question about the biomodels:", key="user_query")
+        @st.cache_resource
+        def get_messages():
+            if "messages" not in st.session_state:
+                st.session_state.messages = []
+            return st.session_state.messages
+        st.session_state.messages = get_messages()
 
-        if user_query:
-            if 'previous_context' not in st.session_state:
-                st.session_state.previous_context = ""
-            
-            # Placeholder for indicating that response generation has started
-            response_placeholder = st.empty()
-            response_placeholder.write("Response generation is beginning...")
+        for message in st.session_state.messages:
+            with st.chat_message(message["role"]):
+                st.markdown(message["content"])
 
-            # Generate and display response
-            response = generate_response(db, user_query, st.session_state.previous_context)
-            st.write(f"Final Response: {response}")
-            
-            st.session_state.previous_context += f"{response}\n"
-            # Clear the "Response generation is beginning..." placeholder after generation
-            response_placeholder.empty()
+        if prompt := st.chat_input(query_text):
+            st.chat_message("user").markdown(prompt)
+            st.session_state.messages.append({"role": "user", "content":prompt})
+            response = generate_response(db, query_text, st.session_state)
+
+            with st.chat_message("assistant"):
+                st.markdown(response)
+
+            st.session_state.messages.append({"role":"assistant","content":response})
+
+if __name__ == "__main__":
+    streamlit_app()