diff options
Diffstat (limited to 'g4f/Provider/Blackbox.py')
| -rw-r--r-- | g4f/Provider/Blackbox.py | 37 |
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 |