Update api/utils.py

api/utils.py

@@ -5,6 +5,7 @@ import asyncio
 import random
 import string
 from typing import Any, Dict, Optional
+import re
 
 import httpx
 from fastapi import HTTPException
@@ -80,15 +81,14 @@ def strip_model_prefix(content: str, model_prefix: Optional[str] = None) -> str:
         return content[len(model_prefix):].strip()
     return content
 
-# Helper function to remove message between special tags
+# Function to remove content between '$~~~$' tags
 def remove_message_between_special_tags(content: str) -> str:
-    """Remove any message that starts with '$~~~$' and ends with '$~~~$'."""
-    if content.startswith('$~~~$') and content.endswith('$~~~$'):
-        # Remove the entire message between the tags
-        return ""
+    """Remove any content that starts with '$~~~$' and ends with '$~~~$', including the tags."""
+    # Use regex to remove the entire block between '$~~~$' tags
+    content = re.sub(r'\$~~~\$.*?\$~~~\$', '', content, flags=re.DOTALL)
     return content
 
-# Modify the streaming response processing
+# Process streaming response with headers from config.py
 async def process_streaming_response(request: ChatRequest):
     # Generate a unique ID for this request
     request_id = f"chatcmpl-{uuid.uuid4()}"
@@ -187,7 +187,7 @@ async def process_streaming_response(request: ChatRequest):
             )
             raise HTTPException(status_code=500, detail=str(e))
 
-# Modify the non-streaming response processing
+# Process non-streaming response with headers from config.py
 async def process_non_streaming_response(request: ChatRequest):
     # Generate a unique ID for this request
     request_id = f"chatcmpl-{uuid.uuid4()}"