| |
| use super::models::*; |
| use crate::proxy::model_specs; |
| use crate::proxy::token_manager::ProxyToken; |
|
|
| use serde_json::{json, Value}; |
|
|
| pub fn transform_openai_request( |
| request: &OpenAIRequest, |
| project_id: &str, |
| mapped_model: &str, |
| token: Option<&ProxyToken>, |
| ) -> (Value, String, usize) { |
| let session_id = crate::proxy::session_manager::SessionManager::extract_openai_session_id(request); |
| let message_count = request.messages.len(); |
| |
| let tools_val = request |
| .tools |
| .as_ref() |
| .map(|list| list.iter().map(|v| v.clone()).collect::<Vec<_>>()); |
|
|
| let mapped_model_lower = mapped_model.to_lowercase(); |
|
|
| |
| let config = crate::proxy::mappers::common_utils::resolve_request_config( |
| &request.model, |
| &mapped_model_lower, |
| &tools_val, |
| request.size.as_deref(), |
| request.quality.as_deref(), |
| request.image_size.as_deref(), |
| None, |
| ); |
|
|
| |
| |
| |
| |
| let is_gemini_3_thinking = mapped_model_lower.contains("gemini") |
| && ( |
| mapped_model_lower.contains("-thinking") |
| || mapped_model_lower.contains("gemini-2.0-pro") |
| || mapped_model_lower.contains("gemini-3-pro") |
| || mapped_model_lower.contains("gemini-3.1-pro") |
| ) |
| && !mapped_model_lower.contains("claude"); |
| |
| |
| let is_gemini_flash_thinking = (mapped_model_lower.contains("gemini-3-flash") |
| || mapped_model_lower.contains("gemini-3.1-flash")) |
| && !mapped_model_lower.contains("claude"); |
| let is_claude_thinking = mapped_model_lower.ends_with("-thinking"); |
| let is_thinking_model = is_gemini_3_thinking || is_claude_thinking || is_gemini_flash_thinking; |
|
|
|
|
| |
| let user_enabled_thinking = request.thinking.as_ref() |
| .map(|t| t.thinking_type.as_deref() == Some("enabled")) |
| .unwrap_or(false); |
| let user_thinking_budget = request.thinking.as_ref() |
| .and_then(|t| t.budget_tokens); |
|
|
| |
| let has_incompatible_assistant_history = request.messages.iter().any(|msg| { |
| msg.role == "assistant" |
| && msg |
| .reasoning_content |
| .as_ref() |
| .map(|s| s.is_empty()) |
| .unwrap_or(true) |
| }); |
| let has_tool_history = request.messages.iter().any(|msg| { |
| msg.role == "tool" || msg.role == "function" || msg.tool_calls.is_some() |
| }); |
|
|
|
|
|
|
| |
| |
| |
| |
| let mut actual_include_thinking = is_thinking_model || user_enabled_thinking; |
| |
| |
| let session_thought_sig = crate::proxy::SignatureCache::global().get_session_signature(&session_id); |
| |
| if is_claude_thinking && has_incompatible_assistant_history && session_thought_sig.is_none() { |
| tracing::warn!("[OpenAI-Thinking] Incompatible assistant history detected for Claude thinking model without session signature. Disabling thinking for this request to avoid 400 error. (sid: {})", session_id); |
| actual_include_thinking = false; |
| } |
| |
| |
| if user_enabled_thinking { |
| tracing::info!( |
| "[OpenAI-Thinking] User explicitly enabled thinking with budget: {:?}", |
| user_thinking_budget |
| ); |
| } |
|
|
| tracing::debug!( |
| "[Debug] OpenAI Request: original='{}', mapped='{}', type='{}', has_image_config={}", |
| request.model, |
| mapped_model, |
| config.request_type, |
| config.image_config.is_some() |
| ); |
|
|
| |
| let mut system_instructions: Vec<String> = request |
| .messages |
| .iter() |
| .filter(|msg| msg.role == "system" || msg.role == "developer") |
| .filter_map(|msg| { |
| msg.content.as_ref().map(|c| match c { |
| OpenAIContent::String(s) => s.clone(), |
| OpenAIContent::Array(blocks) => blocks |
| .iter() |
| .filter_map(|b| { |
| if let OpenAIContentBlock::Text { text } = b { |
| Some(text.clone()) |
| } else { |
| None |
| } |
| }) |
| .collect::<Vec<_>>() |
| .join("\n"), |
| }) |
| }) |
| .collect(); |
|
|
| |
| if let Some(inst) = &request.instructions { |
| if !inst.is_empty() { |
| system_instructions.insert(0, inst.clone()); |
| } |
| } |
|
|
| |
| let mut tool_id_to_name = std::collections::HashMap::new(); |
| for msg in &request.messages { |
| if let Some(tool_calls) = &msg.tool_calls { |
| for call in tool_calls { |
| let name = &call.function.name; |
| let final_name = if name == "local_shell_call" { |
| "shell" |
| } else { |
| name |
| }; |
| tool_id_to_name.insert(call.id.clone(), final_name.to_string()); |
| } |
| } |
| } |
|
|
| |
| let thought_sig = session_thought_sig; |
| if thought_sig.is_some() { |
| tracing::debug!( |
| "[OpenAI-Request] Using session signature (sid: {}, len: {})", |
| session_id, |
| thought_sig.as_ref().unwrap().len() |
| ); |
| } |
|
|
| |
| let mut tool_name_to_schema = std::collections::HashMap::new(); |
| if let Some(tools) = &request.tools { |
| for tool in tools { |
| if let (Some(name), Some(params)) = ( |
| tool.get("function") |
| .and_then(|f| f.get("name")) |
| .and_then(|v| v.as_str()), |
| tool.get("function").and_then(|f| f.get("parameters")), |
| ) { |
| tool_name_to_schema.insert(name.to_string(), params.clone()); |
| } else if let (Some(name), Some(params)) = ( |
| tool.get("name").and_then(|v| v.as_str()), |
| tool.get("parameters"), |
| ) { |
| |
| tool_name_to_schema.insert(name.to_string(), params.clone()); |
| } |
| } |
| } |
|
|
| |
| let contents: Vec<Value> = request |
| .messages |
| .iter() |
| .filter(|msg| msg.role != "system" && msg.role != "developer") |
| .map(|msg| { |
| let role = match msg.role.as_str() { |
| "assistant" => "model", |
| "tool" | "function" => "user", |
| _ => &msg.role, |
| }; |
|
|
| let mut parts = Vec::new(); |
|
|
| |
| if let Some(reasoning) = &msg.reasoning_content { |
| |
| let is_invalid_placeholder = reasoning == "[undefined]" || reasoning.is_empty(); |
| |
| if !is_invalid_placeholder { |
| let thought_part = json!({ |
| "text": reasoning, |
| "thought": true, |
| }); |
| parts.push(thought_part); |
| } |
| } else if actual_include_thinking && role == "model" { |
| |
| |
| |
| tracing::debug!("[OpenAI-Thinking] Injecting placeholder thinking block for assistant message"); |
| let mut thought_part = json!({ |
| "text": "Applying tool decisions and generating response...", |
| "thought": true, |
| }); |
| |
| |
| |
| if is_gemini_3_thinking { |
| thought_part["thoughtSignature"] = json!("skip_thought_signature_validator"); |
| } |
| |
| parts.push(thought_part); |
| } |
|
|
| |
| |
| |
| let is_tool_role = msg.role == "tool" || msg.role == "function"; |
| if let (Some(content), false) = (&msg.content, is_tool_role) { |
| match content { |
| OpenAIContent::String(s) => { |
| if !s.is_empty() { |
| parts.push(json!({"text": s})); |
| } |
| } |
| OpenAIContent::Array(blocks) => { |
| for block in blocks { |
| match block { |
| OpenAIContentBlock::Text { text } => { |
| parts.push(json!({"text": text})); |
| } |
| OpenAIContentBlock::ImageUrl { image_url } => { |
| if image_url.url.starts_with("data:") { |
| if let Some(pos) = image_url.url.find(",") { |
| let mime_part = &image_url.url[5..pos]; |
| let mime_type = mime_part.split(';').next().unwrap_or("image/jpeg"); |
| let data = &image_url.url[pos + 1..]; |
| |
| parts.push(json!({ |
| "inlineData": { "mimeType": mime_type, "data": data } |
| })); |
| } |
| } else if image_url.url.starts_with("http") { |
| parts.push(json!({ |
| "fileData": { "fileUri": &image_url.url, "mimeType": "image/jpeg" } |
| })); |
| } else { |
| |
| let file_path = if image_url.url.starts_with("file://") { |
| |
| #[cfg(target_os = "windows")] |
| { image_url.url.trim_start_matches("file:///").replace('/', "\\") } |
| #[cfg(not(target_os = "windows"))] |
| { image_url.url.trim_start_matches("file://").to_string() } |
| } else { |
| image_url.url.clone() |
| }; |
| |
| tracing::debug!("[OpenAI-Request] Reading local image: {}", file_path); |
| |
| |
| if let Ok(file_bytes) = std::fs::read(&file_path) { |
| use base64::Engine as _; |
| let b64 = base64::engine::general_purpose::STANDARD.encode(&file_bytes); |
| |
| |
| let mime_type = if file_path.to_lowercase().ends_with(".png") { |
| "image/png" |
| } else if file_path.to_lowercase().ends_with(".gif") { |
| "image/gif" |
| } else if file_path.to_lowercase().ends_with(".webp") { |
| "image/webp" |
| } else { |
| "image/jpeg" |
| }; |
| |
| parts.push(json!({ |
| "inlineData": { "mimeType": mime_type, "data": b64 } |
| })); |
| tracing::debug!("[OpenAI-Request] Successfully loaded image: {} ({} bytes)", file_path, file_bytes.len()); |
| } else { |
| tracing::debug!("[OpenAI-Request] Failed to read local image: {}", file_path); |
| } |
| } |
| } |
| OpenAIContentBlock::AudioUrl { audio_url: _ } => { |
| |
| |
| |
| tracing::debug!("[OpenAI-Request] Skipping audio_url (not yet implemented in v3.3.16)"); |
| } |
| } |
| } |
| } |
| } |
| } |
|
|
| |
| if let Some(tool_calls) = &msg.tool_calls { |
| for (_index, tc) in tool_calls.iter().enumerate() { |
| |
| |
| |
| |
| |
| |
| |
|
|
|
|
| let mut args = serde_json::from_str::<Value>(&tc.function.arguments).unwrap_or(json!({})); |
| |
| |
| if let Some(original_schema) = tool_name_to_schema.get(&tc.function.name) { |
| crate::proxy::common::json_schema::fix_tool_call_args(&mut args, original_schema); |
| } |
|
|
| let mut func_call_part = json!({ |
| "functionCall": { |
| "name": if tc.function.name == "local_shell_call" { "shell" } else { &tc.function.name }, |
| "args": args, |
| "id": &tc.id, |
| } |
| }); |
|
|
| |
| crate::proxy::common::json_schema::clean_json_schema(&mut func_call_part); |
|
|
| if let Some(ref sig) = thought_sig { |
| func_call_part["thoughtSignature"] = json!(sig); |
| } else if is_thinking_model || is_gemini_flash_thinking { |
| |
| |
| |
| tracing::debug!("[OpenAI-Signature] Adding GEMINI_SKIP_SIGNATURE for tool_use: {}", tc.id); |
| func_call_part["thoughtSignature"] = json!("skip_thought_signature_validator"); |
| } |
|
|
| parts.push(func_call_part); |
| } |
| } |
|
|
| |
| if msg.role == "tool" || msg.role == "function" { |
| let name = msg.name.as_deref().unwrap_or("unknown"); |
| let final_name = if name == "local_shell_call" { "shell" } |
| else if let Some(id) = &msg.tool_call_id { tool_id_to_name.get(id).map(|s| s.as_str()).unwrap_or(name) } |
| else { name }; |
|
|
| let mut extra_parts = Vec::new(); |
|
|
| let content_val = match &msg.content { |
| Some(OpenAIContent::String(s)) => s.clone(), |
| Some(OpenAIContent::Array(blocks)) => { |
| let mut texts = Vec::new(); |
| for block in blocks { |
| match block { |
| OpenAIContentBlock::Text { text } => texts.push(text.clone()), |
| OpenAIContentBlock::ImageUrl { image_url } => { |
| if image_url.url.starts_with("data:") { |
| if let Some(pos) = image_url.url.find(',') { |
| let mime_part = &image_url.url[5..pos]; |
| let mime_type = mime_part.split(';').next().unwrap_or("image/jpeg"); |
| let data = &image_url.url[pos + 1..]; |
| |
| extra_parts.push(json!({ |
| "inlineData": { "mimeType": mime_type, "data": data } |
| })); |
| } |
| } else { |
| texts.push("[image link]".to_string()); |
| } |
| } |
| _ => {} |
| } |
| } |
| texts.join("\n") |
| }, |
| None => "".to_string() |
| }; |
|
|
| parts.push(json!({ |
| "functionResponse": { |
| "name": final_name, |
| "response": { "result": content_val }, |
| "id": msg.tool_call_id.clone().unwrap_or_default() |
| } |
| })); |
|
|
| for extra in extra_parts { |
| parts.push(extra); |
| } |
| } |
|
|
| json!({ "role": role, "parts": parts }) |
| }) |
| .filter(|msg| !msg["parts"].as_array().map(|a| a.is_empty()).unwrap_or(true)) |
| .collect(); |
|
|
| |
| |
| let mut contents = contents; |
| if actual_include_thinking && has_tool_history { |
| tracing::debug!("[OpenAI-Thinking] Applied thinking recovery (stripping old thought blocks) for tool history"); |
| contents = super::thinking_recovery::strip_all_thinking_blocks(contents); |
| } |
|
|
| |
| let mut merged_contents: Vec<Value> = Vec::new(); |
| for msg in contents { |
| if let Some(last) = merged_contents.last_mut() { |
| if last["role"] == msg["role"] { |
| |
| if let (Some(last_parts), Some(msg_parts)) = |
| (last["parts"].as_array_mut(), msg["parts"].as_array()) |
| { |
| last_parts.extend(msg_parts.iter().cloned()); |
| continue; |
| } |
| } |
| } |
| merged_contents.push(msg); |
| } |
| let contents = merged_contents; |
|
|
| |
|
|
| let mut gen_config = json!({ |
| "temperature": request.temperature.unwrap_or(1.0), |
| |
| "topP": request.top_p.unwrap_or(1.0), |
| |
| "topK": 40, |
| }); |
|
|
| |
| if let Some(max_tokens) = request.max_tokens { |
| gen_config["maxOutputTokens"] = json!(max_tokens); |
| } else { |
| |
| let limit = model_specs::get_max_output_tokens(mapped_model, token); |
| gen_config["maxOutputTokens"] = json!(limit); |
| } |
|
|
| |
| if let Some(n) = request.n { |
| gen_config["candidateCount"] = json!(n); |
| } |
|
|
| |
| if actual_include_thinking { |
| |
| let image_thinking_mode = crate::proxy::config::get_image_thinking_mode(); |
| |
| let is_image_gen_disabled = config.request_type == "image_gen" && image_thinking_mode == "disabled"; |
|
|
| if is_image_gen_disabled { |
| tracing::debug!("[OpenAI-Request] Image thinking mode disabled: enforcing includeThoughts=false for {}", mapped_model); |
| gen_config["thinkingConfig"] = json!({ |
| "includeThoughts": false |
| }); |
| } else { |
| |
| let tb_config = crate::proxy::config::get_thinking_budget_config(); |
| |
| let default_budget = model_specs::get_thinking_budget(mapped_model, token); |
| let user_budget: i64 = user_thinking_budget.map(|b| b as i64).unwrap_or(default_budget as i64); |
| |
| let budget = match tb_config.mode { |
| crate::proxy::config::ThinkingBudgetMode::Passthrough => { |
| user_budget |
| } |
| crate::proxy::config::ThinkingBudgetMode::Custom => { |
| let mut custom_value = tb_config.custom_value as i64; |
| |
| if custom_value > default_budget as i64 { |
| tracing::warn!( |
| "[OpenAI-Request] Custom budget {} exceeds model spec limit {}, capping.", |
| custom_value, default_budget |
| ); |
| custom_value = default_budget as i64; |
| } |
| custom_value |
| } |
| crate::proxy::config::ThinkingBudgetMode::Auto => { |
| |
| if user_budget > default_budget as i64 { |
| default_budget as i64 |
| } else { |
| user_budget |
| } |
| } |
| crate::proxy::config::ThinkingBudgetMode::Adaptive => { |
| user_budget |
| } |
| }; |
|
|
| gen_config["thinkingConfig"] = json!({ |
| "includeThoughts": true, |
| "thinkingBudget": budget |
| }); |
|
|
| |
| |
| let overhead = if config.request_type == "image_gen" { 2048 } else { 32768 }; |
| let min_overhead = if config.request_type == "image_gen" { 1024 } else { 8192 }; |
|
|
| if let Some(max_tokens) = request.max_tokens { |
| if (max_tokens as i64) <= budget { |
| gen_config["maxOutputTokens"] = json!(budget + min_overhead); |
| } |
| } else { |
| |
| gen_config["maxOutputTokens"] = json!(budget + overhead); |
| } |
| |
| let new_max = gen_config["maxOutputTokens"].as_i64().unwrap_or(0); |
| tracing::debug!( |
| "[OpenAI-Request] Adjusted maxOutputTokens to {} for thinking model (budget={})", |
| new_max, budget |
| ); |
| |
| tracing::debug!( |
| "[OpenAI-Request] Injected thinkingConfig for model {}: thinkingBudget={} (mode={:?})", |
| mapped_model, budget, tb_config.mode |
| ); |
| } |
| } |
|
|
| if let Some(stop) = &request.stop { |
| if stop.is_string() { |
| gen_config["stopSequences"] = json!([stop]); |
| } else if stop.is_array() { |
| gen_config["stopSequences"] = stop.clone(); |
| } |
| } |
|
|
| if let Some(fmt) = &request.response_format { |
| if fmt.r#type == "json_object" { |
| gen_config["responseMimeType"] = json!("application/json"); |
| } |
| } |
|
|
| let mut inner_request = json!({ |
| "contents": contents, |
| "generationConfig": gen_config, |
| "safetySettings": [ |
| { "category": "HARM_CATEGORY_HARASSMENT", "threshold": "OFF" }, |
| { "category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "OFF" }, |
| { "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "OFF" }, |
| { "category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "OFF" }, |
| { "category": "HARM_CATEGORY_CIVIC_INTEGRITY", "threshold": "OFF" }, |
| ] |
| }); |
|
|
| |
| crate::proxy::mappers::common_utils::deep_clean_undefined(&mut inner_request, 0); |
|
|
| |
| if let Some(tools) = &request.tools { |
| let mut function_declarations: Vec<Value> = Vec::new(); |
| for tool in tools.iter() { |
| let mut gemini_func = if let Some(func) = tool.get("function") { |
| func.clone() |
| } else { |
| let mut func = tool.clone(); |
| if let Some(obj) = func.as_object_mut() { |
| obj.remove("type"); |
| obj.remove("strict"); |
| obj.remove("additionalProperties"); |
| } |
| func |
| }; |
|
|
| let name_opt = gemini_func.get("name").and_then(|v| v.as_str()).map(|s| s.to_string()); |
|
|
| if let Some(name) = &name_opt { |
| |
| if name == "web_search" |
| || name == "google_search" |
| || name == "web_search_20250305" |
| || name == "builtin_web_search" |
| { |
| continue; |
| } |
|
|
| if name == "local_shell_call" { |
| if let Some(obj) = gemini_func.as_object_mut() { |
| obj.insert("name".to_string(), json!("shell")); |
| } |
| } |
| } else { |
| |
| tracing::warn!("[OpenAI-Request] Skipping tool without name: {:?}", gemini_func); |
| continue; |
| } |
|
|
| |
| if let Some(obj) = gemini_func.as_object_mut() { |
| obj.remove("format"); |
| obj.remove("strict"); |
| obj.remove("additionalProperties"); |
| obj.remove("type"); |
| obj.remove("external_web_access"); |
| } |
|
|
| if let Some(params) = gemini_func.get_mut("parameters") { |
| |
| crate::proxy::common::json_schema::clean_json_schema(params); |
|
|
| |
| |
| |
| if let Some(params_obj) = params.as_object_mut() { |
| if !params_obj.contains_key("type") { |
| params_obj.insert("type".to_string(), json!("OBJECT")); |
| } |
| } |
|
|
| |
| enforce_uppercase_types(params); |
| } else { |
| |
| |
| tracing::debug!( |
| "[OpenAI-Request] Injecting default schema for custom tool: {}", |
| gemini_func |
| .get("name") |
| .and_then(|v| v.as_str()) |
| .unwrap_or("unknown") |
| ); |
|
|
| gemini_func.as_object_mut().unwrap().insert( |
| "parameters".to_string(), |
| json!({ |
| "type": "OBJECT", |
| "properties": { |
| "content": { |
| "type": "STRING", |
| "description": "The raw content or patch to be applied" |
| } |
| }, |
| "required": ["content"] |
| }), |
| ); |
| } |
| function_declarations.push(gemini_func); |
| } |
|
|
| if !function_declarations.is_empty() { |
| inner_request["tools"] = json!([{ "functionDeclarations": function_declarations }]); |
| |
| inner_request["toolConfig"] = json!({ |
| "functionCallingConfig": { "mode": "VALIDATED" } |
| }); |
| } |
| } |
|
|
| |
| let antigravity_identity = "You are Antigravity, a powerful agentic AI coding assistant designed by the Google Deepmind team working on Advanced Agentic Coding.\n\ |
| You are pair programming with a USER to solve their coding task. The task may require creating a new codebase, modifying or debugging an existing codebase, or simply answering a question.\n\ |
| **Absolute paths only**\n\ |
| **Proactiveness**"; |
|
|
| |
| let user_has_antigravity = system_instructions |
| .iter() |
| .any(|s| s.contains("You are Antigravity")); |
|
|
| let mut parts = Vec::new(); |
|
|
| |
| if !user_has_antigravity { |
| parts.push(json!({"text": antigravity_identity})); |
| } |
|
|
| |
| let global_prompt_config = crate::proxy::config::get_global_system_prompt(); |
| if global_prompt_config.enabled && !global_prompt_config.content.trim().is_empty() { |
| parts.push(json!({"text": global_prompt_config.content})); |
| } |
|
|
| |
| for inst in system_instructions { |
| parts.push(json!({"text": inst})); |
| } |
|
|
| inner_request["systemInstruction"] = json!({ |
| "role": "user", |
| "parts": parts |
| }); |
|
|
| if config.inject_google_search { |
| crate::proxy::mappers::common_utils::inject_google_search_tool(&mut inner_request, Some(mapped_model)); |
| } |
|
|
| if let Some(image_config) = config.image_config { |
| if let Some(obj) = inner_request.as_object_mut() { |
| obj.remove("tools"); |
| obj.remove("systemInstruction"); |
| let gen_config = obj.entry("generationConfig").or_insert_with(|| json!({})); |
| if let Some(gen_obj) = gen_config.as_object_mut() { |
| |
| |
| gen_obj.remove("responseMimeType"); |
| gen_obj.remove("responseModalities"); |
| gen_obj.insert("imageConfig".to_string(), image_config); |
| } |
| } |
| } |
|
|
| |
| if let Some(t) = token { |
| inner_request["sessionId"] = json!(crate::proxy::common::session::derive_session_id(&t.account_id)); |
| } |
|
|
| let final_body = json!({ |
| "project": project_id, |
| |
| "requestId": format!("agent/antigravity/{}/{}", &session_id[..session_id.len().min(8)], message_count), |
| "request": inner_request, |
| "model": config.final_model, |
| "userAgent": "antigravity", |
| |
| "requestType": if config.request_type == "image_gen" { "image_gen" } else { "agent" } |
| }); |
|
|
| (final_body, session_id, message_count) |
| } |
|
|
| fn enforce_uppercase_types(value: &mut Value) { |
| if let Value::Object(map) = value { |
| if let Some(type_val) = map.get_mut("type") { |
| if let Value::String(ref mut s) = type_val { |
| *s = s.to_uppercase(); |
| } |
| } |
| if let Some(properties) = map.get_mut("properties") { |
| if let Value::Object(ref mut props) = properties { |
| for v in props.values_mut() { |
| enforce_uppercase_types(v); |
| } |
| } |
| } |
| if let Some(items) = map.get_mut("items") { |
| enforce_uppercase_types(items); |
| } |
| } else if let Value::Array(arr) = value { |
| for item in arr { |
| enforce_uppercase_types(item); |
| } |
| } |
| } |
|
|
| #[cfg(test)] |
| mod tests { |
| use super::*; |
| use crate::proxy::mappers::openai::models::*; |
|
|
| #[test] |
| #[test] |
| fn test_issue_1592_gemini_3_pro_budget_capping() { |
| |
| let req = OpenAIRequest { |
| model: "gemini-3-pro".to_string(), |
| messages: vec![OpenAIMessage { |
| role: "user".to_string(), |
| content: Some(OpenAIContent::String("test".into())), |
| reasoning_content: None, |
| tool_calls: None, |
| tool_call_id: None, |
| name: None, |
| }], |
| ..Default::default() |
| }; |
|
|
| |
| let (result, _sid, _msg_count) = transform_openai_request(&req, "test-v", "gemini-3-pro", None); |
| let budget = result["request"]["generationConfig"]["thinkingConfig"]["thinkingBudget"] |
| .as_i64() |
| .unwrap(); |
| assert_eq!(budget, 24576, "Gemini-3-pro budget must be capped to 24576 in Auto mode"); |
| } |
|
|
| #[test] |
| fn test_issue_1602_custom_mode_gemini_capping() { |
| |
| use crate::proxy::config::{ThinkingBudgetConfig, ThinkingBudgetMode, update_thinking_budget_config}; |
| |
| |
| update_thinking_budget_config(ThinkingBudgetConfig { |
| mode: ThinkingBudgetMode::Custom, |
| custom_value: 32000, |
| effort: None, |
| }); |
|
|
| let req = OpenAIRequest { |
| model: "gemini-2.0-flash-thinking".to_string(), |
| messages: vec![OpenAIMessage { |
| role: "user".to_string(), |
| content: Some(OpenAIContent::String("test".into())), |
| reasoning_content: None, |
| tool_calls: None, |
| tool_call_id: None, |
| name: None, |
| }], |
| stream: false, |
| n: None, |
| max_tokens: None, |
| temperature: None, |
| top_p: None, |
| stop: None, |
| response_format: None, |
| tools: None, |
| tool_choice: None, |
| parallel_tool_calls: None, |
| ..Default::default() |
| }; |
|
|
| |
| let (result, _sid, _msg_count) = transform_openai_request(&req, "test-v", "gemini-2.0-flash-thinking", None); |
| let budget = result["request"]["generationConfig"]["thinkingConfig"]["thinkingBudget"] |
| .as_i64() |
| .unwrap(); |
| assert_eq!(budget, 24576, "Gemini custom budget must be capped to 24576"); |
|
|
| |
| |
| let (result_claude, _, _) = transform_openai_request(&req, "test-v", "claude-3-7-sonnet", None); |
| let budget_claude = result_claude["request"]["generationConfig"]["thinkingConfig"]["thinkingBudget"] |
| .as_i64(); |
| |
| |
| |
|
|
| |
| update_thinking_budget_config(ThinkingBudgetConfig::default()); |
| } |
|
|
| #[test] |
| fn test_transform_openai_request_multimodal() { |
| let req = OpenAIRequest { |
| model: "gpt-4-vision".to_string(), |
| messages: vec![OpenAIMessage { |
| role: "user".to_string(), |
| content: Some(OpenAIContent::Array(vec![ |
| OpenAIContentBlock::Text { text: "What is in this image?".to_string() }, |
| OpenAIContentBlock::ImageUrl { image_url: OpenAIImageUrl { |
| url: "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg==".to_string(), |
| detail: None |
| } } |
| ])), |
| reasoning_content: None, |
| tool_calls: None, |
| tool_call_id: None, |
| name: None, |
| }], |
| stream: false, |
| n: None, |
| max_tokens: None, |
| temperature: None, |
| top_p: None, |
| stop: None, |
| response_format: None, |
| tools: None, |
| tool_choice: None, |
| parallel_tool_calls: None, |
| ..Default::default() |
| }; |
|
|
| let (result, _sid, _msg_count) = transform_openai_request(&req, "test-v", "gemini-1.5-flash", None); |
| let parts = &result["request"]["contents"][0]["parts"]; |
| assert_eq!(parts.as_array().unwrap().len(), 2); |
| assert_eq!(parts[0]["text"].as_str().unwrap(), "What is in this image?"); |
| assert_eq!( |
| parts[1]["inlineData"]["mimeType"].as_str().unwrap(), |
| "image/png" |
| ); |
| } |
| |
| #[test] |
| fn test_gemini_pro_thinking_injection() { |
| let req = OpenAIRequest { |
| model: "gemini-3-pro-preview".to_string(), |
| messages: vec![OpenAIMessage { |
| role: "user".to_string(), |
| content: Some(OpenAIContent::String("Thinking test".to_string())), |
| reasoning_content: None, |
| tool_calls: None, |
| tool_call_id: None, |
| name: None, |
| }], |
| stream: false, |
| n: None, |
| |
| thinking: Some(ThinkingConfig { |
| thinking_type: Some("enabled".to_string()), |
| budget_tokens: Some(16000), |
| effort: None, |
| }), |
| max_tokens: None, |
| temperature: None, |
| top_p: None, |
| stop: None, |
| response_format: None, |
| tools: None, |
| tool_choice: None, |
| parallel_tool_calls: None, |
| ..Default::default() |
| }; |
|
|
| |
| let (result, _sid, _msg_count) = transform_openai_request(&req, "test-p", "gemini-3-pro-preview", None); |
| let gen_config = &result["request"]["generationConfig"]; |
| |
| |
| assert!(gen_config.get("thinkingConfig").is_some(), "thinkingConfig should be injected for gemini-3-pro"); |
| |
| let budget = gen_config["thinkingConfig"]["thinkingBudget"].as_u64().unwrap(); |
| |
| assert_eq!(budget, 16000); |
| } |
| #[test] |
| fn test_gemini_3_pro_image_not_thinking() { |
| let req = OpenAIRequest { |
| model: "gemini-3-pro-image-4k".to_string(), |
| messages: vec![OpenAIMessage { |
| role: "user".to_string(), |
| content: Some(OpenAIContent::String("Generate a cat".to_string())), |
| reasoning_content: None, |
| tool_calls: None, |
| tool_call_id: None, |
| name: None, |
| }], |
| ..Default::default() |
| }; |
|
|
| |
| let (result, _sid, _msg_count) = transform_openai_request(&req, "test-p", "gemini-3-pro-image", None); |
| let gen_config = &result["request"]["generationConfig"]; |
| |
| |
| assert!(gen_config.get("thinkingConfig").is_some(), "thinkingConfig SHOULD be injected for gemini-3-pro-image"); |
| |
| |
| assert!(gen_config.get("imageConfig").is_some(), "imageConfig should be present for image models"); |
| assert_eq!(gen_config["imageConfig"]["imageSize"], "4K"); |
| } |
|
|
| #[test] |
| fn test_default_max_tokens_openai() { |
| let req = OpenAIRequest { |
| model: "gpt-4".to_string(), |
| messages: vec![OpenAIMessage { |
| role: "user".to_string(), |
| content: Some(OpenAIContent::String("Hello".to_string())), |
| reasoning_content: None, |
| tool_calls: None, |
| tool_call_id: None, |
| name: None, |
| }], |
| stream: false, |
| n: None, |
| max_tokens: None, |
| temperature: None, |
| top_p: None, |
| stop: None, |
| response_format: None, |
| tools: None, |
| tool_choice: None, |
| parallel_tool_calls: None, |
| ..Default::default() |
| }; |
|
|
| let (result, _sid, _msg_count) = transform_openai_request(&req, "test-p", "gemini-3-pro-high-thinking", None); |
| let gen_config = &result["request"]["generationConfig"]; |
| let max_output_tokens = gen_config["maxOutputTokens"].as_i64().unwrap(); |
| |
| assert_eq!(max_output_tokens, 57344); |
| |
| |
| let budget = gen_config["thinkingConfig"]["thinkingBudget"].as_i64().unwrap(); |
| |
| assert_eq!(budget, 24576); |
| } |
|
|
| #[test] |
| fn test_flash_thinking_budget_capping() { |
| let req = OpenAIRequest { |
| model: "gpt-4".to_string(), |
| messages: vec![OpenAIMessage { |
| role: "user".to_string(), |
| content: Some(OpenAIContent::String("Hello".to_string())), |
| reasoning_content: None, |
| tool_calls: None, |
| tool_call_id: None, |
| name: None, |
| }], |
| stream: false, |
| n: None, |
| |
| thinking: Some(ThinkingConfig { |
| thinking_type: Some("enabled".to_string()), |
| budget_tokens: Some(32768), |
| effort: None, |
| }), |
| max_tokens: None, |
| temperature: None, |
| top_p: None, |
| stop: None, |
| response_format: None, |
| tools: None, |
| tool_choice: None, |
| parallel_tool_calls: None, |
| ..Default::default() |
| }; |
|
|
| |
| let (result, _sid, _msg_count) = transform_openai_request(&req, "test-p", "gemini-2.0-flash-thinking-exp", None); |
| let gen_config = &result["request"]["generationConfig"]; |
| |
| |
| let budget = gen_config["thinkingConfig"]["thinkingBudget"].as_i64().unwrap(); |
| assert_eq!(budget, 24576); |
|
|
| |
| |
| let max_output_tokens = gen_config["maxOutputTokens"].as_i64().unwrap(); |
| assert_eq!(max_output_tokens, 57344); |
| } |
| #[test] |
| fn test_vertex_ai_sentinel_injection() { |
| |
| let req = OpenAIRequest { |
| model: "claude-3-7-sonnet-thinking".to_string(), |
| messages: vec![OpenAIMessage { |
| role: "assistant".to_string(), |
| content: None, |
| reasoning_content: Some("Thinking...".to_string()), |
| tool_calls: Some(vec![ToolCall { |
| id: "call_123".to_string(), |
| r#type: "function".to_string(), |
| function: ToolFunction { |
| name: "test_tool".to_string(), |
| arguments: "{}".to_string(), |
| }, |
| }]), |
| tool_call_id: None, |
| name: None, |
| }], |
| person_generation: None, |
| ..Default::default() |
| }; |
|
|
| |
| let mapped_model = "projects/my-project/locations/us-central1/publishers/google/models/gemini-2.0-flash-thinking-exp"; |
| |
| let (result, _sid, _msg_count) = transform_openai_request(&req, "test-v", mapped_model, None); |
| |
| |
| let contents = result["contents"].as_array().unwrap(); |
| |
| let parts = contents[0]["parts"].as_array().unwrap(); |
| let tool_part = parts.iter().find(|p| p.get("functionCall").is_some()).expect("Should find functionCall part"); |
| |
| |
| assert_eq!(tool_part["thoughtSignature"].as_str(), Some("skip_thought_signature_validator")); |
| } |
|
|
| #[test] |
| fn test_issue_2167_gemini_flash_thinking_signature() { |
| |
| for model in &["gemini-3-flash", "gemini-3.1-flash"] { |
| let req = OpenAIRequest { |
| model: model.to_string(), |
| messages: vec![OpenAIMessage { |
| role: "assistant".to_string(), |
| content: None, |
| reasoning_content: None, |
| tool_calls: Some(vec![ToolCall { |
| id: "call_flash_test".to_string(), |
| r#type: "function".to_string(), |
| function: ToolFunction { |
| name: "get_weather".to_string(), |
| arguments: "{\"location\":\"Beijing\"}".to_string(), |
| }, |
| }]), |
| tool_call_id: None, |
| name: None, |
| }], |
| ..Default::default() |
| }; |
|
|
| let (result, _sid, _msg_count) = transform_openai_request(&req, "test-proj", model, None); |
|
|
| let contents = result["request"]["contents"].as_array().expect("Should have request.contents"); |
| |
| let model_msg = contents.iter().find(|c| c["role"] == "model").expect("Should find model role message"); |
| let parts = model_msg["parts"].as_array().expect("Should have parts"); |
| let tool_part = parts |
| .iter() |
| .find(|p| p.get("functionCall").is_some()) |
| .expect(&format!("[{model}] Should find functionCall part")); |
|
|
| assert_eq!( |
| tool_part["thoughtSignature"].as_str(), |
| Some("skip_thought_signature_validator"), |
| "[{model}] gemini-3-flash functionCall must contain thoughtSignature sentinel" |
| ); |
| } |
| } |
|
|
| #[test] |
| fn test_openai_image_thinking_mode_disabled() { |
| |
| crate::proxy::config::update_image_thinking_mode(Some("disabled".to_string())); |
|
|
| let req = OpenAIRequest { |
| model: "gemini-3-pro-image".to_string(), |
| messages: vec![OpenAIMessage { |
| role: "user".to_string(), |
| content: Some(OpenAIContent::String("Draw a cat".to_string())), |
| name: None, |
| tool_calls: None, |
| tool_call_id: None, |
| reasoning_content: None, |
| }], |
| tools: None, |
| tool_choice: None, |
| parallel_tool_calls: None, |
| person_generation: None, |
| ..Default::default() |
| }; |
|
|
| |
| let (result, _sid, _msg_count) = transform_openai_request(&req, "test-proj", "gemini-3-pro-image", None); |
|
|
| |
| let gen_config = result["request"]["generationConfig"].as_object().expect("Should have generationConfig in request payload"); |
| let thinking_config = gen_config["thinkingConfig"].as_object().unwrap(); |
| |
| assert_eq!(thinking_config["includeThoughts"], false); |
| |
| |
| crate::proxy::config::update_image_thinking_mode(Some("enabled".to_string())); |
| } |
|
|
| #[test] |
| fn test_mixed_tools_injection_openai() { |
| |
| let req = OpenAIRequest { |
| model: "gpt-4o-online".to_string(), |
| messages: vec![OpenAIMessage { |
| role: "user".to_string(), |
| content: Some(OpenAIContent::String("Hello".to_string())), |
| reasoning_content: None, |
| tool_calls: None, |
| tool_call_id: None, |
| name: None, |
| }], |
| tools: Some(vec![json!({ |
| "type": "function", |
| "function": { |
| "name": "get_weather", |
| "parameters": { |
| "type": "object", |
| "properties": { |
| "location": {"type": "string"} |
| } |
| } |
| } |
| })]), |
| ..Default::default() |
| }; |
|
|
| |
| let (result, _, _) = transform_openai_request(&req, "proj", "gemini-2.0-flash", None); |
| |
| let tools = result["request"]["tools"].as_array().expect("Should have tools"); |
| |
| let has_functions = tools.iter().any(|t| t.get("functionDeclarations").is_some()); |
| let has_google_search = tools.iter().any(|t| t.get("googleSearch").is_some()); |
| |
| assert!(has_functions, "Should contain functionDeclarations"); |
| assert!(has_google_search, "Should contain googleSearch (Gemini 2.0+ supports mixed tools)"); |
| } |
| } |
|
|