diff --git a/src/providers/grok/count_tokens.rs b/src/providers/grok/count_tokens.rs index b456e6e..c6cbc16 100644 --- a/src/providers/grok/count_tokens.rs +++ b/src/providers/grok/count_tokens.rs @@ -1,7 +1,11 @@ -use super::translate::request::{GrokContentPart, GrokInputItem, GrokResponsesRequest}; +use super::translate::request::{ + GrokContentPart, GrokFunctionCallOutput, GrokInputItem, GrokResponsesRequest, +}; const MESSAGE_OVERHEAD_TOKENS: u64 = 4; const TOOL_OVERHEAD_TOKENS: u64 = 4; +// Rough per-image cost for token estimation; actual usage is billed by Grok. +const IMAGE_TOKENS: u64 = 1024; pub fn count_tokens(request: &GrokResponsesRequest) -> u64 { let instructions = request @@ -37,18 +41,25 @@ pub fn count_tokens(request: &GrokResponsesRequest) -> u64 { fn count_input_item(item: &GrokInputItem) -> u64 { match item { - GrokInputItem::Message { content, .. } => content - .iter() - .map(|part| match part { - GrokContentPart::InputText { text } | GrokContentPart::OutputText { text } => { - approx_token_count(text) - } - }) - .sum(), + GrokInputItem::Message { content, .. } => content.iter().map(count_content_part).sum(), GrokInputItem::FunctionCall { name, arguments, .. } => approx_token_count(name) + approx_token_count(arguments), - GrokInputItem::FunctionCallOutput { output, .. } => approx_token_count(output), + GrokInputItem::FunctionCallOutput { output, .. } => match output { + GrokFunctionCallOutput::Text(text) => approx_token_count(text), + GrokFunctionCallOutput::Content(content) => { + content.iter().map(count_content_part).sum() + } + }, + } +} + +fn count_content_part(part: &GrokContentPart) -> u64 { + match part { + GrokContentPart::InputText { text } | GrokContentPart::OutputText { text } => { + approx_token_count(text) + } + GrokContentPart::InputImage { .. } => IMAGE_TOKENS, } } @@ -112,6 +123,21 @@ mod tests { assert!(count_tokens(&long) > count_tokens(&short)); } + #[test] + fn count_tokens_includes_tool_result_images() { + let request = translated_request(json!({ + "model": "grok-4.5", + "messages": [ + {"role":"assistant","content":[{"type":"tool_use","id":"call_1","name":"screenshot","input":{}}]}, + {"role":"user","content":[{"type":"tool_result","tool_use_id":"call_1","content":[ + {"type":"image","source":{"type":"base64","media_type":"image/png","data":"aGVsbG8="}} + ]}]} + ] + })); + + assert!(count_tokens(&request) >= IMAGE_TOKENS); + } + #[test] fn count_tokens_is_deterministic() { let request = translated_request(json!({ diff --git a/src/providers/grok/translate/request.rs b/src/providers/grok/translate/request.rs index d706c6e..b1032da 100644 --- a/src/providers/grok/translate/request.rs +++ b/src/providers/grok/translate/request.rs @@ -36,7 +36,17 @@ pub enum GrokInputItem { arguments: String, }, #[serde(rename = "function_call_output")] - FunctionCallOutput { call_id: String, output: String }, + FunctionCallOutput { + call_id: String, + output: GrokFunctionCallOutput, + }, +} + +#[derive(Debug, Clone, Serialize)] +#[serde(untagged)] +pub enum GrokFunctionCallOutput { + Text(String), + Content(Vec), } #[derive(Debug, Clone, Serialize)] @@ -46,6 +56,12 @@ pub enum GrokContentPart { InputText { text: String }, #[serde(rename = "output_text")] OutputText { text: String }, + #[serde(rename = "input_image")] + InputImage { + image_url: String, + #[serde(skip_serializing_if = "Option::is_none")] + detail: Option, + }, } #[derive(Debug, Clone, Serialize)] @@ -364,6 +380,62 @@ fn parse_tool_choice( } } +fn parse_image(object: &serde_json::Map) -> anyhow::Result { + if object + .keys() + .any(|key| !["type", "source", "cache_control"].contains(&key.as_str())) + || !valid_cache_control(object.get("cache_control")) + { + anyhow::bail!("unsupported image block field"); + } + let source = object + .get("source") + .and_then(Value::as_object) + .ok_or_else(|| anyhow::anyhow!("image source must be an object"))?; + let image_url = match source.get("type").and_then(Value::as_str) { + Some("base64") => { + if source + .keys() + .any(|key| !["type", "media_type", "data"].contains(&key.as_str())) + { + anyhow::bail!("unsupported base64 image source field"); + } + let media_type = source + .get("media_type") + .and_then(Value::as_str) + .ok_or_else(|| anyhow::anyhow!("image media_type is invalid"))?; + if !matches!(media_type, "image/jpeg" | "image/png") { + anyhow::bail!("unsupported image media_type: {media_type}"); + } + let data = source + .get("data") + .and_then(Value::as_str) + .filter(|data| !data.is_empty()) + .ok_or_else(|| anyhow::anyhow!("image data is invalid"))?; + format!("data:{media_type};base64,{data}") + } + Some("url") => { + if source + .keys() + .any(|key| !["type", "url"].contains(&key.as_str())) + { + anyhow::bail!("unsupported URL image source field"); + } + source + .get("url") + .and_then(Value::as_str) + .filter(|url| !url.is_empty()) + .ok_or_else(|| anyhow::anyhow!("image url is invalid"))? + .to_string() + } + _ => anyhow::bail!("unsupported image source"), + }; + Ok(GrokContentPart::InputImage { + image_url, + detail: None, + }) +} + fn parse_message( message: &Message, out: &mut Vec, @@ -406,6 +478,7 @@ fn parse_message( GrokContentPart::InputText { text: text.into() } }); } + ("user", "image") => content.push(parse_image(object)?), ("assistant", "server_tool_use") => { let name = object.get("name").and_then(Value::as_str); if !matches!(name, Some("web_search" | "x_search")) { @@ -476,28 +549,48 @@ fn parse_message( .get("content") .ok_or_else(|| anyhow::anyhow!("tool result content is required"))?; let output = match value { - Value::String(text) => text.clone(), - Value::Array(parts) => parts - .iter() - .map(|part| { + Value::String(text) => GrokFunctionCallOutput::Text(text.clone()), + Value::Array(parts) => { + let mut text = String::new(); + let mut content = Vec::new(); + let mut has_image = false; + for part in parts { let part = part.as_object().ok_or_else(|| { anyhow::anyhow!("tool result child must be an object") })?; - if part.get("type").and_then(Value::as_str) != Some("text") - || part.keys().any(|key| { - !["type", "text", "cache_control"].contains(&key.as_str()) - }) - || !valid_cache_control(part.get("cache_control")) - { - anyhow::bail!("tool result supports text children only"); + match part.get("type").and_then(Value::as_str) { + Some("text") => { + if part.keys().any(|key| { + !["type", "text", "cache_control"].contains(&key.as_str()) + }) || !valid_cache_control(part.get("cache_control")) + { + anyhow::bail!("unsupported tool result text child"); + } + let child_text = + part.get("text").and_then(Value::as_str).ok_or_else( + || anyhow::anyhow!("tool result text is invalid"), + )?; + text.push_str(child_text); + content.push(GrokContentPart::InputText { + text: child_text.into(), + }); + } + Some("image") => { + has_image = true; + content.push(parse_image(part)?); + } + _ => anyhow::bail!( + "tool result supports text and image children only" + ), } - part.get("text") - .and_then(Value::as_str) - .ok_or_else(|| anyhow::anyhow!("tool result text is invalid")) - }) - .collect::>>()? - .join(""), - _ => anyhow::bail!("tool result supports text only"), + } + if has_image { + GrokFunctionCallOutput::Content(content) + } else { + GrokFunctionCallOutput::Text(text) + } + } + _ => anyhow::bail!("tool result supports text and image content only"), }; out.push(GrokInputItem::FunctionCallOutput { call_id: id.into(), @@ -513,6 +606,9 @@ fn parse_message( fn valid_cache_control(value: Option<&Value>) -> bool { let Some(value) = value else { return true }; + if value.is_null() { + return true; + } let Some(object) = value.as_object() else { return false; }; @@ -577,6 +673,82 @@ mod tests { assert_eq!(value["input"][2]["type"], "function_call_output"); assert_eq!(value["tool_choice"]["type"], "function"); } + + #[test] + fn grok_translation_maps_base64_image_to_input_image() { + let request = request_with_blocks(serde_json::json!([ + {"type":"image","source":{"type":"base64","media_type":"image/png","data":"aGVsbG8="},"cache_control":null}, + {"type":"text","text":"describe it"} + ])); + let value = + serde_json::to_value(translate_request(&request, "grok-4.5".into()).unwrap()).unwrap(); + let content = &value["input"][1]["content"]; + assert_eq!(content[0]["type"], "input_image"); + assert_eq!(content[0]["image_url"], "data:image/png;base64,aGVsbG8="); + assert_eq!(content[1]["type"], "input_text"); + } + + #[test] + fn grok_translation_maps_url_image_to_input_image() { + let request = request_with_blocks(serde_json::json!([ + {"type":"image","source":{"type":"url","url":"https://example.com/cat.png"}} + ])); + let value = + serde_json::to_value(translate_request(&request, "grok-4.5".into()).unwrap()).unwrap(); + assert_eq!(value["input"][1]["content"][0]["type"], "input_image"); + assert_eq!( + value["input"][1]["content"][0]["image_url"], + "https://example.com/cat.png" + ); + } + + #[test] + fn grok_translation_maps_tool_result_images_to_function_output_content() { + let request = request_with_blocks(serde_json::json!([{ + "type":"tool_result", + "tool_use_id":"call_1", + "content":[ + {"type":"text","text":"screenshot"}, + {"type":"image","source":{"type":"base64","media_type":"image/png","data":"aGVsbG8="}} + ] + }])); + let value = + serde_json::to_value(translate_request(&request, "grok-4.5".into()).unwrap()).unwrap(); + let output = &value["input"][1]["output"]; + assert_eq!( + output[0], + serde_json::json!({"type":"input_text","text":"screenshot"}) + ); + assert_eq!(output[1]["type"], "input_image"); + assert_eq!(output[1]["image_url"], "data:image/png;base64,aGVsbG8="); + } + + #[test] + fn grok_translation_rejects_unknown_image_source_and_fields() { + for image in [ + serde_json::json!({"type":"image","source":{"type":"file","file_id":"f_1"}}), + serde_json::json!({"type":"image","source":{"type":"url","url":"https://x/y.png"},"unknown":true}), + serde_json::json!({"type":"image","source":{"type":"url","url":"https://x/y.png","unknown":true}}), + serde_json::json!({"type":"image","source":{"type":"url","url":"https://x/y.png","data":"aGVsbG8="}}), + ] { + let request = request_with_blocks(serde_json::json!([image])); + assert!(translate_request(&request, "grok-4.5".into()).is_err()); + } + } + + #[test] + fn grok_translation_rejects_unsupported_or_empty_image_values() { + for image in [ + serde_json::json!({"type":"image","source":{"type":"base64","media_type":"image/gif","data":"aGVsbG8="}}), + serde_json::json!({"type":"image","source":{"type":"base64","media_type":"image/webp","data":"aGVsbG8="}}), + serde_json::json!({"type":"image","source":{"type":"base64","media_type":"image/png","data":""}}), + serde_json::json!({"type":"image","source":{"type":"url","url":""}}), + ] { + let request = request_with_blocks(serde_json::json!([image])); + assert!(translate_request(&request, "grok-4.5".into()).is_err()); + } + } + #[test] fn grok_translation_maps_claude_web_search_to_hosted_web_search() { let request: MessagesRequest = serde_json::from_value(serde_json::json!({