backend/llm/providers/databricks.py

from __future__ import annotations

import os

from .openai_compatible import OpenAICompatibleProvider


class DatabricksProvider(OpenAICompatibleProvider):
    """
    Databricks Model Serving provider through the
    OpenAI-compatible MLflow gateway.
    """

    name = "databricks"

    api_base = os.getenv(
        "DATABRICKS_BASE_URL",
        "https://dbc-76aba418-4d3d.cloud.databricks.com/ai-gateway/mlflow/v1",
    )

    api_key_env = "DATABRICKS_TOKEN"

    default_model = os.getenv(
        "DATABRICKS_MODEL",
        "databricks-meta-llama-3-3-70b-instruct",
    )