Update api/utils.py
Browse files- api/utils.py +7 -16
api/utils.py
CHANGED
@@ -1,10 +1,10 @@
|
|
|
|
1 |
from datetime import datetime
|
2 |
import json
|
3 |
import uuid
|
4 |
import asyncio
|
5 |
import random
|
6 |
import string
|
7 |
-
import re
|
8 |
from typing import Any, Dict, Optional
|
9 |
import httpx
|
10 |
from fastapi import HTTPException
|
@@ -22,10 +22,8 @@ from api.models import ChatRequest
|
|
22 |
from api.logger import setup_logger
|
23 |
from api.validate import getHid # Import the asynchronous getHid function
|
24 |
|
25 |
-
# Setup logger
|
26 |
logger = setup_logger(__name__)
|
27 |
|
28 |
-
# Define the blocked message
|
29 |
BLOCKED_MESSAGE = "Generated by BLACKBOX.AI, try unlimited chat https://www.blackbox.ai"
|
30 |
|
31 |
# Helper function to create chat completion data
|
@@ -53,7 +51,6 @@ def message_to_dict(message, model_prefix: Optional[str] = None):
|
|
53 |
if model_prefix:
|
54 |
content = f"{model_prefix} {content}"
|
55 |
if isinstance(message.content, list) and len(message.content) == 2 and "image_url" in message.content[1]:
|
56 |
-
# Ensure base64 images are always included for all models
|
57 |
image_base64 = message.content[1]["image_url"]["url"]
|
58 |
return {
|
59 |
"role": message.role,
|
@@ -86,16 +83,16 @@ def remove_message_between_special_tags(content: str) -> str:
|
|
86 |
content = re.sub(r'\$~~~\$.*?\$~~~\$', '', content, flags=re.DOTALL)
|
87 |
return content
|
88 |
|
89 |
-
#
|
90 |
def remove_search_results(content: str) -> str:
|
91 |
"""Remove search result sections, keeping the URLs intact."""
|
92 |
-
# Regex to match search result JSON (
|
93 |
-
search_result_pattern =
|
94 |
|
95 |
-
# Remove
|
96 |
cleaned_content = re.sub(search_result_pattern, '', content)
|
97 |
|
98 |
-
#
|
99 |
cleaned_content = cleaned_content.strip()
|
100 |
|
101 |
return cleaned_content
|
@@ -212,16 +209,10 @@ async def process_non_streaming_response(request: ChatRequest):
|
|
212 |
trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
|
213 |
model_prefix = MODEL_PREFIXES.get(request.model, "")
|
214 |
|
215 |
-
# Adjust headers_api_chat and headers_chat since referer_url is removed
|
216 |
headers_api_chat = get_headers_api_chat(BASE_URL)
|
217 |
-
headers_chat = get_headers_chat(
|
218 |
-
BASE_URL,
|
219 |
-
next_action=str(uuid.uuid4()),
|
220 |
-
next_router_state_tree=json.dumps([""]),
|
221 |
-
)
|
222 |
|
223 |
if request.model == 'o1-preview':
|
224 |
-
delay_seconds = random.randint(
|
225 |
logger.info(
|
226 |
f"Introducing a delay of {delay_seconds} seconds for model 'o1-preview' "
|
227 |
f"(Request ID: {request_id})"
|
|
|
1 |
+
import re
|
2 |
from datetime import datetime
|
3 |
import json
|
4 |
import uuid
|
5 |
import asyncio
|
6 |
import random
|
7 |
import string
|
|
|
8 |
from typing import Any, Dict, Optional
|
9 |
import httpx
|
10 |
from fastapi import HTTPException
|
|
|
22 |
from api.logger import setup_logger
|
23 |
from api.validate import getHid # Import the asynchronous getHid function
|
24 |
|
|
|
25 |
logger = setup_logger(__name__)
|
26 |
|
|
|
27 |
BLOCKED_MESSAGE = "Generated by BLACKBOX.AI, try unlimited chat https://www.blackbox.ai"
|
28 |
|
29 |
# Helper function to create chat completion data
|
|
|
51 |
if model_prefix:
|
52 |
content = f"{model_prefix} {content}"
|
53 |
if isinstance(message.content, list) and len(message.content) == 2 and "image_url" in message.content[1]:
|
|
|
54 |
image_base64 = message.content[1]["image_url"]["url"]
|
55 |
return {
|
56 |
"role": message.role,
|
|
|
83 |
content = re.sub(r'\$~~~\$.*?\$~~~\$', '', content, flags=re.DOTALL)
|
84 |
return content
|
85 |
|
86 |
+
# Improved function to remove search results but keep the links intact
|
87 |
def remove_search_results(content: str) -> str:
|
88 |
"""Remove search result sections, keeping the URLs intact."""
|
89 |
+
# Regex pattern to match search result JSON structures (looking for "link" and capturing it)
|
90 |
+
search_result_pattern = r'\$~~~\$.*?"link":\s*"(https?://[^\"]+).*?\$~~~\$'
|
91 |
|
92 |
+
# Remove content inside `$~~~$` tags that contains search result data, keeping only the URLs
|
93 |
cleaned_content = re.sub(search_result_pattern, '', content)
|
94 |
|
95 |
+
# Remove unnecessary newlines or extra spaces left by the removal
|
96 |
cleaned_content = cleaned_content.strip()
|
97 |
|
98 |
return cleaned_content
|
|
|
209 |
trending_agent_mode = TRENDING_AGENT_MODE.get(request.model, {})
|
210 |
model_prefix = MODEL_PREFIXES.get(request.model, "")
|
211 |
|
|
|
212 |
headers_api_chat = get_headers_api_chat(BASE_URL)
|
|
|
|
|
|
|
|
|
|
|
213 |
|
214 |
if request.model == 'o1-preview':
|
215 |
+
delay_seconds = random.randint(1, 60)
|
216 |
logger.info(
|
217 |
f"Introducing a delay of {delay_seconds} seconds for model 'o1-preview' "
|
218 |
f"(Request ID: {request_id})"
|