Update main.py

main.py

@@ -8,12 +8,12 @@ import logging
 import asyncio
 import time
 from collections import defaultdict
-from typing import List, Dict, Any, Optional, AsyncGenerator, Union
+from typing import List, Dict, Any, Optional, Union
 from datetime import datetime
 
-from aiohttp import ClientSession, ClientTimeout, ClientError, ClientResponseError
+from aiohttp import ClientSession, ClientResponseError
 from fastapi import FastAPI, HTTPException, Request, Depends, Header
-from fastapi.responses import StreamingResponse, JSONResponse, RedirectResponse
+from fastapi.responses import JSONResponse
 from pydantic import BaseModel
 
 # Configure logging
@@ -27,18 +27,11 @@ 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
-AVAILABLE_MODELS = os.getenv('AVAILABLE_MODELS', '')  # Comma-separated available models
 
 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.")
 
-# Process available models
-if AVAILABLE_MODELS:
-    AVAILABLE_MODELS = [model.strip() for model in AVAILABLE_MODELS.split(',') if model.strip()]
-else:
-    AVAILABLE_MODELS = []  # If empty, all models are available
-
 # Simple in-memory rate limiter based solely on IP addresses
 rate_limit_store = defaultdict(lambda: {"count": 0, "timestamp": time.time()})
 
@@ -46,27 +39,19 @@ rate_limit_store = defaultdict(lambda: {"count": 0, "timestamp": time.time()})
 CLEANUP_INTERVAL = 60  # seconds
 RATE_LIMIT_WINDOW = 60  # seconds
 
-class ImageResponse:
-    def __init__(self, images: str, alt: str):
-        self.images = images
-        self.alt = alt
-
 class Blackbox:
     label = "Blackbox AI"
     url = "https://www.blackbox.ai"
     api_endpoint = "https://www.blackbox.ai/api/chat"
     working = True
     supports_gpt_4 = True
-    supports_stream = True
     supports_system_message = True
     supports_message_history = True
 
     default_model = 'blackboxai'
-    image_models = ['ImageGeneration']
     models = [
         default_model,
         'blackboxai-pro',
-        *image_models,
         "llama-3.1-8b",
         'llama-3.1-70b',
         'llama-3.1-405b',
@@ -74,25 +59,9 @@ class Blackbox:
         'gemini-pro',
         'gemini-1.5-flash',
         'claude-sonnet-3.5',
-        'PythonAgent',
-        'JavaAgent',
-        'JavaScriptAgent',
-        'HTMLAgent',
-        'GoogleCloudAgent',
-        'AndroidDeveloper',
-        'SwiftDeveloper',
-        'Next.jsAgent',
-        'MongoDBAgent',
-        'PyTorchAgent',
-        'ReactAgent',
-        'XcodeAgent',
-        'AngularJSAgent',
     ]
 
