Update main.py

main.py

@@ -10,18 +10,14 @@ import time
 from collections import defaultdict
 from typing import List, Dict, Any, Optional, Union, AsyncGenerator
 
-from datetime import datetime  # Required for timestamping
-
 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
 
-# ----------------------------- Configuration -----------------------------
-
 # Configure logging
 logging.basicConfig(
-    level=logging.DEBUG,  # Set to DEBUG for detailed logs; change to INFO in production
+    level=logging.INFO,
     format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
     handlers=[logging.StreamHandler()]
 )
@@ -42,50 +38,19 @@ rate_limit_store = defaultdict(lambda: {"count": 0, "timestamp": time.time()})
 CLEANUP_INTERVAL = 60  # seconds
 RATE_LIMIT_WINDOW = 60  # seconds
 
-# ----------------------------- Pydantic Models -----------------------------
-
+# Define ImageResponse for image models
 class ImageResponseModel(BaseModel):
-    images: str  # URL of the generated image
+    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
-    stream: Optional[bool] = False
-    webSearchMode: Optional[bool] = False  # Added based on old code
-
-# ----------------------------- Helper Functions -----------------------------
-
-def create_response(content: str, model: str) -> Dict[str, Any]:
-    """
-    Formats the response chunk.
-    """
-    return {
-        "model": model,
-        "content": content
-    }
-
-# ----------------------------- Blackbox Class -----------------------------
-
+# Updated Blackbox class with new models and functionalities
 class Blackbox:
     label = "Blackbox AI"
     url = "https://www.blackbox.ai"
     api_endpoint = "https://www.blackbox.ai/api/chat"
     working = True
     supports_gpt_4 = True
-    supports_stream = True  # New attribute for streaming support
+    supports_stream = True
     supports_system_message = True
     supports_message_history = True
 
@@ -94,7 +59,7 @@ class Blackbox:
     models = [
         default_model,
         'blackboxai-pro',
-        *image_models,  # Incorporate image models
+        *image_models,
         "llama-3.1-8b",
         'llama-3.1-70b',
         'llama-3.1-405b',
@@ -234,15 +199,8 @@ class Blackbox:
         model: str,
         messages: List[Dict[str, str]],
         proxy: Optional[str] = None,
-        websearch: bool = False,
         **kwargs
-    ) -> Union[str, ImageResponseModel]:
-        """
-        Generates a response by interacting with the Blackbox API.
-        Performs two POST requests:
-        1. Initiates the chat and obtains a chat_id.
-        2. Retrieves the chat response using the chat_id.
-        """
+    ) -> str:
         model = cls.get_model(model)
         chat_id = cls.generate_random_string()
         next_action = cls.generate_next_action()
@@ -316,25 +274,12 @@ class Blackbox:
             "clickedForceWebSearch": False,
             "visitFromDelta": False,
             "mobileClient": False,
-            "webSearchMode": websearch,
+            "webSearchMode": False,
             "userSelectedModel": cls.userSelectedModel.get(model, model)
         }
 
-        headers_chat = {
-            'Accept': 'text/x-component',
-            'Content-Type': 'text/plain;charset=UTF-8',
-            'Referer': f'{cls.url}/chat/{chat_id}?model={model}',
-            'next-action': next_action,
-            'next-router-state-tree': next_router_state_tree,
-            'next-url': '/'
-        }
-        headers_chat_combined = {**common_headers, **headers_chat}
-
-        data_chat = '[]'
-
         async with ClientSession(headers=common_headers) as session:
             try:
