Update main.py

main.py

@@ -11,12 +11,12 @@ 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
 from fastapi.responses import JSONResponse, StreamingResponse
 from pydantic import BaseModel
-from datetime import datetime  # Fix for 'datetime' not defined
 
 # Configure logging
 logging.basicConfig(
@@ -148,20 +148,6 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
         "gpt-4o": "/?model=gpt-4o",
         "gemini-pro": "/?model=gemini-pro",
         "claude-sonnet-3.5": "/?model=claude-sonnet-3.5",
-        "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",
-        "ImageGeneration": "/?model=ImageGeneration",
     }
 
     model_aliases = {
@@ -346,6 +332,7 @@ 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()
@@ -438,6 +425,7 @@ 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,
@@ -447,15 +435,19 @@ 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
+                            yield image_response.json() + "\n"
                         else:
-                            yield cleaned_response
+                            yield cleaned_response + "\n"
                     else:
                         if websearch:
                             match = re.search(r'\$~~~\$(.*?)\$~~~\$', cleaned_response, re.DOTALL)
@@ -481,22 +473,29 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
                             else:
                                 final_response = cleaned_response
 
-                        yield final_response
+                        yield final_response + "\n"
+
             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 error_text
+                yield f"{error_text}\n"
             except Exception as e:
-                yield f"Unexpected error during /api/chat request: {str(e)}"
+                error_text = f"Unexpected error during /api/chat request: {str(e)}"
+                logger.error("Exception: %s", error_text)
+                yield f"{error_text}\n"
 
-            chat_url = f'{cls.url}/chat/{chat_id}?model={model}'
+            # Test yield after API call
+            yield "Streaming response ended.\n"
 
+            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,
@@ -504,18 +503,15 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
                     proxy=proxy
                 ) as response_chat:
                     response_chat.raise_for_status()
-                    pass
+                    logger.debug("Chat POST request successful.")
             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
+                logger.error("ClientResponseError on chat POST: %s", error_text)
+                yield f"{error_text}\n"
             except Exception as e:
-                yield f"Unexpected error during /chat/{chat_id} request: {str(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"
 
 # Custom exception for model not working
 class ModelNotWorkingException(Exception):
@@ -633,8 +629,8 @@ async def chat_completions(request: ChatRequest, req: Request, api_key: str = De
         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
+            proxy=request.headers.get('Proxy'),  # Assuming proxy info is passed via headers
+            websearch=request.query_params.get('websearch', 'false').lower() == 'true'
         )
 
         logger.info(f"Started streaming response for API key: {api_key} | IP: {client_ip}")