from datetime import datetime import json import uuid from typing import Any, Dict, Optional import httpx from fastapi import HTTPException, Depends from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer from . import validate from .config import ( MODEL_MAPPING, AGENT_MODE, TRENDING_AGENT_MODE, MODEL_PREFIXES, BASE_URL, APP_SECRET, ) from .models import ChatRequest from .logger import setup_logger logger = setup_logger(__name__) security = HTTPBearer() def create_chat_completion_data( content: str, model: str, timestamp: int, finish_reason: Optional[str] = None ) -> Dict[str, Any]: 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(security)): if credentials.credentials != APP_SECRET: raise HTTPException(status_code=403, detail="Invalid APP_SECRET") return credentials.credentials def message_to_dict(message, model_prefix: Optional[str] = None): content = 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} async def process_streaming_response(request: ChatRequest): model_prefix = MODEL_PREFIXES.get(request.model, "") json_data = { "messages": [message_to_dict(msg, model_prefix=model_prefix) for msg in request.messages], "previewToken": None, "userId": None, "codeModelMode": True, "agentMode": AGENT_MODE.get(request.model, {}), "trendingAgentMode": TRENDING_AGENT_MODE.get(request.model, {}), "isMicMode": False, "userSystemPrompt": None, "maxTokens": request.max_tokens, "playgroundTopP": request.top_p, "playgroundTemperature": request.temperature, "isChromeExt": False, "githubToken": None, "clickedAnswer2": False, "clickedAnswer3": False, "clickedForceWebSearch": False, "visitFromDelta": False, "mobileClient": False, "userSelectedModel": MODEL_MAPPING.get(request.model, request.model), "validated": validate.get_hid(), } headers = { "Content-Type": "application/json", "User-Agent": ( "Mozilla/5.0 (Windows NT 10.0; Win64; x64) " "AppleWebKit/537.36 (KHTML, like Gecko) " "Chrome/91.0.4472.124 Safari/537.36" ), } async with httpx.AsyncClient() as client: try: async with client.stream( "POST", f"{BASE_URL}/api/chat", headers=headers, json=json_data, timeout=100, ) as response: response.raise_for_status() async for line in response.aiter_lines(): timestamp = int(datetime.now().timestamp()) if line: content = line if "https://www.blackbox.ai" in content: validate.get_hid(True) content = "HID refreshed, please start a new conversation.\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:] yield f"data: {json.dumps(create_chat_completion_data(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: {e}") raise HTTPException(status_code=e.response.status_code, detail=str(e)) except httpx.RequestError as e: logger.error(f"Error occurred during request: {e}") raise HTTPException(status_code=500, detail=str(e)) async def process_non_streaming_response(request: ChatRequest): model_prefix = MODEL_PREFIXES.get(request.model, "") json_data = { "messages": [message_to_dict(msg, model_prefix=model_prefix) for msg in request.messages], "previewToken": None, "userId": None, "codeModelMode": True, "agentMode": AGENT_MODE.get(request.model, {}), "trendingAgentMode": TRENDING_AGENT_MODE.get(request.model, {}), "isMicMode": False, "userSystemPrompt": None, "maxTokens": request.max_tokens, "playgroundTopP": request.top_p, "playgroundTemperature": request.temperature, "isChromeExt": False, "githubToken": None, "clickedAnswer2": False, "clickedAnswer3": False, "clickedForceWebSearch": False, "visitFromDelta": False, "mobileClient": False, "userSelectedModel": MODEL_MAPPING.get(request.model, request.model), "validated": validate.get_hid(), } headers = { "Content-Type": "application/json", "User-Agent": ( "Mozilla/5.0 (Windows NT 10.0; Win64; x64) " "AppleWebKit/537.36 (KHTML, like Gecko) " "Chrome/91.0.4472.124 Safari/537.36" ), } full_response = "" async with httpx.AsyncClient() as client: try: response = await client.post(f"{BASE_URL}/api/chat", headers=headers, json=json_data) response.raise_for_status() full_response = response.text except httpx.HTTPStatusError as e: logger.error(f"HTTP error occurred: {e}") raise HTTPException(status_code=e.response.status_code, detail=str(e)) except httpx.RequestError as e: logger.error(f"Error occurred during request: {e}") raise HTTPException(status_code=500, detail=str(e)) if "https://www.blackbox.ai" in full_response: validate.get_hid(True) full_response = "HID refreshed, please start a new conversation." if full_response.startswith("$@$v=undefined-rv1$@$"): full_response = full_response[21:] return { "id": f"chatcmpl-{uuid.uuid4()}", "object": "chat.completion", "created": int(datetime.now().timestamp()), "model": request.model, "choices": [ { "index": 0, "message": {"role": "assistant", "content": full_response}, "finish_reason": "stop", } ], "usage": None, }