litellm_rs/core/providers/databricks/
mod.rs

1//! Databricks Model Serving OpenAI-compatible contract.
2
3use crate::core::net::ProviderEndpointAccess;
4use crate::core::providers::ProviderError;
5use crate::core::providers::enterprise::{
6    EnterpriseOpenAiProvider, EnterpriseOpenAiSettings, normalize_enterprise_base_url,
7    validate_request_header_value,
8};
9use serde::{Deserialize, Serialize};
10
11pub type DatabricksProvider = EnterpriseOpenAiProvider;
12
13#[derive(Clone, Serialize, Deserialize)]
14#[serde(deny_unknown_fields)]
15pub struct DatabricksConfig {
16    pub workspace_url: String,
17    pub api_key: String,
18    #[serde(default)]
19    pub endpoint_access: ProviderEndpointAccess,
20    #[serde(default = "default_timeout")]
21    pub timeout: u64,
22    #[serde(default = "default_retries")]
23    pub max_retries: u32,
24    #[serde(default)]
25    pub models: Vec<String>,
26}
27
28impl std::fmt::Debug for DatabricksConfig {
29    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
30        f.debug_struct("DatabricksConfig")
31            .field("workspace_url", &self.workspace_url)
32            .field("api_key", &"[REDACTED]")
33            .field("endpoint_access", &self.endpoint_access)
34            .field("timeout", &self.timeout)
35            .field("max_retries", &self.max_retries)
36            .field("models", &self.models)
37            .finish()
38    }
39}
40
41impl DatabricksConfig {
42    pub fn api_base(&self) -> Result<String, ProviderError> {
43        let workspace = normalize_enterprise_base_url("databricks", &self.workspace_url, true)?;
44        Ok(format!("{workspace}/serving-endpoints"))
45    }
46
47    pub async fn build(self) -> Result<DatabricksProvider, ProviderError> {
48        if self.api_key.trim().is_empty() {
49            return Err(ProviderError::configuration(
50                "databricks",
51                "api_key is required",
52            ));
53        }
54        validate_request_header_value(
55            "databricks",
56            "api_key",
57            &format!("Bearer {}", self.api_key),
58        )?;
59        EnterpriseOpenAiProvider::new(
60            "databricks",
61            EnterpriseOpenAiSettings {
62                api_base: self.api_base()?,
63                api_key: self.api_key,
64                model_prefix: "databricks/",
65                endpoint_access: self.endpoint_access,
66                timeout: self.timeout,
67                max_retries: self.max_retries,
68                headers: Default::default(),
69                models: self.models,
70            },
71        )
72        .await
73    }
74}
75
76const fn default_timeout() -> u64 {
77    60
78}
79const fn default_retries() -> u32 {
80    2
81}
82
83#[cfg(test)]
84mod tests {
85    use super::*;
86
87    #[test]
88    fn constructs_serving_endpoint_base_and_rejects_pathful_workspace() {
89        let config = DatabricksConfig {
90            workspace_url: "https://dbc.example.cloud.databricks.com".to_string(),
91            api_key: "test".to_string(),
92            endpoint_access: ProviderEndpointAccess::PublicOnly,
93            timeout: 60,
94            max_retries: 2,
95            models: Vec::new(),
96        };
97        assert_eq!(
98            config.api_base().expect("valid workspace"),
99            "https://dbc.example.cloud.databricks.com/serving-endpoints"
100        );
101        let invalid = DatabricksConfig {
102            workspace_url: "https://dbc.example/x".to_string(),
103            ..config
104        };
105        assert!(invalid.api_base().is_err());
106        for workspace_url in [
107            "https://user:password@dbc.example.com",
108            "https://dbc.example.com?tenant=other",
109            "https://dbc.example.com#fragment",
110        ] {
111            assert!(
112                DatabricksConfig {
113                    workspace_url: workspace_url.to_string(),
114                    ..invalid.clone()
115                }
116                .api_base()
117                .is_err(),
118                "workspace origin must reject {workspace_url}"
119            );
120        }
121    }
122
123    #[tokio::test]
124    async fn runtime_identity_and_capabilities_are_platform_specific() {
125        use crate::core::traits::provider::llm_provider::trait_definition::LLMProvider;
126        use crate::core::types::model::ProviderCapability;
127        let provider = DatabricksConfig {
128            workspace_url: "https://dbc.example.cloud.databricks.com".to_string(),
129            api_key: "test".to_string(),
130            endpoint_access: ProviderEndpointAccess::PublicOnly,
131            timeout: 60,
132            max_retries: 2,
133            models: vec!["served-model".to_string()],
134        }
135        .build()
136        .await
137        .expect("valid config should build");
138        assert_eq!(provider.name(), "databricks");
139        assert!(provider.supports_capability(&ProviderCapability::ChatCompletionStream));
140        assert!(!provider.supports_capability(&ProviderCapability::Embeddings));
141        assert!(provider.supports_model("databricks/served-model"));
142        assert!(!provider.supports_model("other-model"));
143    }
144}
litellm_rs/core/providers/databricks/mod.rs

litellm_rs/core/providers/databricks/
mod.rs
