Update app.py

app.py

@@ -100,9 +100,44 @@ def generate_chat_responses(chat_completion) -> Generator[str, None, None]:
         if chunk.choices[0].delta.content:
             yield chunk.choices[0].delta.content
 
-selected_prompt_text = prompts.get(prompt_selection)
+if prompt := st.chat_input("Enter your prompt here..."):
+    st.session_state.messages.append({"role": "user", "content": prompt})
+
+    with st.chat_message("user", avatar="❓"):  
+        st.markdown(prompt)
+
+    # Fetch response from Groq API
+    try:
+        chat_completion = client.chat.completions.create(
+            model=model_option,
+            messages=[
+                {"role": m["role"], "content": m["content"]}
+                for m in st.session_state.messages
+            ],
+            max_tokens=max_tokens,
+            stream=True,
+        )
+
+        # Use the generator function with st.write_stream
+        with st.chat_message("assistant", avatar="🧠"):
+            chat_responses_generator = generate_chat_responses(chat_completion)
+            full_response = st.write_stream(chat_responses_generator)
+    except Exception as e:
+        st.error(e, icon="🚨")
 
-if prompt := st.chat_input("Enter your prompt here...") or prompt := selected_prompt_text:
+    # Append the full response to session_state.messages
+    if isinstance(full_response, str):
+        st.session_state.messages.append(
+            {"role": "assistant", "content": full_response}
+        )
+    else:
+        # Handle the case where full_response is not a string
+        combined_response = "\n".join(str(item) for item in full_response)
+        st.session_state.messages.append(
+            {"role": "assistant", "content": combined_response}
+        )
+        
+if prompt := prompts.get(prompt_selection):
     st.session_state.messages.append({"role": "user", "content": prompt})
 
     with st.chat_message("user", avatar="❓"):