From 49f99c3d696135dd9d68ea3f380c46a92587f1f4 Mon Sep 17 00:00:00 2001 From: Alex Hancock Date: Fri, 20 Mar 2026 14:16:06 -0400 Subject: [PATCH 1/2] fix: gemini models via databricks --- crates/goose/src/providers/databricks.rs | 27 +++++++---- crates/goose/src/providers/formats/openai.rs | 51 ++++++++++++++++++-- 2 files changed, 64 insertions(+), 14 deletions(-) diff --git a/crates/goose/src/providers/databricks.rs b/crates/goose/src/providers/databricks.rs index 9428be4a8bf5..5264ab588ab3 100644 --- a/crates/goose/src/providers/databricks.rs +++ b/crates/goose/src/providers/databricks.rs @@ -257,6 +257,11 @@ impl DatabricksProvider { normalized.contains("codex") } + fn supports_stream_options(model_name: &str) -> bool { + let normalized = model_name.to_ascii_lowercase(); + !normalized.contains("gemini") + } + fn get_endpoint_path(&self, model_name: &str, is_embedding: bool) -> String { if is_embedding { "serving-endpoints/text-embedding-3-small/invocations".to_string() @@ -389,16 +394,18 @@ impl Provider for DatabricksProvider { .unwrap() .insert("stream".to_string(), Value::Bool(true)); - if let Some(opts) = payload - .get_mut("stream_options") - .and_then(|v| v.as_object_mut()) - { - opts.entry("include_usage").or_insert(json!(true)); - } else { - payload - .as_object_mut() - .unwrap() - .insert("stream_options".to_string(), json!({"include_usage": true})); + if Self::supports_stream_options(&model_config.model_name) { + if let Some(opts) = payload + .get_mut("stream_options") + .and_then(|v| v.as_object_mut()) + { + opts.entry("include_usage").or_insert(json!(true)); + } else { + payload + .as_object_mut() + .unwrap() + .insert("stream_options".to_string(), json!({"include_usage": true})); + } } let mut log = RequestLog::start(model_config, &payload)?; diff --git a/crates/goose/src/providers/formats/openai.rs b/crates/goose/src/providers/formats/openai.rs index 3d84ce3272a8..8c23d56dfda6 100644 --- a/crates/goose/src/providers/formats/openai.rs +++ b/crates/goose/src/providers/formats/openai.rs @@ -48,9 +48,25 @@ struct DeltaToolCall { extra: Option>, } +#[derive(Serialize, Deserialize, Debug)] +#[serde(untagged)] +enum DeltaContent { + String(String), + Array(Vec), +} + +#[derive(Serialize, Deserialize, Debug)] +struct ContentPart { + r#type: String, + text: String, + #[serde(rename = "thoughtSignature")] + thought_signature: Option, +} + #[derive(Serialize, Deserialize, Debug)] struct Delta { - content: Option, + #[serde(default)] + content: Option, role: Option, tool_calls: Option>, reasoning_details: Option>, @@ -74,6 +90,32 @@ struct StreamingChunk { model: Option, } +fn extract_content_and_signature( + delta_content: Option<&DeltaContent>, +) -> (Option, Option) { + 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::(); + + 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 { let mut messages_spec = Vec::new(); for message in messages { @@ -735,9 +777,11 @@ where } } - 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(text) = text_content { if !text.is_empty() { - content.push(MessageContent::text(text)); + content.push(MessageContent::text(&text)); } } @@ -748,7 +792,6 @@ where content, ); - // Add ID if present if let Some(id) = chunk.id { msg = msg.with_id(id); } From 852929dc53388bfc2aae8b59f0a943a6b2774544 Mon Sep 17 00:00:00 2001 From: Alex Hancock Date: Fri, 20 Mar 2026 15:55:14 -0400 Subject: [PATCH 2/2] fix: reacting to PR feedback --- crates/goose/src/providers/databricks.rs | 61 +++++++++++++------- crates/goose/src/providers/formats/openai.rs | 25 ++++++-- 2 files changed, 59 insertions(+), 27 deletions(-) diff --git a/crates/goose/src/providers/databricks.rs b/crates/goose/src/providers/databricks.rs index 5264ab588ab3..5b68accde867 100644 --- a/crates/goose/src/providers/databricks.rs +++ b/crates/goose/src/providers/databricks.rs @@ -257,11 +257,6 @@ impl DatabricksProvider { normalized.contains("codex") } - fn supports_stream_options(model_name: &str) -> bool { - let normalized = model_name.to_ascii_lowercase(); - !normalized.contains("gemini") - } - fn get_endpoint_path(&self, model_name: &str, is_embedding: bool) -> String { if is_embedding { "serving-endpoints/text-embedding-3-small/invocations".to_string() @@ -394,18 +389,16 @@ impl Provider for DatabricksProvider { .unwrap() .insert("stream".to_string(), Value::Bool(true)); - if Self::supports_stream_options(&model_config.model_name) { - if let Some(opts) = payload - .get_mut("stream_options") - .and_then(|v| v.as_object_mut()) - { - opts.entry("include_usage").or_insert(json!(true)); - } else { - payload - .as_object_mut() - .unwrap() - .insert("stream_options".to_string(), json!({"include_usage": true})); - } + if let Some(opts) = payload + .get_mut("stream_options") + .and_then(|v| v.as_object_mut()) + { + opts.entry("include_usage").or_insert(json!(true)); + } else { + payload + .as_object_mut() + .unwrap() + .insert("stream_options".to_string(), json!({"include_usage": true})); } let mut log = RequestLog::start(model_config, &payload)?; @@ -419,16 +412,40 @@ impl Provider for DatabricksProvider { let status = resp.status(); 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::(&error_text).ok(); return Err(map_http_error_to_provider_error(status, json_payload)); } Ok(resp) }) - .await - .inspect_err(|e| { - let _ = log.error(e); - })?; + .await; + + 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::(&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) } diff --git a/crates/goose/src/providers/formats/openai.rs b/crates/goose/src/providers/formats/openai.rs index 8c23d56dfda6..db733fc4dfc3 100644 --- a/crates/goose/src/providers/formats/openai.rs +++ b/crates/goose/src/providers/formats/openai.rs @@ -606,6 +606,7 @@ where let mut accumulated_reasoning: Vec = Vec::new(); let mut accumulated_reasoning_content = String::new(); + let mut last_signature: Option = None; 'outer: while let Some(response) = stream.next().await { let response_str = response?; @@ -727,14 +728,23 @@ where serde_json::from_str::(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 { Ok(params) => { MessageContent::tool_request_with_metadata( id.clone(), Ok(CallToolRequestParams::new(function_name.clone()).with_arguments(object(params))), - metadata, + metadata.as_ref(), ) }, Err(e) => { @@ -746,7 +756,7 @@ where )), 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); @@ -773,11 +783,16 @@ where if let Some(reasoning) = &chunk.choices[0].delta.reasoning_content { if !reasoning.is_empty() { - content.push(MessageContent::thinking(reasoning, "")); + let signature = last_signature.as_deref().unwrap_or(""); + content.push(MessageContent::thinking(reasoning, signature)); } } - let (text_content, _thought_signature) = extract_content_and_signature(chunk.choices[0].delta.content.as_ref()); + 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() {