Update main.py

main.py

@@ -10,7 +10,7 @@ import logging
 import asyncio
 import time
 from collections import defaultdict
-from typing import List, Dict, Any, Optional, AsyncGenerator, Union, Tuple
+from typing import List, Dict, Any, Optional, Union, Tuple
 
 from datetime import datetime
 
@@ -453,272 +453,8 @@ async def security_middleware(request: Request, call_next):
     response = await call_next(request)
     return response
 
-# Request 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: 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
-
-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
-    }
-
-def extract_all_images_from_content(content_items: List[ContentItem]) -> List[Tuple[str, str]]:
-    """
-    Extracts all images from the content list.
-    Returns a list of tuples containing (alt_text, image_data_uri).
-    """
-    images = []
-    for item in content_items:
-        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
-
-# 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(len(" ".join([item.text if isinstance(item, TextContent) else item.image_url['url'] for item in msg.content]).split()) for msg in request.messages)
-        completion_tokens = len(assistant_content.split())
-        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):
-    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')
+# Define other endpoints as needed
+# ... [Insert other endpoints here] ...
 
 # Custom exception handler to match OpenAI's error format
 @app.exception_handler(HTTPException)