Update app.py

app.py

@@ -201,15 +201,23 @@ def generate_response(db, query_text, previous_context):
         filename="unsloth.BF16.gguf",
     )
     
-    # Directly get the output without streaming
-    output = llm(
+    output_stream = llm(
         prompt_template,
+        stream=True,
         temperature=0.1,
         top_p=0.9,
         top_k=20
     )
     
-    full_response = output["choices"][0]["text"]
+    full_response = ""
+    
+    response_placeholder = st.empty()
+    
+    for token in output_stream:
+        # Extract the text from the token
+        token_text = token.get("choices", [{}])[0].get("text", "")
+        full_response += token_text
+        response_placeholder.text(full_response)  # Print token output in real-time
 
     return full_response
 
@@ -219,10 +227,6 @@ 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:")
@@ -255,30 +259,35 @@ def streamlit_app():
                 
                 if final_items:
                     st.session_state.db = create_vector_db(final_items)
-                    st.success("Database initialized and models analyzed!")
+                    st.write("Models have been processed and added to the database.")
                 else:
-                    st.warning("No valid models found for the selected IDs.")
+                    st.error("No items found in the models. Check if the Antimony files were generated correctly.")
 
-    # Chat interface
-    query = st.text_area("Chat with the model:", value="", height=150)
-    
-    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
+    # 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
+
+    st.session_state.messages = get_messages()
 
-    # Display conversation history
     for message in st.session_state.messages:
-        if message["role"] == "user":
-            st.write(f"**You:** {message['content']}")
+        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 st.session_state.db is None:
+            st.error("Database is not initialized. Please process the models first.")
         else:
-            st.write(f"**Assistant:** {message['content']}")
+            response = generate_response(st.session_state.db, prompt, st.session_state.messages)
+
+            st.chat_message("assistant").markdown(response)  # Directly display the final response
+            st.session_state.messages.append({"role": "assistant", "content": response})
 
 if __name__ == "__main__":
     streamlit_app()