api/utils.py

# utils.py

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__)

# Initialize the HTTPBearer security scheme
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  # Assuming 'headers' from config.py is suitable

    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"

                # Indicate the end of the stream
                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  # Assuming 'headers' from config.py is suitable

    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,
    }