fix(databricks): update Unity Gateway model mappings

litellm/llms/databricks/common_utils.py

@@ -10,13 +10,20 @@
 3. Databricks SDK automatic auth - Fallback (uses unified auth)
 """
 
+from __future__ import annotations
+
+import copy
 import os
 import re
-from typing import Final, Literal
+from typing import TYPE_CHECKING, Final, Literal
 from urllib.parse import urlsplit, urlunsplit
 
+from litellm.llms.base_llm.base_utils import BaseLLMModelInfo
 from litellm.llms.base_llm.chat.transformation import BaseLLMException
 
+if TYPE_CHECKING:
+    from litellm.types.llms.openai import AllMessageValues
+
 
 class DatabricksException(BaseLLMException):
     pass
@@ -398,3 +405,39 @@ def databricks_validate_environment(
             api_base = f"{api_base}/embeddings"
 
         return api_base, headers
+
+
+class DatabricksModelInfo(BaseLLMModelInfo):
+    def get_model_cost_key(self, model: str) -> str | None:
+        from litellm.llms.databricks.cost_calculator import get_databricks_cost_key
+
+        return get_databricks_cost_key(model)
+
+    def get_models(self, api_key: str | None = None, api_base: str | None = None) -> list[str]:
+        import litellm
+
+        return copy.deepcopy(litellm.models_by_provider.get("databricks", []))
+
+    @staticmethod
+    def get_api_key(api_key: str | None = None) -> str | None:
+        return api_key or os.getenv("DATABRICKS_API_KEY")
+
+    @staticmethod
+    def get_api_base(api_base: str | None = None) -> str | None:
+        return api_base or os.getenv("DATABRICKS_API_BASE")
+
+    def validate_environment(
+        self,
+        headers: dict,
+        model: str,
+        messages: list[AllMessageValues],
+        optional_params: dict,
+        litellm_params: dict,
+        api_key: str | None = None,
+        api_base: str | None = None,
+    ) -> dict:
+        return headers
+
+    @staticmethod
+    def get_base_model(model: str) -> str | None:
+        return model

litellm/llms/databricks/cost_calculator.py

@@ -9,6 +9,58 @@
 from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
 from litellm.types.utils import Usage
 
+_DATABRICKS_LEGACY_TO_CANONICAL: Final = MappingProxyType(
+    {
+        "databricks-claude-fable-5": "system.ai.claude-fable-5",
+        "databricks-claude-fable-5-1": "system.ai.claude-fable-5-1",
+        "databricks-claude-haiku-4-5": "system.ai.claude-haiku-4-5",
+        "databricks-claude-opus-4-6": "system.ai.claude-opus-4-6",
+        "databricks-claude-opus-4-7": "system.ai.claude-opus-4-7",
+        "databricks-claude-opus-4-8": "system.ai.claude-opus-4-8",
+        "databricks-claude-opus-5": "system.ai.claude-opus-5",
+        "databricks-claude-sonnet-4": "system.ai.claude-sonnet-4",
+        "databricks-claude-sonnet-4-5": "system.ai.claude-sonnet-4-5",
+        "databricks-claude-sonnet-4-6": "system.ai.claude-sonnet-4-6",
+        "databricks-claude-sonnet-5": "system.ai.claude-sonnet-5",
+        "databricks-deepseek-v4-flash-0731": "system.ai.deepseek-v4-flash-0731",
+        "databricks-deepseek-v4-pro-0813": "system.ai.deepseek-v4-pro-0813",
+        "databricks-gemini-2-5-flash": "system.ai.gemini-2-5-flash",
+        "databricks-gemini-2-5-pro": "system.ai.gemini-2-5-pro",
+        "databricks-gemini-3-1-flash-lite": "system.ai.gemini-3-1-flash-lite",
+        "databricks-gemini-3-1-pro": "system.ai.gemini-3-1-pro",
+        "databricks-gemini-3-5-flash": "system.ai.gemini-3-5-flash",
+        "databricks-gemini-3-6-flash": "system.ai.gemini-3-6-flash",
+        "databricks-gemini-3-flash": "system.ai.gemini-3-flash",
+        "databricks-gemma-3-12b": "system.ai.gemma-3-12b",
+        "databricks-glm-5-2": "system.ai.glm-5-2",
+        "databricks-glm-5-3": "system.ai.glm-5-3",
+        "databricks-glm-5-3-flash": "system.ai.glm-5-3-flash",
+        "databricks-gpt-5": "system.ai.gpt-5",
+        "databricks-gpt-5-1": "system.ai.gpt-5-1",
+        "databricks-gpt-5-2": "system.ai.gpt-5-2",
+        "databricks-gpt-5-4": "system.ai.gpt-5-4",
+        "databricks-gpt-5-4-mini": "system.ai.gpt-5-4-mini",
+        "databricks-gpt-5-4-nano": "system.ai.gpt-5-4-nano",
+        "databricks-gpt-5-5": "system.ai.gpt-5-5",
+        "databricks-gpt-5-6-luna": "system.ai.gpt-5-6-luna",
+        "databricks-gpt-5-6-sol": "system.ai.gpt-5-6-sol",
+        "databricks-gpt-5-mini": "system.ai.gpt-5-mini",
+        "databricks-gpt-5-nano": "system.ai.gpt-5-nano",
+        "databricks-gpt-oss-120b": "system.ai.gpt-oss-120b",
+        "databricks-gpt-oss-20b": "system.ai.gpt-oss-20b",
+        "databricks-grok-4-6": "system.ai.grok-4-6",
+        "databricks-gte-large-en": "system.ai.gte-large-en",
+        "databricks-inkling": "system.ai.inkling",
+        "databricks-kimi-k3": "system.ai.kimi-k3",
+        "databricks-llama-4-maverick": "system.ai.llama-4-maverick",
+        "databricks-meta-llama-3-1-8b-instruct": "system.ai.meta-llama-3-1-8b-instruct",
+        "databricks-meta-llama-3-3-70b-instruct": "system.ai.meta-llama-3-3-70b-instruct",
+        "databricks-qwen3-embedding-0-6b": "system.ai.qwen3-embedding-0-6b",
+        "databricks-qwen3-next-80b-a3b-instruct": "system.ai.qwen3-next-80b-a3b-instruct",
+        "databricks-qwen35-122b-a10b": "system.ai.qwen35-122b-a10b",
+    }
+)
+
 _LEGACY_ENDPOINT_NAMES: Final = MappingProxyType(
     {
         "dbrx-instruct": "databricks-dbrx-instruct",
@@ -22,12 +74,30 @@
 )
 
 
+def get_databricks_cost_key(model: str) -> str | None:
+    name: Final = model.removeprefix("databricks/")
+    canonical_slug: Final = _DATABRICKS_LEGACY_TO_CANONICAL.get(name)
+    if canonical_slug is not None:
+        return f"databricks/{canonical_slug}"
+    legacy_endpoint: Final = _LEGACY_ENDPOINT_NAMES.get(name)
+    if legacy_endpoint is not None:
+        resolved_canonical: Final = _DATABRICKS_LEGACY_TO_CANONICAL.get(legacy_endpoint)
+        if resolved_canonical is not None:
+            return f"databricks/{resolved_canonical}"
+        return f"databricks/{legacy_endpoint}"
+    return None
+
+
 def _registry_key(model: str) -> str:
     name: Final = model.removeprefix("databricks/")
-    return next(
+    canonical_slug: Final = _DATABRICKS_LEGACY_TO_CANONICAL.get(name)
+    if canonical_slug is not None:
+        return canonical_slug
+    endpoint: Final = next(
         (key for prefix, key in _LEGACY_ENDPOINT_NAMES.items() if name.startswith(prefix)),
         name,
     )
+    return _DATABRICKS_LEGACY_TO_CANONICAL.get(endpoint, endpoint)
 
 
 def cost_per_token(model: str, usage: Usage) -> tuple[float, float]: