1 files changed, 30 insertions, 35 deletions

diff --git a/g4f/Provider/Blackbox.py b/g4f/Provider/Blackbox.py
index 75abb183..97466c04 100644
--- a/g4f/Provider/Blackbox.py
+++ b/g4f/Provider/Blackbox.py
@@ -10,6 +10,7 @@ import aiohttp
 from ..typing import AsyncResult, Messages, ImageType
 from .base_provider import AsyncGeneratorProvider, ProviderModelMixin
 from ..image import ImageResponse, to_data_uri
+from .helper import get_random_string
 
 class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
     label = "Blackbox AI"
@@ -22,11 +23,13 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
     _last_validated_value = None
 
     default_model = 'blackboxai'
+    default_vision_model = default_model
     default_image_model = 'generate_image'  
     image_models = [default_image_model, 'repomap']
     text_models = [default_model, 'gpt-4o', 'gemini-pro', 'claude-sonnet-3.5', 'blackboxai-pro']
+    vision_models = [default_model, 'gpt-4o', 'gemini-pro', 'blackboxai-pro']
     agentMode = {
-        'Image Generation': {'mode': True, 'id': "ImageGenerationLV45LJp", 'name': "Image Generation"},
+        default_image_model: {'mode': True, 'id': "ImageGenerationLV45LJp", 'name': "Image Generation"},
     }
     trendingAgentMode = {
         "gemini-1.5-flash": {'mode': True, 'id': 'Gemini'},
@@ -111,11 +114,6 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
 
         return cls._last_validated_value
 
-    @staticmethod
-    def generate_id(length=7):
-        characters = string.ascii_letters + string.digits
-        return ''.join(random.choice(characters) for _ in range(length))
-
     @classmethod
     def add_prefix_to_messages(cls, messages: Messages, model: str) -> Messages:
         prefix = cls.model_prefixes.get(model, "")
@@ -143,12 +141,12 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
         **kwargs
     ) -> AsyncResult:
         model = cls.get_model(model)
-        message_id = cls.generate_id()
-        messages_with_prefix = cls.add_prefix_to_messages(messages, model)
+        message_id = get_random_string(7)
+        messages = cls.add_prefix_to_messages(messages, model)
         validated_value = await cls.fetch_validated()
 
         if image is not None:
-            messages_with_prefix[-1]['data'] = {
+            messages[-1]['data'] = {
                 'fileText': '',
                 'imageBase64': to_data_uri(image),
                 'title': image_name
@@ -171,9 +169,9 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
             'sec-fetch-site': 'same-origin',
             'user-agent': 'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/130.0.0.0 Safari/537.36'
         }
-        
+
         data = {
-            "messages": messages_with_prefix,
+            "messages": messages,
             "id": message_id,
             "previewToken": None,
             "userId": None,
@@ -200,27 +198,24 @@ class Blackbox(AsyncGeneratorProvider, ProviderModelMixin):
         async with ClientSession(headers=headers) as session:
             async with session.post(cls.api_endpoint, json=data, proxy=proxy) as response:
                 response.raise_for_status()
-                response_text = await response.text()
-                
-                if model in cls.image_models:
-                    image_matches = re.findall(r'!\[.*?\]\((https?://[^\)]+)\)', response_text)
-                    if image_matches:
-                        image_url = image_matches[0]
-                        image_response = ImageResponse(images=[image_url], alt="Generated Image")
-                        yield image_response
-                        return
-
-                response_text = re.sub(r'Generated by BLACKBOX.AI, try unlimited chat https://www.blackbox.ai', '', response_text, flags=re.DOTALL)
-                
-                json_match = re.search(r'\$~~~\$(.*?)\$~~~\$', response_text, re.DOTALL)
-                if json_match:
-                    search_results = json.loads(json_match.group(1))
-                    answer = response_text.split('$~~~$')[-1].strip()
-                    
-                    formatted_response = f"{answer}\n\n**Source:**"
-                    for i, result in enumerate(search_results, 1):
-                        formatted_response += f"\n{i}. {result['title']}: {result['link']}"
-                    
-                    yield formatted_response
-                else:
-                    yield response_text.strip()
+                async for chunk in response.content.iter_any():
+                    text_chunk = chunk.decode(errors="ignore")
+                    if model in cls.image_models:
+                        image_matches = re.findall(r'!\[.*?\]\((https?://[^\)]+)\)', text_chunk)
+                        if image_matches:
+                            image_url = image_matches[0]
+                            image_response = ImageResponse(images=[image_url])
+                            yield image_response
+                            continue
+
+                    text_chunk = re.sub(r'Generated by BLACKBOX.AI, try unlimited chat https://www.blackbox.ai', '', text_chunk, flags=re.DOTALL)
+                    json_match = re.search(r'\$~~~\$(.*?)\$~~~\$', text_chunk, re.DOTALL)
+                    if json_match:
+                        search_results = json.loads(json_match.group(1))
+                        answer = text_chunk.split('$~~~$')[-1].strip()
+                        formatted_response = f"{answer}\n\n**Source:**"
+                        for i, result in enumerate(search_results, 1):
+                            formatted_response += f"\n{i}. {result['title']}: {result['link']}"
+                        yield formatted_response
+                    else:
+                        yield text_chunk.strip()