Update main.py

main.py

@@ -10,11 +10,10 @@ import time
 from collections import defaultdict
 from typing import List, Dict, Any, Optional, Union, AsyncGenerator
 
-from aiohttp import ClientSession, ClientResponseError, ClientTimeout
+from aiohttp import ClientSession, ClientResponseError
 from fastapi import FastAPI, HTTPException, Request, Depends, Header
-from fastapi.responses import JSONResponse, StreamingResponse
+from fastapi.responses import JSONResponse
 from pydantic import BaseModel
-from datetime import datetime
 
 # Configure logging
 logging.basicConfig(
@@ -39,32 +38,14 @@ rate_limit_store = defaultdict(lambda: {"count": 0, "timestamp": time.time()})
 CLEANUP_INTERVAL = 60  # seconds
 RATE_LIMIT_WINDOW = 60  # seconds
 
-# Define ImageResponseModel
+# Define ImageResponse for handling image outputs
 class ImageResponseModel(BaseModel):
     images: str
     alt: str
 
-def strip_markdown(text: str) -> str:
-    """
-    Strips markdown syntax from the given text to ensure plain text.
-    """
-    # Remove bold (**text** or __text__)
-    text = re.sub(r'(\*\*|__)(.*?)\1', r'\2', text)
-    # Remove italic (*text* or _text_)
-    text = re.sub(r'(\*|_)(.*?)\1', r'\2', text)
-    # Remove inline code (`code`)
-    text = re.sub(r'`(.*?)`', r'\1', text)
-    # Remove links [text](url)
-    text = re.sub(r'\[(.*?)\]\((.*?)\)', r'\1', text)
-    # Remove images ![alt](url)
-    text = re.sub(r'!\[(.*?)\]\((.*?)\)', r'\1', text)
-    # Remove headers (# Header)
-    text = re.sub(r'#+\s+(.*)', r'\1', text)
-    # Remove any remaining markdown characters
-    text = re.sub(r'[*_`>#]', '', text)
-    return text
-
-# Updated Blackbox Class
+# Define Messages type for better type hinting
+Messages = List[Dict[str, Any]]
+
 class Blackbox:
     label = "Blackbox AI"
     url = "https://www.blackbox.ai"
@@ -160,21 +141,7 @@ class Blackbox:
         "blackboxai": "/?model=blackboxai",
         "gpt-4o": "/?model=gpt-4o",
         "gemini-pro": "/?model=gemini-pro",
-        "claude-sonnet-3.5": "/?model=claude-sonnet-3.5",
-        "ImageGeneration": "/?model=ImageGeneration",
-        "PythonAgent": "/?model=PythonAgent",
-        "JavaAgent": "/?model=JavaAgent",
-        "JavaScriptAgent": "/?model=JavaScriptAgent",
-        "HTMLAgent": "/?model=HTMLAgent",
-        "GoogleCloudAgent": "/?model=GoogleCloudAgent",
-        "AndroidDeveloper": "/?model=AndroidDeveloper",
-        "SwiftDeveloper": "/?model=SwiftDeveloper",
-        "Next.jsAgent": "/?model=Next.jsAgent",
-        "MongoDBAgent": "/?model=MongoDBAgent",
-        "PyTorchAgent": "/?model=PyTorchAgent",
-        "ReactAgent": "/?model=ReactAgent",
-        "XcodeAgent": "/?model=XcodeAgent",
-        "AngularJSAgent": "/?model=AngularJSAgent",
+        "claude-sonnet-3.5": "/?model=claude-sonnet-3.5"
     }
 
     model_aliases = {
@@ -226,8 +193,6 @@ class Blackbox:
     def clean_response(text: str) -> str:
         pattern = r'^\$\@\$v=undefined-rv1\$\@\$'
         cleaned_text = re.sub(pattern, '', text)
-        # Strip markdown syntax to prevent bold text
-        cleaned_text = strip_markdown(cleaned_text)
         return cleaned_text
 
     @classmethod
@@ -235,9 +200,12 @@ class Blackbox:
         cls,
         model: str,
         messages: List[Dict[str, str]],
+        temperature: Optional[float] = 1.0,
+        max_tokens: Optional[int] = 1024,
+        websearch: bool = False,
         proxy: Optional[str] = None,
         **kwargs
-    ) -> str:
+    ) -> Union[str, ImageResponseModel]:
         model = cls.get_model(model)
         chat_id = cls.generate_random_string()
         next_action = cls.generate_next_action()
@@ -301,9 +269,9 @@ class Blackbox:
             "trendingAgentMode": trending_agent_mode,
             "isMicMode": False,
             "userSystemPrompt": None,
-            "maxTokens": 1024,
+            "maxTokens": max_tokens,
             "playgroundTopP": 0.9,
