Update app.py

app.py

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