api/provider/blackboxai.py

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', '')