-            "playgroundTemperature": 0.5,
+            "playgroundTemperature": temperature,
             "isChromeExt": False,
             "githubToken": None,
             "clickedAnswer2": False,
@@ -311,11 +279,11 @@ class Blackbox:
             "clickedForceWebSearch": False,
             "visitFromDelta": False,
             "mobileClient": False,
-            "webSearchMode": False,
+            "webSearchMode": websearch,
             "userSelectedModel": cls.userSelectedModel.get(model, model)
         }
 
-        async with ClientSession(headers=common_headers, timeout=ClientTimeout(total=60)) as session:
+        async with ClientSession(headers=common_headers) as session:
             try:
                 async with session.post(
                     cls.api_endpoint,
@@ -324,25 +292,43 @@ class Blackbox:
                     proxy=proxy
                 ) as response_api_chat:
                     response_api_chat.raise_for_status()
-                    # Instead of reading the entire response, iterate over chunks
-                    async for data in response_api_chat.content.iter_chunked(1024):
-                        decoded_data = data.decode('utf-8')
-                        cleaned_response = cls.clean_response(decoded_data)
-                        if model in cls.image_models:
-                            match = re.search(r'!\[.*?\]\((https?://[^\)]+)\)', cleaned_response)
+                    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 = ImageResponseModel(images=image_url, alt="Generated Image")
+                            return image_response
+                        else:
+                            return cleaned_response
+                    else:
+                        if websearch:
+                            match = re.search(r'\$~~~\$(.*?)\$~~~\$', cleaned_response, re.DOTALL)
                             if match:
-                                image_url = match.group(1)
-                                image_response = ImageResponseModel(images=image_url, alt="Generated Image")
-                                return image_response.dict()
+                                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:
-                                return cleaned_response
+                                final_response = cleaned_response
                         else:
                             if '$~~~$' in cleaned_response:
                                 final_response = cleaned_response.split('$~~~$')[0].strip()
                             else:
                                 final_response = cleaned_response
 
-                            return final_response
+                        return final_response
             except ClientResponseError as e:
                 error_text = f"Error {e.status}: {e.message}"
                 try:
@@ -355,158 +341,6 @@ class Blackbox:
             except Exception as e:
                 return f"Unexpected error during /api/chat request: {str(e)}"
 