-                # First POST request to initiate the chat and get chat_id
                 async with session.post(
                     cls.api_endpoint,
                     headers=headers_api_chat_combined,
@@ -343,108 +288,29 @@ 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 (initiate chat): {text}")  # Added logging
                     cleaned_response = cls.clean_response(text)
-                    logger.debug(f"Cleaned response (initiate chat): {cleaned_response}")  # Added logging
-
-                # Second POST request to retrieve the chat response using chat_id
-                chat_url = f'{cls.url}/chat/{chat_id}?model={model}'
-                async with session.post(
-                    chat_url,
-                    headers=headers_chat_combined,
-                    data=data_chat,
-                    proxy=proxy
-                ) as response_chat:
-                    response_chat.raise_for_status()
-                    chat_response_text = await response_chat.text()
-                    logger.debug(f"Raw response from Blackbox API (retrieve chat): {chat_response_text}")  # Added logging
-                    cleaned_chat_response = cls.clean_response(chat_response_text)
-                    logger.debug(f"Cleaned response (retrieve chat): {cleaned_chat_response}")  # Added logging
-
-                    if model in cls.image_models:
-                        match = re.search(r'!\[.*?\]\((https?://[^\)]+)\)', cleaned_chat_response)
-                        if match:
-                            image_url = match.group(1)
-                            image_response = ImageResponseModel(images=image_url, alt="Generated Image")
-                            logger.debug(f"Image URL extracted: {image_url}")  # Added logging
-                            return image_response
-                        else:
-                            logger.debug("No image URL found in the response.")  # Added logging
-                            return cleaned_chat_response
-                    else:
-                        if websearch:
-                            match = re.search(r'\$~~~\$(.*?)\$~~~\$', cleaned_chat_response, re.DOTALL)
-                            if match:
-                                source_part = match.group(1).strip()
-                                answer_part = cleaned_chat_response[match.end():].strip()
-                                try:
-                                    sources = json.loads(source_part)
-                                    source_formatted = "**Source:**\n"
-                                    for item in sources:
-                                        title = item.get('title', 'No Title')
-                                        link = item.get('link', '#')
-                                        position = item.get('position', '')
-                                        source_formatted += f"{position}. [{title}]({link})\n"
-                                    final_response = f"{answer_part}\n\n{source_formatted}"
-                                except json.JSONDecodeError:
-                                    final_response = f"{answer_part}\n\nSource information is unavailable."
-                            else:
-                                final_response = cleaned_chat_response
-                        else:
-                            if '$~~~$' in cleaned_chat_response:
-                                final_response = cleaned_chat_response.split('$~~~$')[0].strip()
-                            else:
-                                final_response = cleaned_chat_response
-
-                        logger.debug(f"Final response to return: {final_response}")  # Added logging
-                        return final_response
+                    return cleaned_response
             except ClientResponseError as e:
                 error_text = f"Error {e.status}: {e.message}"
                 try:
                     error_response = await e.response.text()
                     cleaned_error = cls.clean_response(error_response)
                     error_text += f" - {cleaned_error}"
-                    logger.error(f"ClientResponseError: {error_text}")  # Added logging
                 except Exception:
                     pass
                 return error_text
             except Exception as e:
-                logger.exception(f"Unexpected error during /api/chat request: {str(e)}")  # Added logging
                 return f"Unexpected error during /api/chat request: {str(e)}"
 
-    @classmethod
-    async def create_async_generator(
-        cls,
-        model: str,
-        messages: List[Dict[str, str]],
-        proxy: Optional[str] = None,
-        websearch: bool = False,
-        **kwargs
-    ) -> AsyncGenerator[Union[str, ImageResponseModel], None]:
-        """
-        Creates an asynchronous generator for streaming responses from Blackbox AI.
-        """
-        try:
-            response = await cls.generate_response(
-                model=model,
-                messages=messages,
-                proxy=proxy,
-                websearch=websearch,
-                **kwargs
-            )
-            if isinstance(response, ImageResponseModel):
-                yield response
-            else:
-                yield response
-        except Exception as e:
-            logger.exception(f"Error in create_async_generator: {str(e)}")
-            yield f"Unexpected error: {str(e)}"
-
-# ----------------------------- FastAPI App Setup -----------------------------
+    # If needed, you can integrate create_async_generator here
+    # ...
 
-app = FastAPI()
-
-# ----------------------------- Rate Limiter Cleanup Task -----------------------------
+# 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():
     """
@@ -458,14 +324,46 @@ async def cleanup_rate_limit_stores():
             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
+
+# 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.")
 
-# ----------------------------- Security Middleware -----------------------------
-
+# 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
@@ -488,41 +386,22 @@ async def security_middleware(request: Request, call_next):
     response = await call_next(request)
     return response
 
-# ----------------------------- API Key Dependency -----------------------------
-
-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
-
-# ----------------------------- Rate Limiter Dependency -----------------------------
-
-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
+# Request Models
+class Message(BaseModel):
+    role: str
+    content: str
 
-# ----------------------------- Chat Completions Endpoint -----------------------------
+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)):
@@ -538,132 +417,36 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
             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.")
 
-        # Check if the model is an image generation model
-        is_image_model = request.model in Blackbox.image_models
-
-        if request.stream:
-            # Streaming response
-            async_generator = Blackbox.create_async_generator(
-                model=request.model,
-                messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],
-                websearch=request.webSearchMode
-            )
-
-            async def generate():
-                try:
-                    async for chunk in async_generator:
-                        if isinstance(chunk, ImageResponseModel):
-                            image_markdown = f"![image]({chunk.images})"
-                            response_chunk = create_response(image_markdown, request.model)
-                        else:
-                            response_chunk = create_response(chunk, request.model)
-                        
-                        yield f"data: {json.dumps(response_chunk)}\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(), media_type="text/event-stream")
-        else:
-            # Non-streaming response
-            async_generator = Blackbox.create_async_generator(
-                model=request.model,
-                messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],
-                websearch=request.webSearchMode
-            )
-
-            response_content = ""
-            async for chunk in async_generator:
-                if isinstance(chunk, ImageResponseModel):
-                    response_content += f"![image]({chunk.images})\n"
-                else:
-                    response_content += chunk
-
-            logger.info(f"Completed non-streaming response generation for API key: {api_key} | IP: {client_ip}")
-            return {
-                "id": f"chatcmpl-{uuid.uuid4()}",
-                "object": "chat.completion",
-                "created": int(datetime.now().timestamp()),
-                "model": request.model,
-                "choices": [
-                    {
-                        "message": {
-                            "role": "assistant",
-                            "content": response_content
-                        },
-                        "finish_reason": "stop",
-                        "index": 0
-                    }
-                ],
-                "usage": {
-                    "prompt_tokens": sum(len(msg.content.split()) for msg in request.messages),
-                    "completion_tokens": len(response_content.split()),
-                    "total_tokens": sum(len(msg.content.split()) for msg in request.messages) + len(response_content.split())
-                },
-            }
-    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))
-
-# ----------------------------- Streaming Endpoint (Optional) -----------------------------
-
-@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)):
-    """
-    Optional endpoint for streaming responses. Can be removed if not needed.
-    """
-    client_ip = req.client.host
-    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.")
-
-        # Check if the model is an image generation model
-        is_image_model = request.model in Blackbox.image_models
-
-        # Create an asynchronous generator
-        async_gen = Blackbox.create_async_generator(
+        # Process the request with actual message content, but don't log it
+        response_content = await Blackbox.generate_response(
             model=request.model,
             messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],
-            websearch=request.webSearchMode
+            temperature=request.temperature,
+            max_tokens=request.max_tokens
         )
 
-        async def stream_response() -> AsyncGenerator[bytes, None]:
-            async for chunk in async_gen:
-                if isinstance(chunk, ImageResponseModel):
-                    # For image responses, send the URL directly
-                    yield json.dumps({
-                        "role": "assistant",
-                        "content": chunk.images
-                    }).encode('utf-8') + b'\n'
-                else:
-                    yield json.dumps({
+        logger.info(f"Completed response generation for API key: {api_key} | IP: {client_ip}")
+        return {
+            "id": f"chatcmpl-{uuid.uuid4()}",
+            "object": "chat.completion",
+            "created": int(datetime.now().timestamp()),
+            "model": request.model,
+            "choices": [
+                {
+                    "index": 0,
+                    "message": {
                         "role": "assistant",
-                        "content": chunk
-                    }).encode('utf-8') + b'\n'
-
-        logger.info(f"Streaming response started for API key: {api_key} | IP: {client_ip}")
-        return StreamingResponse(
-            stream_response(),
-            media_type='text/event-stream'
-        )
+                        "content": response_content
+                    },
+                    "finish_reason": "stop"
+                }
+            ],
+            "usage": {
+                "prompt_tokens": sum(len(msg.content.split()) for msg in request.messages),
+                "completion_tokens": len(response_content.split()),
+                "total_tokens": sum(len(msg.content.split()) for msg in request.messages) + len(response_content.split())
+            },
+        }
     except ModelNotWorkingException as e:
         logger.warning(f"Model not working: {e} | IP: {client_ip}")
         raise HTTPException(status_code=503, detail=str(e))
@@ -671,25 +454,24 @@ async def chat_completions_stream(request: ChatRequest, req: Request, api_key: s
         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}.")
+        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))
 
-# ----------------------------- Additional Endpoints -----------------------------
-
+# 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"}
 
-# ----------------------------- Exception Handlers -----------------------------
-
+# 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
@@ -706,8 +488,6 @@ async def http_exception_handler(request: Request, exc: HTTPException):
         },
     )
 
-# ----------------------------- Main Entry Point -----------------------------
-
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=8000)