Upload 22 files

.dockerignore

@@ -0,0 +1,9 @@
+__pycache__
+*.pyc
+*.pyo
+*.pyd
+*.swp
+*.swo
+*.log
+.env
+.git

.env

@@ -0,0 +1 @@
+APP_SECRET=123456

Dockerfile

@@ -1,33 +1,18 @@
-# Use the official Python 3.11 slim image as the base
-FROM python:3.11-slim
+# Use an official Python runtime as a parent image
+FROM python:3.10-slim
 
-# Install Tesseract OCR and its development libraries
-RUN apt-get update && apt-get install -y --no-install-recommends \
-    tesseract-ocr \
-    libtesseract-dev \
-    && rm -rf /var/lib/apt/lists/*
-
-# Set environment variables to prevent Python from writing .pyc files and to buffer stdout/stderr
-ENV PYTHONDONTWRITEBYTECODE=1
-ENV PYTHONUNBUFFERED=1
-
-# Set the working directory inside the container
+# Set the working directory in the container
 WORKDIR /app
 
-# Copy the requirements file into the container
-COPY requirements.txt .
-
-# Upgrade pip to the latest version
-RUN pip install --upgrade pip
-
-# Install Python dependencies from requirements.txt
-RUN pip install -r requirements.txt
+# Copy the current directory contents into the container
+COPY . /app
 
-# Copy the rest of your application code into the container
-COPY . .
+# Install any needed packages specified in requirements.txt
+RUN pip install --no-cache-dir --upgrade pip && \
+    pip install --no-cache-dir -r requirements.txt
 
-# Expose port 8000 to the outside world
-EXPOSE 8000
+# Expose the port that the FastAPI app runs on
+EXPOSE 8001
 
-# Define the default command to run your application using Uvicorn
-CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]
+# Command to run the app with Gunicorn and Uvicorn workers
+CMD ["gunicorn", "-k", "uvicorn.workers.UvicornWorker", "--workers", "4", "--bind", "0.0.0.0:8001", "main:app"]

api/app.py

@@ -0,0 +1,34 @@
+from fastapi import FastAPI, Request
+from starlette.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse
+from api.logger import setup_logger
+from api.routes import router  # 导入router而不是单独的函数
+
+logger = setup_logger(__name__)
+
+def create_app():
+    app = FastAPI()
+
+    # 配置CORS
+    app.add_middleware(
+        CORSMiddleware,
+        allow_origins=["*"],
+        allow_credentials=True,
+        allow_methods=["*"],
+        allow_headers=["*"],
+    )
+
+    # 添加路由
+    app.include_router(router)
+
+    @app.exception_handler(Exception)
+    async def global_exception_handler(request: Request, exc: Exception):
+        logger.error(f"An error occurred: {str(exc)}")
+        return JSONResponse(
+            status_code=500,
+            content={"message": "An internal server error occurred."},
+        )
+
+    return app
+
+app = create_app()

api/auth.py

@@ -0,0 +1,10 @@
+from fastapi import Depends, HTTPException
+from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
+from api.config import APP_SECRET
+
+security = HTTPBearer()
+
+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

api/config.py

@@ -0,0 +1,36 @@
+import os
+from dotenv import load_dotenv
+
+load_dotenv()
+
+BASE_URL = "https://www.blackbox.ai"
+headers = {
+    'accept': '*/*',
+    'accept-language': 'zh-CN,zh;q=0.9',
+    'origin': 'https://www.blackbox.ai',
+    'priority': 'u=1, i',
+    'sec-ch-ua': '"Google Chrome";v="129", "Not=A?Brand";v="8", "Chromium";v="129"',
+    'sec-ch-ua-mobile': '?0',
+    'sec-ch-ua-platform': '"Windows"',
+    'sec-fetch-dest': 'empty',
+    'sec-fetch-mode': 'cors',
+    'sec-fetch-site': 'same-origin',
+    'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/129.0.0.0 Safari/537.36',
+}
+APP_SECRET = os.getenv("APP_SECRET")
+ALLOWED_MODELS = [
+    {"id": "gpt-4o", "name": "gpt-4o"},
+    {"id": "gemini-1.5-pro-latest", "name": "gemini-pro"},
+    {"id": "gemini-1.5-pro", "name": "gemini-pro"},
+    {"id": "gemini-pro", "name": "gemini-pro"},
+    {"id": "claude-3-5-sonnet-20240620", "name": "claude-sonnet-3.5"},
+    {"id": "claude-3-5-sonnet", "name": "claude-sonnet-3.5"},
+]
+MODEL_MAPPING = {
+    "gpt-4o":"gpt-4o",
+    "gemini-1.5-pro-latest": "gemini-pro",
+    "gemini-1.5-pro":"gemini-1.5-pro",
+    "gemini-pro":"gemini-pro",
+    "claude-3-5-sonnet-20240620":"claude-sonnet-3.5",
+    "claude-3-5-sonnet":"claude-sonnet-3.5",
+}

