Merge 3a32e92126 into a3236ba947

litellm/llms/custom_httpx/llm_http_handler.py

@@ -1604,6 +1604,7 @@ class BaseLLMHTTPHandler:
             model=model,
             messages=[],
             optional_params=optional_params,
+            api_base=api_base,
             litellm_params=litellm_params,
         )
 

litellm/llms/databricks/audio_transcription/__init__.py

@@ -0,0 +1 @@
+"""Databricks audio transcription implementation."""

litellm/llms/databricks/audio_transcription/transformation.py

@@ -0,0 +1,186 @@
+"""
+Translates from OpenAI's `/v1/audio/transcriptions` to Databricks Model
+Serving's MLflow invoke route `/serving-endpoints/{endpoint}/invocations`.
+
+The api_base (parameter or DATABRICKS_API_BASE) is accepted either as a
+bare workspace host or as a base that already carries `/serving-endpoints`
+(the litellm docs convention); a bare host is normalized by appending the
+path. The OpenAI-compatible audio routes require a token with `all-apis`
+scope; the invoke route accepts any workspace token, so it is the default
+target.
+"""
+
+import base64
+import os
+import re
+from collections.abc import Mapping, Sequence
+from typing import Final
+
+from httpx import Headers, Response
+from pydantic import TypeAdapter, ValidationError
+
+from litellm.litellm_core_utils.audio_utils.utils import process_audio_file
+from litellm.llms.base_llm.chat.transformation import BaseLLMException
+from litellm.types.llms.openai import (
+    AllMessageValues,
+    OpenAIAudioTranscriptionOptionalParams,
+)
+from litellm.types.utils import FileTypes, TranscriptionResponse
+
+from ...base_llm.audio_transcription.transformation import (
+    AudioTranscriptionRequestData,
+    BaseAudioTranscriptionConfig,
+)
+from ..common_utils import DatabricksBase, DatabricksException
+
+_OBJECT_LIST: Final = TypeAdapter(list[object])
+_STRING_OBJECT_DICT: Final = TypeAdapter(dict[str, object])
+# dots are excluded so path-shaping inputs such as ".." cannot reach the URL
+_ENDPOINT_NAME: Final = re.compile(r"[A-Za-z0-9_-]+")
+
+
+class DatabricksAudioTranscriptionConfig(BaseAudioTranscriptionConfig, DatabricksBase):
+    def get_supported_openai_params(
+        self, model: str
+    ) -> list[OpenAIAudioTranscriptionOptionalParams]:  # mutable-ok: base class signature returns list
+        # accept-and-ignore: advertising these avoids UnsupportedParamsError; the invoke route has no parameter surface
+        return [  # mutable-ok: fixed provider param set, never grown
+            "language",
+            "prompt",
+            "response_format",
+            "temperature",
+            "timestamp_granularities",
+        ]
+
+    def map_openai_params(
+        self,
+        non_default_params: Mapping[str, object],
+        optional_params: Mapping[str, object],
+        model: str,
+        drop_params: bool,
+    ) -> dict[str, object]:  # mutable-ok: base class signature returns dict
+        # accept-and-ignore, see get_supported_openai_params
+        return dict(optional_params)  # mutable-ok: base class signature returns dict
+
+    def get_error_class(
+        self,
+        error_message: str,
+        status_code: int,
+        headers: dict[str, object] | Headers,  # mutable-ok: base class signature takes dict
+    ) -> BaseLLMException:
+        return DatabricksException(message=error_message, status_code=status_code, headers=headers)
+
+    def validate_environment(
+        self,
+        headers: dict[str, object],  # mutable-ok: base class signature takes and returns dict
+        model: str,
+        messages: Sequence[AllMessageValues],
+        optional_params: Mapping[str, object],
+        litellm_params: Mapping[str, object],
+        api_key: str | None = None,
+        api_base: str | None = None,
+    ) -> dict[str, object]:  # mutable-ok: base class signature returns dict
+        # endpoint_type="chat_completions" reuses the shared Databricks auth; the
+        # /chat/completions suffix it appends is discarded (get_complete_url builds the URL)
+        auth_headers: Final = _STRING_OBJECT_DICT.validate_python(
+            self.databricks_validate_environment(  # pyright: ignore[reportUnknownMemberType]  # common_utils auth predates strict typing; re-annotating it is out of scope here
+                api_key=api_key,
+                api_base=api_base,
+                endpoint_type="chat_completions",
+                custom_endpoint=False,
+                headers=headers,
+            )[1]
+        )
+        auth_headers["Content-Type"] = "application/json"
+        return auth_headers
+
+    def get_complete_url(
+        self,
+        api_base: str | None,
+        api_key: str | None,
+        model: str,
+        optional_params: Mapping[str, object],
+        litellm_params: Mapping[str, object],
+        stream: bool | None = None,
+    ) -> str:
+        endpoint: Final = model.removeprefix("databricks/")
+        if _ENDPOINT_NAME.fullmatch(endpoint) is None:
+            raise DatabricksException(
+                status_code=400,
+                message=(
+                    f"Invalid Databricks endpoint name {endpoint!r}: names may only contain "
+                    "letters, digits, underscores, and hyphens."
+                ),
+            )
+        resolved_base: Final = self._get_api_base(api_base or os.getenv("DATABRICKS_API_BASE")).rstrip("/")
+        # _get_api_base appends /serving-endpoints only in its SDK-fallback branch
+        serving_base: Final = (
+            resolved_base if resolved_base.endswith("/serving-endpoints") else f"{resolved_base}/serving-endpoints"
+        )
+        return f"{serving_base}/{endpoint}/invocations"
+
+    def transform_audio_transcription_request(
+        self,
+        model: str,
+        audio_file: FileTypes,
+        optional_params: Mapping[str, object],
+        litellm_params: Mapping[str, object],
+    ) -> AudioTranscriptionRequestData:
+        processed_audio: Final = process_audio_file(audio_file)
+        encoded_audio: Final = base64.b64encode(processed_audio.file_content).decode("ascii")
+        return AudioTranscriptionRequestData(
+            data={"inputs": [encoded_audio]},  # mutable-ok: AudioTranscriptionRequestData.data is a dict payload
+            files=None,
+            content_type="application/json",
+        )
+
+    def transform_audio_transcription_response(
+        self,
+        raw_response: Response,
+    ) -> TranscriptionResponse:
+        if raw_response.status_code >= 400:
+            raise self.get_error_class(
+                error_message=raw_response.text,
+                status_code=raw_response.status_code,
+                headers=raw_response.headers,
+            )
+        try:
+            payload: Final = _STRING_OBJECT_DICT.validate_python(raw_response.json())
+        except ValueError as e:
+            raise DatabricksException(
+                status_code=500,
+                message=f"Error parsing Databricks transcription response: {e}\nResponse: {raw_response.text}",
+            )
+        text: Final = self._extract_prediction(payload)
+        if text is None:  # "" is a valid transcript (e.g. silent audio); only a missing prediction is an error
+            raise DatabricksException(
+                status_code=500,
+                message=(
+                    f"Databricks transcription response carried no prediction text. Response: {raw_response.text}"
+                ),
+            )
+        response: Final = TranscriptionResponse(text=text)
+        response["task"] = "transcribe"
+        response._hidden_params = payload  # pyright: ignore[reportPrivateUsage]  # TranscriptionResponse exposes no public hidden-params setter
+        return response
+
+    @staticmethod
+    def _extract_prediction(payload: Mapping[str, object]) -> str | None:
+        predictions: Final[object] = payload.get("predictions")
+        if isinstance(predictions, str):
+            return predictions
+        try:
+            prediction_items: Final = _OBJECT_LIST.validate_python(predictions)
+        except ValidationError:
+            return None
+        if not prediction_items:
+            return None
+        first: Final[object] = prediction_items[0]
+        if isinstance(first, str):
+            return first
+        try:
+            first_dict: Final = _STRING_OBJECT_DICT.validate_python(first)
+        except ValidationError:
+            return None
+        text: Final[object] = first_dict.get("text")
+        return text if isinstance(text, str) else None

