now the ai message is not replaced by [Assistant reply streamed above] in the chat ui

app.py

@@ -113,11 +113,13 @@ class StreamHandler:
     def __init__(self, client: Any):
         self.client = client
         self.text_processor = TextProcessor()
+        self.complete_response = []
     
     def stream_data(self) -> Generator[Any, None, None]:
         """Stream data from the assistant run."""
         st.toast("Thinking...", icon="🤔")
         content_produced = False
+        self.complete_response = []  # Reset for new stream
 
         try:
             for event in st.session_state.run_stream:
@@ -131,6 +133,8 @@ class StreamHandler:
                         raise StreamingError(f"Assistant run failed: {event}")
                     
             st.toast("Completed", icon="✅")
+            # Return the complete response for storage
+            return "".join(self.complete_response)
         except Exception as e:
             logger.error(f"Streaming error: {e}")
             st.error(f"An error occurred while streaming: {str(e)}")
@@ -141,7 +145,9 @@ class StreamHandler:
         content = event.data.delta.content[0]
         match content.type:
             case "text":
-                yield self.text_processor.remove_citations(content.text.value)
+                processed_text = self.text_processor.remove_citations(content.text.value)
+                self.complete_response.append(processed_text)  # Store the chunk
+                yield processed_text
             case "image_file":
                 image_content = io.BytesIO(self.client.files.content(content.image_file.file_id).read())
                 yield Image.open(image_content)
@@ -199,7 +205,7 @@ class AssistantManager:
         self.tool_handler = ToolRequestHandler()
     
     @traceable
-    def generate_reply(self, user_input: str) -> None:
+    def generate_reply(self, user_input: str) -> str:
         """Generate and stream assistant's reply."""
         # Ensure thread exists
         if not st.session_state.thread:
@@ -212,24 +218,28 @@ class AssistantManager:
             content=user_input
         )
         
+        complete_response = ""
+        
         # Stream initial response
         with self.client.beta.threads.runs.stream(
             thread_id=st.session_state.thread.id,
             assistant_id=self.assistant_id,
         ) as run_stream:
-            self._display_stream(run_stream)
+            complete_response = self._display_stream(run_stream)
         
         # Handle any tool requests
         self._process_tool_requests()
+        
+        return complete_response
     
-    def _display_stream(self, run_stream: Any, create_context: bool = True) -> None:
+    def _display_stream(self, run_stream: Any, create_context: bool = True) -> str:
         """Display streaming content."""
         st.session_state.run_stream = run_stream
         if create_context:
             with st.chat_message("assistant"):
-                st.write_stream(self.stream_handler.stream_data)
+                return st.write_stream(self.stream_handler.stream_data)
         else:
-            st.write_stream(self.stream_handler.stream_data)
+            return st.write_stream(self.stream_handler.stream_data)
     
     def _process_tool_requests(self) -> None:
         """Process any pending tool requests."""
@@ -292,10 +302,10 @@ class ChatApplication:
             
             # Generate and display assistant reply
             try:
-                self.assistant_manager.generate_reply(user_input)
+                complete_response = self.assistant_manager.generate_reply(user_input)
                 self.state_manager.add_message(
                     "assistant",
-                    "[Assistant reply streamed above]"
+                    complete_response
                 )
             except Exception as e:
                 st.error(f"Error generating response: {str(e)}")