Update main.py

main.py

@@ -9,16 +9,19 @@ import asyncio
 import time
 from collections import defaultdict
 from typing import List, Dict, Any, Optional, Union, AsyncGenerator
-from datetime import datetime  # <-- Added import
+
+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
+from fastapi.responses import JSONResponse, StreamingResponse
 from pydantic import BaseModel
 
+# ----------------------------- Configuration -----------------------------
+
 # Configure logging
 logging.basicConfig(
-    level=logging.DEBUG,  # Changed to DEBUG for detailed logs
+    level=logging.INFO,  # Set to DEBUG for more detailed logs
     format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
     handlers=[logging.StreamHandler()]
 )
@@ -39,19 +42,43 @@ rate_limit_store = defaultdict(lambda: {"count": 0, "timestamp": time.time()})
 CLEANUP_INTERVAL = 60  # seconds
 RATE_LIMIT_WINDOW = 60  # seconds
 
-# Define the ImageResponse model (as used in the new Blackbox class)
+# ----------------------------- Pydantic Models -----------------------------
+
 class ImageResponseModel(BaseModel):
     images: str  # URL of the generated image
     alt: str
 
-# 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)
+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 Functionality
 class Blackbox:
     label = "Blackbox AI"
     url = "https://www.blackbox.ai"
@@ -364,6 +391,8 @@ class Blackbox:
                 logger.exception(f"Unexpected error during /api/chat request: {str(e)}")  # Added logging
                 return f"Unexpected error during /api/chat request: {str(e)}"
 
+        # The following block is unreachable due to the return statements above
+        # It can be removed unless it's intended for future use
         chat_url = f'{cls.url}/chat/{chat_id}?model={model}'
 
         try:
@@ -562,6 +591,8 @@ class Blackbox:
             except Exception as e:
                 yield f"Unexpected error during /api/chat request: {str(e)}"
 
+        # The following block is unreachable due to the yield statements above
+        # It can be removed unless it's intended for future use
         chat_url = f'{cls.url}/chat/{chat_id}?model={model}'
 
         try:
@@ -585,10 +616,12 @@ class Blackbox:
         except Exception as e:
             yield f"Unexpected error during /chat/{chat_id} request: {str(e)}"
 
-# FastAPI app setup
+# ----------------------------- FastAPI App Setup -----------------------------
+
 app = FastAPI()
 
-# Rate Limiter Cleanup Task
+# ----------------------------- Rate Limiter Cleanup Task -----------------------------
+
 async def cleanup_rate_limit_stores():
     """
     Periodically cleans up stale entries in the rate_limit_store to prevent memory bloat.
@@ -607,7 +640,8 @@ 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
+# ----------------------------- Security Middleware -----------------------------
+
 @app.middleware("http")
 async def security_middleware(request: Request, call_next):
     client_ip = request.client.host
@@ -630,24 +664,24 @@ async def security_middleware(request: Request, call_next):
     response = await call_next(request)
     return response
 
-# Request Models
-class Message(BaseModel):
-    role: str
-    content: str
+# ----------------------------- API Key Dependency -----------------------------
 
-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
+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 -----------------------------
 
-# Rate Limiter Dependency
 async def rate_limiter_per_ip(request: Request):
     """
     Rate limiter that enforces a limit based on the client's IP address.
@@ -664,22 +698,8 @@ async def rate_limiter_per_ip(request: Request):
             raise HTTPException(status_code=429, detail='Rate limit exceeded for IP address | NiansuhAI')
         rate_limit_store[client_ip]["count"] += 1
 
-# 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
+# ----------------------------- Chat Completions Endpoint -----------------------------
 
-# Endpoint: POST /v1/chat/completions
 @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
@@ -697,61 +717,72 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
         # Check if the model is an image generation model
         is_image_model = request.model in Blackbox.image_models
 
-        # Generate response
-        response_content = await Blackbox.generate_response(
-            model=request.model,
-            messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],
-            temperature=request.temperature,
-            max_tokens=request.max_tokens
-        )
+        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
 
-        # If the model is for image generation, handle accordingly
-        if is_image_model and isinstance(response_content, ImageResponseModel):
-            logger.info(f"Completed image generation for API key: {api_key} | IP: {client_ip}")
+            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()),  # Fixed: datetime is now imported
+                "created": int(datetime.now().timestamp()),
                 "model": request.model,
                 "choices": [
                     {
-                        "index": 0,
                         "message": {
                             "role": "assistant",
-                            "content": response_content.images  # Return the image URL
+                            "content": response_content
                         },
-                        "finish_reason": "stop"
+                        "finish_reason": "stop",
+                        "index": 0
                     }
                 ],
                 "usage": {
                     "prompt_tokens": sum(len(msg.content.split()) for msg in request.messages),
-                    "completion_tokens": len(response_content.images.split()),
-                    "total_tokens": sum(len(msg.content.split()) for msg in request.messages) + len(response_content.images.split())
+                    "completion_tokens": len(response_content.split()),
+                    "total_tokens": sum(len(msg.content.split()) for msg in request.messages) + len(response_content.split())
                 },
             }
-
-        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()),  # Fixed: datetime is now imported
-            "model": request.model,
-            "choices": [
-                {
-                    "index": 0,
-                    "message": {
-                        "role": "assistant",
-                        "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))
@@ -762,12 +793,13 @@ 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))
 
-# Optional: Endpoint for Streaming Responses (Requires Client Support)
-# If you wish to support streaming, you can implement an endpoint that leverages the asynchronous generator.
-# This requires clients to handle streaming responses appropriately.
+# ----------------------------- 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]
 
@@ -786,14 +818,13 @@ async def chat_completions_stream(request: ChatRequest, req: Request, api_key: s
         async_gen = 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
+            websearch=request.webSearchMode
         )
 
         async def stream_response() -> AsyncGenerator[bytes, None]:
             async for chunk in async_gen:
                 if isinstance(chunk, ImageResponseModel):
-                    # For image responses, you might want to send the URL directly
+                    # For image responses, send the URL directly
                     yield json.dumps({
                         "role": "assistant",
                         "content": chunk.images
@@ -805,10 +836,9 @@ async def chat_completions_stream(request: ChatRequest, req: Request, api_key: s
                     }).encode('utf-8') + b'\n'
 
         logger.info(f"Streaming response started for API key: {api_key} | IP: {client_ip}")
-        return JSONResponse(
-            content=None,  # The actual streaming is handled by the generator
-            media_type='text/event-stream',
-            background=stream_response()
+        return StreamingResponse(
+            stream_response(),
+            media_type='text/event-stream'
         )
     except ModelNotWorkingException as e:
         logger.warning(f"Model not working: {e} | IP: {client_ip}")
@@ -820,21 +850,22 @@ async def chat_completions_stream(request: ChatRequest, req: Request, api_key: s
         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))
 
-# Endpoint: GET /v1/models
+# ----------------------------- Additional Endpoints -----------------------------
+
 @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"}
 
-# Custom exception handler to match OpenAI's error format
+# ----------------------------- Exception Handlers -----------------------------
+
 @app.exception_handler(HTTPException)
 async def http_exception_handler(request: Request, exc: HTTPException):
     client_ip = request.client.host
@@ -851,6 +882,8 @@ 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)