litellm/utils.py

@@ -9093,6 +9093,12 @@ class ProviderConfigManager:
             )
 
             return AzureSpeechAudioTranscriptionConfig()
+        elif litellm.LlmProviders.DATABRICKS == provider:
+            from litellm.llms.databricks.audio_transcription.transformation import (
+                DatabricksAudioTranscriptionConfig,
+            )
+
+            return DatabricksAudioTranscriptionConfig()
         elif litellm.LlmProviders.DEEPGRAM == provider:
             return litellm.DeepgramAudioTranscriptionConfig()
         elif litellm.LlmProviders.ELEVENLABS == provider:

tests/unit/llms/databricks/audio_transcription/test_databricks_audio_transcription_transformation.py

@@ -0,0 +1,411 @@
+import base64
+import io
+import json
+import sys
+from unittest.mock import patch
+
+import httpx
+import pytest
+
+import litellm
+from litellm.llms.base_llm.audio_transcription.transformation import (
+    AudioTranscriptionRequestData,
+)
+from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
+from litellm.llms.databricks.audio_transcription.transformation import (
+    DatabricksAudioTranscriptionConfig,
+)
+from litellm.llms.databricks.common_utils import DatabricksException
+from litellm.types.utils import LlmProviders
+from litellm.utils import ProviderConfigManager
+
+CONFIG = DatabricksAudioTranscriptionConfig()
+
+WAV_BYTES = b"RIFF" + b"\x00" * 64
+
+API_BASE = "https://adb-123.azuredatabricks.net/serving-endpoints"
+
+BARE_HOST = "https://adb-123.azuredatabricks.net"
+
+INVOKE_URL = f"{API_BASE}/whisper-t/invocations"
+
+TRANSCRIPT = "The quick brown fox jumps over the lazy dog."
+
+
+@pytest.fixture(autouse=True)
+def clean_databricks_env(monkeypatch):
+    for var in (
+        "DATABRICKS_API_BASE",
+        "DATABRICKS_API_KEY",
+        "DATABRICKS_CLIENT_ID",
+        "DATABRICKS_CLIENT_SECRET",
+        "DATABRICKS_HOST",
+        "DATABRICKS_TOKEN",
+    ):
+        monkeypatch.delenv(var, raising=False)
+
+
+def test_transform_request_builds_invoke_body():
+    result = CONFIG.transform_audio_transcription_request(
+        model="databricks/whisper-t",
+        audio_file=WAV_BYTES,
+        optional_params={},
+        litellm_params={},
+    )
+
+    assert isinstance(result, AudioTranscriptionRequestData)
+    assert result.files is None
+    assert result.content_type == "application/json"
+    assert list(result.data) == ["inputs"]
+    assert base64.b64decode(result.data["inputs"][0]) == WAV_BYTES
+
+
+@pytest.mark.parametrize(
+    "audio_file",
+    [
+        WAV_BYTES,
+        io.BytesIO(WAV_BYTES),
+        ("sample.wav", WAV_BYTES, "audio/wav"),
+    ],
+    ids=["bytes", "file-like", "tuple"],
+)
+def test_transform_request_file_variants(audio_file):
+    result = CONFIG.transform_audio_transcription_request(
+        model="databricks/whisper-t",
+        audio_file=audio_file,
+        optional_params={},
+        litellm_params={},
+    )
+
+    assert base64.b64decode(result.data["inputs"][0]) == WAV_BYTES
+    assert result.files is None
+
+
+@pytest.mark.parametrize(
+    "api_base, model, expected",
+    [
+        (API_BASE, "databricks/whisper-t", INVOKE_URL),
+        (f"{API_BASE}/", "databricks/whisper-t", INVOKE_URL),
+        (BARE_HOST, "whisper-t", INVOKE_URL),
+        (f"{BARE_HOST}/", "whisper-t", INVOKE_URL),
+    ],
+    ids=[
+        "serving-endpoints-base",
+        "serving-endpoints-trailing-slash",
+        "bare-host-normalized",
+        "bare-host-trailing-slash",
+    ],
+)
+def test_get_complete_url(api_base, model, expected):
+    url = CONFIG.get_complete_url(
+        api_base=api_base,
+        api_key=None,
+        model=model,
+        optional_params={},
+        litellm_params={},
+    )
+
+    assert url == expected
+
+
+def test_get_complete_url_from_env(monkeypatch):
+    monkeypatch.setenv("DATABRICKS_API_BASE", API_BASE)
+    url = CONFIG.get_complete_url(
+        api_base=None,
+        api_key=None,
+        model="databricks/whisper-t",
+        optional_params={},
+        litellm_params={},
+    )
+
+    assert url == INVOKE_URL
+
+
+@pytest.mark.parametrize(
+    "model",
+    ["databricks/../x", "databricks/.."],
+    ids=["dotdot-path", "dotdot-only"],
+)
+def test_get_complete_url_rejects_path_shaping(model):
+    with pytest.raises(DatabricksException, match="Invalid Databricks endpoint name") as exc_info:
+        CONFIG.get_complete_url(
+            api_base=API_BASE,
+            api_key=None,
+            model=model,
+            optional_params={},
+            litellm_params={},
+        )
+
+    assert exc_info.value.status_code == 400
+
+
+def test_get_complete_url_no_base_raises(monkeypatch):
+    monkeypatch.delenv("DATABRICKS_API_BASE", raising=False)
+    with patch.dict(sys.modules, {"databricks": None}):
+        with pytest.raises(DatabricksException, match="DATABRICKS_API_BASE") as exc_info:
+            CONFIG.get_complete_url(
+                api_base=None,
+                api_key=None,
+                model="databricks/whisper-t",
+                optional_params={},
+                litellm_params={},
+            )
+
+    assert exc_info.value.status_code == 400
+
+
+def test_validate_environment_sets_bearer_and_json():
+    headers = CONFIG.validate_environment(
+        headers={},
+        model="databricks/whisper-t",
+        messages=[],
+        optional_params={},
+        litellm_params={},
+        api_key="dapi-test",
+        api_base=API_BASE,
+    )
+
+    assert headers["Authorization"] == "Bearer dapi-test"
+    assert headers["Content-Type"] == "application/json"
+
+
+def test_validate_environment_resolves_base_from_env(monkeypatch):
+    monkeypatch.setenv("DATABRICKS_API_BASE", API_BASE)
+    headers = CONFIG.validate_environment(
+        headers={},
+        model="databricks/whisper-t",
+        messages=[],
+        optional_params={},
+        litellm_params={},
+        api_key="dapi-test",
+        api_base=None,
+    )
+
+    assert headers["Authorization"] == "Bearer dapi-test"
+    assert headers["Content-Type"] == "application/json"
+
+
+def test_validate_environment_m2m(monkeypatch):
+    monkeypatch.setenv("DATABRICKS_CLIENT_ID", "client-id")
+    monkeypatch.setenv("DATABRICKS_CLIENT_SECRET", "client-secret")
+    monkeypatch.setenv("DATABRICKS_API_BASE", API_BASE)
+    with patch.object(CONFIG, "_get_oauth_m2m_token", return_value="m2m-token") as mock_token:
+        headers = CONFIG.validate_environment(
+            headers={},
+            model="databricks/whisper-t",
+            messages=[],
+            optional_params={},
+            litellm_params={},
+        )
+
+    mock_token.assert_called_once_with(API_BASE, "client-id", "client-secret")
+    assert headers["Authorization"] == "Bearer m2m-token"
+    assert headers["Content-Type"] == "application/json"
+
+
+def test_validate_environment_no_credentials_raises():
+    with patch.dict(sys.modules, {"databricks": None}):
+        with pytest.raises(DatabricksException, match="databricks-sdk") as exc_info:
+            CONFIG.validate_environment(
+                headers={},
+                model="databricks/whisper-t",
+                messages=[],
+                optional_params={},
+                litellm_params={},
+            )
+
+    assert exc_info.value.status_code == 400
+
+
+def test_transform_response_predictions_list():
+    raw = httpx.Response(
+        200,
+        json={"predictions": [TRANSCRIPT]},
+        request=httpx.Request("POST", INVOKE_URL),
+    )
+    response = CONFIG.transform_audio_transcription_response(raw_response=raw)
+
+    assert response.text == TRANSCRIPT
+    assert response["task"] == "transcribe"
+    assert response._hidden_params == {"predictions": [TRANSCRIPT]}
+
+
+@pytest.mark.parametrize(
+    "payload, expected_text",
+    [
+        ({"predictions": TRANSCRIPT}, TRANSCRIPT),
+        ({"predictions": [{"text": TRANSCRIPT}]}, TRANSCRIPT),
+    ],
+    ids=["bare-string-predictions", "dict-element-predictions"],
+)
+def test_transform_response_predictions_variants(payload, expected_text):
+    raw = httpx.Response(200, json=payload, request=httpx.Request("POST", INVOKE_URL))
+    response = CONFIG.transform_audio_transcription_response(raw_response=raw)
+
+    assert response.text == expected_text
+
+
+def test_transform_response_empty_transcript_is_valid():
+    raw = httpx.Response(200, json={"predictions": [""]}, request=httpx.Request("POST", INVOKE_URL))
+    response = CONFIG.transform_audio_transcription_response(raw_response=raw)
+
+    assert response.text == ""
+
+
+@pytest.mark.parametrize(
+    "payload",
+    [{"predictions": []}, {}, {"predictions": [42]}],
+    ids=["empty-list", "missing-key", "non-text-element"],
+)
+def test_transform_response_empty_predictions_raises(payload):
+    raw = httpx.Response(200, json=payload, request=httpx.Request("POST", INVOKE_URL))
+    with pytest.raises(DatabricksException, match="no prediction text") as exc_info:
+        CONFIG.transform_audio_transcription_response(raw_response=raw)
+
+    assert exc_info.value.status_code == 500
+
+
+def test_transform_response_non_json_body_raises():
+    raw = httpx.Response(
+        200,
+        content=b"<html>bad gateway</html>",
+        request=httpx.Request("POST", INVOKE_URL),
+    )
+    with pytest.raises(DatabricksException, match="Error parsing Databricks transcription response") as exc_info:
+        CONFIG.transform_audio_transcription_response(raw_response=raw)
+
+    assert exc_info.value.status_code == 500
+
+
+@pytest.mark.parametrize("status_code", [401, 404])
+def test_transform_response_error_status_keeps_code(status_code):
+    raw = httpx.Response(
+        status_code,
+        json={"error_code": "ENDPOINT_NOT_FOUND", "message": "Endpoint not found"},
+        request=httpx.Request("POST", INVOKE_URL),
+    )
+    with pytest.raises(DatabricksException) as exc_info:
+        CONFIG.transform_audio_transcription_response(raw_response=raw)
+
+    assert exc_info.value.status_code == status_code
+    assert "Endpoint not found" in exc_info.value.message
+    # the provider error_code rides in the message (no error_code field on BaseLLMException)
+    assert "ENDPOINT_NOT_FOUND" in exc_info.value.message
+
+
+def test_supported_params_advertise_whisper_set():
+    assert CONFIG.get_supported_openai_params(model="whisper-t") == [
+        "language",
+        "prompt",
+        "response_format",
+        "temperature",
+        "timestamp_granularities",
+    ]
+    mapped = CONFIG.map_openai_params(
+        non_default_params={"response_format": "json"},
+        optional_params={},
+        model="whisper-t",
+        drop_params=False,
+    )
+
+    assert mapped == {}
+
+
+def test_transcription_dispatches_to_databricks_invoke(monkeypatch):
+    monkeypatch.setenv("DATABRICKS_API_BASE", API_BASE)
+    captured: dict = {}
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        captured["request"] = request
+        return httpx.Response(
+            200,
+            json={"predictions": ["transcribed text"]},
+            request=request,
+        )
+
+    http_handler = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler)))
+    response = litellm.transcription(
+        model="databricks/whisper-t",
+        file=("sample.wav", WAV_BYTES, "audio/wav"),
+        api_key="dapi-test",
+        api_base=API_BASE,
+        language="en",
+        response_format="json",
+        client=http_handler,
+    )
+
+    request = captured["request"]
+    assert str(request.url) == INVOKE_URL
+    assert request.headers["Authorization"] == "Bearer dapi-test"
+    assert request.headers["Content-Type"] == "application/json"
+    body = json.loads(request.content)
+    assert base64.b64decode(body["inputs"][0]) == WAV_BYTES
+    assert response.text == "transcribed text"
+
+
+def test_transcription_uses_per_call_api_base_for_auth(monkeypatch):
+    # regression: auth must see the per-call base, else it silently falls into the SDK fallback
+    captured: dict = {}
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        captured["request"] = request
+        return httpx.Response(
+            200,
+            json={"predictions": ["transcribed text"]},
+            request=request,
+        )
+
+    http_handler = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler)))
+    response = litellm.transcription(
+        model="databricks/whisper-t",
+        file=("sample.wav", WAV_BYTES, "audio/wav"),
+        api_key="dapi-test",
+        api_base=API_BASE,
+        client=http_handler,
+    )
+
+    request = captured["request"]
+    assert str(request.url) == INVOKE_URL
+    assert request.headers["Authorization"] == "Bearer dapi-test"
+    assert response.text == "transcribed text"
+
+
+async def test_atranscription_dispatches_to_databricks_invoke(monkeypatch):
+    monkeypatch.setenv("DATABRICKS_API_BASE", API_BASE)
+    captured: dict = {}
+
+    async def handler(request: httpx.Request) -> httpx.Response:
+        captured["request"] = request
+        return httpx.Response(
+            200,
+            json={"predictions": ["transcribed text"]},
+            request=request,
+        )
+
+    http_handler = AsyncHTTPHandler(transport=httpx.MockTransport(handler))
+    response = await litellm.atranscription(
+        model="databricks/whisper-t",
+        file=("sample.wav", WAV_BYTES, "audio/wav"),
+        api_key="dapi-test",
+        api_base=API_BASE,
+        language="en",
+        response_format="json",
+        client=http_handler,
+    )
+
+    request = captured["request"]
+    assert str(request.url) == INVOKE_URL
+    assert request.headers["Authorization"] == "Bearer dapi-test"
+    assert request.headers["Content-Type"] == "application/json"
+    body = json.loads(request.content)
+    assert base64.b64decode(body["inputs"][0]) == WAV_BYTES
+    assert response.text == "transcribed text"
+
+
+def test_provider_config_manager_returns_databricks_config():
+    config = ProviderConfigManager.get_provider_audio_transcription_config(
+        model="whisper-t",
+        provider=LlmProviders.DATABRICKS,
+    )
+
+    assert isinstance(config, DatabricksAudioTranscriptionConfig)