0.3.2.2

1 files changed, 25 insertions, 12 deletions

diff --git a/g4f/Provider/Blackbox.py b/g4f/Provider/Blackbox.py
index 6e1e3949..a86471f2 100644
--- a/g4f/Provider/Blackbox.py
+++ b/g4f/Provider/Blackbox.py
@@ -2,31 +2,35 @@ from __future__ import annotations
 
 import uuid
 import secrets
-from aiohttp import ClientSession
+import re
+from aiohttp import ClientSession, ClientResponse
+from typing import AsyncGenerator, Optional
 
 from ..typing import AsyncResult, Messages, ImageType
 from ..image import to_data_uri
-from .base_provider import AsyncGeneratorProvider
+from .base_provider import AsyncGeneratorProvider, ProviderModelMixin
 
-class Blackbox(AsyncGeneratorProvider):
+class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
     url = "https://www.blackbox.ai"
     working = True
+    default_model = 'blackbox'
 
     @classmethod
     async def create_async_generator(
         cls,
         model: str,
         messages: Messages,
-        proxy: str = None,
-        image: ImageType = None,
-        image_name: str = None,
+        proxy: Optional[str] = None,
+        image: Optional[ImageType] = None,
+        image_name: Optional[str] = None,
         **kwargs
-    ) -> AsyncResult:
+    ) -> AsyncGenerator[str, None]:
         if image is not None:
             messages[-1]["data"] = {
-                "fileText":	image_name,
+                "fileText": image_name,
                 "imageBase64": to_data_uri(image)
             }
+
         headers = {
             "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/122.0.0.0 Safari/537.36",
             "Accept": "*/*",
@@ -40,9 +44,11 @@ class Blackbox(AsyncGeneratorProvider):
             "Alt-Used": "www.blackbox.ai",
             "Connection": "keep-alive",
         }
+
         async with ClientSession(headers=headers) as session:
             random_id = secrets.token_hex(16)
             random_user_id = str(uuid.uuid4())
+
             data = {
                 "messages": messages,
                 "id": random_id,
@@ -55,10 +61,17 @@ class Blackbox(AsyncGeneratorProvider):
                 "playgroundMode": False,
                 "webSearchMode": False,
                 "userSystemPrompt": "",
-                "githubToken": None
+                "githubToken": None,
+                "maxTokens": None
             }
-            async with session.post(f"{cls.url}/api/chat", json=data, proxy=proxy) as response:
+
+            async with session.post(
+                f"{cls.url}/api/chat", json=data, proxy=proxy
+            ) as response:  # type: ClientResponse
                 response.raise_for_status()
-                async for chunk in response.content:
+                async for chunk in response.content.iter_any():
                     if chunk:
-                        yield chunk.decode()
+                        # Decode the chunk and clean up unwanted prefixes using a regex
+                        decoded_chunk = chunk.decode()
+                        cleaned_chunk = re.sub(r'\$@\$.+?\$@\$|\$@\$', '', decoded_chunk)
+                        yield cleaned_chunk