Class OpenAIAPI

datastudio/models/openai_api.py:22–633

Async OpenAI-compatible API wrapper. Uses a shared :class:`aiohttp.ClientSession` with connection pooling and ``asyncio.Semaphore`` for concurrency control. Each request handles retries independently with exponential backoff.

@MODELS.register_module()
class OpenAIAPI(BaseAPI):
    """Async OpenAI-compatible API wrapper.

    Uses a shared :class:`aiohttp.ClientSession` with connection pooling
    and ``asyncio.Semaphore`` for concurrency control.  Each request handles
    retries independently with exponential backoff.
    """

    def __init__(
        self,
        model: str,
        key: str = None,
        temperature: float = 0.6,
        top_p: float = 0.95,
        top_k: int = 20,
        min_p: float = 0.0,
        presence_penalty: float = 1.0,
        repetition_penalty: float = 1.0,
        api_base: str = None,
        port: int = None,
        retry: int = 10,
        wait: int = 3,
        timeout: tuple = (30, 1800),
        max_tokens: int = 16384,
        thread_num: int = 8192,
        return_dict: bool = False,
        logger=None,
        max_connections: int = None,
        enable_thinking: bool = False,
        **kwargs,
    ):
        """Initialize OpenAI-compatible API wrapper.

        Args:
            model: Model name, e.g., gpt-4, Qwen2-VL-72B.
            key: API key (uses OPENAI_API_KEY env var if not provided).
            temperature: Generation temperature.
            top_p: Nucleus sampling threshold (0~1).
            top_k: Top-K sampling (0 to disable). vLLM/SGLang extra param.
            min_p: Minimum probability threshold. vLLM/SGLang extra param.
            presence_penalty: Penalize tokens already present in the output.
            repetition_penalty: Penalize repeated tokens. vLLM/SGLang extra param.
            api_base: API base URL (full URL to chat/completions endpoint).
            port: Port number for local deployments.
            retry: Number of retry attempts on API failure.
            wait: Max wait time between retries (seconds).
            timeout: Request timeout as (connect_timeout, read_timeout) tuple.
            max_tokens: Maximum tokens in response.
            thread_num: Maximum number of concurrent requests (semaphore limit).
            return_dict: Whether to parse response as dict via load_str_to_dict.
            logger: Logger instance.
            max_connections: Max TCP connections in the pool. Defaults to
                min(thread_num, 16384).
            enable_thinking: Whether to enable thinking mode for models that support it
                (e.g., Qwen3 on vLLM/SGLang). Default False.
            **kwargs: Additional API parameters passed to payload.
        """
        super().__init__(