Merge pull request #209 from Utkarsh-Patel-13/fix/chat-optimization

Updated Streaming Service to efficiently stream content

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
-            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))
+                langchain_chat_history.append(AIMessage(content=message['content']))
 
-    response = StreamingResponse(stream_connector_search_results(
-        user_query,
-        user.id,
-        search_space_id,  # Already converted to int in lines 32-37
-        session,
-        research_mode,
-        selected_connectors,
-        langchain_chat_history,
-        search_mode_str,
-        document_ids_to_add_in_context
-    ))
-    response.headers['x-vercel-ai-data-stream'] = 'v1'
+    response = StreamingResponse(
+        stream_connector_search_results(
+            user_query,
+            user.id,
+            search_space_id,
+            session,
+            research_mode,
+            selected_connectors,
+            langchain_chat_history,
+            search_mode_str,
+            document_ids_to_add_in_context,
+        )
+    )
+    
+    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
-        # Note: there's a typo in 'yeild_value' in the code, but we need to match it
-        if isinstance(chunk, dict) and 'yeild_value' in chunk:
-            yield chunk['yeild_value']
-    
-    yield streaming_service.format_completion()
+        if isinstance(chunk, dict):
+            if "yield_value" in chunk:
+                yield chunk["yield_value"]
+
+    yield streaming_service.format_completion()

surfsense_web/app/dashboard/[search_space_id]/researcher/[chat_id]/page.tsx

@@ -981,19 +981,16 @@ const ChatPage = () => {
 	const renderTerminalContent = (message: any) => {
 		if (!message.annotations) return null;
 
-		// Get all TERMINAL_INFO annotations
-		const terminalInfoAnnotations = (message.annotations as any[]).filter(
-			(a) => a.type === "TERMINAL_INFO",
-		);
-
-		// Get the latest TERMINAL_INFO annotation
-		const latestTerminalInfo =
-			terminalInfoAnnotations.length > 0
-				? terminalInfoAnnotations[terminalInfoAnnotations.length - 1]
-				: null;
+		// Get all TERMINAL_INFO annotations content
+		const terminalInfoAnnotations = (message.annotations as any[]).map(item => {
+            if(item.type === "TERMINAL_INFO") {
+                return item.content.map((a: any) => a.text)
+                
+            }
+        }).flat().filter(Boolean)
 
 		// Render the content of the latest TERMINAL_INFO annotation
-		return latestTerminalInfo?.content.map((item: any, idx: number) => (
+		return terminalInfoAnnotations.map((item: any, idx: number) => (
 			<div key={idx} className="py-0.5 flex items-start text-gray-300">
 				<span className="text-gray-500 text-xs mr-2 w-10 flex-shrink-0">
 					[{String(idx).padStart(2, "0")}:
@@ -1008,7 +1005,7 @@ const ChatPage = () => {
           ${item.type === "warning" ? "text-yellow-300" : ""}
         `}
 				>
-					{item.text}
+					{item}
 				</span>
 			</div>
 		));