Update main.py

main.py

@@ -11,7 +11,6 @@ import asyncio
 import time
 from collections import defaultdict
 from typing import List, Dict, Any, Optional, Union, AsyncGenerator
-from datetime import datetime  # Essential for timestamping
 
 from aiohttp import ClientSession, ClientResponseError
 from fastapi import FastAPI, HTTPException, Request, Depends, Header
@@ -41,18 +40,13 @@ rate_limit_store = defaultdict(lambda: {"count": 0, "timestamp": time.time()})
 CLEANUP_INTERVAL = 60  # seconds
 RATE_LIMIT_WINDOW = 60  # seconds
 
-# Define ImageResponse and base classes if not defined elsewhere
+# Define the ImageResponse model
 class ImageResponse(BaseModel):
     images: str
     alt: str
 
-class AsyncGeneratorProvider:
-    pass
-
-class ProviderModelMixin:
-    pass
-
-class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
+# Updated Blackbox Provider
+class Blackbox:
     label = "Blackbox AI"
     url = "https://www.blackbox.ai"
     api_endpoint = "https://www.blackbox.ai/api/chat"
@@ -147,7 +141,7 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
         "blackboxai": "/?model=blackboxai",
         "gpt-4o": "/?model=gpt-4o",
         "gemini-pro": "/?model=gemini-pro",
-        "claude-sonnet-3.5": "/?model=claude-sonnet-3.5",
+        "claude-sonnet-3.5": "/?model=claude-sonnet-3.5"
     }
 
     model_aliases = {
@@ -201,115 +195,6 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
         cleaned_text = re.sub(pattern, '', text)
         return cleaned_text
 
-    @classmethod
-    async def generate_response(
-        cls,
-        model: str,
-        messages: List[Dict[str, str]],
-        proxy: Optional[str] = None,
-        **kwargs
-    ) -> str:
-        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": False,
-            "userSelectedModel": cls.userSelectedModel.get(model, model)
-        }
-
-        async with ClientSession(headers=common_headers) 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()
-                    text = await response_api_chat.text()
-                    cleaned_response = cls.clean_response(text)
-                    return cleaned_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
-                return error_text
-            except Exception as e:
-                return f"Unexpected error during /api/chat request: {str(e)}"
-
     @classmethod
     async def create_async_generator(
         cls,
@@ -332,7 +217,6 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
         Yields:
             Union[str, ImageResponse]: Segments of the generated response or ImageResponse objects.
         """
-        logger.debug("Starting async generator for model: %s", model)
         model = cls.get_model(model)
 
         chat_id = cls.generate_random_string()
@@ -425,7 +309,6 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
 
         async with ClientSession(headers=common_headers) as session:
             try:
-                logger.debug("Sending POST request to Blackbox API at %s", cls.api_endpoint)
                 async with session.post(
                     cls.api_endpoint,
                     headers=headers_api_chat_combined,
@@ -435,19 +318,15 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
                     response_api_chat.raise_for_status()
                     text = await response_api_chat.text()
                     cleaned_response = cls.clean_response(text)
-                    logger.debug("Received response from Blackbox API: %s", cleaned_response)
-
-                    # Test yield to verify streaming works
-                    yield "Streaming response started...\n"
 
                     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.json() + "\n"
+                            yield image_response
                         else:
-                            yield cleaned_response + "\n"
+                            yield cleaned_response
                     else:
                         if websearch:
                             match = re.search(r'\$~~~\$(.*?)\$~~~\$', cleaned_response, re.DOTALL)
@@ -473,29 +352,22 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
                             else:
                                 final_response = cleaned_response
 
-                        yield final_response + "\n"
-
+                        yield final_response
             except ClientResponseError as e:
                 error_text = f"Error {e.status}: {e.message}"
-                logger.error("ClientResponseError: %s", error_text)
                 try:
                     error_response = await e.response.text()
                     cleaned_error = cls.clean_response(error_response)
                     error_text += f" - {cleaned_error}"
                 except Exception:
                     pass
-                yield f"{error_text}\n"
+                yield error_text
             except Exception as e:
-                error_text = f"Unexpected error during /api/chat request: {str(e)}"
-                logger.error("Exception: %s", error_text)
-                yield f"{error_text}\n"
-
-            # Test yield after API call
-            yield "Streaming response ended.\n"
+                yield f"Unexpected error during /api/chat request: {str(e)}"
 
             chat_url = f'{cls.url}/chat/{chat_id}?model={model}'
+
             try:
-                logger.debug("Sending POST request to Chat URL: %s", chat_url)
                 async with session.post(
                     chat_url,
                     headers=headers_chat_combined,
@@ -503,15 +375,18 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
                     proxy=proxy
                 ) as response_chat:
                     response_chat.raise_for_status()
-                    logger.debug("Chat POST request successful.")
+                    pass
             except ClientResponseError as e:
                 error_text = f"Error {e.status}: {e.message}"
-                logger.error("ClientResponseError on chat POST: %s", error_text)
-                yield f"{error_text}\n"
+                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:
-                error_text = f"Unexpected error during /chat/{chat_id} request: {str(e)}"
-                logger.error("Exception on chat POST: %s", error_text)
-                yield f"{error_text}\n"
+                yield f"Unexpected error during /chat/{chat_id} request: {str(e)}"
 
 # Custom exception for model not working
 class ModelNotWorkingException(Exception):
@@ -612,7 +487,11 @@ class ChatRequest(BaseModel):
     user: Optional[str] = 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)):
+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]
@@ -625,16 +504,42 @@ 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.")
 
-        # Create an asynchronous generator for streaming response
-        generator = Blackbox.create_async_generator(
+        # 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],
-            proxy=req.headers.get('Proxy'),  # Corrected usage
-            websearch=req.query_params.get('websearch', 'false').lower() == 'true'
+            temperature=request.temperature,
+            max_tokens=request.max_tokens
         )
 
         logger.info(f"Started streaming response for API key: {api_key} | IP: {client_ip}")
-        return StreamingResponse(generator, media_type="text/event-stream")
+
+        # Define a generator function to yield the streamed data in the OpenAI format
+        async def stream_response():
+            async for chunk in async_generator:
+                if isinstance(chunk, ImageResponse):
+                    # Handle image responses if necessary
+                    # For simplicity, converting ImageResponse to a string representation
+                    # You might want to handle it differently based on your requirements
+                    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"
+
+        return StreamingResponse(stream_response(), media_type="text/event-stream")
 
     except ModelNotWorkingException as e:
         logger.warning(f"Model not working: {e} | IP: {client_ip}")