Update app.py

app.py

@@ -219,6 +219,10 @@ def streamlit_app():
     # Initialize db in session state if not already present
     if "db" not in st.session_state:
         st.session_state.db = None
+    if "messages" not in st.session_state:
+        st.session_state.messages = []
+    if "previous_context" not in st.session_state:
+        st.session_state.previous_context = ""
 
     # Search query input
     search_str = st.text_input("Enter search query:")
@@ -251,26 +255,30 @@ def streamlit_app():
                 
                 if final_items:
                     st.session_state.db = create_vector_db(final_items)
-                    st.write("Models have been processed and added to the database.")
+                    st.success("Database initialized and models analyzed!")
                 else:
-                    st.error("No items found in the models. Check if the Antimony files were generated correctly.")
+                    st.warning("No valid models found for the selected IDs.")
 
-    # Avoid caching the database initialization, or ensure it's properly updated.
-    @st.cache_resource
-    def get_messages():
-        if "messages" not in st.session_state:
-            st.session_state.messages = []
-        return st.session_state.messages
-
-    # Query input and response generation
-    query = st.text_area("Ask a question about the models:")
+    # Chat interface
+    query = st.text_area("Chat with the model:", value="", height=150)
     
-    if st.button("Get Response"):
-        if st.session_state.db:
-            response = generate_response(st.session_state.db, query, "")
-            st.write("Response: ", response)
+    if query:
+        # Update messages with user query
+        st.session_state.messages.append({"role": "user", "content": query})
+        
+        # Generate response based on previous context and new query
+        response = generate_response(st.session_state.db, query, st.session_state.previous_context)
+        st.session_state.messages.append({"role": "assistant", "content": response})
+        
+        # Update previous context with the latest response
+        st.session_state.previous_context = response
+
+    # Display conversation history
+    for message in st.session_state.messages:
+        if message["role"] == "user":
+            st.write(f"**You:** {message['content']}")
         else:
-            st.warning("Database is not initialized. Please search and analyze models first.")
+            st.write(f"**Assistant:** {message['content']}")
 
 if __name__ == "__main__":
     streamlit_app()