Skip to content

Support image input for Grok #58

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
hassaans wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Support image input for Grok #58

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a4d77b5
Support image input for Grok
hassaans Jul 17, 2026
e25df19
grok: validate and preserve image results
raine Jul 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 36 additions & 10 deletions src/providers/grok/count_tokens.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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,
}
}

Expand Down Expand Up @@ -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!({
Expand Down
210 changes: 191 additions & 19 deletions src/providers/grok/translate/request.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<GrokContentPart>),
}

#[derive(Debug, Clone, Serialize)]
Expand All @@ -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<String>,
},
}

#[derive(Debug, Clone, Serialize)]
Expand Down Expand Up @@ -364,6 +380,62 @@ fn parse_tool_choice(
}
}

fn parse_image(object: &serde_json::Map<String, Value>) -> anyhow::Result<GrokContentPart> {
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<GrokInputItem>,
Expand Down Expand Up @@ -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")) {
Expand Down Expand Up @@ -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::<anyhow::Result<Vec<_>>>()?
.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(),
Expand All @@ -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;
};
Expand Down Expand Up @@ -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!({
Expand Down