Update main.py

main.py

@@ -12,23 +12,19 @@ from typing import List, Dict, Any, Optional, Union, AsyncGenerator
 
 from aiohttp import ClientSession, ClientResponseError
 from fastapi import FastAPI, HTTPException, Request, Depends, Header
-from fastapi.responses import JSONResponse, StreamingResponse
+from fastapi.responses import JSONResponse
 from pydantic import BaseModel
 from datetime import datetime
 
-# =====================
-# 1. Configure Logging
-# =====================
+# Configure logging
 logging.basicConfig(
-    level=logging.DEBUG,  # Set to DEBUG for detailed logs during development
+    level=logging.INFO,
     format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
     handlers=[logging.StreamHandler()]
 )
 logger = logging.getLogger(__name__)
 
-# ============================
-# 2. Load Environment Variables
-# ============================
+# 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
 
@@ -36,50 +32,19 @@ 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.")
 
-# ====================================
-# 3. Define Rate Limiting Structures
-# ====================================
+# 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
 
-# ========================
-# 4. Define Pydantic Models
-# ========================
+# Define ImageResponse if needed
 class ImageResponseModel(BaseModel):
     images: str
     alt: str
 
-class Message(BaseModel):
-    role: str
-    content: str
-
-class ChatRequest(BaseModel):
-    model: str
-    messages: List[Message]
-    temperature: Optional[float] = 1.0
-    top_p: Optional[float] = 1.0
-    n: Optional[int] = 1
-    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
-
-# ===============================
-# 5. Define Custom Exceptions
-# ===============================
-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)
-
-# =======================
-# 6. Define the Blackbox
-# =======================
+# Updated Blackbox Class
 class Blackbox:
     label = "Blackbox AI"
     url = "https://www.blackbox.ai"
@@ -175,7 +140,21 @@ class Blackbox:
         "blackboxai": "/?model=blackboxai",
         "gpt-4o": "/?model=gpt-4o",
         "gemini-pro": "/?model=gemini-pro",
-        "claude-sonnet-3.5": "/?model=claude-sonnet-3.5"
+        "claude-sonnet-3.5": "/?model=claude-sonnet-3.5",
+        "ImageGeneration": "/?model=ImageGeneration",
+        "PythonAgent": "/?model=PythonAgent",
+        "JavaAgent": "/?model=JavaAgent",
+        "JavaScriptAgent": "/?model=JavaScriptAgent",
+        "HTMLAgent": "/?model=HTMLAgent",
+        "GoogleCloudAgent": "/?model=GoogleCloudAgent",
+        "AndroidDeveloper": "/?model=AndroidDeveloper",
+        "SwiftDeveloper": "/?model=SwiftDeveloper",
+        "Next.jsAgent": "/?model=Next.jsAgent",
+        "MongoDBAgent": "/?model=MongoDBAgent",
+        "PyTorchAgent": "/?model=PyTorchAgent",
+        "ReactAgent": "/?model=ReactAgent",
+        "XcodeAgent": "/?model=XcodeAgent",
+        "AngularJSAgent": "/?model=AngularJSAgent",
     }
 
     model_aliases = {
@@ -227,12 +206,7 @@ class Blackbox:
     def clean_response(text: str) -> str:
         pattern = r'^\$\@\$v=undefined-rv1\$\@\$'
         cleaned_text = re.sub(pattern, '', text)
-        try:
-            response_json = json.loads(cleaned_text)
-            # Adjust based on actual response structure
-            return response_json.get("response", response_json.get("data", cleaned_text))
-        except json.JSONDecodeError:
-            return cleaned_text.strip()
+        return cleaned_text
 
     @classmethod
     async def generate_response(
@@ -252,7 +226,6 @@ class Blackbox:
 
         prefix = cls.model_prefixes.get(model, "")
         
-        # Construct the prompt
         formatted_prompt = ""
         for message in messages:
             role = message.get('role', '').capitalize()
@@ -298,12 +271,30 @@ class Blackbox:
                     "role": "user"
                 }
             ],
-            "model": model  # Simplified payload for testing
+            "id": chat_id,
+            "previewToken": None,
+            "userId": None,
+            "codeModelMode": True,
+            "agentMode": agent_mode,
+            "trendingAgentMode": trending_agent_mode,
+            "isMicMode": False,
+            "userSystemPrompt": None,
+            "maxTokens": 1024,
+            "playgroundTopP": 0.9,
+            "playgroundTemperature": 0.5,
+            "isChromeExt": False,
+            "githubToken": None,
+            "clickedAnswer2": False,
+            "clickedAnswer3": False,
+            "clickedForceWebSearch": False,
+            "visitFromDelta": False,
+            "mobileClient": False,
+            "webSearchMode": False,
+            "userSelectedModel": cls.userSelectedModel.get(model, model)
         }
 
         async with ClientSession(headers=common_headers) as session:
             try:
-                logger.debug(f"Payload sent to Blackbox API: {json.dumps(payload_api_chat)}")
                 async with session.post(
                     cls.api_endpoint,
                     headers=headers_api_chat_combined,
@@ -312,9 +303,7 @@ class Blackbox:
                 ) as response_api_chat:
                     response_api_chat.raise_for_status()
                     text = await response_api_chat.text()
-                    logger.debug(f"Raw response from Blackbox API: {text}")  # Log raw response
                     cleaned_response = cls.clean_response(text)
-                    logger.debug(f"Cleaned response: {cleaned_response}")  # Log cleaned response
                     return cleaned_response
             except ClientResponseError as e:
                 error_text = f"Error {e.status}: {e.message}"
@@ -322,43 +311,18 @@ class Blackbox:
                     error_response = await e.response.text()
                     cleaned_error = cls.clean_response(error_response)
                     error_text += f" - {cleaned_error}"
-                    logger.error(f"Blackbox API ClientResponseError: {error_text}")
                 except Exception:
                     pass
                 return error_text
             except Exception as e:
-                logger.exception(f"Unexpected error during /api/chat request: {str(e)}")
                 return f"Unexpected error during /api/chat request: {str(e)}"
 
-# ============================
-# 7. Initialize FastAPI App
-# ============================
-app = FastAPI()
-
-# ====================================
-# 8. Define Middleware and Dependencies
-# ====================================
-@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
+# 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)
 
 async def cleanup_rate_limit_stores():
     """
@@ -402,17 +366,55 @@ async def get_api_key(request: Request, authorization: str = Header(None)) -> st
         raise HTTPException(status_code=401, detail='Invalid API key')
     return api_key
 
-# =====================================
-# 9. Define FastAPI Event Handlers
-# =====================================
+# 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.")
 
-# ==========================================
-# 10. Define FastAPI Endpoints
-# ==========================================
+# 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: str
+
+class ChatRequest(BaseModel):
+    model: str
+    messages: List[Message]
+    temperature: Optional[float] = 1.0
+    top_p: Optional[float] = 1.0
+    n: Optional[int] = 1
+    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
+
 @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
@@ -467,21 +469,21 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
         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: 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/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"}
 
-# ========================================
-# 11. Define Custom Exception Handler
-# ========================================
+# 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
@@ -498,46 +500,6 @@ async def http_exception_handler(request: Request, exc: HTTPException):
         },
     )
 
-# ============================
-# 12. Optional: Streaming Endpoint
-# ============================
-@app.post("/v1/chat/completions/stream", dependencies=[Depends(rate_limiter_per_ip)])
-async def chat_completions_stream(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 streaming chat completions request from API key: {api_key} | IP: {client_ip} | Model: {request.model} | Messages: {redacted_messages}")
-
-    try:
-        # Validate that the requested model is available
-        if request.model not in Blackbox.models and request.model not in Blackbox.model_aliases:
-            logger.warning(f"Attempt to use unavailable model: {request.model} from IP: {client_ip}")
-            raise HTTPException(status_code=400, detail="Requested model is not available.")
-
-        # Create an asynchronous generator for the response
-        async_generator = Blackbox.create_async_generator(
-            model=request.model,
-            messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],
-            temperature=request.temperature,
-            max_tokens=request.max_tokens
-        )
-
-        logger.info(f"Started streaming response for API key: {api_key} | IP: {client_ip}")
-        return StreamingResponse(async_generator, media_type="text/event-stream")
-    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 streaming chat completions request from IP: {client_ip}.")
-        raise HTTPException(status_code=500, detail=str(e))
-
-# ========================================
-# 13. Run the Application with Uvicorn
-# ========================================
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=8000)