|
|
|
|
|
|
|
|
from datetime import datetime |
|
|
import json |
|
|
import uuid |
|
|
import asyncio |
|
|
import random |
|
|
import string |
|
|
from typing import Any, Dict, Optional |
|
|
|
|
|
import httpx |
|
|
from fastapi import Depends, HTTPException |
|
|
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials |
|
|
|
|
|
from api import validate |
|
|
from api.config import ( |
|
|
MODEL_MAPPING, |
|
|
headers, |
|
|
AGENT_MODE, |
|
|
TRENDING_AGENT_MODE, |
|
|
MODEL_PREFIXES, |
|
|
BASE_URL |
|
|
) |
|
|
from api.models import ChatRequest |
|
|
|
|
|
from api.logger import setup_logger |
|
|
|
|
|
logger = setup_logger(__name__) |
|
|
|
|
|
|
|
|
bearer_scheme = HTTPBearer() |
|
|
|
|
|
|
|
|
def generate_chat_id(length: int = 7) -> str: |
|
|
"""Generate a random alphanumeric chat ID.""" |
|
|
characters = string.ascii_letters + string.digits |
|
|
return ''.join(random.choices(characters, k=length)) |
|
|
|
|
|
|
|
|
def create_chat_completion_data( |
|
|
content: str, model: str, timestamp: int, finish_reason: Optional[str] = None |
|
|
) -> Dict[str, Any]: |
|
|
"""Create a dictionary representing a chat completion chunk.""" |
|
|
return { |
|
|
"id": f"chatcmpl-{uuid.uuid4()}", |
|
|
"object": "chat.completion.chunk", |
|
|
"created": timestamp, |
|
|
"model": model, |
|
|
"choices": [ |
|
|
{ |
|
|
"index": 0, |
|
|
"delta": {"content": content, "role": "assistant"}, |
|
|
"finish_reason": finish_reason, |
|
|
} |
|
|
], |
|
|
"usage": None, |
|
|
} |
|
|
|
|
|
|
|
|
def verify_app_secret(credentials: HTTPAuthorizationCredentials = Depends(bearer_scheme)): |
|
|
"""Verify the APP_SECRET from the authorization credentials.""" |
|
|
if credentials.credentials != APP_SECRET: |
|
|
logger.warning("Invalid APP_SECRET provided.") |
|
|
raise HTTPException(status_code=403, detail="Invalid APP_SECRET") |
|
|
logger.debug("APP_SECRET verified successfully.") |
|
|
return credentials.credentials |
|
|
|
|
|
|
|
|
def message_to_dict(message, model_prefix: Optional[str] = None): |
|
|
""" |
|
|
Convert a message object to a dictionary suitable for the API request. |
|
|
Ensures base64 data and optional model prefix are handled. |
|
|
""" |
|
|
if isinstance(message.content, str): |
|
|
content = message.content |
|
|
elif isinstance(message.content, list): |
|
|
content = message.content[0].get("text", "") |
|
|
else: |
|
|
content = str(message.content) |
|
|
|
|
|
if model_prefix: |
|
|
content = f"{model_prefix} {content}" |
|
|
|
|
|
if isinstance(message.content, list) and len(message.content) == 2 and "image_url" in message.content[1]: |
|
|
return { |
|
|
"role": message.role, |
|
|
"content": content, |
|
|
"data": { |
|
|
"imageBase64": message.content[1]["image_url"]["url"], |
|
|
"fileText": "", |
|
|
"title": "snapshot", |
|
|
}, |
|
|
} |
|
|
|
|
|
return {"role": message.role, "content": content} |
|
|
|
|
|
|
|
|
def strip_model_prefix(content: str, model_prefix: Optional[str] = None) -> str: |
|
|
"""Remove the model prefix from the response content if present.""" |
|
|
if model_prefix and content.startswith(model_prefix): |
|
|
logger.debug(f"Stripping prefix '{model_prefix}' from content.") |
|
|
return content[len(model_prefix):].strip() |
|
|
return content |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def get_agent_mode_config(model: str) -> Dict[str, Any]: |
|
|
""" |
|
|
Retrieve the agent mode configuration. |
|
|
Logs a warning if the configuration is missing. |
|
|
""" |
|
|
agent_mode = AGENT_MODE.get(model) |
|
|
if not agent_mode: |
|
|
logger.warning(f"No AGENT_MODE configuration found for model: {model}") |
|
|
return agent_mode or {} |
|
|
|
|
|
|
|
|
def get_trending_agent_mode_config(model: str) -> Dict[str, Any]: |
|
|
""" |
|
|
Retrieve the trending agent mode configuration. |
|
|
Logs a warning if the configuration is missing. |
|
|
""" |
|
|
trending_agent_mode = TRENDING_AGENT_MODE.get(model) |
|
|
if not trending_agent_mode: |
|
|
logger.warning(f"No TRENDING_AGENT_MODE configuration found for model: {model}") |
|
|
return trending_agent_mode or {} |
|
|
|
|
|
|
|
|
async def process_streaming_response(request: ChatRequest): |
|
|
"""Process a streaming response for a chat completion request.""" |
|
|
chat_id = generate_chat_id() |
|
|
referer_url = get_referer_url(chat_id, request.model) |
|
|
logger.info(f"Generated Chat ID: {chat_id} - Model: {request.model} - URL: {referer_url}") |
|
|
|
|
|
agent_mode = get_agent_mode_config(request.model) |
|
|
trending_agent_mode = get_trending_agent_mode_config(request.model) |
|
|
model_prefix = MODEL_PREFIXES.get(request.model, "") |
|
|
|
|
|
headers_api_chat = headers |
|
|
|
|
|
if request.model == 'o1-preview': |
|
|
delay_seconds = random.randint(1, 60) |
|
|
logger.info(f"Introducing a delay of {delay_seconds} seconds for model 'o1-preview' (Chat ID: {chat_id})") |
|
|
await asyncio.sleep(delay_seconds) |
|
|
|
|
|
json_data = { |
|
|
"agentMode": agent_mode, |
|
|
"clickedAnswer2": False, |
|
|
"clickedAnswer3": False, |
|
|
"clickedForceWebSearch": False, |
|
|
"codeModelMode": True, |
|
|
"githubToken": None, |
|
|
"id": chat_id, |
|
|
"isChromeExt": False, |
|
|
"isMicMode": False, |
|
|
"maxTokens": request.max_tokens, |
|
|
"messages": [message_to_dict(msg, model_prefix=model_prefix) for msg in request.messages], |
|
|
"mobileClient": False, |
|
|
"playgroundTemperature": request.temperature, |
|
|
"playgroundTopP": request.top_p, |
|
|
"previewToken": None, |
|
|
"trendingAgentMode": trending_agent_mode, |
|
|
"userId": None, |
|
|
"userSelectedModel": MODEL_MAPPING.get(request.model, request.model), |
|
|
"userSystemPrompt": None, |
|
|
"validated": validate.getHid(), |
|
|
"visitFromDelta": False, |
|
|
} |
|
|
|
|
|
async with httpx.AsyncClient() as client: |
|
|
try: |
|
|
async with client.stream( |
|
|
"POST", |
|
|
f"{BASE_URL}/api/chat", |
|
|
headers=headers_api_chat, |
|
|
json=json_data, |
|
|
timeout=100.0, |
|
|
) as response: |
|
|
response.raise_for_status() |
|
|
timestamp = int(datetime.now().timestamp()) |
|
|
async for line in response.aiter_lines(): |
|
|
if line: |
|
|
content = line.strip() + "\n" |
|
|
if "https://www.blackbox.ai" in content: |
|
|
validate.getHid(True) |
|
|
content = "hid已刷新,重新对话即可\n" |
|
|
yield f"data: {json.dumps(create_chat_completion_data(content, request.model, timestamp))}\n\n" |
|
|
break |
|
|
if content.startswith("$@$v=undefined-rv1$@$"): |
|
|
content = content[21:] |
|
|
cleaned_content = strip_model_prefix(content, model_prefix) |
|
|
yield f"data: {json.dumps(create_chat_completion_data(cleaned_content, request.model, timestamp))}\n\n" |
|
|
|
|
|
|
|
|
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: |
|
|
logger.error(f"HTTP error occurred for Chat ID {chat_id}: {e}") |
|
|
raise HTTPException(status_code=e.response.status_code, detail=str(e)) |
|
|
except httpx.RequestError as e: |
|
|
logger.error(f"Request error occurred for Chat ID {chat_id}: {e}") |
|
|
raise HTTPException(status_code=500, detail=str(e)) |
|
|
except Exception as e: |
|
|
logger.error(f"Unexpected error for Chat ID {chat_id}: {e}") |
|
|
raise HTTPException(status_code=500, detail=str(e)) |
|
|
|
|
|
|
|
|
async def process_non_streaming_response(request: ChatRequest): |
|
|
"""Process a non-streaming response for a chat completion request.""" |
|
|
chat_id = generate_chat_id() |
|
|
referer_url = get_referer_url(chat_id, request.model) |
|
|
logger.info(f"Generated Chat ID: {chat_id} - Model: {request.model} - URL: {referer_url}") |
|
|
|
|
|
agent_mode = get_agent_mode_config(request.model) |
|
|
trending_agent_mode = get_trending_agent_mode_config(request.model) |
|
|
model_prefix = MODEL_PREFIXES.get(request.model, "") |
|
|
|
|
|
headers_api_chat = headers |
|
|
|
|
|
if request.model == 'o1-preview': |
|
|
delay_seconds = random.randint(20, 60) |
|
|
logger.info(f"Introducing a delay of {delay_seconds} seconds for model 'o1-preview' (Chat ID: {chat_id})") |
|
|
await asyncio.sleep(delay_seconds) |
|
|
|
|
|
json_data = { |
|
|
"agentMode": agent_mode, |
|
|
"clickedAnswer2": False, |
|
|
"clickedAnswer3": False, |
|
|
"clickedForceWebSearch": False, |
|
|
"codeModelMode": True, |
|
|
"githubToken": None, |
|
|
"id": chat_id, |
|
|
"isChromeExt": False, |
|
|
"isMicMode": False, |
|
|
"maxTokens": request.max_tokens, |
|
|
"messages": [message_to_dict(msg, model_prefix=model_prefix) for msg in request.messages], |
|
|
"mobileClient": False, |
|
|
"playgroundTemperature": request.temperature, |
|
|
"playgroundTopP": request.top_p, |
|
|
"previewToken": None, |
|
|
"trendingAgentMode": trending_agent_mode, |
|
|
"userId": None, |
|
|
"userSelectedModel": MODEL_MAPPING.get(request.model, request.model), |
|
|
"userSystemPrompt": None, |
|
|
"validated": validate.getHid(), |
|
|
"visitFromDelta": False, |
|
|
} |
|
|
|
|
|
full_response = "" |
|
|
async with httpx.AsyncClient() as client: |
|
|
try: |
|
|
async with client.stream( |
|
|
method="POST", url=f"{BASE_URL}/api/chat", headers=headers_api_chat, json=json_data |
|
|
) as response: |
|
|
response.raise_for_status() |
|
|
async for chunk in response.aiter_text(): |
|
|
full_response += chunk |
|
|
except httpx.HTTPStatusError as e: |
|
|
logger.error(f"HTTP error occurred for Chat ID {chat_id}: {e}") |
|
|
raise HTTPException(status_code=e.response.status_code, detail=str(e)) |
|
|
except httpx.RequestError as e: |
|
|
logger.error(f"Request error occurred for Chat ID {chat_id}: {e}") |
|
|
raise HTTPException(status_code=500, detail=str(e)) |
|
|
|
|
|
if full_response.startswith("$@$v=undefined-rv1$@$"): |
|
|
full_response = full_response[21:] |
|
|
|
|
|
cleaned_full_response = strip_model_prefix(full_response, model_prefix) |
|
|
|
|
|
return { |
|
|
"id": f"chatcmpl-{uuid.uuid4()}", |
|
|
"object": "chat.completion", |
|
|
"created": int(datetime.now().timestamp()), |
|
|
"model": request.model, |
|
|
"choices": [ |
|
|
{ |
|
|
"index": 0, |
|
|
"message": {"role": "assistant", "content": cleaned_full_response}, |
|
|
"finish_reason": "stop", |
|
|
} |
|
|
], |
|
|
"usage": None, |
|
|
} |
|
|
|
