Update main.py

main.py

@@ -23,7 +23,7 @@ import base64
 
 from dotenv import load_dotenv
 
-# Load environment variables
+# Load environment variables from .env file
 load_dotenv()
 
 # Configure logging
@@ -145,7 +145,7 @@ class ImageResponseCustom:
         self.url = url
         self.alt = alt
 
-# Blackbox AI Integration (Placeholder for actual implementation)
+# Placeholder for Blackbox AI Integration
 class Blackbox:
     url = "https://www.blackbox.ai"
     api_endpoint = "https://www.blackbox.ai/api/chat"  # Placeholder endpoint
@@ -421,7 +421,7 @@ class Blackbox:
                 if attempt == retry_attempts - 1:
                     raise HTTPException(status_code=500, detail=str(e))
 
-# FastAPI app setup
+# Initialize FastAPI app
 app = FastAPI()
 
 # Add the cleanup task when the app starts
@@ -453,8 +453,308 @@ async def security_middleware(request: Request, call_next):
     response = await call_next(request)
     return response
 
-# Define other endpoints as needed
-# ... [Insert other endpoints here] ...
+# Pydantic Models
+
+class TextContent(BaseModel):
+    type: str = "text"
+    text: str
+
+    @validator('type')
+    def type_must_be_text(cls, v):
+        if v != "text":
+            raise ValueError("Type must be 'text'")
+        return v
+
+class ImageContent(BaseModel):
+    type: str = "image_url"
+    image_url: Dict[str, str]
+
+    @validator('type')
+    def type_must_be_image_url(cls, v):
+        if v != "image_url":
+            raise ValueError("Type must be 'image_url'")
+        return v
+
+ContentItem = Union[TextContent, ImageContent]
+
+class Message(BaseModel):
+    role: str
+    content: Union[str, List[ContentItem]]
+
+    @validator('role')
+    def role_must_be_valid(cls, v):
+        if v not in {"system", "user", "assistant"}:
+            raise ValueError("Role must be 'system', 'user', or 'assistant'")
+        return v
+
+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
+
+class TokenizerRequest(BaseModel):
+    text: str
+
+# Utility Functions
+
+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 count_tokens(text: str) -> int:
+    """
+    Counts the number of tokens in a given text using tiktoken.
+    """
+    try:
+        import tiktoken
+        encoding = tiktoken.get_encoding("cl100k_base")
+        return len(encoding.encode(text))
+    except ImportError:
+        # Fallback if tiktoken is not installed
+        return len(text.split())
+
+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
+    }
+
+def extract_all_images_from_content(content: Union[str, List[ContentItem]]) -> List[Tuple[str, str]]:
+    """
+    Extracts all images from the content.
+    Returns a list of tuples containing (alt_text, image_data_uri).
+    """
+    images = []
+    if isinstance(content, list):
+        for item in content:
+            if isinstance(item, ImageContent):
+                alt_text = item.image_url.get('alt', '')  # Optional alt text
+                image_data_uri = item.image_url.get('url', '')
+                if image_data_uri:
+                    images.append((alt_text, image_data_uri))
+    return images
+
+# Image Analysis Function (Placeholder)
+async def analyze_image(image_data_uri: str) -> str:
+    """
+    Placeholder function to analyze the image.
+    Replace this with actual image analysis logic or API calls.
+    """
+    try:
+        # Extract base64 data
+        image_data = image_data_uri.split(",")[1]
+        # Decode the image
+        image_bytes = base64.b64decode(image_data)
+        
+        # Here, integrate with an image analysis API or implement your own logic
+        # For demonstration, we'll simulate analysis with a dummy response.
+        await asyncio.sleep(1)  # Simulate processing delay
+        return "Image analysis result: The image depicts a beautiful sunset over the mountains."
+    except Exception as e:
+        logger.error(f"Failed to analyze image: {e}")
+        raise HTTPException(status_code=400, detail="Failed to process the provided image.")
+
+# 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
+    # 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}")
+
+    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.")
+
+        # Initialize response content
+        assistant_content = ""
+
+        # Iterate through messages to find and process images
+        for msg in request.messages:
+            if msg.role == "user":
+                # Extract all images from the message content
+                images = extract_all_images_from_content(msg.content)
+                for alt_text, image_data_uri in images:
+                    # Analyze the image
+                    analysis_result = await analyze_image(image_data_uri)
+                    assistant_content += analysis_result + "\n"
+
+        # Example response content
+        assistant_content += "Based on the image you provided, here are the insights..."
+
+        # Calculate token usage (simple approximation)
+        prompt_tokens = sum(count_tokens(" ".join([item.text if isinstance(item, TextContent) else item.image_url['url'] for item in msg.content]) ) for msg in request.messages)
+        completion_tokens = count_tokens(assistant_content)
+        total_tokens = prompt_tokens + completion_tokens
+        estimated_cost = calculate_estimated_cost(prompt_tokens, completion_tokens)
+
+        logger.info(f"Completed response generation for API key: {api_key} | IP: {client_ip}")
+
+        if request.stream:
+            async def generate():
+                try:
+                    for msg in request.messages:
+                        if msg.role == "user":
+                            images = extract_all_images_from_content(msg.content)
+                            for alt_text, image_data_uri in images:
+                                analysis_result = await analyze_image(image_data_uri)
+                                response_chunk = {
+                                    "id": f"chatcmpl-{uuid.uuid4()}",
+                                    "object": "chat.completion.chunk",
+                                    "created": int(datetime.now().timestamp()),
+                                    "model": request.model,
+                                    "choices": [
+                                        {
+                                            "index": 0,
+                                            "delta": {"content": analysis_result + "\n", "role": "assistant"},
+                                            "finish_reason": None,
+                                        }
+                                    ],
+                                    "usage": None,
+                                }
+                                yield f"data: {json.dumps(response_chunk)}\n\n"
+
+                    # Final message
+                    final_response = {
+                        "id": f"chatcmpl-{uuid.uuid4()}",
+                        "object": "chat.completion",
+                        "created": int(datetime.now().timestamp()),
+                        "model": request.model,
+                        "choices": [
+                            {
+                                "message": {
+                                    "role": "assistant",
+                                    "content": assistant_content.strip()
+                                },
+                                "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(), media_type="text/event-stream")
+        else:
+            return {
+                "id": f"chatcmpl-{uuid.uuid4()}",
+                "object": "chat.completion",
+                "created": int(datetime.now().timestamp()),
+                "model": request.model,
+                "choices": [
+                    {
+                        "message": {
+                            "role": "assistant",
+                            "content": assistant_content.strip()
+                        },
+                        "finish_reason": "stop",
+                        "index": 0
+                    }
+                ],
+                "usage": {
+                    "prompt_tokens": prompt_tokens,
+                    "completion_tokens": completion_tokens,
+                    "total_tokens": total_tokens,
+                    "estimated_cost": estimated_cost
+                },
+            }
+    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, api_key: str = Depends(get_api_key)):
+    client_ip = req.client.host
+    text = request.text
+    token_count = count_tokens(text)
+    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, api_key: str = Depends(get_api_key)):
+    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, api_key: str = Depends(get_api_key)):
+    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, api_key: str = Depends(get_api_key)):
+    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, api_key: str = Depends(get_api_key)):
+    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)