Update api/utils.py

api/utils.py

@@ -80,7 +80,15 @@ def strip_model_prefix(content: str, model_prefix: Optional[str] = None) -> str:
         return content[len(model_prefix):].strip()
     return content
 
-# Process streaming response with headers from config.py
+# Helper function to remove message between special 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 ""
+    return content
+
+# Modify the streaming response processing
 async def process_streaming_response(request: ChatRequest):
     # Generate a unique ID for this request
     request_id = f"chatcmpl-{uuid.uuid4()}"
@@ -161,6 +169,10 @@ async def process_streaming_response(request: ChatRequest):
                             content = content.replace(BLOCKED_MESSAGE, '').strip()
                             if not content:
                                 continue  # Skip if content is empty after removal
+                        
+                        # Remove content between special tags
+                        content = remove_message_between_special_tags(content)
+
                         cleaned_content = strip_model_prefix(content, model_prefix)
                         yield f"data: {json.dumps(create_chat_completion_data(cleaned_content, request.model, timestamp))}\n\n"
 
@@ -175,7 +187,7 @@ async def process_streaming_response(request: ChatRequest):
             )
             raise HTTPException(status_code=500, detail=str(e))
 
-# Process non-streaming response with headers from config.py
+# Modify the non-streaming response processing
 async def process_non_streaming_response(request: ChatRequest):
     # Generate a unique ID for this request
     request_id = f"chatcmpl-{uuid.uuid4()}"
@@ -272,6 +284,9 @@ async def process_non_streaming_response(request: ChatRequest):
                 status_code=500, detail="Blocked message detected in response."
             )
 
+    # Remove content between special tags
+    full_response = remove_message_between_special_tags(full_response)
+
     cleaned_full_response = strip_model_prefix(full_response, model_prefix)
 
     return {