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, validator from typing import List, Dict, Any, Optional, Union, 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, Any]], 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: if isinstance(message['content'], list): for content_part in message['content']: if content_part.get('type') == 'text': tokens += len(encoding.encode(content_part['text'])) elif content_part.get('type') == 'image_url': tokens += len(encoding.encode(content_part['image_url']['url'])) else: 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, Any]], 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 messages and isinstance(messages[0]['content'], list): # Prepend prefix to the first text message for content_part in messages[0]['content']: if content_part.get('type') == 'text' and not content_part['text'].startswith(prefix): logger.debug(f"Adding prefix '{prefix}' to the first text message.") content_part['text'] = f"{prefix} {content_part['text']}" break elif messages and isinstance(messages[0]['content'], str) and not messages[0]['content'].startswith(prefix): messages[0]['content'] = f"{prefix} {messages[0]['content']}" random_id = ''.join(random.choices(string.ascii_letters + string.digits, k=7)) # Assuming the last message is from the user if messages: last_message = messages[-1] if isinstance(last_message['content'], list): for content_part in last_message['content']: if content_part.get('type') == 'text': content_part['role'] = 'user' else: last_message['id'] = random_id last_message['role'] = 'user' if image is not None: # Process image if required # This implementation assumes that image URLs are handled by the external service pass # Implement as needed 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: 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(): yield decoded_chunk 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 TextContent(BaseModel): type: str = Field(..., description="Type of content, e.g., 'text'.") text: str = Field(..., description="The text content.") class ImageURLContent(BaseModel): type: str = Field(..., description="Type of content, e.g., 'image_url'.") image_url: Dict[str, str] = Field(..., description="Dictionary containing the image URL.") Content = Union[TextContent, ImageURLContent] class Message(BaseModel): role: str = Field(..., description="The role of the message author.") content: Union[str, List[Content]] = Field(..., description="The content of the message. Can be a string or a list of content parts.") @validator('content', pre=True) def validate_content(cls, v): if isinstance(v, list): return [Content(**item) for item in v] elif isinstance(v, str): return v else: raise ValueError("Content must be either a string or a list of content parts.") 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( chat_request: ChatRequest, # Renamed from 'request' to 'chat_request' request: Request, # Added 'request: Request' parameter api_key: str = Depends(get_api_key) ): logger.info(f"Received chat completions request: {chat_request}") try: # Process messages for token counting and sending to Blackbox processed_messages = [] for msg in chat_request.messages: if isinstance(msg.content, list): # Convert list of content parts to a structured format combined_content = [] for part in msg.content: if isinstance(part, TextContent): combined_content.append({"type": part.type, "text": part.text}) elif isinstance(part, ImageURLContent): combined_content.append({"type": part.type, "image_url": part.image_url}) processed_messages.append({"role": msg.role, "content": combined_content}) else: processed_messages.append({"role": msg.role, "content": msg.content}) prompt_tokens = count_tokens(processed_messages, chat_request.model) async_generator = Blackbox.create_async_generator( model=chat_request.model, messages=processed_messages, image=None, # Adjust if image handling is required image_name=None, webSearchMode=chat_request.webSearchMode ) if chat_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, chat_request.model) yield f"data: {json.dumps(response_chunk)}\n\n" completion_tokens += len(image_markdown.split()) else: response_chunk = create_response(chunk, chat_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=chat_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, # Ensure 'request: Request' parameter is present 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, # Ensure 'request: Request' parameter is present 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 # Ensure 'request: Request' parameter is present ): """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)