api/logger.py

@@ -0,0 +1,20 @@
+import logging
+
+def setup_logger(name):
+    logger = logging.getLogger(name)
+    if not logger.handlers:
+        logger.setLevel(logging.INFO)
+        formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+        
+        # 控制台处理器
+        console_handler = logging.StreamHandler()
+        console_handler.setFormatter(formatter)
+        logger.addHandler(console_handler)
+        
+        # 文件处理器 - 错误级别
+        # error_file_handler = logging.FileHandler('error.log')
+        # error_file_handler.setFormatter(formatter)
+        # error_file_handler.setLevel(logging.ERROR)
+        # logger.addHandler(error_file_handler)
+    
+    return logger

api/models.py

@@ -0,0 +1,16 @@
+from typing import List, Optional
+from pydantic import BaseModel
+
+
+class Message(BaseModel):
+    role: str
+    content: str | list
+
+
+class ChatRequest(BaseModel):
+    model: str
+    messages: List[Message]
+    stream: Optional[bool] = False
+    temperature: Optional[float] = 0.7
+    top_p: Optional[float] = 0.9
+    max_tokens: Optional[int] = 8192

api/routes.py

@@ -0,0 +1,62 @@
+import json
+from fastapi import APIRouter, Depends, HTTPException, Request, Response
+from fastapi.responses import StreamingResponse
+from api.auth import verify_app_secret
+from api.config import ALLOWED_MODELS
+from api.models import ChatRequest
+from api.utils import process_non_streaming_response, process_streaming_response
+from api.logger import setup_logger
+
+logger = setup_logger(__name__)
+
+router = APIRouter()
+
+@router.options("/v1/chat/completions")
+@router.options("/api/v1/chat/completions")
+async def chat_completions_options():
+    return Response(
+        status_code=200,
+        headers={
+            "Access-Control-Allow-Origin": "*",
+            "Access-Control-Allow-Methods": "POST, OPTIONS",
+            "Access-Control-Allow-Headers": "Content-Type, Authorization",
+        },
+    )
+
+@router.get("/v1/models")
+@router.get("/api/v1/models")
+async def list_models():
+    return {"object": "list", "data": ALLOWED_MODELS}
+
+@router.post("/v1/chat/completions")
+@router.post("/api/v1/chat/completions")
+async def chat_completions(
+    request: ChatRequest, app_secret: str = Depends(verify_app_secret)
+):
+    logger.info("Entering chat_completions route")
+    logger.info(f"Received request: {request}")
+    logger.info(f"App secret: {app_secret}")
+    logger.info(f"Received chat completion request for model: {request.model}")
+
+    if request.model not in [model["id"] for model in ALLOWED_MODELS]:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Model {request.model} is not allowed. Allowed models are: {', '.join(model['id'] for model in ALLOWED_MODELS)}",
+        )
+
+    if request.stream:
+        logger.info("Streaming response")
+        return StreamingResponse(process_streaming_response(request), media_type="text/event-stream")
+    else:
+        logger.info("Non-streaming response")
+        return await process_non_streaming_response(request)
+
+
+@router.route('/')  
+@router.route('/healthz')  
+@router.route('/ready')  
+@router.route('/alive')  
+@router.route('/status') 
+@router.get("/health")
+def health_check(request: Request):  
+    return Response(content=json.dumps({"status": "ok"}), media_type="application/json")  

api/utils.py

