fix: gemini models via databricks (#8042)
Co-authored-by: Michael Neale <michael.neale@gmail.com>
This commit is contained in:
@@ -412,16 +412,40 @@ impl Provider for DatabricksProvider {
|
|||||||
let status = resp.status();
|
let status = resp.status();
|
||||||
let error_text = resp.text().await.unwrap_or_default();
|
let error_text = resp.text().await.unwrap_or_default();
|
||||||
|
|
||||||
// Parse as JSON if possible to pass to map_http_error_to_provider_error
|
|
||||||
let json_payload = serde_json::from_str::<Value>(&error_text).ok();
|
let json_payload = serde_json::from_str::<Value>(&error_text).ok();
|
||||||
return Err(map_http_error_to_provider_error(status, json_payload));
|
return Err(map_http_error_to_provider_error(status, json_payload));
|
||||||
}
|
}
|
||||||
Ok(resp)
|
Ok(resp)
|
||||||
})
|
})
|
||||||
.await
|
.await;
|
||||||
.inspect_err(|e| {
|
|
||||||
let _ = log.error(e);
|
let response = match response {
|
||||||
})?;
|
Err(e) if e.to_string().contains("stream_options") => {
|
||||||
|
payload.as_object_mut().unwrap().remove("stream_options");
|
||||||
|
self.with_retry(|| async {
|
||||||
|
let resp = self
|
||||||
|
.api_client
|
||||||
|
.response_post(Some(session_id), &path, &payload)
|
||||||
|
.await?;
|
||||||
|
if !resp.status().is_success() {
|
||||||
|
let status = resp.status();
|
||||||
|
let error_text = resp.text().await.unwrap_or_default();
|
||||||
|
let json_payload = serde_json::from_str::<Value>(&error_text).ok();
|
||||||
|
return Err(map_http_error_to_provider_error(status, json_payload));
|
||||||
|
}
|
||||||
|
Ok(resp)
|
||||||
|
})
|
||||||
|
.await
|
||||||
|
.inspect_err(|e| {
|
||||||
|
let _ = log.error(e);
|
||||||
|
})?
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
let _ = log.error(&e);
|
||||||
|
return Err(e);
|
||||||
|
}
|
||||||
|
Ok(resp) => resp,
|
||||||
|
};
|
||||||
|
|
||||||
stream_openai_compat(response, log)
|
stream_openai_compat(response, log)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -48,9 +48,25 @@ struct DeltaToolCall {
|
|||||||
extra: Option<serde_json::Map<String, Value>>,
|
extra: Option<serde_json::Map<String, Value>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Serialize, Deserialize, Debug)]
|
||||||
|
#[serde(untagged)]
|
||||||
|
enum DeltaContent {
|
||||||
|
String(String),
|
||||||
|
Array(Vec<ContentPart>),
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Serialize, Deserialize, Debug)]
|
||||||
|
struct ContentPart {
|
||||||
|
r#type: String,
|
||||||
|
text: String,
|
||||||
|
#[serde(rename = "thoughtSignature")]
|
||||||
|
thought_signature: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
#[derive(Serialize, Deserialize, Debug)]
|
#[derive(Serialize, Deserialize, Debug)]
|
||||||
struct Delta {
|
struct Delta {
|
||||||
content: Option<String>,
|
#[serde(default)]
|
||||||
|
content: Option<DeltaContent>,
|
||||||
role: Option<String>,
|
role: Option<String>,
|
||||||
tool_calls: Option<Vec<DeltaToolCall>>,
|
tool_calls: Option<Vec<DeltaToolCall>>,
|
||||||
reasoning_details: Option<Vec<Value>>,
|
reasoning_details: Option<Vec<Value>>,
|
||||||
@@ -74,6 +90,32 @@ struct StreamingChunk {
|
|||||||
model: Option<String>,
|
model: Option<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn extract_content_and_signature(
|
||||||
|
delta_content: Option<&DeltaContent>,
|
||||||
|
) -> (Option<String>, Option<String>) {
|
||||||
|
match delta_content {
|
||||||
|
Some(DeltaContent::String(s)) => (Some(s.clone()), None),
|
||||||
|
Some(DeltaContent::Array(parts)) => {
|
||||||
|
let text_parts: Vec<_> = parts.iter().filter(|p| p.r#type == "text").collect();
|
||||||
|
|
||||||
|
let text = text_parts
|
||||||
|
.iter()
|
||||||
|
.map(|p| p.text.as_str())
|
||||||
|
.collect::<String>();
|
||||||
|
|
||||||
|
let signature = text_parts
|
||||||
|
.iter()
|
||||||
|
.find_map(|p| p.thought_signature.as_ref())
|
||||||
|
.cloned();
|
||||||
|
|
||||||
|
let text = if text.is_empty() { None } else { Some(text) };
|
||||||
|
|
||||||
|
(text, signature)
|
||||||
|
}
|
||||||
|
None => (None, None),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
pub fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec<Value> {
|
pub fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec<Value> {
|
||||||
let mut messages_spec = Vec::new();
|
let mut messages_spec = Vec::new();
|
||||||
for message in messages {
|
for message in messages {
|
||||||
@@ -564,6 +606,7 @@ where
|
|||||||
|
|
||||||
let mut accumulated_reasoning: Vec<Value> = Vec::new();
|
let mut accumulated_reasoning: Vec<Value> = Vec::new();
|
||||||
let mut accumulated_reasoning_content = String::new();
|
let mut accumulated_reasoning_content = String::new();
|
||||||
|
let mut last_signature: Option<String> = None;
|
||||||
|
|
||||||
'outer: while let Some(response) = stream.next().await {
|
'outer: while let Some(response) = stream.next().await {
|
||||||
let response_str = response?;
|
let response_str = response?;
|
||||||
@@ -685,14 +728,23 @@ where
|
|||||||
serde_json::from_str::<Value>(arguments)
|
serde_json::from_str::<Value>(arguments)
|
||||||
};
|
};
|
||||||
|
|
||||||
let metadata = extra_fields.as_ref().filter(|m| !m.is_empty());
|
let metadata = if let Some(sig) = &last_signature {
|
||||||
|
let mut combined = extra_fields.clone().unwrap_or_default();
|
||||||
|
combined.insert(
|
||||||
|
crate::providers::formats::google::THOUGHT_SIGNATURE_KEY.to_string(),
|
||||||
|
json!(sig)
|
||||||
|
);
|
||||||
|
Some(combined)
|
||||||
|
} else {
|
||||||
|
extra_fields.as_ref().filter(|m| !m.is_empty()).cloned()
|
||||||
|
};
|
||||||
|
|
||||||
let content = match parsed {
|
let content = match parsed {
|
||||||
Ok(params) => {
|
Ok(params) => {
|
||||||
MessageContent::tool_request_with_metadata(
|
MessageContent::tool_request_with_metadata(
|
||||||
id.clone(),
|
id.clone(),
|
||||||
Ok(CallToolRequestParams::new(function_name.clone()).with_arguments(object(params))),
|
Ok(CallToolRequestParams::new(function_name.clone()).with_arguments(object(params))),
|
||||||
metadata,
|
metadata.as_ref(),
|
||||||
)
|
)
|
||||||
},
|
},
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
@@ -704,7 +756,7 @@ where
|
|||||||
)),
|
)),
|
||||||
data: None,
|
data: None,
|
||||||
};
|
};
|
||||||
MessageContent::tool_request_with_metadata(id.clone(), Err(error), metadata)
|
MessageContent::tool_request_with_metadata(id.clone(), Err(error), metadata.as_ref())
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
contents.push(content);
|
contents.push(content);
|
||||||
@@ -731,13 +783,20 @@ where
|
|||||||
|
|
||||||
if let Some(reasoning) = &chunk.choices[0].delta.reasoning_content {
|
if let Some(reasoning) = &chunk.choices[0].delta.reasoning_content {
|
||||||
if !reasoning.is_empty() {
|
if !reasoning.is_empty() {
|
||||||
content.push(MessageContent::thinking(reasoning, ""));
|
let signature = last_signature.as_deref().unwrap_or("");
|
||||||
|
content.push(MessageContent::thinking(reasoning, signature));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
if let Some(text) = &chunk.choices[0].delta.content {
|
let (text_content, thought_signature) = extract_content_and_signature(chunk.choices[0].delta.content.as_ref());
|
||||||
|
|
||||||
|
if let Some(sig) = thought_signature {
|
||||||
|
last_signature = Some(sig);
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(text) = text_content {
|
||||||
if !text.is_empty() {
|
if !text.is_empty() {
|
||||||
content.push(MessageContent::text(text));
|
content.push(MessageContent::text(&text));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -748,7 +807,6 @@ where
|
|||||||
content,
|
content,
|
||||||
);
|
);
|
||||||
|
|
||||||
// Add ID if present
|
|
||||||
if let Some(id) = chunk.id {
|
if let Some(id) = chunk.id {
|
||||||
msg = msg.with_id(id);
|
msg = msg.with_id(id);
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user