use crate::error::{Result, ShimError};
use crate::provider::{Provider, ProviderRequest};
use crate::vision;
use serde_json::{json, Value};
pub struct Gemini {
pub api_key: String,
pub base_url: String,
}
impl Gemini {
pub fn new(api_key: String) -> Self {
Self {
api_key,
base_url: "https://generativelanguage.googleapis.com/v1beta".to_string(),
}
}
pub fn with_base_url(mut self, url: String) -> Self {
self.base_url = url;
self
}
}
fn transform_messages(messages: &[Value]) -> (Option<Value>, Vec<Value>) {
let mut system_parts: Vec<String> = Vec::new();
let mut contents: Vec<Value> = Vec::new();
let mut call_id_to_name: std::collections::HashMap<String, String> =
std::collections::HashMap::new();
for msg in messages {
let role = msg.get("role").and_then(|r| r.as_str()).unwrap_or("");
match role {
"system" | "developer" => {
if let Some(text) = msg.get("content").and_then(|c| c.as_str()) {
system_parts.push(text.to_string());
}
}
"assistant" => {
if let Some(tool_calls) = msg.get("tool_calls").and_then(|t| t.as_array()) {
for tc in tool_calls {
if let (Some(id), Some(name)) = (
tc.get("id").and_then(|v| v.as_str()),
tc.get("function")
.and_then(|f| f.get("name"))
.and_then(|n| n.as_str()),
) {
call_id_to_name.insert(id.to_string(), name.to_string());
}
}
}
let mut parts = build_parts(msg);
sanitize_parts(&mut parts);
contents.push(json!({
"role": "model",
"parts": parts,
}));
}
"tool" => {
let call_id = msg
.get("tool_call_id")
.and_then(|v| v.as_str())
.unwrap_or("");
let name = msg
.get("name")
.and_then(|n| n.as_str())
.map(|s| s.to_string())
.or_else(|| call_id_to_name.get(call_id).cloned())
.unwrap_or_else(|| "function".to_string());
let content = msg.get("content").and_then(|c| c.as_str()).unwrap_or("");
let parsed: Value =
serde_json::from_str(content).unwrap_or_else(|_| json!({"result": content}));
let response = if parsed.is_object() {
parsed
} else {
json!({"result": parsed})
};
contents.push(json!({
"role": "user",
"parts": [{"functionResponse": {"name": name, "response": response}}]
}));
}
_ => {
let parts = build_parts(msg);
contents.push(json!({
"role": "user",
"parts": parts,
}));
}
}
}
let system_instruction = if system_parts.is_empty() {
None
} else {
Some(json!({
"parts": [{"text": system_parts.join("\n\n")}]
}))
};
let contents = enforce_gemini_turn_order(contents);
(system_instruction, contents)
}
fn strip_function_calls(turn: &mut Value) {
if let Some(parts) = turn.get_mut("parts").and_then(|p| p.as_array_mut()) {
parts.retain(|p| p.get("functionCall").is_none() && p.get("thoughtSignature").is_none());
if parts.is_empty() {
parts.push(json!({"text": ""}));
}
}
}
fn enforce_gemini_turn_order(contents: Vec<Value>) -> Vec<Value> {
if contents.is_empty() {
return contents;
}
let mut merged = merge_same_role(contents);
for _ in 0..20 {
let mut changed = false;
for i in 0..merged.len() {
let has_fc = merged[i]
.get("parts")
.and_then(|p| p.as_array())
.map(|parts| parts.iter().any(|p| p.get("functionCall").is_some()))
.unwrap_or(false);
if !has_fc {
continue;
}
let prev_ok =
i > 0 && merged[i - 1].get("role").and_then(|r| r.as_str()) == Some("user");
let next_ok = merged.get(i + 1).is_some_and(|next| {
next.get("parts")
.and_then(|p| p.as_array())
.map(|parts| parts.iter().any(|p| p.get("functionResponse").is_some()))
.unwrap_or(false)
});
let missing_sig = merged[i]
.get("parts")
.and_then(|p| p.as_array())
.map(|parts| {
parts.iter().any(|p| {
p.get("functionCall").is_some() && p.get("thoughtSignature").is_none()
})
})
.unwrap_or(false);
if !prev_ok || !next_ok || missing_sig {
strip_function_calls(&mut merged[i]);
changed = true;
}
}
let before = merged.len();
let mut i = 0;
while i < merged.len() {
let has_fr = merged[i]
.get("parts")
.and_then(|p| p.as_array())
.map(|parts| parts.iter().any(|p| p.get("functionResponse").is_some()))
.unwrap_or(false);
if has_fr {
let prev_fc = i > 0
&& merged[i - 1]
.get("parts")
.and_then(|p| p.as_array())
.map(|parts| parts.iter().any(|p| p.get("functionCall").is_some()))
.unwrap_or(false);
if !prev_fc {
merged.remove(i);
continue;
}
}
i += 1;
}
if merged.len() != before {
changed = true;
}
let before = merged.len();
merged.retain(|turn| {
turn.get("parts")
.and_then(|p| p.as_array())
.map(|parts| {
parts.iter().any(|p| {
p.get("text")
.and_then(|t| t.as_str())
.map(|s| !s.is_empty())
.unwrap_or(true)
})
})
.unwrap_or(true)
});
if merged.len() != before {
changed = true;
}
let before = merged.len();
merged = merge_same_role(merged);
if merged.len() != before {
changed = true;
}
if !changed {
break;
}
}
merged
}
fn merge_same_role(turns: Vec<Value>) -> Vec<Value> {
let mut merged: Vec<Value> = Vec::new();
for turn in turns {
let role = turn.get("role").and_then(|r| r.as_str()).unwrap_or("");
let last_role = merged
.last()
.and_then(|t| t.get("role"))
.and_then(|r| r.as_str())
.unwrap_or("");
if role == last_role {
if let Some(new_parts) = turn.get("parts").and_then(|p| p.as_array()) {
if let Some(last) = merged.last_mut() {
if let Some(existing) = last.get_mut("parts").and_then(|p| p.as_array_mut()) {
existing.extend(new_parts.clone());
}
}
}
} else {
merged.push(turn);
}
}
merged
}
fn build_parts(msg: &Value) -> Vec<Value> {
let mut parts = Vec::new();
match msg.get("content") {
Some(Value::String(text)) if !text.is_empty() => {
parts.push(json!({"text": text}));
}
Some(Value::Array(blocks)) => {
for block in blocks {
match block.get("type").and_then(|t| t.as_str()) {
Some("text") => {
if let Some(text) = block.get("text").and_then(|t| t.as_str()) {
parts.push(json!({"text": text}));
}
}
Some("image_url" | "input_image" | "image") => {
if let Some(gemini_part) = vision::to_gemini(block) {
parts.push(gemini_part);
}
}
_ => {} }
}
}
_ => {}
}
if let Some(tool_calls) = msg.get("tool_calls").and_then(|t| t.as_array()) {
for tc in tool_calls {
if let Some(func) = tc.get("function") {
let name = func.get("name").and_then(|n| n.as_str()).unwrap_or("");
let args: Value = func
.get("arguments")
.and_then(|a| a.as_str())
.and_then(|s| serde_json::from_str(s).ok())
.unwrap_or(json!({}));
let mut fc_part = json!({"functionCall": {"name": name, "args": args}});
if let Some(sig) = tc.get("thought_signature") {
fc_part["thoughtSignature"] = sig.clone();
}
parts.push(fc_part);
}
}
}
if parts.is_empty() {
parts.push(json!({"text": ""}));
}
parts
}
fn sanitize_parts(parts: &mut Vec<Value>) {
parts.retain(|p| {
if p.get("text").is_none() && p.get("functionCall").is_none() {
return true;
}
true
});
}
fn transform_tools(tools: &[Value]) -> Value {
let empty = json!("");
let default_params = json!({"type": "object", "properties": {}});
let declarations: Vec<Value> = tools
.iter()
.filter_map(|tool| {
let source = tool.get("function").unwrap_or(tool);
let name = source.get("name")?;
let description = source.get("description").unwrap_or(&empty);
let parameters = source.get("parameters").unwrap_or(&default_params).clone();
let sanitized = sanitize_schema(parameters);
Some(json!({
"name": name,
"description": description,
"parameters": sanitized,
}))
})
.collect();
json!([{ "functionDeclarations": declarations }])
}
fn sanitize_schema(mut schema: Value) -> Value {
let defs = schema
.as_object_mut()
.and_then(|obj| obj.remove("$defs").or_else(|| obj.remove("definitions")));
if let Some(defs_val) = &defs {
resolve_refs(&mut schema, defs_val);
}
clean_schema_fields(&mut schema);
schema
}
fn resolve_refs(value: &mut Value, defs: &Value) {
match value {
Value::Object(obj) => {
if let Some(ref_val) = obj.remove("$ref") {
if let Some(ref_str) = ref_val.as_str() {
let def_name = ref_str
.strip_prefix("#/$defs/")
.or_else(|| ref_str.strip_prefix("#/definitions/"));
if let Some(name) = def_name {
if let Some(def) = defs.get(name) {
let mut inlined = def.clone();
resolve_refs(&mut inlined, defs);
clean_schema_fields(&mut inlined);
*value = inlined;
return;
}
}
}
}
for v in obj.values_mut() {
resolve_refs(v, defs);
}
}
Value::Array(arr) => {
for v in arr.iter_mut() {
resolve_refs(v, defs);
}
}
_ => {}
}
}
fn clean_schema_fields(value: &mut Value) {
match value {
Value::Object(obj) => {
obj.remove("$schema");
obj.remove("$defs");
obj.remove("definitions");
obj.remove("$ref"); obj.remove("additionalProperties");
obj.remove("default");
if let Some(type_val) = obj.get_mut("type") {
if let Some(arr) = type_val.as_array().cloned() {
let non_null: Vec<&Value> =
arr.iter().filter(|v| v.as_str() != Some("null")).collect();
if non_null.len() == 1 {
*type_val = non_null[0].clone();
}
}
}
for v in obj.values_mut() {
clean_schema_fields(v);
}
}
Value::Array(arr) => {
for v in arr.iter_mut() {
clean_schema_fields(v);
}
}
_ => {}
}
}
fn translate_tool_choice(tc: &Value) -> Option<Value> {
let mode = if let Some(s) = tc.as_str() {
match s {
"auto" => "AUTO",
"required" => "ANY",
"none" => "NONE",
_ => return None,
}
} else {
let obj = tc.as_object()?;
match obj.get("type").and_then(|t| t.as_str()) {
Some("auto") => "AUTO",
Some("any" | "required") => "ANY",
Some("none") => "NONE",
_ => return None,
}
};
Some(json!({"functionCallingConfig": {"mode": mode}}))
}
fn transform_response_to_openai(model: &str, resp: &Value) -> Result<Value> {
let candidate = resp
.get("candidates")
.and_then(|c| c.as_array())
.and_then(|a| a.first())
.ok_or_else(|| ShimError::ProviderError {
status: 500,
body: format!("no candidates in response: {}", resp),
})?;
let parts = candidate
.pointer("/content/parts")
.and_then(|p| p.as_array())
.cloned()
.unwrap_or_default();
let mut text_parts: Vec<String> = Vec::new();
let mut thought_parts: Vec<String> = Vec::new();
let mut tool_calls: Vec<Value> = Vec::new();
for part in &parts {
let is_thought = part
.get("thought")
.and_then(|t| t.as_bool())
.unwrap_or(false);
if let Some(text) = part.get("text").and_then(|t| t.as_str()) {
if !text.is_empty() {
if is_thought {
thought_parts.push(text.to_string());
} else {
text_parts.push(text.to_string());
}
}
}
if let Some(fc) = part.get("functionCall") {
let name = fc.get("name").and_then(|n| n.as_str()).unwrap_or("");
if name.is_empty() {
continue;
}
let id = fc
.get("id")
.and_then(|v| v.as_str())
.map(|s| s.to_string())
.unwrap_or_else(|| format!("call_{}", tool_calls.len()));
let args_str = fc
.get("args")
.filter(|a| !a.is_null())
.map(|a| serde_json::to_string(a).unwrap_or_else(|_| "{}".to_string()))
.unwrap_or_else(|| "{}".to_string());
let mut tc = json!({
"id": id,
"type": "function",
"function": {
"name": name,
"arguments": args_str,
}
});
if let Some(sig) = part.get("thoughtSignature") {
tc["thought_signature"] = sig.clone();
}
tool_calls.push(tc);
}
}
let content = if text_parts.is_empty() {
Value::Null
} else {
json!(text_parts.join(""))
};
let finish_reason = candidate
.get("finishReason")
.and_then(|f| f.as_str())
.map(|f| match f {
"STOP" => "stop",
"MAX_TOKENS" => "length",
"SAFETY" => "content_filter",
_ => "stop",
})
.unwrap_or("stop");
let usage = resp.get("usageMetadata").cloned().unwrap_or(json!({}));
let mut message = json!({
"role": "assistant",
"content": content,
});
if !tool_calls.is_empty() {
message["tool_calls"] = json!(tool_calls);
}
if !thought_parts.is_empty() {
message["reasoning_content"] = json!(thought_parts.join("\n"));
}
Ok(json!({
"id": resp.get("responseId").cloned().unwrap_or(json!("")),
"object": "chat.completion",
"model": model,
"choices": [{
"index": 0,
"message": message,
"finish_reason": finish_reason,
}],
"usage": {
"prompt_tokens": usage.get("promptTokenCount").cloned().unwrap_or(json!(0)),
"completion_tokens": usage.get("candidatesTokenCount").cloned().unwrap_or(json!(0)),
"total_tokens": usage.get("totalTokenCount").cloned().unwrap_or(json!(0)),
}
}))
}
fn cannot_disable_thinking(model: &str) -> bool {
model.to_lowercase().contains("3.1-pro")
}
impl Provider for Gemini {
fn name(&self) -> &str {
"gemini"
}
fn transform_request(&self, model: &str, request: &Value) -> Result<ProviderRequest> {
let obj = request.as_object().ok_or(ShimError::MissingModel)?;
let messages = obj
.get("messages")
.and_then(|m| m.as_array())
.ok_or_else(|| {
ShimError::Json(serde_json::Error::io(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"missing messages array",
)))
})?;
let (system_instruction, contents) = transform_messages(messages);
let mut body = json!({"contents": contents});
let body_obj = body.as_object_mut().unwrap();
if let Some(si) = system_instruction {
body_obj.insert("systemInstruction".to_string(), si);
}
let mut gen_config = json!({});
let gc = gen_config.as_object_mut().unwrap();
if let Some(v) = obj.get("temperature") {
gc.insert("temperature".to_string(), v.clone());
}
if let Some(v) = obj.get("top_p") {
gc.insert("topP".to_string(), v.clone());
}
if let Some(v) = obj.get("top_k") {
gc.insert("topK".to_string(), v.clone());
}
if let Some(v) = obj.get("max_tokens").or(obj.get("max_completion_tokens")) {
gc.insert("maxOutputTokens".to_string(), v.clone());
}
if let Some(v) = obj.get("stop") {
gc.insert("stopSequences".to_string(), v.clone());
}
let effort = obj
.get("reasoning_effort")
.and_then(|e| e.as_str())
.or_else(|| {
obj.get("output_config")
.and_then(|oc| oc.get("effort"))
.and_then(|e| e.as_str())
});
let pro = obj
.get("reasoning_mode")
.and_then(|m| m.as_str())
.map(|m| m == "pro")
.unwrap_or(false);
let level = effort.map(|e| {
if e == "none" {
return if cannot_disable_thinking(model) {
"low"
} else {
"minimal"
};
}
let base = match e {
"minimal" | "low" => "low",
"medium" => "medium",
"high" | "xhigh" | "max" => "high",
_ => "medium",
};
if pro {
match base {
"low" => "medium",
_ => "high",
}
} else {
base
}
});
{
let mut thinking_config = json!({"includeThoughts": true});
if let Some(lvl) = level {
thinking_config["thinkingLevel"] = json!(lvl);
}
gc.insert("thinkingConfig".to_string(), thinking_config);
}
if let Some(ext) = obj.get("x-gemini").and_then(|e| e.as_object()) {
if let Some(tc) = ext.get("thinkingConfig") {
gc.insert("thinkingConfig".to_string(), tc.clone());
}
for (k, v) in ext {
if k != "thinkingConfig" {
body_obj.insert(k.clone(), v.clone());
}
}
}
if !gc.is_empty() {
body_obj.insert("generationConfig".to_string(), gen_config);
}
if let Some(tools) = obj.get("tools").and_then(|t| t.as_array()) {
body_obj.insert("tools".to_string(), transform_tools(tools));
}
if let Some(tc) = obj.get("tool_choice") {
if let Some(config) = translate_tool_choice(tc) {
body_obj.insert("toolConfig".to_string(), config);
}
}
let is_stream = obj.get("stream").and_then(|s| s.as_bool()).unwrap_or(false);
let method = if is_stream {
"streamGenerateContent"
} else {
"generateContent"
};
let mut url = format!(
"{}/models/{}:{}?key={}",
self.base_url, model, method, self.api_key
);
if is_stream {
url.push_str("&alt=sse");
}
Ok(ProviderRequest {
url,
headers: vec![("Content-Type".into(), "application/json".into())],
body,
})
}
fn transform_response(&self, model: &str, response: Value) -> Result<Value> {
if let Some(err) = response.get("error") {
let msg = err
.get("message")
.and_then(|m| m.as_str())
.unwrap_or("unknown error");
let code = err.get("code").and_then(|c| c.as_u64()).unwrap_or(400) as u16;
return Err(ShimError::ProviderError {
status: code,
body: msg.to_string(),
});
}
transform_response_to_openai(model, &response)
}
fn transform_stream_chunk(&self, model: &str, chunk: &str) -> Result<Option<String>> {
let trimmed = chunk.trim();
if trimmed.is_empty() {
return Ok(None);
}
let parsed: Value = serde_json::from_str(trimmed)?;
if parsed.get("error").is_some() {
return self.transform_response(model, parsed).map(|_| None);
}
let candidate = match parsed
.get("candidates")
.and_then(|c| c.as_array())
.and_then(|a| a.first())
{
Some(c) => c,
None => return Ok(None),
};
let parts = candidate
.pointer("/content/parts")
.and_then(|p| p.as_array())
.cloned()
.unwrap_or_default();
let mut text = String::new();
let mut thought_text = String::new();
let mut has_function_call = false;
let mut tool_calls: Vec<Value> = Vec::new();
for part in &parts {
let is_thought = part
.get("thought")
.and_then(|t| t.as_bool())
.unwrap_or(false);
if let Some(t) = part.get("text").and_then(|t| t.as_str()) {
if !t.is_empty() {
if is_thought {
thought_text.push_str(t);
} else {
text.push_str(t);
}
}
}
if let Some(fc) = part.get("functionCall") {
let name = fc.get("name").and_then(|n| n.as_str()).unwrap_or("");
if name.is_empty() {
continue;
}
has_function_call = true;
let id = fc
.get("id")
.and_then(|v| v.as_str())
.map(|s| s.to_string())
.unwrap_or_else(|| format!("call_{}", tool_calls.len()));
let args_str = fc
.get("args")
.filter(|a| !a.is_null())
.map(|a| serde_json::to_string(a).unwrap_or_else(|_| "{}".to_string()))
.unwrap_or_else(|| "{}".to_string());
let mut tc = json!({
"index": tool_calls.len(),
"id": id,
"type": "function",
"function": {
"name": name,
"arguments": args_str,
}
});
if let Some(sig) = part.get("thoughtSignature") {
tc["thought_signature"] = sig.clone();
}
tool_calls.push(tc);
}
}
let finish_reason =
candidate
.get("finishReason")
.and_then(|f| f.as_str())
.map(|f| match f {
"STOP" => "stop",
"MAX_TOKENS" => "length",
"SAFETY" => "content_filter",
_ => "stop",
});
let mut delta = json!({});
if !thought_text.is_empty() {
delta["reasoning_content"] = json!(thought_text);
}
if !text.is_empty() {
delta["content"] = json!(text);
}
if has_function_call {
delta["tool_calls"] = json!(tool_calls);
}
if delta.as_object().map(|o| o.is_empty()).unwrap_or(true) && finish_reason.is_none() {
return Ok(None);
}
let mut chunk_json = json!({
"object": "chat.completion.chunk",
"model": model,
"choices": [{
"index": 0,
"delta": delta,
"finish_reason": finish_reason,
}]
});
if finish_reason.is_some() {
if let Some(usage) = parsed.get("usageMetadata") {
chunk_json["usage"] = json!({
"prompt_tokens": usage.get("promptTokenCount").cloned().unwrap_or(json!(0)),
"completion_tokens": usage.get("candidatesTokenCount").cloned().unwrap_or(json!(0)),
});
}
}
Ok(Some(serde_json::to_string(&chunk_json)?))
}
}