Update main.py

main.py

@@ -1,53 +1,22 @@
 from __future__ import annotations
 
-import os
-import re
+import asyncio
+import aiohttp
 import random
 import string
-import uuid
 import json
-import logging
-import asyncio
-import time
-from collections import defaultdict
-from typing import List, Dict, Any, Optional, Union, AsyncGenerator
+import uuid
+import re
+from typing import Optional, AsyncGenerator, Union
 
 from aiohttp import ClientSession, ClientResponseError
-from fastapi import FastAPI, HTTPException, Request, Depends, Header
-from fastapi.responses import JSONResponse, StreamingResponse
-from pydantic import BaseModel
-from datetime import datetime
-
-# Configure logging
-logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
-    handlers=[logging.StreamHandler()]
-)
-logger = logging.getLogger(__name__)
-
-# Load environment variables
-API_KEYS = os.getenv('API_KEYS', '').split(',')  # Comma-separated API keys
-RATE_LIMIT = int(os.getenv('RATE_LIMIT', '60'))  # Requests per minute
-
-if not API_KEYS or API_KEYS == ['']:
-    logger.error("No API keys found. Please set the API_KEYS environment variable.")
-    raise Exception("API_KEYS environment variable not set.")
-
-# Simple in-memory rate limiter based solely on IP addresses
-rate_limit_store = defaultdict(lambda: {"count": 0, "timestamp": time.time()})
-
-# Define cleanup interval and window
-CLEANUP_INTERVAL = 60  # seconds
-RATE_LIMIT_WINDOW = 60  # seconds
-
-# Define the ImageResponse model
-class ImageResponseModel(BaseModel):
-    images: str
-    alt: str
-
-# Updated Blackbox Provider
-class Blackbox:
+
+from ..typing import AsyncResult, Messages
+from .base_provider import AsyncGeneratorProvider, ProviderModelMixin
+from ..image import ImageResponse
+
+
+class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
     label = "Blackbox AI"
     url = "https://www.blackbox.ai"
     api_endpoint = "https://www.blackbox.ai/api/chat"
@@ -200,23 +169,23 @@ class Blackbox:
     async def create_async_generator(
         cls,
         model: str,
-        messages: List[Dict[str, str]],
+        messages: Messages,
         proxy: Optional[str] = None,
         websearch: bool = False,
         **kwargs
-    ) -> AsyncGenerator[Union[str, ImageResponseModel], None]:
+    ) -> AsyncGenerator[Union[str, ImageResponse], None]:
         """
         Creates an asynchronous generator for streaming responses from Blackbox AI.
 
         Parameters:
             model (str): Model to use for generating responses.
-            messages (List[Dict[str, str]]): Message history.
+            messages (Messages): Message history.
             proxy (Optional[str]): Proxy URL, if needed.
             websearch (bool): Enables or disables web search mode.
             **kwargs: Additional keyword arguments.
 
         Yields:
-            Union[str, ImageResponseModel]: Segments of the generated response or ImageResponseModel objects.
+            Union[str, ImageResponse]: Segments of the generated response or ImageResponse objects.
         """
         model = cls.get_model(model)
 
@@ -317,26 +286,43 @@ class Blackbox:
                     proxy=proxy
                 ) as response_api_chat:
                     response_api_chat.raise_for_status()
-                    # Simulate streaming by breaking the response into chunks
-                    # Since the external API may not support streaming, we'll simulate it
                     text = await response_api_chat.text()
                     cleaned_response = cls.clean_response(text)
 
-                    # For demonstration, we'll split the response into words and yield them with delays
-                    words = cleaned_response.split()
-                    for word in words:
-                        await asyncio.sleep(0.1)  # Simulate delay
-                        yield word + ' '
-                    
-                    # If the model is an image model, handle accordingly
                     if model in cls.image_models:
                         match = re.search(r'!\[.*?\]\((https?://[^\)]+)\)', cleaned_response)
                         if match:
                             image_url = match.group(1)
-                            image_response = ImageResponseModel(images=image_url, alt="Generated Image")
+                            image_response = ImageResponse(images=image_url, alt="Generated Image")
                             yield image_response
                         else:
                             yield cleaned_response
