Updated Streaming Service to efficently stream content\

\
- Earlier for each chunk, whole message (with all annotations included)
  were streamed. Leading to extremely large data length.
- Fixed to only stream new chunk.
- Updated ANSWER part to be streamed as message content (following
  Vercel's Stream Protocol)\
- Fixed yield typo

surfsense_backend/app/routes/chats_routes.py

@@ -54,32 +54,23 @@ async def handle_chat_data(
         if message['role'] == "user":
             langchain_chat_history.append(HumanMessage(content=message['content']))
         elif message['role'] == "assistant":
-            # Find the last "ANSWER" annotation specifically
+                langchain_chat_history.append(AIMessage(content=message['content']))
-            answer_annotation = None
-            for annotation in reversed(message['annotations']):
-                if annotation['type'] == "ANSWER":
-                    answer_annotation = annotation
-                    break
-            
-            if answer_annotation:
-                answer_text = answer_annotation['content']
-                # If content is a list, join it into a single string
-                if isinstance(answer_text, list):
-                    answer_text = "\n".join(answer_text)
-                langchain_chat_history.append(AIMessage(content=answer_text))
 
-    response = StreamingResponse(stream_connector_search_results(
+    response = StreamingResponse(
-        user_query,
+        stream_connector_search_results(
-        user.id,
+            user_query,
-        search_space_id,  # Already converted to int in lines 32-37
+            user.id,
-        session,
+            search_space_id,
-        research_mode,
+            session,
-        selected_connectors,
+            research_mode,
-        langchain_chat_history,
+            selected_connectors,
-        search_mode_str,
+            langchain_chat_history,
-        document_ids_to_add_in_context
+            search_mode_str,
-    ))
+            document_ids_to_add_in_context,
-    response.headers['x-vercel-ai-data-stream'] = 'v1'
+        )
+    )
+    
+    response.headers["x-vercel-ai-data-stream"] = "v1"
     return response
 
 

surfsense_backend/app/services/streaming_service.py

@@ -23,17 +23,138 @@ class StreamingService:
                 "content": []
             }
         ]
-    # It is used to send annotations to the frontend
+
+    # DEPRECATED: This sends the full annotation array every time (inefficient)
     def _format_annotations(self) -> str:
         """
         Format the annotations as a string
-        
+
+        DEPRECATED: This method sends the full annotation state every time.
+        Use the delta formatters instead for optimal streaming.
+
         Returns:
             str: The formatted annotations string
         """
         return f'8:{json.dumps(self.message_annotations)}\n'