@@ -0,0 +1,158 @@
+from datetime import datetime
+from http.client import HTTPException
+import json
+from typing import Any, Dict, Optional
+import uuid
+
+import httpx
+from api.config import MODEL_MAPPING, headers
+from fastapi import Depends, security
+from fastapi.security import HTTPAuthorizationCredentials
+
+from api.config import APP_SECRET, BASE_URL
+from api.models import ChatRequest
+
+from api.logger import setup_logger
+
+logger = setup_logger(__name__)
+
+
+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):
+    if isinstance(message.content, str):
+        return {"role": message.role, "content": message.content}
+    elif isinstance(message.content, list) and len(message.content) == 2:
+        return {
+            "role": message.role,
+            "content": message.content[0]["text"],
+            "data": {
+                "imageBase64": message.content[1]["image_url"]["url"],
+                "fileText": "",
+                "title": "snapshoot",
+            },
+        }
+    else:
+        return {"role": message.role, "content": message.content}
+
+
+async def process_streaming_response(request: ChatRequest):
+    json_data = {
+        "messages": [message_to_dict(msg) for msg in request.messages],
+        "previewToken": None,
+        "userId": None,
+        "codeModelMode": True,
+        "agentMode": {},
+        "trendingAgentMode": {},
+        "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),
+    }
+
+    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 + "\n"
+                        if content.startswith("$@$v=undefined-rv1$@$"):
+                            yield f"data: {json.dumps(create_chat_completion_data(content[21:], request.model, timestamp))}\n\n"
+                        else:
+                            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):
+    json_data = {
+        "messages": [message_to_dict(msg) for msg in request.messages],
+        "previewToken": None,
+        "userId": None,
+        "codeModelMode": True,
+        "agentMode": {},
+        "trendingAgentMode": {},
+        "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),
+    }
+    full_response = ""
+    async with httpx.AsyncClient() as client:
+        async with client.stream(
+            method="POST", url=f"{BASE_URL}/api/chat", headers=headers, json=json_data
+        ) as response:
+            async for chunk in response.aiter_text():
+                full_response += chunk
+    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,
+    }

main.py

