diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 4a9b8fee2e8..ad7571d4b06 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/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, ) diff --git a/litellm/llms/databricks/audio_transcription/__init__.py b/litellm/llms/databricks/audio_transcription/__init__.py new file mode 100644 index 00000000000..a42b3737127 --- /dev/null +++ b/litellm/llms/databricks/audio_transcription/__init__.py @@ -0,0 +1 @@ +"""Databricks audio transcription implementation.""" diff --git a/litellm/llms/databricks/audio_transcription/transformation.py b/litellm/llms/databricks/audio_transcription/transformation.py new file mode 100644 index 00000000000..02d0a38b884 --- /dev/null +++ b/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 diff --git a/litellm/utils.py b/litellm/utils.py index dec57245ab2..3f5dccef1f9 100644 --- a/litellm/utils.py +++ b/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: diff --git a/tests/unit/llms/databricks/audio_transcription/__init__.py b/tests/unit/llms/databricks/audio_transcription/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/databricks/audio_transcription/test_databricks_audio_transcription_transformation.py b/tests/unit/llms/databricks/audio_transcription/test_databricks_audio_transcription_transformation.py new file mode 100644 index 00000000000..a78e501b505 --- /dev/null +++ b/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"bad gateway", + 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)