-    
+
-    # It is used to end Streaming
+    def format_terminal_info_delta(self, text: str, message_type: str = "info") -> str:
+        """
+        Format a single terminal info message as a delta annotation
+
+        Args:
+            text: The terminal message text
+            message_type: The message type (info, error, success, etc.)
+
+        Returns:
+            str: The formatted annotation delta string
+        """
+        message = {"id": self.terminal_idx, "text": text, "type": message_type}
+        self.terminal_idx += 1
+
+        # Update internal state for reference
+        self.message_annotations[0]["content"].append(message)
+
+        # Return only the delta annotation
+        annotation = {"type": "TERMINAL_INFO", "content": [message]}
+        return f"8:[{json.dumps(annotation)}]\n"
+
+    def format_sources_delta(self, sources: List[Dict[str, Any]]) -> str:
+        """
+        Format sources as a delta annotation
+
+        Args:
+            sources: List of source objects
+
+        Returns:
+            str: The formatted annotation delta string
+        """
+        # Update internal state
+        self.message_annotations[1]["content"] = sources
+
+        # Return only the delta annotation
+        annotation = {"type": "SOURCES", "content": sources}
+        return f"8:[{json.dumps(annotation)}]\n"
+
+    def format_answer_delta(self, answer_chunk: str) -> str:
+        """
+        Format a single answer chunk as a delta annotation
+
+        Args:
+            answer_chunk: The new answer chunk to add
+
+        Returns:
+            str: The formatted annotation delta string
+        """
+        # Update internal state by appending the chunk
+        if isinstance(self.message_annotations[2]["content"], list):
+            self.message_annotations[2]["content"].append(answer_chunk)
+        else:
+            self.message_annotations[2]["content"] = [answer_chunk]
+
+        # Return only the delta annotation with the new chunk
+        annotation = {"type": "ANSWER", "content": [answer_chunk]}
+        return f"8:[{json.dumps(annotation)}]\n"
+
+    def format_answer_annotation(self, answer_lines: List[str]) -> str:
+        """
+        Format the complete answer as a replacement annotation
+
+        Args:
+            answer_lines: Complete list of answer lines
+
+        Returns:
+            str: The formatted annotation string
+        """
+        # Update internal state
+        self.message_annotations[2]["content"] = answer_lines
+
+        # Return the full answer annotation
+        annotation = {"type": "ANSWER", "content": answer_lines}
+        return f"8:[{json.dumps(annotation)}]\n"
+
+    def format_further_questions_delta(
+        self, further_questions: List[Dict[str, Any]]
+    ) -> str:
+        """
+        Format further questions as a delta annotation
+
+        Args:
+            further_questions: List of further question objects
+
+        Returns:
+            str: The formatted annotation delta string
+        """
+        # Update internal state
+        self.message_annotations[3]["content"] = further_questions
+
+        # Return only the delta annotation
+        annotation = {"type": "FURTHER_QUESTIONS", "content": further_questions}
+        return f"8:[{json.dumps(annotation)}]\n"
+
+    def format_text_chunk(self, text: str) -> str:
+        """
+        Format a text chunk using the text stream part
+
+        Args:
+            text: The text chunk to stream
+
+        Returns:
+            str: The formatted text part string
+        """
+        return f"0:{json.dumps(text)}\n"
+
+    def format_error(self, error_message: str) -> str:
+        """
+        Format an error using the error stream part
+
+        Args:
+            error_message: The error message
+
+        Returns:
+            str: The formatted error part string
+        """
+        return f"3:{json.dumps(error_message)}\n"
+
     def format_completion(self, prompt_tokens: int = 156, completion_tokens: int = 204) -> str:
         """
         Format a completion message
@@ -56,7 +177,12 @@ class StreamingService:
         }
         return f'd:{json.dumps(completion_data)}\n' 
     
+
+    # DEPRECATED METHODS: Keep for backward compatibility but mark as deprecated
     def only_update_terminal(self, text: str, message_type: str = "info") -> str:
+        """
+        DEPRECATED: Use format_terminal_info_delta() instead for optimal streaming
+        """
         self.message_annotations[0]["content"].append({
             "id": self.terminal_idx,
             "text": text,
@@ -66,17 +192,23 @@ class StreamingService:
         return self.message_annotations
 
     def only_update_sources(self, sources: List[Dict[str, Any]]) -> str:
+        """
+        DEPRECATED: Use format_sources_delta() instead for optimal streaming
+        """
         self.message_annotations[1]["content"] = sources
         return self.message_annotations
     
     def only_update_answer(self, answer: List[str]) -> str:
+        """
+        DEPRECATED: Use format_answer_delta() or format_answer_annotation() instead for optimal streaming
+        """
         self.message_annotations[2]["content"] = answer
         return self.message_annotations
-    
+
     def only_update_further_questions(self, further_questions: List[Dict[str, Any]]) -> str:
         """
-        Update the further questions annotation
+        DEPRECATED: Use format_further_questions_delta() instead for optimal streaming
-        
+
         Args:
             further_questions: List of further question objects with id and question fields
             

surfsense_backend/app/tasks/stream_connector_search_results.py

@@ -83,9 +83,8 @@ async def stream_connector_search_results(
         config=config,
         stream_mode="custom",
     ):
-        # If the chunk contains a 'yeild_value' key, print its value
+        if isinstance(chunk, dict):
-        # Note: there's a typo in 'yeild_value' in the code, but we need to match it
+            if "yield_value" in chunk:
-        if isinstance(chunk, dict) and 'yeild_value' in chunk:
+                yield chunk["yield_value"]
-            yield chunk['yeild_value']
+
-    
+    yield streaming_service.format_completion()
-    yield streaming_service.format_completion()