-    @classmethod
-    async def create_async_generator(
-        cls,
-        model: str,
-        messages: List[Dict[str, str]],
-        proxy: Optional[str] = None,
-        websearch: bool = False,
-        **kwargs
-    ) -> AsyncGenerator[str, 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.
-            websearch (bool): Enables or disables web search mode.
-            **kwargs: Additional keyword arguments.
-
-        Yields:
-            str: Segments of the generated response.
-        """
-        model = cls.get_model(model)
-
-        chat_id = cls.generate_random_string()
-        next_action = cls.generate_next_action()
-        next_router_state_tree = cls.generate_next_router_state_tree()
-
-        agent_mode = cls.agentMode.get(model, {})
-        trending_agent_mode = cls.trendingAgentMode.get(model, {})
-
-        prefix = cls.model_prefixes.get(model, "")
-        
-        formatted_prompt = ""
-        for message in messages:
-            role = message.get('role', '').capitalize()
-            content = message.get('content', '')
-            if role and content:
-                formatted_prompt += f"{role}: {content}\n"
-        
-        if prefix:
-            formatted_prompt = f"{prefix} {formatted_prompt}".strip()
-
-        referer_path = cls.model_referers.get(model, f"/?model={model}")
-        referer_url = f"{cls.url}{referer_path}"
-
-        common_headers = {
-            'accept': '*/*',
-            'accept-language': 'en-US,en;q=0.9',
-            'cache-control': 'no-cache',
-            'origin': cls.url,
-            'pragma': 'no-cache',
-            'priority': 'u=1, i',
-            'sec-ch-ua': '"Chromium";v="129", "Not=A?Brand";v="8"',
-            'sec-ch-ua-mobile': '?0',
-            'sec-ch-ua-platform': '"Linux"',
-            'sec-fetch-dest': 'empty',
-            'sec-fetch-mode': 'cors',
-            'sec-fetch-site': 'same-origin',
-            'user-agent': 'Mozilla/5.0 (X11; Linux x86_64) '
-                          'AppleWebKit/537.36 (KHTML, like Gecko) '
-                          'Chrome/129.0.0.0 Safari/537.36'
-        }
-
-        headers_api_chat = {
-            'Content-Type': 'application/json',
-            'Referer': referer_url
-        }
-        headers_api_chat_combined = {**common_headers, **headers_api_chat}
-
-        payload_api_chat = {
-            "messages": [
-                {
-                    "id": chat_id,
-                    "content": formatted_prompt,
-                    "role": "user"
-                }
-            ],
-            "id": chat_id,
-            "previewToken": None,
-            "userId": None,
-            "codeModelMode": True,
-            "agentMode": agent_mode,
-            "trendingAgentMode": trending_agent_mode,
-            "isMicMode": False,
-            "userSystemPrompt": None,
-            "maxTokens": 1024,
-            "playgroundTopP": 0.9,
-            "playgroundTemperature": 0.5,
-            "isChromeExt": False,
-            "githubToken": None,
-            "clickedAnswer2": False,
-            "clickedAnswer3": False,
-            "clickedForceWebSearch": False,
-            "visitFromDelta": False,
-            "mobileClient": False,
-            "webSearchMode": websearch,
-            "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, timeout=ClientTimeout(total=60)) as session:
-            try:
-                async with session.post(
-                    cls.api_endpoint,
-                    headers=headers_api_chat_combined,
-                    json=payload_api_chat,
-                    proxy=proxy
-                ) as response_api_chat:
-                    response_api_chat.raise_for_status()
-                    # Iterate over the response in chunks
-                    async for data in response_api_chat.content.iter_any():
-                        decoded_data = data.decode('utf-8')
-                        cleaned_response = cls.clean_response(decoded_data)
-                        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")
-                                yield f"Image URL: {image_response.images}\n"
-                            else:
-                                yield cleaned_response
-                        else:
-                            if '$~~~$' in cleaned_response:
-                                final_response = cleaned_response.split('$~~~$')[0].strip()
-                            else:
-                                final_response = cleaned_response
-
-                            yield f"{final_response}\n"
-            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)}"
-
 # Custom exception for model not working
 class ModelNotWorkingException(Exception):
     def __init__(self, model: str):
@@ -599,12 +433,12 @@ class ChatRequest(BaseModel):
     temperature: Optional[float] = 1.0
     top_p: Optional[float] = 1.0
     n: Optional[int] = 1
-    max_tokens: Optional[int] = None
+    max_tokens: Optional[int] = 1024
     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
+    websearch: Optional[bool] = False  # Added websearch 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)):
@@ -612,7 +446,7 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
     # 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}")
+    logger.info(f"Received chat completions request from API key: {api_key} | IP: {client_ip} | Model: {request.model} | Messages: {redacted_messages} | Websearch: {request.websearch}")
 
     try:
         # Validate that the requested model is available
@@ -620,36 +454,41 @@ 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.")
 
-        if request.stream:
-            # Streaming response
-            async def event_generator():
-                async for chunk in Blackbox.create_async_generator(
-                    model=request.model,
-                    messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],
-                    proxy=None,  # Add proxy if needed
-                    websearch=False  # Modify if websearch is needed
-                ):
-                    if isinstance(chunk, ImageResponseModel):
-                        # Yield image URLs as plain text in SSE format
-                        yield f"data: Image URL: {chunk.images}\n\n"
-                    else:
-                        # Ensure chunk is a string and yield as plain text in SSE format
-                        if isinstance(chunk, str):
-                            yield f"data: {chunk}\n\n"
-                        else:
-                            yield f"data: {str(chunk)}\n\n"
-
-            logger.info(f"Initiating streaming response for API key: {api_key} | IP: {client_ip}")
-            return StreamingResponse(event_generator(), media_type='text/event-stream')
+        # Generate response using the updated Blackbox class
+        response_content = await Blackbox.generate_response(
+            model=request.model,
+            messages=[{"role": msg.role, "content": msg.content} for msg in request.messages],
+            temperature=request.temperature,
+            max_tokens=request.max_tokens,
+            websearch=request.websearch
+        )
+
+        # Handle image responses
+        if isinstance(response_content, ImageResponseModel):
+            logger.info(f"Generated image 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.images,
+                            "alt": response_content.alt
+                        },
+                        "finish_reason": "stop"
+                    }
+                ],
+                "usage": {
+                    "prompt_tokens": sum(len(msg.content.split()) for msg in request.messages),
+                    "completion_tokens": len(response_content.images.split()),
+                    "total_tokens": sum(len(msg.content.split()) for msg in request.messages) + len(response_content.images.split())
+                },
+            }
         else:
-            # Non-streaming response
-            response_content = await Blackbox.generate_response(
-                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"Completed response generation for API key: {api_key} | IP: {client_ip}")
             return {
                 "id": f"chatcmpl-{uuid.uuid4()}",
@@ -684,14 +523,14 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
 
 # Endpoint: GET /v1/models
 @app.get("/v1/models", dependencies=[Depends(rate_limiter_per_ip)])
-async def get_models(req: Request):
+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/health
 @app.get("/v1/health", dependencies=[Depends(rate_limiter_per_ip)])
-async def health_check(req: Request):
+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"}