From c9e699b1910ad1a21d9061653d7dd2f632bc3f33 Mon Sep 17 00:00:00 2001 From: narmstrong-berry Date: Thu, 15 Oct 2026 02:22:53 +0000 Subject: [PATCH] feat(rust): serve Databricks GPT responses natively The Responses route now resolves databricks. The Databricks config sends the OpenAI Responses payload to the serving-endpoints base with the personal access token, as Python does for GPT endpoints. Non-GPT endpoints take Python's Chat Completions bridge, and OAuth M2M and Databricks SDK auth decline to Python. The parity harness now compares the user agent too --- .../core-utils/src/get_llm_provider_logic.rs | 2 + .../crates/inference-responses/src/prepare.rs | 8 + litellm-rust/crates/llms/AGENTS.md | 2 +- .../crates/llms/src/databricks/mod.rs | 1 + .../llms/src/databricks/responses/AGENTS.md | 10 + .../llms/src/databricks/responses/mod.rs | 1 + .../databricks/responses/transformation.rs | 177 ++++++++++++++++++ litellm-rust/crates/llms/src/lib.rs | 1 + .../python-bridge/src/routes/responses/mod.rs | 2 +- litellm/rust_bridge/host/routes/responses.py | 24 +++ .../sdk/responses/test_native_parity.py | 10 +- .../rust_bridge/host/routes/test_responses.py | 56 ++++++ 12 files changed, 291 insertions(+), 3 deletions(-) create mode 100644 litellm-rust/crates/llms/src/databricks/mod.rs create mode 100644 litellm-rust/crates/llms/src/databricks/responses/AGENTS.md create mode 100644 litellm-rust/crates/llms/src/databricks/responses/mod.rs create mode 100644 litellm-rust/crates/llms/src/databricks/responses/transformation.rs diff --git a/litellm-rust/crates/core-utils/src/get_llm_provider_logic.rs b/litellm-rust/crates/core-utils/src/get_llm_provider_logic.rs index 2e88afdd8ff..1454795fb09 100644 --- a/litellm-rust/crates/core-utils/src/get_llm_provider_logic.rs +++ b/litellm-rust/crates/core-utils/src/get_llm_provider_logic.rs @@ -22,6 +22,8 @@ pub enum LlmProviders { BedrockMantle, #[strum(serialize = "cohere")] Cohere, + #[strum(serialize = "databricks")] + Databricks, #[strum(serialize = "deepseek")] Deepseek, #[strum(serialize = "edenai")] diff --git a/litellm-rust/crates/inference-responses/src/prepare.rs b/litellm-rust/crates/inference-responses/src/prepare.rs index c7f72372b51..a846b64e7db 100644 --- a/litellm-rust/crates/inference-responses/src/prepare.rs +++ b/litellm-rust/crates/inference-responses/src/prepare.rs @@ -5,6 +5,7 @@ use litellm_host::interceptors::{ProviderIdentity, RequestContext}; use litellm_llms::{ azure::responses::transformation::AZURE_RESPONSES_CONFIG, base_llm::responses::transformation::BaseResponsesApiConfig, + databricks::responses::transformation::DATABRICKS_RESPONSES_CONFIG, fireworks_ai::responses::transformation::FIREWORKS_AI_RESPONSES_CONFIG, hosted_vllm::responses::transformation::HOSTED_VLLM_RESPONSES_CONFIG, openai::responses::transformation::OpenAiResponsesApiConfig, @@ -99,6 +100,7 @@ pub(super) fn resolve_provider( Ok(LlmProviders::Openrouter) => (&OPENROUTER_RESPONSES_CONFIG, resolved.into()), Ok(LlmProviders::Azure) => (&AZURE_RESPONSES_CONFIG, resolved.into()), Ok(LlmProviders::Xai) => (&XAI_RESPONSES_CONFIG, resolved.into()), + Ok(LlmProviders::Databricks) => (&DATABRICKS_RESPONSES_CONFIG, resolved.into()), Ok(LlmProviders::FireworksAi) => (&FIREWORKS_AI_RESPONSES_CONFIG, resolved.into()), Ok(LlmProviders::HostedVllm) => (&HOSTED_VLLM_RESPONSES_CONFIG, resolved.into()), _ => return Err(Error::Unsupported("native HTTP responses provider")), @@ -150,6 +152,12 @@ mod tests { #[case::openrouter_native_id_implicit("openrouter/auto", None, "openrouter", "auto")] #[case::azure_prefix("azure/my-deployment", None, "azure", "my-deployment")] #[case::azure_explicit("my-deployment", Some("azure"), "azure", "my-deployment")] + #[case::databricks_prefix( + "databricks/databricks-gpt-5-nano", + None, + "databricks", + "databricks-gpt-5-nano" + )] #[case::fireworks_prefix( "fireworks_ai/accounts/fireworks/models/qwen3", None, diff --git a/litellm-rust/crates/llms/AGENTS.md b/litellm-rust/crates/llms/AGENTS.md index 245b6756a50..93b3d354fb0 100644 --- a/litellm-rust/crates/llms/AGENTS.md +++ b/litellm-rust/crates/llms/AGENTS.md @@ -58,7 +58,7 @@ - `AnthropicConfig`, `AmazonConverseConfig` and the non-OCR base traits are partial. Preserve their acceptance gates, passthrough and host fallback contracts - `OpenAiResponsesApiConfig` implements WebSocket transformations and direct HTTP Responses, without Python's model-specific param rewriting or Responses-to-Chat emulation -- OpenAI-compatible Responses hosts (`openrouter`, `azure`, `xai`, `hosted_vllm`, `fireworks_ai`) reuse the OpenAI request and response transforms as free functions and own only their credentials, URL and host-specific item rules +- OpenAI-compatible Responses hosts (`openrouter`, `azure`, `xai`, `hosted_vllm`, `fireworks_ai`, `databricks`) reuse the OpenAI request and response transforms as free functions and own only their credentials, URL and host-specific item rules - Chat `supported_openai_param_mappings` lists OpenAI/provider name pairs, unlike Python's name-only list. Audio `map_transcription_params` is a Rust filtering helper - Auth resolution is split between configs and route preparation in `inference-` diff --git a/litellm-rust/crates/llms/src/databricks/mod.rs b/litellm-rust/crates/llms/src/databricks/mod.rs new file mode 100644 index 00000000000..b396b037bc5 --- /dev/null +++ b/litellm-rust/crates/llms/src/databricks/mod.rs @@ -0,0 +1 @@ +pub mod responses; diff --git a/litellm-rust/crates/llms/src/databricks/responses/AGENTS.md b/litellm-rust/crates/llms/src/databricks/responses/AGENTS.md new file mode 100644 index 00000000000..b97b48f39bd --- /dev/null +++ b/litellm-rust/crates/llms/src/databricks/responses/AGENTS.md @@ -0,0 +1,10 @@ +# rules + +- Send the OpenAI Responses payload to `{api_base}/responses`, where the base is the workspace `/serving-endpoints` URL from `api_base` or `DATABRICKS_API_BASE` +- Authenticate with the personal access token from `api_key` or `DATABRICKS_API_KEY`. OAuth M2M and Databricks SDK auth stay on Python +- Only OpenAI GPT endpoints serve Responses; other models take Python's Chat Completions bridge +- No native WebSocket mode + +# references + +- https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/api-reference diff --git a/litellm-rust/crates/llms/src/databricks/responses/mod.rs b/litellm-rust/crates/llms/src/databricks/responses/mod.rs new file mode 100644 index 00000000000..f239b6921fa --- /dev/null +++ b/litellm-rust/crates/llms/src/databricks/responses/mod.rs @@ -0,0 +1 @@ +pub mod transformation; diff --git a/litellm-rust/crates/llms/src/databricks/responses/transformation.rs b/litellm-rust/crates/llms/src/databricks/responses/transformation.rs new file mode 100644 index 00000000000..4230e923ddf --- /dev/null +++ b/litellm-rust/crates/llms/src/databricks/responses/transformation.rs @@ -0,0 +1,177 @@ +//! Databricks Responses for its OpenAI GPT serving endpoints, matching +//! `litellm/llms/databricks/responses/transformation.py` + +use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_llms_types::formats::responses::ResponsesApiResponse; +use serde_json::{Map, Value}; + +use crate::{ + Error, + base_llm::{ + auth::{AuthScheme, ValidatedEnvironment}, + responses::transformation::BaseResponsesApiConfig, + }, + openai::responses::{ + request::ResponsesHost, + transformation::{responses_request, responses_response}, + }, +}; + +const API_KEY_ENV: &str = "DATABRICKS_API_KEY"; +const API_BASE_ENV: &str = "DATABRICKS_API_BASE"; + +pub struct DatabricksResponsesApiConfig; +pub const DATABRICKS_RESPONSES_CONFIG: DatabricksResponsesApiConfig = DatabricksResponsesApiConfig; + +impl BaseResponsesApiConfig for DatabricksResponsesApiConfig { + fn secret_names( + &self, + api_key: Option<&str>, + api_base: Option<&str>, + ) -> &'static [&'static str] { + match ( + api_key.is_some_and(|key| !key.is_empty()), + api_base.is_some_and(|base| !base.is_empty()), + ) { + (true, true) => &[], + (true, false) => &[API_BASE_ENV], + (false, true) => &[API_KEY_ENV], + (false, false) => &[API_KEY_ENV, API_BASE_ENV], + } + } + + fn validate_environment( + &self, + headers: Vec<(String, String)>, + api_key: Option<&str>, + lookup: &dyn Fn(&str) -> Option, + ) -> Result { + let key = + resolved(api_key, lookup, API_KEY_ENV).ok_or(litellm_auth::Error::MissingApiKey { + provider: "Databricks", + environment_variable: API_KEY_ENV, + })?; + Ok(ValidatedEnvironment { + headers: litellm_http::request::with_default_headers( + headers, + &[("content-type", "application/json")], + ), + auth: AuthScheme::Credential { + placement: CredentialPlacement::Bearer, + secret: SecretValue::new(key), + }, + }) + } + + fn get_complete_url( + &self, + api_base: Option<&str>, + _api_version: Option<&str>, + lookup: &dyn Fn(&str) -> Option, + ) -> Result { + let api_base = resolved(api_base, lookup, API_BASE_ENV).ok_or( + litellm_auth::Error::MissingApiBase { + provider: "Databricks", + guidance: "Set `api_base` or the DATABRICKS_API_BASE environment variable", + }, + )?; + Ok(format!("{}/responses", api_base.trim_end_matches('/'))) + } + + fn transform_responses_api_request( + &self, + model: &str, + input: Value, + params: Map, + ) -> Result { + let model = model.strip_prefix("databricks/").unwrap_or(model); + responses_request(model, input, params, ResponsesHost::Compatible) + } + + fn transform_response_api_response(&self, body: Value) -> Result { + responses_response(body) + } +} + +/// Python's `value or os.getenv(name)`: an empty value falls through, and nothing is trimmed +fn resolved( + value: Option<&str>, + lookup: &dyn Fn(&str) -> Option, + name: &str, +) -> Option { + value + .filter(|value| !value.is_empty()) + .map(str::to_owned) + .or_else(|| lookup(name).filter(|value| !value.is_empty())) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + fn env(pairs: &'static [(&'static str, &'static str)]) -> impl Fn(&str) -> Option { + move |name| { + pairs + .iter() + .find(|(key, _)| *key == name) + .map(|(_, value)| (*value).to_owned()) + } + } + + #[rstest] + #[case::explicit(Some("https://w.cloud.databricks.com/serving-endpoints/"), &[], "https://w.cloud.databricks.com/serving-endpoints/responses")] + #[case::secret(None, &[("DATABRICKS_API_BASE", "https://s.cloud.databricks.com/serving-endpoints")], "https://s.cloud.databricks.com/serving-endpoints/responses")] + fn url_is_the_serving_endpoints_base_plus_responses( + #[case] api_base: Option<&str>, + #[case] secrets: &'static [(&'static str, &'static str)], + #[case] expected: &str, + ) { + assert_eq!( + DATABRICKS_RESPONSES_CONFIG + .get_complete_url(api_base, None, &env(secrets)) + .unwrap(), + expected + ); + } + + #[test] + fn url_needs_an_api_base() { + assert!(matches!( + DATABRICKS_RESPONSES_CONFIG.get_complete_url(None, None, &env(&[])), + Err(Error::Auth(litellm_auth::Error::MissingApiBase { .. })) + )); + } + + #[rstest] + #[case::explicit(Some("dapi-explicit"), &[("DATABRICKS_API_KEY", "dapi-env")], "dapi-explicit")] + #[case::secret(None, &[("DATABRICKS_API_KEY", "dapi-env")], "dapi-env")] + fn bearer_is_the_personal_access_token( + #[case] api_key: Option<&str>, + #[case] secrets: &'static [(&'static str, &'static str)], + #[case] expected: &str, + ) { + let environment = DATABRICKS_RESPONSES_CONFIG + .validate_environment(vec![], api_key, &env(secrets)) + .unwrap(); + let AuthScheme::Credential { placement, secret } = environment.auth else { + panic!("expected a credential"); + }; + assert_eq!(placement, CredentialPlacement::Bearer); + assert_eq!(secret.expose(), expected); + } + + #[test] + fn a_repeated_provider_prefix_is_stripped() { + let body = DATABRICKS_RESPONSES_CONFIG + .transform_responses_api_request( + "databricks/databricks-gpt-5-nano", + json!("hi"), + Map::new(), + ) + .unwrap(); + assert_eq!(body["model"], json!("databricks-gpt-5-nano")); + } +} diff --git a/litellm-rust/crates/llms/src/lib.rs b/litellm-rust/crates/llms/src/lib.rs index d4c3344ab1e..6fb5e1fa6fa 100644 --- a/litellm-rust/crates/llms/src/lib.rs +++ b/litellm-rust/crates/llms/src/lib.rs @@ -6,6 +6,7 @@ pub mod base_llm; pub mod bedrock; pub mod bedrock_mantle; pub mod cohere; +pub mod databricks; pub mod deepseek; pub mod edenai; mod error; diff --git a/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs b/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs index 38ca89f712b..41973795ca1 100644 --- a/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs @@ -47,7 +47,7 @@ fn run_responses(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> Py None => !model.contains('/'), Some(resolved) => match resolved.custom_llm_provider { "openai" => !resolved.model.contains('/'), - "azure" | "xai" | "hosted_vllm" | "fireworks_ai" => true, + "azure" | "xai" | "hosted_vllm" | "fireworks_ai" | "databricks" => true, _ => false, }, }; diff --git a/litellm/rust_bridge/host/routes/responses.py b/litellm/rust_bridge/host/routes/responses.py index 1d354b3be9d..f975cca69d5 100644 --- a/litellm/rust_bridge/host/routes/responses.py +++ b/litellm/rust_bridge/host/routes/responses.py @@ -1,5 +1,6 @@ from __future__ import annotations +import os from collections.abc import Mapping from types import MappingProxyType from typing import Final, cast # noqa: TID251 # narrows caller json the checker cannot see into @@ -7,6 +8,7 @@ from typing import Final, cast # noqa: TID251 # narrows caller json the checke import litellm from litellm import get_llm_provider from litellm.llms.azure.responses.transformation import AzureOpenAIResponsesAPIConfig +from litellm.llms.databricks.responses.transformation import DatabricksResponsesAPIConfig from litellm.llms.fireworks_ai.common_utils import get_fireworks_session_id from litellm.llms.fireworks_ai.responses.transformation import FireworksAIResponsesAPIConfig from litellm.llms.hosted_vllm.responses.transformation import HostedVLLMResponsesAPIConfig @@ -27,6 +29,7 @@ _CONFIGS: Final[Mapping[str, OpenAIResponsesAPIConfig]] = MappingProxyType( "xai": XAIResponsesAPIConfig(), "hosted_vllm": HostedVLLMResponsesAPIConfig(), "fireworks_ai": FireworksAIResponsesAPIConfig(), + "databricks": DatabricksResponsesAPIConfig(), } ) _AZURE_KEY_ENVIRONMENT: Final = ("AZURE_OPENAI_API_KEY", "AZURE_API_KEY") @@ -37,6 +40,8 @@ def body_params(request: Mapping[str, object]) -> Mapping[str, object]: def connection_defaults(provider: str) -> tuple[str | None, str | None]: + if provider == "databricks": + return None, None provider_key: Final = {"openai": litellm.openai_key, "azure": litellm.azure_key}.get(provider) return litellm.api_key or provider_key, litellm.api_base @@ -92,6 +97,8 @@ def decline_reason(request: Mapping[str, object]) -> str | None: def _provider_decline_reason(provider: str, model: str, request: Mapping[str, object]) -> str | None: if provider == "azure": return _azure_decline_reason(model, request) + if provider == "databricks": + return _databricks_decline_reason(model, request) if provider not in ("xai", "hosted_vllm", "fireworks_ai"): return None provider_key, provider_base = provider_connection( @@ -111,6 +118,23 @@ def _provider_decline_reason(provider: str, model: str, request: Mapping[str, ob return None +def _databricks_decline_reason(model: str, request: Mapping[str, object]) -> str | None: + if "gpt" not in model.lower(): + return "Databricks serves Responses only on GPT endpoints" + api_base: Final = ( + optional_str(request.get("api_base")) + or optional_str(request.get("base_url")) + or os.getenv("DATABRICKS_API_BASE") + ) + if not api_base: + return "native Databricks Responses does not implement Databricks SDK auth" + if os.getenv("DATABRICKS_CLIENT_ID") and os.getenv("DATABRICKS_CLIENT_SECRET"): + return "native Databricks Responses does not implement OAuth M2M" + if not (optional_str(request.get("api_key")) or os.getenv("DATABRICKS_API_KEY")): + return "native Databricks Responses does not implement Databricks SDK auth" + return None + + def _azure_decline_reason(model: str, request: Mapping[str, object]) -> str | None: if "o_series/" in model: return "native Azure Responses does not implement the o-series config" diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/responses/test_native_parity.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/responses/test_native_parity.py index f8db38f03bc..65c265030f2 100644 --- a/tests/rust-python-harness/strategies/e2e_parity/sdk/responses/test_native_parity.py +++ b/tests/rust-python-harness/strategies/e2e_parity/sdk/responses/test_native_parity.py @@ -25,7 +25,14 @@ pytestmark = pytest.mark.requires_rust_extension MODEL: Final = "openai/gpt-6.1-sol" _OBJECT: Final = TypeAdapter(dict[str, JsonValue]) _LOGGED: Final = TypeAdapter(dict[str, object]) -_COMPARED_HEADERS: Final = ("api-key", "authorization", "content-type", "openai-organization", "openai-project") +_COMPARED_HEADERS: Final = ( + "api-key", + "authorization", + "content-type", + "openai-organization", + "openai-project", + "user-agent", +) _COMPARED_LOG_FIELDS: Final = ( "call_type", "custom_llm_provider", @@ -374,6 +381,7 @@ CASES: Final = ( {"model": FIREWORKS_MODEL, "input": "hi", "litellm_metadata": {"session_id": "sess-1"}}, falls_back=True, ), + Case("databricks_gpt_endpoint", {"model": "databricks/databricks-gpt-5-nano", "input": "hi"}), Case("hosted_vllm_placeholder_key", {"model": HOSTED_VLLM_MODEL, "input": "hi", "api_key": None}), Case( "hosted_vllm_keeps_replayed_reasoning_and_ids", diff --git a/tests/unit/rust_bridge/host/routes/test_responses.py b/tests/unit/rust_bridge/host/routes/test_responses.py index a096331a9a0..b0c2d83b03a 100644 --- a/tests/unit/rust_bridge/host/routes/test_responses.py +++ b/tests/unit/rust_bridge/host/routes/test_responses.py @@ -260,3 +260,59 @@ def test_openai_file_search_stays_native() -> None: "tools": [{"type": "file_search", "vector_store_ids": ["vs_1"]}], } assert decline_reason(request) is None + + +@pytest.mark.parametrize( + ("request_fields", "environment", "expected"), + [ + pytest.param({}, {}, None, id="personal_access_token"), + pytest.param( + {"model": "databricks/databricks-claude-sonnet"}, + {}, + "Databricks serves Responses only on GPT endpoints", + id="non_gpt_endpoint", + ), + pytest.param( + {"api_key": None}, + {}, + "native Databricks Responses does not implement Databricks SDK auth", + id="no_token", + ), + pytest.param( + {"api_key": None, "api_base": None}, + {"DATABRICKS_API_KEY": "dapi-env"}, + "native Databricks Responses does not implement Databricks SDK auth", + id="no_api_base", + ), + pytest.param( + {}, + {"DATABRICKS_CLIENT_ID": "id", "DATABRICKS_CLIENT_SECRET": "secret"}, + "native Databricks Responses does not implement OAuth M2M", + id="oauth_m2m", + ), + ], +) +def test_databricks_requests_outside_the_native_config_stay_on_python( + monkeypatch: pytest.MonkeyPatch, + request_fields: dict[str, object], + environment: dict[str, str], + expected: str | None, +) -> None: + for name in ("DATABRICKS_API_KEY", "DATABRICKS_API_BASE", "DATABRICKS_CLIENT_ID", "DATABRICKS_CLIENT_SECRET"): + monkeypatch.delenv(name, raising=False) + for name, value in environment.items(): + monkeypatch.setenv(name, value) + request: Final = { + "model": "databricks/databricks-gpt-5-nano", + "input": "hi", + "api_key": "dapi-explicit", + "api_base": "https://w.cloud.databricks.com/serving-endpoints", + **request_fields, + } + assert decline_reason(request) == expected + + +def test_databricks_connection_defaults_ignore_the_global_settings(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "api_key", "global") + monkeypatch.setattr(litellm, "api_base", "https://global.invalid") + assert connection_defaults("databricks") == (None, None)