Update api/utils.py

api/utils.py

@@ -120,7 +120,10 @@ async def process_streaming_response(request: ChatRequest):
 
     json_data = build_json_data(request, h_value, model_prefix)
 
+    # Initialize buffer to handle BLOCKED_MESSAGE that may be split across chunks
     buffer = ""
+    buffer_size = len(BLOCKED_MESSAGE) - 1
+
     async with httpx.AsyncClient() as client:
         try:
             async with client.stream(
@@ -134,39 +137,33 @@ async def process_streaming_response(request: ChatRequest):
                 async for chunk in response.aiter_text():
                     if chunk:
                         buffer += chunk
-                        # Check if buffer contains the blocked message
+
+                        # Remove any occurrence of BLOCKED_MESSAGE in buffer
                         if BLOCKED_MESSAGE in buffer:
                             logger.info("Blocked message detected in response.")
-                            buffer = buffer.replace(BLOCKED_MESSAGE, '').strip()
-                            if not buffer:
-                                continue  # Skip if buffer is empty after removal
-
-                        # Process the buffer to yield content
-                        # Decide on a suitable chunk size or yield the entire buffer
-                        while buffer:
-                            content_to_yield = buffer
-                            buffer = ""
-                            # Remove the prefix if present
+                            buffer = buffer.replace(BLOCKED_MESSAGE, '')
+
+                        # Process the buffer
+                        # Keep last buffer_size characters in buffer to handle partial BLOCKED_MESSAGE
+                        while len(buffer) >= buffer_size:
+                            content_to_yield = buffer[:-buffer_size+1]
+                            buffer = buffer[-buffer_size+1:]
+
+                            # Remove model prefix if present
                             cleaned_content = strip_model_prefix(content_to_yield, model_prefix)
                             timestamp = int(datetime.now().timestamp())
                             yield f"data: {json.dumps(create_chat_completion_data(cleaned_content, request.model, timestamp))}\n\n"
 
-                # After streaming is done, check if any remaining content is in the buffer
+                # Process any remaining content in buffer
                 if buffer:
-                    # Remove the blocked message if present
+                    # Remove any occurrence of BLOCKED_MESSAGE in buffer
                     if BLOCKED_MESSAGE in buffer:
                         logger.info("Blocked message detected in remaining buffer.")
-                        buffer = buffer.replace(BLOCKED_MESSAGE, '').strip()
-                        if not buffer:
-                            pass  # No content to yield
-                        else:
-                            cleaned_content = strip_model_prefix(buffer, model_prefix)
-                            timestamp = int(datetime.now().timestamp())
-                            yield f"data: {json.dumps(create_chat_completion_data(cleaned_content, request.model, timestamp))}\n\n"
-                    else:
-                        cleaned_content = strip_model_prefix(buffer, model_prefix)
-                        timestamp = int(datetime.now().timestamp())
-                        yield f"data: {json.dumps(create_chat_completion_data(cleaned_content, request.model, timestamp))}\n\n"
+                        buffer = buffer.replace(BLOCKED_MESSAGE, '')
+
+                    cleaned_content = strip_model_prefix(buffer, model_prefix)
+                    timestamp = int(datetime.now().timestamp())
+                    yield f"data: {json.dumps(create_chat_completion_data(cleaned_content, request.model, timestamp))}\n\n"
 
                 # Signal the end of the streaming
                 timestamp = int(datetime.now().timestamp())