Update app.py

app.py

@@ -97,18 +97,24 @@ for message in st.session_state.messages:
 if "is_streaming" not in st.session_state:
     st.session_state.is_streaming = False
 
-# Chat input handling
+# Session state initialization
+if "is_streaming" not in st.session_state:
+    st.session_state.is_streaming = False  # Indicate if streaming/processing
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+# Prevent input during streaming
 if st.session_state.is_streaming:
-    st.chat_input("The assistant is currently responding. Please wait...")  # Inform the user to wait
+    st.info("Assistant is processing. Please wait...")
 else:
-    # If not streaming, allow user input
     if prompt := st.chat_input("Ask me anything about diabetes"):
-        st.session_state.is_streaming = True  # Set the flag to indicate streaming has started
+        st.session_state.is_streaming = True  # Mark as streaming
 
+        # Display user message
         with st.chat_message("user"):
             st.markdown(prompt)
 
-        # Add the user message to chat history
+        # Append user message to chat history
         st.session_state.messages.append({"role": "user", "content": prompt})
 
         instructions = """
@@ -124,30 +130,25 @@ else:
     You will answer as if you are talking to a patient directly
     """
 
+        # Full prompt
         full_prompt = f"<s>[INST] {prompt} [/INST] {instructions}</s>"
 
-        # Display assistant response in chat message container
+        # Stream assistant's response
         with st.chat_message("assistant"):
-            # Stream the response
             stream = client.chat.completions.create(
                 model=model_links[selected_model],
                 messages=[
-                    {"role": m["role"], "content": full_prompt}
-                    for m in st.session_state.messages
+                    {"role": "assistant", "content": prompt},
+                    {"role": "user", "content": full_prompt},
                 ],
-                temperature=temp_values,
                 stream=True,
+                temperature=temp_values,
                 max_tokens=1024,
             )
-            response = st.write_stream(stream)
-
-            # Process and clean the response
-            response = response.replace('</s>', '').strip()  # Clean unnecessary characters
 
+            response = st.write_stream(stream).replace("</s>", "").strip()
             st.markdown(response)
 
-            # Indicate that streaming is complete
-            st.session_state.is_streaming = False
-
-        # Store the final response
+        # Indicate streaming completion and update chat history
+        st.session_state.is_streaming = False
         st.session_state.messages.append({"role": "assistant", "content": response})