+                    else:
+                        if websearch:
+                            match = re.search(r'\$~~~\$(.*?)\$~~~\$', cleaned_response, re.DOTALL)
+                            if match:
+                                source_part = match.group(1).strip()
+                                answer_part = cleaned_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_response
+                        else:
+                            if '$~~~$' in cleaned_response:
+                                final_response = cleaned_response.split('$~~~$')[0].strip()
+                            else:
+                                final_response = cleaned_response
+
+                        yield final_response
             except ClientResponseError as e:
                 error_text = f"Error {e.status}: {e.message}"
                 try:
@@ -349,7 +335,6 @@ class Blackbox:
             except Exception as e:
                 yield f"Unexpected error during /api/chat request: {str(e)}"
 
-            # Simulate the second API call to /chat/{chat_id}
             chat_url = f'{cls.url}/chat/{chat_id}?model={model}'
 
             try:
@@ -360,7 +345,7 @@ class Blackbox:
                     proxy=proxy
                 ) as response_chat:
                     response_chat.raise_for_status()
-                    # Assuming no streaming from this endpoint
+                    pass
             except ClientResponseError as e:
                 error_text = f"Error {e.status}: {e.message}"
                 try:
@@ -372,258 +357,3 @@ class Blackbox:
                 yield error_text
             except Exception as e:
                 yield f"Unexpected error during /chat/{chat_id} request: {str(e)}"
-
-# 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():
-    """
-    Periodically cleans up stale entries in the rate_limit_store to prevent memory bloat.
-    """
-    while True:
-        current_time = time.time()
-        ips_to_delete = [ip for ip, value in rate_limit_store.items() if current_time - value["timestamp"] > RATE_LIMIT_WINDOW * 2]
-        for ip in ips_to_delete:
-            del rate_limit_store[ip]
-            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.")
-
-# 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: 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  # Added stream parameter
-
-@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} | Stream: {request.stream}")
-
-    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.")
-
-        if request.stream:
-            # Create the asynchronous generator for streaming responses
-            async_generator = 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
-            )
-
-            logger.info(f"Started streaming response for API key: {api_key} | IP: {client_ip}")
-
-            # Define a generator function to yield the streamed data in the OpenAI format
-            async def stream_response():
-                try:
-                    async for chunk in async_generator:
-                        if isinstance(chunk, ImageResponseModel):
-                            # Handle image responses by sending markdown image syntax
-                            data = json.dumps({
-                                "choices": [{
-                                    "delta": {
-                                        "content": f"![Image]({chunk.images})"
-                                    }
-                                }]
-                            })
-                        else:
-                            # Assuming chunk is a string
-                            data = json.dumps({
-                                "choices": [{
-                                    "delta": {
-                                        "content": chunk
-                                    }
-                                }]
-                            })
-                        yield f"data: {data}\n\n"
-                except Exception as e:
-                    # If an error occurs during streaming, send it as an SSE error event
-                    error_data = json.dumps({
-                        "choices": [{
-                            "delta": {
-                                "content": f"Error: {str(e)}"
-                            }
-                        }]
-                    })
-                    yield f"data: {error_data}\n\n"
-
-            return StreamingResponse(stream_response(), media_type="text/event-stream")
-        else:
-            # Non-streaming: Collect all chunks and assemble the final response
-            all_chunks = []
-            async for chunk in 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
-            ):
-                if isinstance(chunk, ImageResponseModel):
-                    # Convert ImageResponseModel to markdown image syntax
-                    content = f"![Image]({chunk.images})"
-                else:
-                    content = chunk
-                all_chunks.append(content)
-
-            # Assemble the full response content
-            full_content = "\n".join(all_chunks)
-
-            # Calculate token usage (approximation)
-            prompt_tokens = sum(len(msg.content.split()) for msg in request.messages)
-            completion_tokens = len(full_content.split())
-            total_tokens = prompt_tokens + completion_tokens
-
-            logger.info(f"Completed non-stream 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": full_content
-                        },
-                        "finish_reason": "stop"
-                    }
-                ],
-                "usage": {
-                    "prompt_tokens": prompt_tokens,
-                    "completion_tokens": completion_tokens,
-                    "total_tokens": total_tokens
-                },
-            }
-    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: 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"}
-
-# 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
-            }
-        },
-    )
-
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)