-    agentMode = {
-        'ImageGeneration': {'mode': True, 'id': "ImageGenerationLV45LJp", 'name': "Image Generation"},
-    }
-
+    agentMode = {}
     trendingAgentMode = {
         "blackboxai": {},
         "gemini-1.5-flash": {'mode': True, 'id': 'Gemini'},
@@ -100,19 +69,6 @@ class Blackbox:
         'llama-3.1-70b': {'mode': True, 'id': "llama-3.1-70b"},
         'llama-3.1-405b': {'mode': True, 'id': "llama-3.1-405b"},
         'blackboxai-pro': {'mode': True, 'id': "BLACKBOXAI-PRO"},
-        'PythonAgent': {'mode': True, 'id': "Python Agent"},
-        'JavaAgent': {'mode': True, 'id': "Java Agent"},
-        'JavaScriptAgent': {'mode': True, 'id': "JavaScript Agent"},
-        'HTMLAgent': {'mode': True, 'id': "HTML Agent"},
-        'GoogleCloudAgent': {'mode': True, 'id': "Google Cloud Agent"},
-        'AndroidDeveloper': {'mode': True, 'id': "Android Developer"},
-        'SwiftDeveloper': {'mode': True, 'id': "Swift Developer"},
-        'Next.jsAgent': {'mode': True, 'id': "Next.js Agent"},
-        'MongoDBAgent': {'mode': True, 'id': "MongoDB Agent"},
-        'PyTorchAgent': {'mode': True, 'id': "PyTorch Agent"},
-        'ReactAgent': {'mode': True, 'id': "React Agent"},
-        'XcodeAgent': {'mode': True, 'id': "Xcode Agent"},
-        'AngularJSAgent': {'mode': True, 'id': "AngularJS Agent"},
     }
 
     userSelectedModel = {
@@ -125,21 +81,7 @@ class Blackbox:
         'gpt-4o': '@GPT-4o',
         'gemini-pro': '@Gemini-PRO',
         'claude-sonnet-3.5': '@Claude-Sonnet-3.5',
-        'PythonAgent': '@Python Agent',
-        'JavaAgent': '@Java Agent',
-        'JavaScriptAgent': '@JavaScript Agent',
-        'HTMLAgent': '@HTML Agent',
-        'GoogleCloudAgent': '@Google Cloud Agent',
-        'AndroidDeveloper': '@Android Developer',
-        'SwiftDeveloper': '@Swift Developer',
-        'Next.jsAgent': '@Next.js Agent',
-        'MongoDBAgent': '@MongoDB Agent',
-        'PyTorchAgent': '@PyTorch Agent',
-        'ReactAgent': '@React Agent',
-        'XcodeAgent': '@Xcode Agent',
-        'AngularJSAgent': '@AngularJS Agent',
         'blackboxai-pro': '@BLACKBOXAI-PRO',
-        'ImageGeneration': '@Image Generation',
     }
 
     model_referers = {
@@ -152,7 +94,6 @@ class Blackbox:
     model_aliases = {
         "gemini-flash": "gemini-1.5-flash",
         "claude-3.5-sonnet": "claude-sonnet-3.5",
-        "flux": "ImageGeneration",
     }
 
     @classmethod
@@ -201,30 +142,14 @@ class Blackbox:
         return cleaned_text
 
     @classmethod
-    async def create_async_generator(
+    async def generate_response(
         cls,
         model: str,
         messages: List[Dict[str, str]],
         proxy: Optional[str] = None,
-        web_search_mode: bool = False,
         **kwargs
-    ) -> 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.
-            proxy (Optional[str]): Proxy URL, if needed.
-            web_search_mode (bool): Enables or disables web search mode.
-            **kwargs: Additional keyword arguments.
-        Yields:
-            Union[str, ImageResponse]: Segments of the generated response or ImageResponse objects.
-        """
+    ) -> str:
         model = cls.get_model(model)
-        if model is None:
-            logger.error(f"Model {model} is not available.")
-            raise ModelNotWorkingException(model)
-
         chat_id = cls.generate_random_string()
         next_action = cls.generate_next_action()
         next_router_state_tree = cls.generate_next_router_state_tree()
@@ -297,22 +222,10 @@ class Blackbox:
             "clickedForceWebSearch": False,
             "visitFromDelta": False,
             "mobileClient": False,
-            "webSearchMode": web_search_mode,
+            "webSearchMode": False,
             "userSelectedModel": cls.userSelectedModel.get(model, model)
         }
 
-        headers_chat = {
-            'Accept': 'text/x-component',
-            'Content-Type': 'text/plain;charset=UTF-8',
-            'Referer': f'{cls.url}/chat/{chat_id}?model={model}',
-            'next-action': next_action,
-            'next-router-state-tree': next_router_state_tree,
-            'next-url': '/'
-        }
-        headers_chat_combined = {**common_headers, **headers_chat}
-
-        data_chat = '[]'
-
         async with ClientSession(headers=common_headers) as session:
             try:
                 async with session.post(
@@ -324,64 +237,7 @@ class Blackbox:
                     response_api_chat.raise_for_status()
                     text = await response_api_chat.text()
                     cleaned_response = cls.clean_response(text)
-
-                    if model in cls.image_models:
-                        match = re.search(r'!\[.*?\]\((https?://[^\)]+)\)', cleaned_response)
-                        if match:
-                            image_url = match.group(1)
-                            image_response = ImageResponse(images=image_url, alt="Generated Image")
-                            yield image_response
-                        else:
-                            yield cleaned_response
-                    else:
-                        if web_search_mode:
-                            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:
-                    error_response = await e.response.text()
-                    cleaned_error = cls.clean_response(error_response)
-                    error_text += f" - {cleaned_error}"
-                except Exception:
-                    pass
-                yield error_text
-            except Exception as e:
-                yield f"Unexpected error during /api/chat request: {str(e)}"
-
-            chat_url = f'{cls.url}/chat/{chat_id}?model={model}'
-
-            try:
-                async with session.post(
-                    chat_url,
-                    headers=headers_chat_combined,
-                    data=data_chat,
-                    proxy=proxy
-                ) as response_chat:
-                    response_chat.raise_for_status()
-                    pass
+                    return cleaned_response
             except ClientResponseError as e:
                 error_text = f"Error {e.status}: {e.message}"
                 try:
@@ -390,9 +246,9 @@ class Blackbox:
                     error_text += f" - {cleaned_error}"
                 except Exception:
                     pass
-                yield error_text
+                return error_text
             except Exception as e:
-                yield f"Unexpected error during /chat/{chat_id} request: {str(e)}"
+                return f"Unexpected error during /api/chat request: {str(e)}"
 
 # Custom exception for model not working
 class ModelNotWorkingException(Exception):
@@ -486,30 +342,11 @@ class ChatRequest(BaseModel):
     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
-    web_search_mode: Optional[bool] = False  # Custom parameter
-
-def create_response(content: str, model: str, finish_reason: Optional[str] = None) -> Dict[str, Any]:
-    return {
-        "id": f"chatcmpl-{uuid.uuid4()}",
-        "object": "chat.completion.chunk",
-        "created": int(datetime.now().timestamp()),
-        "model": model,
-        "choices": [
-            {
-                "index": 0,
-                "delta": {"content": content, "role": "assistant"},
-                "finish_reason": finish_reason,
-            }
-        ],
-        "usage": None,
-    }
 
 @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)):
@@ -525,65 +362,35 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
             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.")
 
-        # Process the request with actual message content, but don't log it
-        async_generator = Blackbox.create_async_generator(
+        # Process the request with actual message content, but don't log itresponse_content = await Blackbox.generate_response(
             model=request.model,
-            messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],  # Actual message content used here
-            web_search_mode=request.web_search_mode
+            messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],
+            temperature=request.temperature,
+            max_tokens=request.max_tokens
         )
 
-        if request.stream:
-            async def generate():
-                try:
-                    async for chunk in async_generator:
-                        if isinstance(chunk, ImageResponse):
-                            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:
-            response_content = ""
-            async for chunk in async_generator:
-                if isinstance(chunk, ImageResponse):
-                    response_content += f"![image]({chunk.images})\n"
-                else:
-                    response_content += chunk
-
-            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()),
-                "model": request.model,
-                "choices": [
-                    {
-                        "message": {
-                            "role": "assistant",
-                            "content": response_content
-                        },
-                        "finish_reason": "stop",
-                        "index": 0
-                    }
-                ],
-                "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())
-                },
-            }
+        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()),
+            "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))
@@ -594,18 +401,6 @@ 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))
 
-# Endpoint: POST /v1/tokenizer
-class TokenizerRequest(BaseModel):
-    text: str
-
-@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):
@@ -613,20 +408,6 @@ async def get_models(req: Request):
     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):
@@ -634,13 +415,6 @@ async def health_check(req: Request):
     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):