From 5aa61e83604c6b0d8ac1807ea811d10e109a3498 Mon Sep 17 00:00:00 2001 From: centdix <40307056+centdix@users.noreply.github.com> Date: Mon, 26 Jan 2026 14:15:18 +0100 Subject: [PATCH] fix(aiagent): fix usage for gemini 3 models (#7682) * fix for gemini3 models * cleaning * cleaning * small opti --- backend/windmill-api/src/bedrock.rs | 14 ++++++----- backend/windmill-common/src/ai_bedrock.rs | 7 +++--- .../src/ai/providers/google_ai.rs | 10 ++++++++ .../src/ai/providers/openai.rs | 4 +--- .../windmill-worker/src/ai/query_builder.rs | 4 +++- backend/windmill-worker/src/ai/sse.rs | 24 ++++++++++++------- backend/windmill-worker/src/ai/tools.rs | 2 +- 7 files changed, 41 insertions(+), 24 deletions(-) diff --git a/backend/windmill-api/src/bedrock.rs b/backend/windmill-api/src/bedrock.rs index 8965c8c377..4ed1f2cdfa 100644 --- a/backend/windmill-api/src/bedrock.rs +++ b/backend/windmill-api/src/bedrock.rs @@ -21,9 +21,11 @@ use serde::Deserialize; use windmill_common::ai_bedrock::{ bedrock_stream_event_is_block_stop, bedrock_stream_event_to_text, bedrock_stream_event_to_tool_delta, bedrock_stream_event_to_tool_start, format_bedrock_error, - BedrockClient, OpenAIMessage, OpenAIToolCall, + BedrockClient }; use windmill_common::error::{Error, Result}; +use windmill_common::ai_types::{OpenAIFunction, OpenAIMessage, OpenAIToolCall, ToolDef, ToolDefFunction}; +use windmill_common::ai_bedrock::build_tool_config; // ============================================================================ // Shared Request Types for SDK-Based Handlers @@ -122,11 +124,11 @@ fn build_tool_config_from_request( tool_choice: Option<&serde_json::Value>, ) -> Result> { if let Some(tools) = tools { - let tool_defs: Vec = tools + let tool_defs: Vec = tools .iter() - .map(|t| windmill_common::ai_bedrock::ToolDef { + .map(|t| ToolDef { r#type: "function".to_string(), - function: windmill_common::ai_bedrock::ToolDefFunction { + function: ToolDefFunction { name: t.function.name.clone(), description: t.function.description.clone(), parameters: Box::from( @@ -152,7 +154,7 @@ fn build_tool_config_from_request( .map(|tc| tc == "required" || tc.as_str() == Some("required")) .unwrap_or(false); - windmill_common::ai_bedrock::build_tool_config(Some(&tool_defs), force_tool_use) + build_tool_config(Some(&tool_defs), force_tool_use) } else { Ok(None) } @@ -668,7 +670,7 @@ pub async fn handle_bedrock_sdk_non_streaming( let input_json = document_to_json(tool_use.input()); tool_calls.push(OpenAIToolCall { id: tool_use.tool_use_id().to_string(), - function: windmill_common::ai_bedrock::OpenAIFunction { + function: OpenAIFunction { name: tool_use.name().to_string(), arguments: serde_json::to_string(&input_json).unwrap_or_default(), }, diff --git a/backend/windmill-common/src/ai_bedrock.rs b/backend/windmill-common/src/ai_bedrock.rs index b7fcdc3aa7..5a1ba8837c 100644 --- a/backend/windmill-common/src/ai_bedrock.rs +++ b/backend/windmill-common/src/ai_bedrock.rs @@ -19,10 +19,9 @@ use serde::{Deserialize, Serialize}; use crate::error::Error; -// Re-export types from ai_types -pub use crate::ai_types::{ - ContentPart, ExtraContent, ImageUrlData, OpenAIContent, OpenAIFunction, OpenAIMessage, - OpenAIToolCall, ToolDef, ToolDefFunction, UrlCitation, +use crate::ai_types::{ + ContentPart, OpenAIContent, OpenAIFunction, OpenAIMessage, + OpenAIToolCall, ToolDef, }; // ============================================================================ diff --git a/backend/windmill-worker/src/ai/providers/google_ai.rs b/backend/windmill-worker/src/ai/providers/google_ai.rs index b2c31f9e35..db17194a40 100644 --- a/backend/windmill-worker/src/ai/providers/google_ai.rs +++ b/backend/windmill-worker/src/ai/providers/google_ai.rs @@ -37,6 +37,9 @@ pub enum GeminiPart { FunctionCall { #[serde(rename = "functionCall")] function_call: GeminiFunctionCall, + /// Thought signature for Gemini 3+ models - required for function calling + #[serde(rename = "thoughtSignature", skip_serializing_if = "Option::is_none")] + thought_signature: Option, }, FunctionResponse { #[serde(rename = "functionResponse")] @@ -369,11 +372,18 @@ impl GoogleAIQueryBuilder { for tc in tool_calls { let args: serde_json::Value = serde_json::from_str(&tc.function.arguments).unwrap_or_default(); + // Extract thought_signature from extra_content if present + let thought_signature = tc + .extra_content + .as_ref() + .and_then(|ec| ec.google.as_ref()) + .and_then(|g| g.thought_signature.clone()); parts.push(GeminiPart::FunctionCall { function_call: GeminiFunctionCall { name: tc.function.name.clone(), args, }, + thought_signature, }); } } diff --git a/backend/windmill-worker/src/ai/providers/openai.rs b/backend/windmill-worker/src/ai/providers/openai.rs index e27d752871..d9f66cd2d4 100644 --- a/backend/windmill-worker/src/ai/providers/openai.rs +++ b/backend/windmill-worker/src/ai/providers/openai.rs @@ -11,9 +11,7 @@ use crate::ai::{ utils::extract_text_content, }; -pub use windmill_common::ai_types::{ - ExtraContent, OpenAIFunction, OpenAIToolCall, -}; +use windmill_common::ai_types::OpenAIToolCall; // Responses API structures #[derive(Deserialize)] diff --git a/backend/windmill-worker/src/ai/query_builder.rs b/backend/windmill-worker/src/ai/query_builder.rs index a33177f74c..9fb2a9f548 100644 --- a/backend/windmill-worker/src/ai/query_builder.rs +++ b/backend/windmill-worker/src/ai/query_builder.rs @@ -9,7 +9,7 @@ use crate::{ providers::{ anthropic::AnthropicQueryBuilder, google_ai::GoogleAIQueryBuilder, - openai::{OpenAIQueryBuilder, OpenAIToolCall}, + openai::{OpenAIQueryBuilder}, openrouter::OpenRouterQueryBuilder, other::OtherQueryBuilder, }, @@ -18,6 +18,8 @@ use crate::{ job_logger::append_result_stream, }; +use windmill_common::ai_types::OpenAIToolCall; + /// Arguments for building an AI request pub struct BuildRequestArgs<'a> { pub messages: &'a [OpenAIMessage], diff --git a/backend/windmill-worker/src/ai/sse.rs b/backend/windmill-worker/src/ai/sse.rs index cbc4da9580..e59898d2fa 100644 --- a/backend/windmill-worker/src/ai/sse.rs +++ b/backend/windmill-worker/src/ai/sse.rs @@ -8,11 +8,12 @@ use tokio_stream::StreamExt; use windmill_common::{error::Error, utils::rd_string}; use crate::ai::{ - providers::openai::{ExtraContent, OpenAIFunction, OpenAIToolCall}, query_builder::StreamEventProcessor, types::{StreamingEvent, UrlCitation}, }; +use windmill_common::ai_types::{ExtraContent, GoogleExtraContent, OpenAIFunction, OpenAIToolCall}; + #[derive(Deserialize)] pub struct OpenAIChoiceDeltaToolCallFunction { pub name: Option, @@ -24,8 +25,6 @@ pub struct OpenAIChoiceDeltaToolCall { pub index: Option, pub id: Option, pub function: Option, - /// Extra content for provider-specific metadata (e.g., Google Gemini thought signatures) - pub extra_content: Option, } #[derive(Deserialize)] @@ -140,10 +139,6 @@ impl SSEParser for OpenAISSEParser { if let Some(arguments) = function.arguments { existing_tool_call.function.arguments += &arguments; } - // Update extra_content if provided in this delta (for thought signatures) - if let Some(extra) = tool_call.extra_content { - existing_tool_call.extra_content = Some(extra); - } } else { let fun_name = function.name.unwrap_or_default(); let call_id = tool_call.id.unwrap_or_else(|| rd_string(24)); @@ -163,7 +158,7 @@ impl SSEParser for OpenAISSEParser { arguments: function.arguments.unwrap_or_default(), }, r#type: "function".to_string(), - extra_content: tool_call.extra_content, + extra_content: None, }, ); } @@ -425,6 +420,9 @@ pub struct GeminiSSEPart { pub text: Option, #[serde(rename = "functionCall")] pub function_call: Option, + /// Thought signature for Gemini 3+ models - required for function calling + #[serde(rename = "thoughtSignature")] + pub thought_signature: Option, } /// Function call in Gemini streaming response @@ -547,6 +545,14 @@ impl SSEParser for GeminiSSEParser { .send(event, &mut self.events_str) .await?; + // Build extra_content with thought_signature if present + let extra_content = + part.thought_signature.map(|sig| ExtraContent { + google: Some(GoogleExtraContent { + thought_signature: Some(sig), + }), + }); + // Store accumulated tool call self.accumulated_tool_calls.insert( idx, @@ -560,7 +566,7 @@ impl SSEParser for GeminiSSEParser { .unwrap_or_else(|_| "{}".to_string()), }, r#type: "function".to_string(), - extra_content: None, + extra_content, }, ); } diff --git a/backend/windmill-worker/src/ai/tools.rs b/backend/windmill-worker/src/ai/tools.rs index 584be8fe2b..e93d1bb782 100644 --- a/backend/windmill-worker/src/ai/tools.rs +++ b/backend/windmill-worker/src/ai/tools.rs @@ -1,4 +1,4 @@ -use crate::ai::providers::openai::OpenAIToolCall; +use windmill_common::ai_types::OpenAIToolCall; use crate::ai::query_builder::StreamEventProcessor; use crate::ai::types::*; use crate::ai::utils::{