claude

app.py

@@ -233,12 +233,18 @@ def main():
         # Generate assistant reply
         generate_assistant_reply(user_input)
 
-        # In a real app, you might keep track of the final text
-        # from the streamed tokens. For simplicity, we store
-        # the entire streamed result as one block in session state:
+        # Track the final text from streamed tokens
+        final_text = ""
+        for msg in st.session_state.get("streamed_messages", []):
+            final_text += msg
+            
+        # Store the complete assistant response in session state
         st.session_state["messages"].append(
-            {"role": "assistant", "content": "[assistant reply streamed above]"}
+            {"role": "assistant", "content": final_text}
         )
+        
+        # Clear streamed messages for next interaction
+        st.session_state["streamed_messages"] = []
 
 if __name__ == "__main__":
     main()