fix: page through all Databricks AI Gateway v2 endpoints when listing models (#9753)
Signed-off-by: Kalvin Chau <kalvin@block.xyz>
This commit is contained in:
@@ -37,6 +37,7 @@ use rmcp::model::Tool;
|
||||
|
||||
const DATABRICKS_V2_PROVIDER_NAME: &str = "databricks_v2";
|
||||
const DATABRICKS_V2_LIST_ENDPOINTS_PATH: &str = "api/ai-gateway/v2/endpoints";
|
||||
const DATABRICKS_V2_LIST_ENDPOINTS_PAGE_SIZE: usize = 100;
|
||||
pub const DATABRICKS_V2_DEFAULT_MODEL: &str = "databricks-gpt-5-5";
|
||||
pub const DATABRICKS_V2_KNOWN_MODELS: &[&str] =
|
||||
&["databricks-gpt-5-5", "databricks-claude-opus-4-7"];
|
||||
@@ -189,7 +190,9 @@ impl DatabricksV2Provider {
|
||||
model_name.contains("claude")
|
||||
}
|
||||
|
||||
fn parse_list_endpoints_response(json: &Value) -> Result<Vec<String>, ProviderError> {
|
||||
fn parse_list_endpoints_response(
|
||||
json: &Value,
|
||||
) -> Result<(Vec<String>, Option<String>), ProviderError> {
|
||||
let endpoints = json
|
||||
.get("endpoints")
|
||||
.and_then(|v| v.as_array())
|
||||
@@ -200,7 +203,7 @@ impl DatabricksV2Provider {
|
||||
)
|
||||
})?;
|
||||
|
||||
let mut models: Vec<String> = endpoints
|
||||
let models: Vec<String> = endpoints
|
||||
.iter()
|
||||
.filter_map(|endpoint| {
|
||||
endpoint
|
||||
@@ -209,8 +212,14 @@ impl DatabricksV2Provider {
|
||||
.map(str::to_string)
|
||||
})
|
||||
.collect();
|
||||
models.sort();
|
||||
Ok(models)
|
||||
|
||||
let next_page_token = json
|
||||
.get("next_page_token")
|
||||
.and_then(|v| v.as_str())
|
||||
.filter(|token| !token.is_empty())
|
||||
.map(str::to_string);
|
||||
|
||||
Ok((models, next_page_token))
|
||||
}
|
||||
|
||||
async fn stream_openai_responses(
|
||||
@@ -417,31 +426,53 @@ impl Provider for DatabricksV2Provider {
|
||||
}
|
||||
|
||||
async fn fetch_supported_models(&self) -> Result<Vec<String>, ProviderError> {
|
||||
let response = self
|
||||
.api_client
|
||||
.response_get(None, DATABRICKS_V2_LIST_ENDPOINTS_PATH)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
let mut models = Vec::new();
|
||||
let mut page_token: Option<String> = None;
|
||||
|
||||
loop {
|
||||
let mut path = format!(
|
||||
"{}?page_size={}",
|
||||
DATABRICKS_V2_LIST_ENDPOINTS_PATH, DATABRICKS_V2_LIST_ENDPOINTS_PAGE_SIZE
|
||||
);
|
||||
if let Some(token) = &page_token {
|
||||
path.push_str(&format!("&page_token={}", urlencoding::encode(token)));
|
||||
}
|
||||
|
||||
let response = self
|
||||
.api_client
|
||||
.response_get(None, &path)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
ProviderError::RequestFailed(format!(
|
||||
"Failed to fetch Databricks AI Gateway endpoints: {e}"
|
||||
))
|
||||
})?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
let status = response.status();
|
||||
let detail = response.text().await.unwrap_or_default();
|
||||
return Err(ProviderError::RequestFailed(format!(
|
||||
"Failed to fetch Databricks AI Gateway endpoints: {status} {detail}"
|
||||
)));
|
||||
}
|
||||
|
||||
let json: Value = response.json().await.map_err(|e| {
|
||||
ProviderError::RequestFailed(format!(
|
||||
"Failed to fetch Databricks AI Gateway endpoints: {e}"
|
||||
"Failed to parse Databricks AI Gateway endpoints response: {e}"
|
||||
))
|
||||
})?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
let status = response.status();
|
||||
let detail = response.text().await.unwrap_or_default();
|
||||
return Err(ProviderError::RequestFailed(format!(
|
||||
"Failed to fetch Databricks AI Gateway endpoints: {status} {detail}"
|
||||
)));
|
||||
let (page_models, next_page_token) = Self::parse_list_endpoints_response(&json)?;
|
||||
models.extend(page_models);
|
||||
|
||||
if next_page_token.is_none() || next_page_token == page_token {
|
||||
break;
|
||||
}
|
||||
page_token = next_page_token;
|
||||
}
|
||||
|
||||
let json: Value = response.json().await.map_err(|e| {
|
||||
ProviderError::RequestFailed(format!(
|
||||
"Failed to parse Databricks AI Gateway endpoints response: {e}"
|
||||
))
|
||||
})?;
|
||||
|
||||
Self::parse_list_endpoints_response(&json)
|
||||
models.sort();
|
||||
Ok(models)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -480,19 +511,22 @@ mod tests {
|
||||
{"name": "databricks-claude-opus-4-7"},
|
||||
{"name": "databricks-gpt-5-5"},
|
||||
{"name": "custom-model"}
|
||||
]
|
||||
],
|
||||
"next_page_token": "tok"
|
||||
});
|
||||
|
||||
let models = DatabricksV2Provider::parse_list_endpoints_response(&json).unwrap();
|
||||
let (models, next_page_token) =
|
||||
DatabricksV2Provider::parse_list_endpoints_response(&json).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
models,
|
||||
vec![
|
||||
"custom-model".to_string(),
|
||||
"databricks-claude-opus-4-7".to_string(),
|
||||
"databricks-gpt-5-5".to_string(),
|
||||
"custom-model".to_string(),
|
||||
]
|
||||
);
|
||||
assert_eq!(next_page_token.as_deref(), Some("tok"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
Reference in New Issue
Block a user