import os import re import random import string import uuid import json import logging import asyncio from aiohttp import ClientSession, ClientTimeout, ClientError from fastapi import FastAPI, HTTPException, Request, Depends, Header, status from fastapi.responses import StreamingResponse, JSONResponse from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel, Field from typing import List, Dict, Any, Optional, AsyncGenerator from datetime import datetime from slowapi import Limiter, _rate_limit_exceeded_handler from slowapi.util import get_remote_address from slowapi.errors import RateLimitExceeded import tiktoken from dotenv import load_dotenv # Load environment variables from .env file load_dotenv() # Configure logging logging.basicConfig( level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", handlers=[ logging.StreamHandler() ] ) logger = logging.getLogger(__name__) # Initialize FastAPI app app = FastAPI(title="OpenAI-Compatible API") # Configure CORS (adjust origins as needed) origins = [ "*", # Allow all origins; replace with specific origins in production ] app.add_middleware( CORSMiddleware, allow_origins=origins, allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # Initialize Rate Limiter from environment variable RATE_LIMIT = os.getenv("RATE_LIMIT", "60/minute") # Default to 60 requests per minute limiter = Limiter(key_func=get_remote_address, default_limits=[RATE_LIMIT]) app.state.limiter = limiter app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler) # API Key Authentication API_KEYS = set(api_key.strip() for api_key in os.getenv("API_KEYS", "").split(",") if api_key.strip()) async def get_api_key(authorization: Optional[str] = Header(None)): """ Dependency to validate API Key from the Authorization header. """ if authorization is None: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Authorization header missing", headers={"WWW-Authenticate": "Bearer"}, ) parts = authorization.split() if parts[0].lower() != "bearer" or len(parts) != 2: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid authorization header format", headers={"WWW-Authenticate": "Bearer"}, ) token = parts[1] if token not in API_KEYS: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid API Key", headers={"WWW-Authenticate": "Bearer"}, ) return token # Custom exception for model not working class ModelNotWorkingException(Exception): def __init__(self, model: str): self.model = model self.message = f"The model '{model}' is currently not working. Please try another model or wait for it to be fixed." super().__init__(self.message) # Mock implementations for ImageResponse and to_data_uri (custom functionality) class ImageResponse: def __init__(self, url: str, alt: str): self.url = url self.alt = alt def to_data_uri(image: Any) -> str: return "data:image/png;base64,..." # Replace with actual base64 data if needed # Token Counting using tiktoken def count_tokens(messages: List[Dict[str, str]], model: str) -> int: """ Counts the number of tokens in the messages using tiktoken. Adjust the encoding based on the model. """ try: encoding = tiktoken.get_encoding("cl100k_base") # Adjust encoding as per model except: encoding = tiktoken.get_encoding("cl100k_base") # Default encoding tokens = 0 for message in messages: tokens += len(encoding.encode(message['content'])) return tokens # Blackbox Class: Handles interaction with the external AI service class Blackbox: url = "https://www.blackbox.ai" api_endpoint = os.getenv("EXTERNAL_API_ENDPOINT", "https://www.blackbox.ai/api/chat") working = True supports_stream = True supports_system_message = True supports_message_history = True default_model = 'blackboxai' image_models = ['ImageGeneration'] models = [ default_model, 'blackboxai-pro', "llama-3.1-8b", 'llama-3.1-70b', 'llama-3.1-405b', 'gpt-4o', 'gemini-pro', 'gemini-1.5-flash', 'claude-sonnet-3.5', 'PythonAgent', 'JavaAgent', 'JavaScriptAgent', 'HTMLAgent', 'GoogleCloudAgent', 'AndroidDeveloper', 'SwiftDeveloper', 'Next.jsAgent', 'MongoDBAgent', 'PyTorchAgent', 'ReactAgent', 'XcodeAgent', 'AngularJSAgent', *image_models, 'Niansuh', ] agentMode = { 'ImageGeneration': {'mode': True, 'id': "ImageGenerationLV45LJp", 'name': "Image Generation"}, 'Niansuh': {'mode': True, 'id': "NiansuhAIk1HgESy", 'name': "Niansuh"}, } trendingAgentMode = { "blackboxai": {}, "gemini-1.5-flash": {'mode': True, 'id': 'Gemini'}, "llama-3.1-8b": {'mode': True, 'id': "llama-3.1-8b"}, 'llama-3.1-70b': {'mode': True, 'id': "llama-3.1-70b"}, 'llama-3.1-405b': {'mode': True, 'id': "llama-3.1-405b"}, 'blackboxai-pro': {'mode': True, 'id': "BLACKBOXAI-PRO"}, 'PythonAgent': {'mode': True, 'id': "Python Agent"}, 'JavaAgent': {'mode': True, 'id': "Java Agent"}, 'JavaScriptAgent': {'mode': True, 'id': "JavaScript Agent"}, 'HTMLAgent': {'mode': True, 'id': "HTML Agent"}, 'GoogleCloudAgent': {'mode': True, 'id': "Google Cloud Agent"}, 'AndroidDeveloper': {'mode': True, 'id': "Android Developer"}, 'SwiftDeveloper': {'mode': True, 'id': "Swift Developer"}, 'Next.jsAgent': {'mode': True, 'id': "Next.js Agent"}, 'MongoDBAgent': {'mode': True, 'id': "MongoDB Agent"}, 'PyTorchAgent': {'mode': True, 'id': "PyTorch Agent"}, 'ReactAgent': {'mode': True, 'id': "React Agent"}, 'XcodeAgent': {'mode': True, 'id': "Xcode Agent"}, 'AngularJSAgent': {'mode': True, 'id': "AngularJS Agent"}, } userSelectedModel = { "gpt-4o": "gpt-4o", "gemini-pro": "gemini-pro", 'claude-sonnet-3.5': "claude-sonnet-3.5", } model_prefixes = { 'gpt-4o': '@GPT-4o', 'gemini-pro': '@Gemini-PRO', 'claude-sonnet-3.5': '@Claude-Sonnet-3.5', 'PythonAgent': '@Python Agent', 'JavaAgent': '@Java Agent', 'JavaScriptAgent': '@JavaScript Agent', 'HTMLAgent': '@HTML Agent', 'GoogleCloudAgent': '@Google Cloud Agent', 'AndroidDeveloper': '@Android Developer', 'SwiftDeveloper': '@Swift Developer', 'Next.jsAgent': '@Next.js Agent', 'MongoDBAgent': '@MongoDB Agent', 'PyTorchAgent': '@PyTorch Agent', 'ReactAgent': '@React Agent', 'XcodeAgent': '@Xcode Agent', 'AngularJSAgent': '@AngularJS Agent', 'blackboxai-pro': '@BLACKBOXAI-PRO', 'ImageGeneration': '@Image Generation', 'Niansuh': '@Niansuh', } model_referers = { "blackboxai": f"{url}/?model=blackboxai", "gpt-4o": f"{url}/?model=gpt-4o", "gemini-pro": f"{url}/?model=gemini-pro", "claude-sonnet-3.5": f"{url}/?model=claude-sonnet-3.5" } model_aliases = { "gemini-flash": "gemini-1.5-flash", "claude-3.5-sonnet": "claude-sonnet-3.5", "flux": "ImageGeneration", "niansuh": "Niansuh", } @classmethod def get_model(cls, model: str) -> str: if model in cls.models: return model elif model in cls.userSelectedModel: return model elif model in cls.model_aliases: return cls.model_aliases[model] else: return cls.default_model @classmethod async def create_async_generator( cls, model: str, messages: List[Dict[str, str]], proxy: Optional[str] = None, image: Any = None, image_name: Optional[str] = None, webSearchMode: bool = False, **kwargs ) -> AsyncGenerator[Any, None]: model = cls.get_model(model) logger.info(f"Selected model: {model}") if not cls.working or model not in cls.models: logger.error(f"Model {model} is not working or not supported.") raise ModelNotWorkingException(model) headers = { "accept": "*/*", "accept-language": "en-US,en;q=0.9", "cache-control": "no-cache", "content-type": "application/json", "origin": cls.url, "pragma": "no-cache", "priority": "u=1, i", "referer": cls.model_referers.get(model, cls.url), "sec-ch-ua": '"Chromium";v="129", "Not=A?Brand";v="8"', "sec-ch-ua-mobile": "?0", "sec-ch-ua-platform": '"Linux"', "sec-fetch-dest": "empty", "sec-fetch-mode": "cors", "sec-fetch-site": "same-origin", "user-agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/129.0.0.0 Safari/537.36", } if model in cls.model_prefixes: prefix = cls.model_prefixes[model] if not messages[0]['content'].startswith(prefix): logger.debug(f"Adding prefix '{prefix}' to the first message.") messages[0]['content'] = f"{prefix} {messages[0]['content']}" random_id = ''.join(random.choices(string.ascii_letters + string.digits, k=7)) messages[-1]['id'] = random_id messages[-1]['role'] = 'user' if image is not None: messages[-1]['data'] = { 'fileText': '', 'imageBase64': to_data_uri(image), 'title': image_name } messages[-1]['content'] = 'FILE:BB\n$#$\n\n$#$\n' + messages[-1]['content'] logger.debug("Image data added to the message.") data = { "messages": messages, "id": random_id, "previewToken": None, "userId": None, "codeModelMode": True, "agentMode": {}, "trendingAgentMode": {}, "isMicMode": False, "userSystemPrompt": None, "maxTokens": int(os.getenv("MAX_TOKENS", "4096")), "playgroundTopP": 0.9, "playgroundTemperature": 0.5, "isChromeExt": False, "githubToken": None, "clickedAnswer2": False, "clickedAnswer3": False, "clickedForceWebSearch": False, "visitFromDelta": False, "mobileClient": False, "userSelectedModel": None, "webSearchMode": webSearchMode, } if model in cls.agentMode: data["agentMode"] = cls.agentMode[model] elif model in cls.trendingAgentMode: data["trendingAgentMode"] = cls.trendingAgentMode[model] elif model in cls.userSelectedModel: data["userSelectedModel"] = cls.userSelectedModel[model] logger.info(f"Sending request to {cls.api_endpoint} with data: {data}") timeout = ClientTimeout(total=60) # Set an appropriate timeout retry_attempts = 10 # Set the number of retry attempts for attempt in range(retry_attempts): try: async with ClientSession(headers=headers, timeout=timeout) as session: async with session.post(cls.api_endpoint, json=data, proxy=proxy) as response: response.raise_for_status() logger.info(f"Received response with status {response.status}") if model == 'ImageGeneration': response_text = await response.text() url_match = re.search(r'https://storage\.googleapis\.com/[^\s\)]+', response_text) if url_match: image_url = url_match.group(0) logger.info(f"Image URL found: {image_url}") yield ImageResponse(image_url, alt=messages[-1]['content']) else: logger.error("Image URL not found in the response.") raise Exception("Image URL not found in the response") else: full_response = "" search_results_json = "" try: async for chunk, _ in response.content.iter_chunks(): if chunk: decoded_chunk = chunk.decode(errors='ignore') decoded_chunk = re.sub(r'\$@\$v=[^$]+\$@\$', '', decoded_chunk) if decoded_chunk.strip(): if '$~~~$' in decoded_chunk: search_results_json += decoded_chunk else: full_response += decoded_chunk yield decoded_chunk logger.info("Finished streaming response chunks.") except Exception as e: logger.exception("Error while iterating over response chunks.") raise e if data["webSearchMode"] and search_results_json: match = re.search(r'\$~~~\$(.*?)\$~~~\$', search_results_json, re.DOTALL) if match: try: search_results = json.loads(match.group(1)) formatted_results = "\n\n**Sources:**\n" for i, result in enumerate(search_results[:5], 1): formatted_results += f"{i}. [{result['title']}]({result['link']})\n" logger.info("Formatted search results.") yield formatted_results except json.JSONDecodeError as je: logger.error("Failed to parse search results JSON.") raise je break # Exit the retry loop if successful except ClientError as ce: logger.error(f"Client error occurred: {ce}. Retrying attempt {attempt + 1}/{retry_attempts}") if attempt == retry_attempts - 1: raise HTTPException(status_code=502, detail="Error communicating with the external API.") except asyncio.TimeoutError: logger.error(f"Request timed out. Retrying attempt {attempt + 1}/{retry_attempts}") if attempt == retry_attempts - 1: raise HTTPException(status_code=504, detail="External API request timed out.") except Exception as e: logger.error(f"Unexpected error: {e}. Retrying attempt {attempt + 1}/{retry_attempts}") if attempt == retry_attempts - 1: raise HTTPException(status_code=500, detail=str(e)) # Pydantic Models class Message(BaseModel): role: str = Field(..., description="The role of the message author.") content: str = Field(..., description="The content of the message.") class ChatRequest(BaseModel): model: str = Field(..., description="ID of the model to use.") messages: List[Message] = Field(..., description="A list of messages comprising the conversation.") stream: Optional[bool] = Field(False, description="Whether to stream the response.") webSearchMode: Optional[bool] = Field(False, description="Whether to enable web search mode.") class ChatCompletionChoice(BaseModel): index: int delta: Dict[str, Any] finish_reason: Optional[str] = None class ChatCompletionResponse(BaseModel): id: str object: str created: int model: str choices: List[ChatCompletionChoice] usage: Optional[Dict[str, int]] = None # Utility Function to Create Response def create_response(content: str, model: str, finish_reason: Optional[str] = None) -> Dict[str, Any]: return { "id": f"chatcmpl-{uuid.uuid4()}", "object": "chat.completion.chunk", "created": int(datetime.now().timestamp()), "model": model, "choices": [ { "index": 0, "delta": {"content": content, "role": "assistant"}, "finish_reason": finish_reason, } ], "usage": None, # To be populated if usage metrics are available } # Endpoint: Chat Completions @app.post("/v1/chat/completions", response_model=ChatCompletionResponse) @limiter.limit("60/minute") # Example: 60 requests per minute per IP async def chat_completions( request: ChatRequest, req: Request, api_key: str = Depends(get_api_key) ): logger.info(f"Received chat completions request: {request}") try: messages = [{"role": msg.role, "content": msg.content} for msg in request.messages] prompt_tokens = count_tokens(messages, request.model) async_generator = Blackbox.create_async_generator( model=request.model, messages=messages, image=None, # Adjust if image handling is required image_name=None, webSearchMode=request.webSearchMode ) if request.stream: async def generate(): try: completion_tokens = 0 async for chunk in async_generator: if isinstance(chunk, ImageResponse): image_markdown = f"![image]({chunk.url})" response_chunk = create_response(image_markdown, request.model) yield f"data: {json.dumps(response_chunk)}\n\n" completion_tokens += len(image_markdown.split()) else: response_chunk = create_response(chunk, request.model) yield f"data: {json.dumps(response_chunk)}\n\n" completion_tokens += len(chunk.split()) # Signal the end of the stream yield "data: [DONE]\n\n" except HTTPException as he: error_response = {"error": he.detail} yield f"data: {json.dumps(error_response)}\n\n" except Exception as e: logger.exception("Error during streaming response generation.") error_response = {"error": str(e)} yield f"data: {json.dumps(error_response)}\n\n" return StreamingResponse(generate(), media_type="text/event-stream") else: response_content = "" completion_tokens = 0 async for chunk in async_generator: if isinstance(chunk, ImageResponse): response_content += f"![image]({chunk.url})\n" completion_tokens += len(f"![image]({chunk.url})\n".split()) else: response_content += chunk completion_tokens += len(chunk.split()) total_tokens = prompt_tokens + completion_tokens logger.info("Completed non-streaming response generation.") return ChatCompletionResponse( id=f"chatcmpl-{uuid.uuid4()}", object="chat.completion", created=int(datetime.now().timestamp()), model=request.model, choices=[ ChatCompletionChoice( index=0, delta={"content": response_content, "role": "assistant"}, finish_reason="stop" ) ], usage={ "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, "total_tokens": total_tokens } ) except ModelNotWorkingException as e: logger.warning(f"Model not working: {e}") raise HTTPException(status_code=503, detail=str(e)) except HTTPException as he: logger.warning(f"HTTPException: {he.detail}") raise he except Exception as e: logger.exception("An unexpected error occurred while processing the chat completions request.") raise HTTPException(status_code=500, detail=str(e)) # Endpoint: List Models @app.get("/v1/models", response_model=Dict[str, List[Dict[str, str]]]) @limiter.limit("60/minute") async def get_models( request: Request, api_key: str = Depends(get_api_key) ): logger.info("Fetching available models.") return {"data": [{"id": model} for model in Blackbox.models]} # Endpoint: Model Status @app.get("/v1/models/{model}/status", response_model=Dict[str, str]) @limiter.limit("60/minute") async def model_status( model: str, request: Request, api_key: str = Depends(get_api_key) ): """Check if a specific model is available.""" if model in Blackbox.models: return {"model": model, "status": "available"} elif model in Blackbox.model_aliases: actual_model = Blackbox.model_aliases[model] return {"model": actual_model, "status": "available via alias"} else: raise HTTPException(status_code=404, detail="Model not found") # Endpoint: Health Check @app.get("/v1/health", response_model=Dict[str, str]) @limiter.limit("60/minute") async def health_check(request: Request): """Health check endpoint to verify the service is running.""" return {"status": "ok"} # Run the application if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=8000)