Update app.py

app.py

@@ -213,11 +213,18 @@ def generate_response(db, query_text, previous_context):
 
     return full_response
 
+import streamlit as st
+
 def streamlit_app():
     st.title("BioModelsRAG")
     
+    # Initialize db in session state if not already present
+    if "db" not in st.session_state:
+        st.session_state.db = None
+
+    # Search query input
     search_str = st.text_input("Enter search query:")
-    
+
     if search_str:
         cached_data = fetch_github_json()
         models = search_models(search_str, cached_data)
@@ -245,13 +252,11 @@ def streamlit_app():
                     final_items.extend(split_biomodels(antimony_file_path))
                 
                 if final_items:
-                    db = create_vector_db(final_items)
+                    st.session_state.db = create_vector_db(final_items)
                     st.write("Models have been processed and added to the database.")
                 else:
                     st.error("No items found in the models. Check if the Antimony files were generated correctly.")
 
-                st.write("Models have been processed and written to the database.")
-
     # Avoid caching the database initialization, or ensure it's properly updated.
     @st.cache_resource
     def get_messages():
@@ -264,15 +269,16 @@ def streamlit_app():
     for message in st.session_state.messages:
         with st.chat_message(message["role"]):
             st.markdown(message["content"])
-    
+
+    # Chat input section
     if prompt := st.chat_input("Ask a question about the models:"):
         st.chat_message("user").markdown(prompt)
         st.session_state.messages.append({"role": "user", "content": prompt})
-        
-        if db is None:
+
+        if st.session_state.db is None:
             st.error("Database is not initialized. Please process the models first.")
         else:
-            response = generate_response(db, prompt, st.session_state.messages)
+            response = generate_response(st.session_state.db, prompt, st.session_state.messages)
             
             with st.chat_message("assistant"):
                 st.markdown(response)