Update api/utils.py

api/utils.py

@@ -122,6 +122,7 @@ async def process_streaming_response(request: ChatRequest):
 
     json_data = build_json_data(request, h_value, model_prefix)
 
+    buffer = ""
     async with httpx.AsyncClient() as client:
         try:
             async with client.stream(
@@ -133,20 +134,44 @@ async def process_streaming_response(request: ChatRequest):
             ) as response:
                 response.raise_for_status()
                 async for chunk in response.aiter_text():
-                    timestamp = int(datetime.now().timestamp())
                     if chunk:
-                        content = chunk
-                        if content.startswith("$@$v=undefined-rv1$@$"):
-                            content = content[21:]
-                        # Remove the blocked message if present
-                        if BLOCKED_MESSAGE in content:
+                        buffer += chunk
+                        # Check if buffer contains the blocked message
+                        if BLOCKED_MESSAGE in buffer:
                             logger.info("Blocked message detected in response.")
-                            content = content.replace(BLOCKED_MESSAGE, '').strip()
-                            if not content:
-                                continue  # Skip if content is empty after removal
-                        cleaned_content = strip_model_prefix(content, model_prefix)
+                            buffer = buffer.replace(BLOCKED_MESSAGE, '').strip()
+                            if not buffer:
+                                continue  # Skip if buffer is empty after removal
+
+                        # Process the buffer to yield content
+                        while buffer:
+                            # Decide on a suitable chunk size or yield the entire buffer
+                            content_to_yield = buffer
+                            buffer = ""
+                            # Remove the 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
+                if buffer:
+                    # Remove the blocked message if present
+                    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"
 
+                # Signal the end of the streaming
+                timestamp = int(datetime.now().timestamp())
                 yield f"data: {json.dumps(create_chat_completion_data('', request.model, timestamp, 'stop'))}\n\n"
                 yield "data: [DONE]\n\n"
         except httpx.HTTPStatusError as e: