Update main.py

main.py

@@ -386,370 +386,3 @@ class Blackbox:
                 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")
-            streaming_response = await 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_response = {
-                        "id": f"chatcmpl-{uuid.uuid4()}",
-                        "object": "chat.completion",
-                        "created": int(datetime.now().timestamp()),
-                        "model": modified_request.model,
-                        "choices": [
-                            {
-                                "message": {
-                                    "role": "assistant",
-                                    "content": assistant_content
-                                },
-                                "finish_reason": "stop",
-                                "index": 0
-                            }
-                        ],
-                        "usage": {
-                            "prompt_tokens": prompt_tokens,
-                            "completion_tokens": completion_tokens,
-                            "total_tokens": total_tokens,
-                            "estimated_cost": estimated_cost
-                        },
-                    }
-
-                    if analysis_result:
-                        final_response["choices"][0]["message"]["content"] += f"\n\n**Image Analysis:** {analysis_result}"
-                    
-                    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")
-            streaming_response = await 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)
-
-            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
-                },
-            }
-
-            if analysis_result:
-                response["choices"][0]["message"]["content"] += f"\n\n**Image Analysis:** {analysis_result}"
-
-            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)