main.py

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)