Create blackboxai.py

api/provider/blackboxai.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+import json
+from aiohttp import ClientSession
+
+from ..typing import AsyncResult, Messages
+from ..image import ImageResponse
+from .base_provider import AsyncGeneratorProvider, ProviderModelMixin
+from .helper import format_prompt
+
+class BlackBoxAI(AsyncGeneratorProvider, ProviderModelMixin):
+    url = "https://api.blackboxai.com/assistant/"
+    api_endpoint = "https://api.blackboxai.com/api/data/users/inferenceServer.infer"
+    working = True
+
+    supports_system_message = True
+    supports_message_history = True
+
+    # Chat models
+    default_model = 'chat-blackboxai-basic'
+    chat_models = [
+        default_model,
+        # Add other chat models as needed
+    ]
+
+    # Image models
+    image_models = [
+        'bbai-sdxl',
+        # Add other image models as needed
+    ]
+
+    models = [*chat_models, *image_models]
+
+    model_aliases = {
+        # Define aliases as needed
+    }
+
+    @classmethod
+    def get_model(cls, model: str) -> str:
+        if model in cls.models:
+            return model
+        elif model in cls.model_aliases:
+            return cls.model_aliases[model]
+        else:
+            return cls.default_model
+
+    @classmethod
+    def is_image_model(cls, model: str) -> bool:
+        return model in cls.image_models
+
+    @classmethod
+    async def create_async_generator(
+        cls,
+        model: str,
+        messages: Messages,
+        proxy: str = None,
+        **kwargs
+    ) -> AsyncResult:
+        model = cls.get_model(model)
+
+        headers = {
+            'Accept': 'application/json, text/plain, */*',
+            'Accept-Language': 'en-US,en;q=0.9',
+            'Cache-Control': 'no-cache',
+            'Connection': 'keep-alive',
+            'Content-Type': 'application/json',
+            'Origin': 'https://api.blackboxai.com',
+            'Pragma': 'no-cache',
+            '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/130.0.0.0 Safari/537.36',
+            'sec-ch-ua': '"Not?A_Brand";v="99", "Chromium";v="130"',
+            'sec-ch-ua-mobile': '?0',
+            'sec-ch-ua-platform': '"Linux"'
+        }
+
+        async with ClientSession() as session:
+            if cls.is_image_model(model):
+                # Image generation
+                prompt = messages[-1]["content"]
+                data = {
+                    "model": model,
+                    "input": {
+                        "width": "1024",
+                        "height": "1024",
+                        "steps": 4,
+                        "output_format": "webp",
+                        "batch_size": 1,
+                        "mode": "plan",
+                        "prompt": prompt
+                    }
+                }
+                async with session.post(
+                    cls.api_endpoint,
+                    headers=headers,
+                    data=json.dumps(data),
+                    proxy=proxy
+                ) as response:
+                    response.raise_for_status()
+                    response_data = await response.json()
+                    if response_data.get('status') == 'completed' and response_data.get('output'):
+                        for url in response_data['output']:
+                            yield ImageResponse(images=url, alt="Generated Image")
+            else:
+                # Chat completion
+                data = {
+                    "model": model,
+                    "input": {
+                        "messages": [
+                            {
+                                "type": "human",
+                                "content": format_prompt(messages)
+                            }
+                        ],
+                        "mode": "plan"
+                    },
+                    "noStream": True
+                }
+                async with session.post(
+                    cls.api_endpoint,
+                    headers=headers,
+                    data=json.dumps(data),
+                    proxy=proxy
+                ) as response:
+                    response.raise_for_status()
+                    result = await response.json()
+                    yield result.get('output', '')