providers/cerebras/client.py

"""Cerebras provider using OpenAI-compatible API."""

from typing import Any

from config.settings import Settings
from core.anthropic import ReasoningReplayMode, build_base_request_body
from providers.base import ProviderConfig
from providers.defaults import CEREBRAS_DEFAULT_BASE
from providers.openai_compat import OpenAIChatTransport


class CerebrasProvider(OpenAIChatTransport):
    """Cerebras provider using OpenAI-compatible /chat/completions."""

    # Mapping of proxy model refs to Cerebras API model IDs.
    CEREBRAS_MODEL_MAP: dict[str, str] = {
        "cerebras/llama3.1-8b": "llama3.1-8b",
        "cerebras/qwen-3-235b-a22b-instruct-2507": "qwen-3-235b-a22b-instruct-2507",
    }

    def __init__(self, config: ProviderConfig, *, settings: Settings):
        base_url = (config.base_url or CEREBRAS_DEFAULT_BASE).rstrip("/")
        if not base_url.endswith("/v1"):
            base_url = base_url + "/v1"
        # Cerebras has generous rate limits
        super().__init__(
            config,
            provider_name="Cerebras",
            base_url=base_url,
            api_key=config.api_key,
            nim_rate_limit=300,
            nim_max_concurrency=80,
        )
        self._settings = settings

    def _build_request_body(
        self, request: Any, thinking_enabled: bool | None = None
    ) -> dict:
        thinking = self._is_thinking_enabled(request, thinking_enabled)
        reasoning_replay = (
            ReasoningReplayMode.REASONING_CONTENT
            if thinking
            else ReasoningReplayMode.DISABLED
        )
        body = build_base_request_body(request, reasoning_replay=reasoning_replay)
        model = body.get("model", "")
        if model in self.CEREBRAS_MODEL_MAP:
            body["model"] = self.CEREBRAS_MODEL_MAP[model]
        elif model.startswith("cerebras/"):
            body["model"] = model[len("cerebras/") :]
        return body