@@ -1,757 +1,5 @@
-import os
-import re
-import random
-import string
-import uuid
-import json
-import logging
-import asyncio
-import time
-from collections import defaultdict
-from typing import List, Dict, Any, Optional, AsyncGenerator, Union
-
-from datetime import datetime
-
-from aiohttp import ClientSession, ClientTimeout, ClientError
-from fastapi import FastAPI, HTTPException, Request, Depends, Header
-from fastapi.responses import StreamingResponse, JSONResponse, RedirectResponse
-from pydantic import BaseModel
-
-from PIL import Image
-import base64
-from io import BytesIO
-
-# Configure logging
-logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
-    handlers=[logging.StreamHandler()]
-)
-logger = logging.getLogger(__name__)
-
-# Load environment variables
-API_KEYS = os.getenv('API_KEYS', '').split(',')  # Comma-separated API keys
-RATE_LIMIT = int(os.getenv('RATE_LIMIT', '60'))  # Requests per minute
-AVAILABLE_MODELS = os.getenv('AVAILABLE_MODELS', '')  # Comma-separated available models
-
-if not API_KEYS or API_KEYS == ['']:
-    logger.error("No API keys found. Please set the API_KEYS environment variable.")
-    raise Exception("API_KEYS environment variable not set.")
-
-# Process available models
-if AVAILABLE_MODELS:
-    AVAILABLE_MODELS = [model.strip() for model in AVAILABLE_MODELS.split(',') if model.strip()]
-else:
-    AVAILABLE_MODELS = []  # If empty, all models are available
-
-# Simple in-memory rate limiter based solely on IP addresses
-rate_limit_store = defaultdict(lambda: {"count": 0, "timestamp": time.time()})
-
-# Define cleanup interval and window
-CLEANUP_INTERVAL = 60  # seconds
-RATE_LIMIT_WINDOW = 60  # seconds
-
-async def cleanup_rate_limit_stores():
-    """
-    Periodically cleans up stale entries in the rate_limit_store to prevent memory bloat.
-    """
-    while True:
-        current_time = time.time()
-        ips_to_delete = [ip for ip, value in rate_limit_store.items() if current_time - value["timestamp"] > RATE_LIMIT_WINDOW * 2]
-        for ip in ips_to_delete:
-            del rate_limit_store[ip]
-            logger.debug(f"Cleaned up rate_limit_store for IP: {ip}")
-        await asyncio.sleep(CLEANUP_INTERVAL)
-
-async def rate_limiter_per_ip(request: Request):
-    """
-    Rate limiter that enforces a limit based on the client's IP address.
-    """
-    client_ip = request.client.host
-    current_time = time.time()
-
-    # Initialize or update the count and timestamp
-    if current_time - rate_limit_store[client_ip]["timestamp"] > RATE_LIMIT_WINDOW:
-        rate_limit_store[client_ip] = {"count": 1, "timestamp": current_time}
-    else:
-        if rate_limit_store[client_ip]["count"] >= RATE_LIMIT:
-            logger.warning(f"Rate limit exceeded for IP address: {client_ip}")
-            raise HTTPException(status_code=429, detail='Rate limit exceeded for IP address | NiansuhAI')
-        rate_limit_store[client_ip]["count"] += 1
-
-async def get_api_key(request: Request, authorization: str = Header(None)) -> str:
-    """
-    Dependency to extract and validate the API key from the Authorization header.
-    """
-    client_ip = request.client.host
-    if authorization is None or not authorization.startswith('Bearer '):
-        logger.warning(f"Invalid or missing authorization header from IP: {client_ip}")
-        raise HTTPException(status_code=401, detail='Invalid authorization header format')
-    api_key = authorization[7:]
-    if api_key not in API_KEYS:
-        logger.warning(f"Invalid API key attempted: {api_key} from IP: {client_ip}")
-        raise HTTPException(status_code=401, detail='Invalid API key')
-    return api_key
-
-# 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
-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
-
-class Blackbox:
-    url = "https://www.blackbox.ai"
-    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',
-    ]
-
-    # Filter models based on AVAILABLE_MODELS
-    if AVAILABLE_MODELS:
-        models = [model for model in models if model in AVAILABLE_MODELS]
-
-    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) -> Optional[str]:
-        if model in cls.models:
-            return model
-        elif model in cls.userSelectedModel and cls.userSelectedModel[model] in cls.models:
-            return cls.userSelectedModel[model]
-        elif model in cls.model_aliases and cls.model_aliases[model] in cls.models:
-            return cls.model_aliases[model]
-        else:
-            return cls.default_model if cls.default_model in cls.models else None
-
-    @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)
-        if model is None:
-            logger.error(f"Model {model} is not available.")
-            raise ModelNotWorkingException(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'
-
-        # Don't log the full message content for privacy
-        logger.debug(f"Generated message ID: {random_id} for model: {model}")
-
-        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": 99999999,
-            "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 (excluding messages).")
-
-        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.")
-                                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
-            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))
-
-# FastAPI app setup
-app = FastAPI()
-
-# Add the cleanup task when the app starts
-@app.on_event("startup")
-async def startup_event():
-    asyncio.create_task(cleanup_rate_limit_stores())
-    logger.info("Started rate limit store cleanup task.")
-
-# Middleware to enhance security and enforce Content-Type for specific endpoints
-@app.middleware("http")
-async def security_middleware(request: Request, call_next):
-    client_ip = request.client.host
-    # Enforce that POST requests to /v1/chat/completions must have Content-Type: application/json
-    if request.method == "POST" and request.url.path == "/v1/chat/completions":
-        content_type = request.headers.get("Content-Type")
-        if content_type != "application/json":
-            logger.warning(f"Invalid Content-Type from IP: {client_ip} for path: {request.url.path}")
-            return JSONResponse(
-                status_code=400,
-                content={
-                    "error": {
-                        "message": "Content-Type must be application/json",
-                        "type": "invalid_request_error",
-                        "param": None,
-                        "code": None
-                    }
-                },
-            )
-    response = await call_next(request)
-    return response
-
-# Request Models
-class Message(BaseModel):
-    role: str
-    content: Union[str, List[Any]]  # content can be a string or a list (for images)
-
-class ChatRequest(BaseModel):
-    model: str
-    messages: List[Message]
-    temperature: Optional[float] = 1.0
-    top_p: Optional[float] = 1.0
-    n: Optional[int] = 1
-    stream: Optional[bool] = False
-    stop: Optional[Union[str, List[str]]] = None
-    max_tokens: Optional[int] = None
-    presence_penalty: Optional[float] = 0.0
-    frequency_penalty: Optional[float] = 0.0
-    logit_bias: Optional[Dict[str, float]] = None
-    user: Optional[str] = None
-    webSearchMode: Optional[bool] = False  # Custom parameter
-    image: Optional[str] = None  # Base64-encoded image
-
-class TokenizerRequest(BaseModel):
-    text: str
-
-def calculate_estimated_cost(prompt_tokens: int, completion_tokens: int) -> float:
-    """
-    Calculate the estimated cost based on the number of tokens.
-    Replace the pricing below with your actual pricing model.
-    """
-    # Example pricing: $0.00000268 per token
-    cost_per_token = 0.00000268
-    return round((prompt_tokens + completion_tokens) * cost_per_token, 8)
-
-def create_response(content: str, model: str, finish_reason: Optional[str] = None) -> Dict[str, Any]:
-    return {
-        "id": f"chatcmpl-{uuid.uuid4()}",
-        "object": "chat.completion",
-        "created": int(datetime.now().timestamp()),
-        "model": model,
-        "choices": [
-            {
-                "index": 0,
-                "message": {
-                    "role": "assistant",
-                    "content": content
-                },
-                "finish_reason": finish_reason
-            }
-        ],
-        "usage": None,  # To be filled in non-streaming responses
-    }
-
-@app.post("/v1/chat/completions", dependencies=[Depends(rate_limiter_per_ip)])
-async def chat_completions(request: ChatRequest, req: Request, api_key: str = Depends(get_api_key)):
-    client_ip = req.client.host
-    # Redact user messages only for logging purposes
-    redacted_messages = [{"role": msg.role, "content": "[redacted]"} for msg in request.messages]
-
-    logger.info(f"Received chat completions request from API key: {api_key} | IP: {client_ip} | Model: {request.model} | Messages: {redacted_messages}")
-
-    analysis_result = None
-    if request.image:
-        try:
-            image = decode_base64_image(request.image)
-            analysis_result = analyze_image(image)
-            logger.info("Image analysis completed successfully.")
-        except HTTPException as he:
-            logger.error(f"Image analysis failed: {he.detail}")
-            raise he
-        except Exception as e:
-            logger.exception("Unexpected error during image analysis.")
-            raise HTTPException(status_code=500, detail="Image analysis failed.") from e
-
-    # Prepare messages to send to the external API, excluding image data
-    processed_messages = []
-    for msg in request.messages:
-        if isinstance(msg.content, list) and len(msg.content) == 2:
-            # Assume the second item is image data, skip it
-            processed_messages.append({
-                "role": msg.role,
-                "content": msg.content[0]["text"]  # Only include the text part
-            })
-        else:
-            processed_messages.append({
-                "role": msg.role,
-                "content": msg.content
-            })
-
-    # Create a modified ChatRequest without the image
-    modified_request = ChatRequest(
-        model=request.model,
-        messages=[msg for msg in processed_messages],
-        stream=request.stream,
-        temperature=request.temperature,
-        top_p=request.top_p,
-        max_tokens=request.max_tokens,
-        presence_penalty=request.presence_penalty,
-        frequency_penalty=request.frequency_penalty,
-        logit_bias=request.logit_bias,
-        user=request.user,
-        webSearchMode=request.webSearchMode,
-        image=None  # Exclude image from external API
-    )
-
-    try:
-        if request.stream:
-            logger.info("Streaming response")
-            # **Removed the 'await' keyword here**
-            streaming_response = Blackbox.create_async_generator(
-                model=modified_request.model,
-                messages=[{"role": msg.role, "content": msg.content} for msg in modified_request.messages],
-                proxy=None,
-                image=None,
-                image_name=None,
-                webSearchMode=modified_request.webSearchMode
-            )
-
-            # Wrap the streaming generator to include image analysis at the end
-            async def generate_with_analysis():
-                assistant_content = ""
-                try:
-                    async for chunk in streaming_response:
-                        if isinstance(chunk, ImageResponse):
-                            # Handle image responses if necessary
-                            image_markdown = f"![image]({chunk.url})\n"
-                            assistant_content += image_markdown
-                            response_chunk = create_response(image_markdown, modified_request.model, finish_reason=None)
-                        else:
-                            assistant_content += chunk
-                            # Yield the chunk as a partial choice
-                            response_chunk = {
-                                "id": f"chatcmpl-{uuid.uuid4()}",
-                                "object": "chat.completion.chunk",
-                                "created": int(datetime.now().timestamp()),
-                                "model": modified_request.model,
-                                "choices": [
-                                    {
-                                        "index": 0,
-                                        "delta": {"content": chunk, "role": "assistant"},
-                                        "finish_reason": None,
-                                    }
-                                ],
-                                "usage": None,  # Usage can be updated if you track tokens in real-time
-                            }
-                        yield f"data: {json.dumps(response_chunk)}\n\n"
-                    
-                    # After all chunks are sent, send the final message with finish_reason
-                    prompt_tokens = sum(len(msg["content"].split()) for msg in modified_request.messages)
-                    completion_tokens = len(assistant_content.split())
-                    total_tokens = prompt_tokens + completion_tokens
-                    estimated_cost = calculate_estimated_cost(prompt_tokens, completion_tokens)
-
-                    final_content = assistant_content
-                    if analysis_result:
-                        final_content += f"\n\n**Image Analysis:** {analysis_result}"
-
-                    final_response = {
-                        "id": f"chatcmpl-{uuid.uuid4()}",
-                        "object": "chat.completion",
-                        "created": int(datetime.now().timestamp()),
-                        "model": modified_request.model,
-                        "choices": [
-                            {
-                                "message": {
-                                    "role": "assistant",
-                                    "content": final_content
-                                },
-                                "finish_reason": "stop",
-                                "index": 0
-                            }
-                        ],
-                        "usage": {
-                            "prompt_tokens": prompt_tokens,
-                            "completion_tokens": completion_tokens,
-                            "total_tokens": total_tokens,
-                            "estimated_cost": estimated_cost
-                        },
-                    }
-
-                    yield f"data: {json.dumps(final_response)}\n\n"
-                    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(f"Error during streaming response generation from IP: {client_ip}.")
-                    error_response = {"error": str(e)}
-                    yield f"data: {json.dumps(error_response)}\n\n"
-
-            return StreamingResponse(generate_with_analysis(), media_type="text/event-stream")
-        else:
-            logger.info("Non-streaming response")
-            # **Removed the 'await' keyword here as well**
-            streaming_response = Blackbox.create_async_generator(
-                model=modified_request.model,
-                messages=[{"role": msg.role, "content": msg.content} for msg in modified_request.messages],
-                proxy=None,
-                image=None,
-                image_name=None,
-                webSearchMode=modified_request.webSearchMode
-            )
-
-            response_content = ""
-            async for chunk in streaming_response:
-                if isinstance(chunk, ImageResponse):
-                    response_content += f"![image]({chunk.url})\n"
-                else:
-                    response_content += chunk
-
-            prompt_tokens = sum(len(msg["content"].split()) for msg in modified_request.messages)
-            completion_tokens = len(response_content.split())
-            total_tokens = prompt_tokens + completion_tokens
-            estimated_cost = calculate_estimated_cost(prompt_tokens, completion_tokens)
-
-            if analysis_result:
-                response_content += f"\n\n**Image Analysis:** {analysis_result}"
-
-            logger.info(f"Completed non-streaming response generation for API key: {api_key} | IP: {client_ip}")
-
-            response = {
-                "id": f"chatcmpl-{uuid.uuid4()}",
-                "object": "chat.completion",
-                "created": int(datetime.now().timestamp()),
-                "model": modified_request.model,
-                "choices": [
-                    {
-                        "message": {
-                            "role": "assistant",
-                            "content": response_content
-                        },
-                        "finish_reason": "stop",
-                        "index": 0
-                    }
-                ],
-                "usage": {
-                    "prompt_tokens": prompt_tokens,
-                    "completion_tokens": completion_tokens,
-                    "total_tokens": total_tokens,
-                    "estimated_cost": estimated_cost
-                },
-            }
-
-            return response
-    except ModelNotWorkingException as e:
-        logger.warning(f"Model not working: {e} | IP: {client_ip}")
-        raise HTTPException(status_code=503, detail=str(e))
-    except HTTPException as he:
-        logger.warning(f"HTTPException: {he.detail} | IP: {client_ip}")
-        raise he
-    except Exception as e:
-        logger.exception(f"An unexpected error occurred while processing the chat completions request from IP: {client_ip}.")
-        raise HTTPException(status_code=500, detail=str(e))
-
-# Endpoint: POST /v1/tokenizer
-@app.post("/v1/tokenizer", dependencies=[Depends(rate_limiter_per_ip)])
-async def tokenizer(request: TokenizerRequest, req: Request):
-    client_ip = req.client.host
-    text = request.text
-    token_count = len(text.split())
-    logger.info(f"Tokenizer requested from IP: {client_ip} | Text length: {len(text)}")
-    return {"text": text, "tokens": token_count}
-
-# Endpoint: GET /v1/models
-@app.get("/v1/models", dependencies=[Depends(rate_limiter_per_ip)])
-async def get_models(req: Request):
-    client_ip = req.client.host
-    logger.info(f"Fetching available models from IP: {client_ip}")
-    return {"data": [{"id": model, "object": "model"} for model in Blackbox.models]}
-
-# Endpoint: GET /v1/models/{model}/status
-@app.get("/v1/models/{model}/status", dependencies=[Depends(rate_limiter_per_ip)])
-async def model_status(model: str, req: Request):
-    client_ip = req.client.host
-    logger.info(f"Model status requested for '{model}' from IP: {client_ip}")
-    if model in Blackbox.models:
-        return {"model": model, "status": "available"}
-    elif model in Blackbox.model_aliases and Blackbox.model_aliases[model] in Blackbox.models:
-        actual_model = Blackbox.model_aliases[model]
-        return {"model": actual_model, "status": "available via alias"}
-    else:
-        logger.warning(f"Model not found: {model} from IP: {client_ip}")
-        raise HTTPException(status_code=404, detail="Model not found")
-
-# Endpoint: GET /v1/health
-@app.get("/v1/health", dependencies=[Depends(rate_limiter_per_ip)])
-async def health_check(req: Request):
-    client_ip = req.client.host
-    logger.info(f"Health check requested from IP: {client_ip}")
-    return {"status": "ok"}
-
-# Endpoint: GET /v1/chat/completions (GET method)
-@app.get("/v1/chat/completions")
-async def chat_completions_get(req: Request):
-    client_ip = req.client.host
-    logger.info(f"GET request made to /v1/chat/completions from IP: {client_ip}, redirecting to 'about:blank'")
-    return RedirectResponse(url='about:blank')
-
-# Custom exception handler to match OpenAI's error format
-@app.exception_handler(HTTPException)
-async def http_exception_handler(request: Request, exc: HTTPException):
-    client_ip = request.client.host
-    logger.error(f"HTTPException: {exc.detail} | Path: {request.url.path} | IP: {client_ip}")
-    return JSONResponse(
-        status_code=exc.status_code,
-        content={
-            "error": {
-                "message": exc.detail,
-                "type": "invalid_request_error",
-                "param": None,
-                "code": None
-            }
-        },
-    )
-
-# Image Processing Utilities
-def decode_base64_image(base64_str: str) -> Image.Image:
-    try:
-        image_data = base64.b64decode(base64_str)
-        image = Image.open(BytesIO(image_data))
-        return image
-    except Exception as e:
-        logger.error("Failed to decode base64 image.")
-        raise HTTPException(status_code=400, detail="Invalid base64 image data.") from e
-
-def analyze_image(image: Image.Image) -> str:
-    """
-    Placeholder for image analysis.
-    Replace this with actual image analysis logic.
-    """
-    # Example: Return image size as analysis
-    width, height = image.size
-    return f"Image analyzed successfully. Width: {width}px, Height: {height}px."
-
-# Run the application
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+import uvicorn
+from api.app import app
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=8001)

requirements.txt

@@ -1,9 +1,7 @@
-fastapi==0.95.1
-uvicorn==0.22.0
-httpx==0.24.0
-python-multipart==0.0.5
-python-dotenv==1.0.0
-Pillow==10.0.0
-aiohttp==3.8.4
-pytesseract==0.3.10
-numpy==1.21.0
+fastapi
+httpx
+pydantic
+pyinstaller
+python-dotenv
+starlette
+uvicorn