use crate::ai::token_budget::TokenBudget;
use crate::ai::{
AiErrorClass, AiProvider, AiRequest, AiResponse, AiRole, AiUsage, ClassifyAiError,
ProviderCapabilities, ToolCall, classify_status_code,
};
use crate::utils::redact_secret;
use anyhow::{Context, Result};
use async_trait::async_trait;
use regex::Regex;
use reqwest::Client;
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use std::time::Duration;
use tokio::sync::RwLock;
#[derive(Debug, Serialize, Deserialize, Clone, Default)]
#[serde(rename_all = "camelCase")]
pub struct Content {
#[serde(default)]
pub role: String,
#[serde(default)]
pub parts: Vec<Part>,
}
#[derive(Debug, Serialize, Deserialize, Clone)]
#[serde(untagged)]
pub enum Part {
Text {
text: String,
#[serde(rename = "thoughtSignature", skip_serializing_if = "Option::is_none")]
thought_signature: Option<String>,
#[serde(default)]
thought: bool,
},
FunctionCall {
#[serde(rename = "functionCall")]
function_call: FunctionCall,
#[serde(rename = "thoughtSignature", skip_serializing_if = "Option::is_none")]
thought_signature: Option<String>,
},
FunctionResponse {
#[serde(rename = "functionResponse")]
function_response: FunctionResponse,
},
}
#[derive(Debug, Serialize, Deserialize, Clone)]
#[serde(rename_all = "camelCase")]
pub struct FunctionCall {
pub name: String,
pub args: Value,
}
#[derive(Debug, Serialize, Deserialize, Clone)]
#[serde(rename_all = "camelCase")]
pub struct FunctionResponse {
pub name: String,
pub response: Value,
}
#[derive(Debug, Serialize, Deserialize, Clone)]
#[serde(rename_all = "camelCase")]
pub struct Tool {
pub function_declarations: Vec<FunctionDeclaration>,
}
#[derive(Debug, Serialize, Deserialize, Clone)]
#[serde(rename_all = "camelCase")]
pub struct FunctionDeclaration {
pub name: String,
pub description: String,
pub parameters: Value, }
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct GenerateContentRequest {
pub contents: Vec<Content>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tools: Option<Vec<Tool>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub system_instruction: Option<Content>,
#[serde(skip_serializing_if = "Option::is_none")]
pub generation_config: Option<GenerationConfig>,
}
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ThinkingConfig {
pub include_thoughts: bool,
}
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct GenerationConfig {
#[serde(skip_serializing_if = "Option::is_none")]
pub response_mime_type: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub response_schema: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub temperature: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub thinking_config: Option<ThinkingConfig>,
}
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct PromptFeedback {
pub block_reason: Option<String>,
}
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct GenerateContentResponse {
pub candidates: Option<Vec<Candidate>>,
pub prompt_feedback: Option<PromptFeedback>,
pub usage_metadata: Option<UsageMetadata>,
}
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct Candidate {
#[serde(default)]
pub content: Content,
pub finish_reason: Option<String>,
}
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct UsageMetadata {
pub prompt_token_count: u32,
pub candidates_token_count: Option<u32>,
pub total_token_count: u32,
#[serde(skip_serializing_if = "Option::is_none")]
pub cached_content_token_count: Option<u32>,
#[serde(flatten)]
pub extra: Option<std::collections::HashMap<String, Value>>,
}
#[async_trait]
pub trait GenAiClient: Send + Sync {
async fn generate_content(
&self,
request: GenerateContentRequest,
) -> Result<GenerateContentResponse>;
}
fn extract_gemini_error_reason(error_text: &str) -> Option<String> {
let json = serde_json::from_str::<serde_json::Value>(error_text).ok()?;
let err = json.get("error")?;
let reason = err
.get("details")
.and_then(|d| d.as_array())
.and_then(|a| a.first())
.and_then(|f| f.get("reason"))
.and_then(|r| r.as_str());
let message = err.get("message").and_then(|m| m.as_str());
if let Some(r) = reason {
return Some(r.to_string());
} else if let Some(m) = message {
return Some(m.to_string());
}
None
}
#[derive(Debug)]
pub enum GeminiError {
QuotaExceeded(Duration),
TransientError(Duration, String),
PermissionDenied(String),
ApiError(reqwest::StatusCode, String),
Other(String),
}
impl std::fmt::Display for GeminiError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
GeminiError::QuotaExceeded(d) => write!(f, "Quota exceeded, retry in {:?}", d),
GeminiError::TransientError(d, msg) => {
write!(f, "Transient error: {}, retry in {:?}", msg, d)
}
GeminiError::PermissionDenied(msg) => {
if let Some(reason) = extract_gemini_error_reason(msg) {
write!(f, "Permission denied (reason: {}): {}", reason, msg)
} else {
write!(f, "Permission denied: {}", msg)
}
}
GeminiError::ApiError(s, msg) => {
if let Some(reason) = extract_gemini_error_reason(msg) {
write!(f, "Gemini API error ({}) (reason: {}): {}", s, reason, msg)
} else {
write!(f, "Gemini API error ({}): {}", s, msg)
}
}
GeminiError::Other(msg) => write!(f, "{}", msg),
}
}
}
impl std::error::Error for GeminiError {}
impl ClassifyAiError for GeminiError {
fn ai_error_class(&self) -> AiErrorClass {
match self {
GeminiError::QuotaExceeded(retry_after) => AiErrorClass::RateLimit {
retry_after: *retry_after,
},
GeminiError::TransientError(retry_after, _) => AiErrorClass::Transient {
retry_after: *retry_after,
},
GeminiError::PermissionDenied(_) => AiErrorClass::Fatal,
GeminiError::ApiError(status, _) => {
classify_status_code(*status).unwrap_or(AiErrorClass::Fatal)
}
GeminiError::Other(_) => AiErrorClass::Fatal,
}
}
}
pub struct GeminiClient {
model: String,
base_url: String,
api_key: String,
client: RwLock<Client>,
}
impl GeminiClient {
pub fn new(model: String) -> Self {
let api_key = std::env::var("GEMINI_API_KEY")
.or_else(|_| std::env::var("LLM_API_KEY"))
.unwrap_or_default();
let base_url = std::env::var("GOOGLE_GEMINI_BASE_URL")
.or_else(|_| std::env::var("GEMINI_BASE_URL"))
.unwrap_or_else(|_| "https://generativelanguage.googleapis.com".to_string());
let client = Self::create_http_client(&api_key);
Self {
model,
base_url,
api_key,
client: RwLock::new(client),
}
}
fn create_http_client(api_key: &str) -> Client {
let mut headers = reqwest::header::HeaderMap::new();
if !api_key.is_empty()
&& let Ok(value) = reqwest::header::HeaderValue::from_str(api_key)
{
headers.insert("x-goog-api-key", value);
}
reqwest::Client::builder()
.default_headers(headers)
.timeout(Duration::from_secs(600))
.tcp_keepalive(Duration::from_secs(60))
.pool_idle_timeout(Duration::from_secs(90))
.pool_max_idle_per_host(5)
.build()
.unwrap_or_else(|_| reqwest::Client::new())
}
async fn refresh_client(&self) {
tracing::warn!("Refreshing Gemini HTTP client due to transport errors...");
let new_client = Self::create_http_client(&self.api_key);
let mut guard = self.client.write().await;
*guard = new_client;
}
pub async fn generate_content_single(
&self,
request: &GenerateContentRequest,
) -> Result<GenerateContentResponse> {
tracing::info!("Sending Gemini request...");
let url = format!(
"{}/v1beta/models/{}:generateContent",
self.base_url, self.model
);
self.post_request(&url, request).await
}
async fn post_request<T: Serialize>(
&self,
url: &str,
body: &T,
) -> Result<GenerateContentResponse> {
let re = Regex::new(r"Please retry in ([0-9.]+)s").unwrap();
let client = self.client.read().await.clone();
let res = match client.post(url).json(body).send().await {
Ok(res) => res,
Err(e) => {
let err_str = redact_secret(&format!("{:#}", anyhow::Error::from(e)));
tracing::error!("Gemini request failed (transport): {}", err_str);
self.refresh_client().await;
return Err(GeminiError::TransientError(Duration::from_secs(30), err_str).into());
}
};
if res.status().is_success() {
let body_text = res.text().await?;
match serde_json::from_str::<GenerateContentResponse>(&body_text) {
Ok(response) => {
if let Some(usage) = &response.usage_metadata {
let cached = usage.cached_content_token_count.unwrap_or(0);
tracing::info!(
"{}Gemini response received. Tokens: in={}, cached={}, out={}",
crate::ai::get_log_prefix(),
usage.prompt_token_count.saturating_sub(cached),
cached,
usage.candidates_token_count.unwrap_or(0)
);
} else {
tracing::info!(
"{}Gemini response received. No usage metadata.",
crate::ai::get_log_prefix()
);
}
return Ok(response);
}
Err(e) => {
tracing::error!("Failed to decode Gemini response: {}", e);
anyhow::bail!("Failed to decode response: {}. Body: {}", e, body_text);
}
}
}
let status = res.status();
let retry_after_duration = res
.headers()
.get(reqwest::header::RETRY_AFTER)
.and_then(|h| h.to_str().ok())
.and_then(|s| s.parse::<f64>().ok());
let error_text = redact_secret(&res.text().await?);
if status == reqwest::StatusCode::TOO_MANY_REQUESTS {
let retry_seconds = if let Some(secs) = retry_after_duration {
secs
} else if let Some(caps) = re.captures(&error_text) {
caps[1].parse::<f64>().unwrap_or(30.0)
} else {
30.0
};
tracing::warn!(
"Gemini 429 Quota Exceeded. Retry suggested in {}s. Body: {}",
retry_seconds,
error_text
);
return Err(
GeminiError::QuotaExceeded(Duration::from_secs_f64(retry_seconds + 1.0)).into(),
);
}
if status == reqwest::StatusCode::FORBIDDEN {
let mut reason_str = String::new();
if let Some(reason) = extract_gemini_error_reason(&error_text) {
reason_str = format!(" (reason: {})", reason);
}
tracing::error!(
"Gemini Permission Denied (403){}: {}",
reason_str,
error_text
);
return Err(GeminiError::PermissionDenied(error_text).into());
}
let is_transient = status.is_server_error() || status.as_u16() == 499;
if is_transient {
let retry_duration = retry_after_duration
.map(Duration::from_secs_f64)
.unwrap_or(Duration::from_secs(0));
tracing::warn!(
"Gemini API Transient Error: status={}, body={}",
status,
error_text
);
return Err(GeminiError::TransientError(retry_duration, error_text).into());
}
let mut reason_str = String::new();
if let Some(reason) = extract_gemini_error_reason(&error_text) {
reason_str = format!(" (reason: {})", reason);
}
tracing::error!(
"Gemini API Error{}: status={}, body={}",
reason_str,
status,
error_text
);
Err(GeminiError::ApiError(status, error_text).into())
}
}
#[async_trait]
impl GenAiClient for GeminiClient {
async fn generate_content(
&self,
request: GenerateContentRequest,
) -> Result<GenerateContentResponse> {
self.generate_content_single(&request).await
}
}
pub struct StdioGeminiClient {
registry: std::sync::Arc<crate::ai::IpcRegistry>,
writer: std::sync::Arc<crate::ai::AtomicWriter>,
reader_started: std::sync::atomic::AtomicBool,
}
impl StdioGeminiClient {
pub fn new() -> Self {
let registry = std::sync::Arc::new(crate::ai::IpcRegistry::new());
let writer = std::sync::Arc::new(crate::ai::AtomicWriter::new());
Self {
registry,
writer,
reader_started: std::sync::atomic::AtomicBool::new(false),
}
}
}
impl Default for StdioGeminiClient {
fn default() -> Self {
Self::new()
}
}
#[async_trait]
impl AiProvider for StdioGeminiClient {
async fn generate_content(&self, request: AiRequest) -> Result<AiResponse> {
if !self
.reader_started
.swap(true, std::sync::atomic::Ordering::SeqCst)
{
crate::ai::start_stdin_reader(std::sync::Arc::downgrade(&self.registry));
}
let tx_id = self.registry.next_id();
let envelope = json!({
"type": "ai_request",
"tx_id": tx_id,
"payload": request
});
let line = serde_json::to_string(&envelope)?;
let (tx, rx) = tokio::sync::oneshot::channel();
self.registry.register(tx_id, tx).await;
self.writer.write_line(&line).await?;
match rx.await {
Ok(Ok(resp)) => Ok(resp),
Ok(Err(remote_err)) => Err(remote_err.into()),
Err(_) => Err(anyhow::anyhow!(
"IPC channel disconnected waiting for response"
)),
}
}
fn estimate_tokens(&self, request: &AiRequest) -> usize {
estimate_tokens_generic(request)
}
fn get_capabilities(&self) -> ProviderCapabilities {
ProviderCapabilities {
model_name: "stdio-gemini".to_string(),
context_window_size: 1_000_000,
}
}
}
fn translate_ai_request(request: AiRequest) -> Result<GenerateContentRequest> {
let mut contents: Vec<Content> = Vec::new();
let mut system_instruction = None;
if let Some(sys_content) = request.system {
system_instruction = Some(Content {
role: "user".to_string(),
parts: vec![Part::Text {
text: sys_content,
thought_signature: None,
thought: false,
}],
});
}
for msg in request.messages {
match msg.role {
AiRole::System => {
if let Some(content) = msg.content {
if system_instruction.is_none() {
system_instruction = Some(Content {
role: "user".to_string(), parts: vec![Part::Text {
text: content,
thought_signature: None,
thought: false,
}],
});
}
}
}
AiRole::User => {
let part = Part::Text {
text: msg.content.unwrap_or_default(),
thought_signature: None,
thought: false,
};
if let Some(last) = contents.last_mut()
&& last.role == "user"
{
last.parts.push(part);
continue;
}
contents.push(Content {
role: "user".to_string(),
parts: vec![part],
});
}
AiRole::Assistant => {
let mut parts = Vec::new();
if let Some(thought) = msg.thought {
parts.push(Part::Text {
text: thought,
thought_signature: None,
thought: true,
});
}
if let Some(text) = msg.content {
parts.push(Part::Text {
text,
thought_signature: None,
thought: false,
});
}
if let Some(tool_calls) = msg.tool_calls {
for call in tool_calls {
parts.push(Part::FunctionCall {
function_call: FunctionCall {
name: call.function_name,
args: call.arguments,
},
thought_signature: call.thought_signature,
});
}
}
contents.push(Content {
role: "model".to_string(),
parts,
});
}
AiRole::Tool => {
contents.push(Content {
role: "function".to_string(),
parts: vec![Part::FunctionResponse {
function_response: FunctionResponse {
name: msg
.tool_call_id
.context("Tool message missing tool_call_id")?,
response: serde_json::from_str(
&msg.content.unwrap_or_else(|| "{}".to_string()),
)
.unwrap_or(json!({})),
},
}],
});
}
}
}
let tools = request.tools.map(|t| {
vec![Tool {
function_declarations: t
.into_iter()
.map(|tool| FunctionDeclaration {
name: tool.name,
description: tool.description,
parameters: normalize_schema(tool.parameters),
})
.collect(),
}]
});
let mut response_mime_type = None;
let mut response_schema = None;
if let Some(format) = request.response_format {
match format {
crate::ai::AiResponseFormat::Text => {
response_mime_type = Some("text/plain".to_string());
}
crate::ai::AiResponseFormat::Json { schema } => {
response_mime_type = Some("application/json".to_string());
response_schema = schema.map(normalize_schema);
}
}
}
Ok(GenerateContentRequest {
contents,
tools,
system_instruction,
generation_config: Some(GenerationConfig {
response_mime_type,
response_schema,
temperature: request.temperature,
thinking_config: Some(ThinkingConfig {
include_thoughts: true,
}),
}),
})
}
fn normalize_schema(mut schema: Value) -> Value {
if let Some(obj) = schema.as_object_mut() {
if let Some(ty) = obj.get_mut("type")
&& let Some(s) = ty.as_str()
{
*ty = Value::String(s.to_uppercase());
}
for (_, val) in obj.iter_mut() {
if val.is_object() || val.is_array() {
*val = normalize_schema(val.clone());
}
}
} else if let Some(arr) = schema.as_array_mut() {
for val in arr.iter_mut() {
if val.is_object() || val.is_array() {
*val = normalize_schema(val.clone());
}
}
}
schema
}
fn translate_ai_response(resp: GenerateContentResponse) -> Result<AiResponse> {
if let Some(reason) = resp.prompt_feedback.and_then(|f| f.block_reason) {
return Err(anyhow::anyhow!(
"Gemini request blocked by prompt feedback (reason: {})",
reason
));
}
let candidate = resp
.candidates
.as_ref()
.and_then(|c| c.first())
.ok_or_else(|| anyhow::anyhow!("No candidates returned from Gemini"))?;
if let Some(finish_reason) = &candidate.finish_reason {
if finish_reason == "SAFETY"
|| finish_reason == "RECITATION"
|| finish_reason == "OTHER"
|| finish_reason == "BLOCKLIST"
|| finish_reason == "PROHIBITED_CONTENT"
{
return Err(anyhow::anyhow!(
"Gemini candidate blocked (finish reason: {})",
finish_reason
));
} else if finish_reason != "STOP" && finish_reason != "MAX_TOKENS" {
tracing::warn!(
"Gemini candidate finished with unexpected reason: {}",
finish_reason
);
}
}
let mut content = String::new();
let mut thought = String::new();
let mut tool_calls = Vec::new();
for part in &candidate.content.parts {
match part {
Part::Text {
text,
thought: is_thought,
..
} => {
if *is_thought {
thought.push_str(text);
} else {
content.push_str(text);
}
}
Part::FunctionCall {
function_call,
thought_signature,
} => {
tool_calls.push(ToolCall {
id: function_call.name.clone(), function_name: function_call.name.clone(),
arguments: function_call.args.clone(),
thought_signature: thought_signature.clone(),
});
}
_ => {}
}
}
let truncated = candidate
.finish_reason
.as_ref()
.map(|r| r == "MAX_TOKENS")
.unwrap_or(false);
if truncated {
tracing::warn!(
"{}Gemini response truncated due to MAX_TOKENS.",
crate::ai::get_log_prefix()
);
}
let usage = resp.usage_metadata.map(|m| AiUsage {
prompt_tokens: m.prompt_token_count as usize,
completion_tokens: m.candidates_token_count.unwrap_or(0) as usize,
total_tokens: m.total_token_count as usize,
cached_tokens: m.cached_content_token_count.map(|c| c as usize),
});
Ok(AiResponse {
content: if content.is_empty() {
None
} else {
Some(content)
},
thought: if thought.is_empty() {
None
} else {
Some(thought)
},
thought_signature: None,
tool_calls: if tool_calls.is_empty() {
None
} else {
Some(tool_calls)
},
usage,
truncated,
})
}
fn estimate_tokens_generic(request: &AiRequest) -> usize {
let mut total = 0;
for msg in &request.messages {
if let Some(content) = &msg.content {
total += TokenBudget::estimate_tokens(content);
}
if let Some(tool_calls) = &msg.tool_calls {
for call in tool_calls {
total += TokenBudget::estimate_tokens(&call.function_name);
total += TokenBudget::estimate_tokens(&call.arguments.to_string());
}
}
}
if let Some(tools) = &request.tools {
for tool in tools {
total += TokenBudget::estimate_tokens(&tool.name);
total += TokenBudget::estimate_tokens(&tool.description);
total += TokenBudget::estimate_tokens(&tool.parameters.to_string());
}
}
total
}
#[async_trait]
impl AiProvider for GeminiClient {
async fn generate_content(&self, request: AiRequest) -> Result<AiResponse> {
let gen_req = translate_ai_request(request)?;
let resp = GenAiClient::generate_content(self, gen_req).await?;
translate_ai_response(resp)
}
fn estimate_tokens(&self, request: &AiRequest) -> usize {
estimate_tokens_generic(request)
}
fn get_capabilities(&self) -> ProviderCapabilities {
ProviderCapabilities {
model_name: self.model.clone(),
context_window_size: 1_000_000, }
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ai::{
AiErrorClass, AiMessage, AiResponseFormat, AiRole, AiTool, ClassifyAiError,
DEFAULT_RETRY_AFTER, ToolCall,
};
use serde_json::json;
#[test]
fn test_quota_exceeded_classifies_as_rate_limit() {
let retry_after = Duration::from_secs(7);
let err = GeminiError::QuotaExceeded(retry_after);
assert_eq!(
err.ai_error_class(),
AiErrorClass::RateLimit { retry_after }
);
}
#[test]
fn test_transient_error_classifies_as_transient() {
let retry_after = Duration::from_secs(11);
let err = GeminiError::TransientError(retry_after, "busy".to_string());
assert_eq!(
err.ai_error_class(),
AiErrorClass::Transient { retry_after }
);
}
#[test]
fn test_permission_denied_classifies_as_fatal() {
let err = GeminiError::PermissionDenied("forbidden".to_string());
assert_eq!(err.ai_error_class(), AiErrorClass::Fatal);
}
#[test]
fn test_api_error_server_status_classifies_as_transient() {
let err = GeminiError::ApiError(
reqwest::StatusCode::SERVICE_UNAVAILABLE,
"unavailable".to_string(),
);
assert_eq!(
err.ai_error_class(),
AiErrorClass::Transient {
retry_after: DEFAULT_RETRY_AFTER,
}
);
}
#[test]
fn test_api_error_client_status_classifies_as_fatal() {
let err =
GeminiError::ApiError(reqwest::StatusCode::BAD_REQUEST, "bad request".to_string());
assert_eq!(err.ai_error_class(), AiErrorClass::Fatal);
}
#[test]
fn test_other_classifies_as_fatal() {
let err = GeminiError::Other("unknown".to_string());
assert_eq!(err.ai_error_class(), AiErrorClass::Fatal);
}
#[test]
fn test_translate_ai_request_system_and_user() -> Result<()> {
let request = AiRequest {
system: None,
messages: vec![
AiMessage {
role: AiRole::System,
content: Some("You are a helpful assistant.".to_string()),
thought: None,
thought_signature: None,
tool_calls: None,
tool_call_id: None,
},
AiMessage {
role: AiRole::User,
content: Some("Hello!".to_string()),
thought: None,
thought_signature: None,
tool_calls: None,
tool_call_id: None,
},
],
tools: None,
temperature: Some(0.7),
response_format: None,
context_tag: None,
};
let gemini_req = translate_ai_request(request)?;
assert!(gemini_req.system_instruction.is_some());
let sys_part = &gemini_req.system_instruction.unwrap().parts[0];
if let Part::Text { text, .. } = sys_part {
assert_eq!(text, "You are a helpful assistant.");
} else {
panic!("Expected Text part in system instruction");
}
assert_eq!(gemini_req.contents.len(), 1);
assert_eq!(gemini_req.contents[0].role, "user");
let user_part = &gemini_req.contents[0].parts[0];
if let Part::Text { text, .. } = user_part {
assert_eq!(text, "Hello!");
} else {
panic!("Expected Text part in user content");
}
assert_eq!(gemini_req.generation_config.unwrap().temperature, Some(0.7));
Ok(())
}
#[test]
fn test_translate_ai_request_assistant_tool_call() -> Result<()> {
let request = AiRequest {
system: None,
messages: vec![AiMessage {
role: AiRole::Assistant,
content: Some("I will use a tool.".to_string()),
thought: None,
thought_signature: None,
tool_calls: Some(vec![ToolCall {
id: "call_123".to_string(),
function_name: "test_tool".to_string(),
arguments: json!({"arg1": "val1"}),
thought_signature: Some("thought_sig_abc".to_string()),
}]),
tool_call_id: None,
}],
tools: None,
temperature: None,
response_format: None,
context_tag: None,
};
let gemini_req = translate_ai_request(request)?;
assert_eq!(gemini_req.contents.len(), 1);
assert_eq!(gemini_req.contents[0].role, "model");
assert_eq!(gemini_req.contents[0].parts.len(), 2);
if let Part::Text { text, .. } = &gemini_req.contents[0].parts[0] {
assert_eq!(text, "I will use a tool.");
} else {
panic!("Expected Text part first");
}
if let Part::FunctionCall {
function_call,
thought_signature,
} = &gemini_req.contents[0].parts[1]
{
assert_eq!(function_call.name, "test_tool");
assert_eq!(function_call.args["arg1"], "val1");
assert_eq!(thought_signature.as_deref(), Some("thought_sig_abc"));
} else {
panic!("Expected FunctionCall part second");
}
Ok(())
}
#[test]
fn test_translate_ai_response_with_thought_signature() -> Result<()> {
let gemini_resp = GenerateContentResponse {
candidates: Some(vec![Candidate {
content: Content {
role: "model".to_string(),
parts: vec![Part::FunctionCall {
function_call: FunctionCall {
name: "test_tool".to_string(),
args: json!({"arg1": "val1"}),
},
thought_signature: Some("thought_sig_xyz".to_string()),
}],
},
finish_reason: Some("STOP".to_string()),
}]),
prompt_feedback: None,
usage_metadata: Some(UsageMetadata {
prompt_token_count: 10,
candidates_token_count: Some(20),
total_token_count: 30,
cached_content_token_count: None,
extra: None,
}),
};
let ai_resp = translate_ai_response(gemini_resp)?;
assert!(ai_resp.content.is_none());
let tool_calls = ai_resp.tool_calls.unwrap();
assert_eq!(tool_calls.len(), 1);
assert_eq!(tool_calls[0].function_name, "test_tool");
assert_eq!(
tool_calls[0].thought_signature.as_deref(),
Some("thought_sig_xyz")
);
Ok(())
}
#[test]
fn test_translate_ai_request_tool_response() -> Result<()> {
let request = AiRequest {
system: None,
messages: vec![AiMessage {
role: AiRole::Tool,
content: Some(json!({"result": "success"}).to_string()),
thought: None,
thought_signature: None,
tool_calls: None,
tool_call_id: Some("call_123".to_string()),
}],
tools: None,
temperature: None,
response_format: None,
context_tag: None,
};
let gemini_req = translate_ai_request(request)?;
assert_eq!(gemini_req.contents.len(), 1);
assert_eq!(gemini_req.contents[0].role, "function");
if let Part::FunctionResponse { function_response } = &gemini_req.contents[0].parts[0] {
assert_eq!(function_response.name, "call_123");
assert_eq!(function_response.response["result"], "success");
} else {
panic!("Expected FunctionResponse part");
}
Ok(())
}
#[test]
fn test_translate_ai_request_json_format() -> Result<()> {
let schema = json!({
"type": "OBJECT",
"properties": {
"score": {"type": "NUMBER"}
}
});
let request = AiRequest {
system: None,
messages: vec![AiMessage {
role: AiRole::User,
content: Some("Score this.".to_string()),
thought: None,
thought_signature: None,
tool_calls: None,
tool_call_id: None,
}],
tools: None,
temperature: None,
response_format: Some(AiResponseFormat::Json {
schema: Some(schema.clone()),
}),
context_tag: None,
};
let gemini_req = translate_ai_request(request)?;
let config = gemini_req.generation_config.unwrap();
assert_eq!(
config.response_mime_type.as_deref(),
Some("application/json")
);
assert_eq!(config.response_schema.as_ref(), Some(&schema));
Ok(())
}
#[test]
fn test_translate_ai_request_conversation_chain() -> Result<()> {
let request = AiRequest {
system: None,
messages: vec![
AiMessage {
role: AiRole::User,
content: Some("Use tool".to_string()),
thought: None,
thought_signature: None,
tool_calls: None,
tool_call_id: None,
},
AiMessage {
role: AiRole::Assistant,
content: None,
thought: None,
thought_signature: None,
tool_calls: Some(vec![ToolCall {
id: "c1".to_string(),
function_name: "t1".to_string(),
arguments: json!({}),
thought_signature: Some("s1".to_string()),
}]),
tool_call_id: None,
},
AiMessage {
role: AiRole::Tool,
content: Some("{\"ok\":true}".to_string()),
thought: None,
thought_signature: None,
tool_calls: None,
tool_call_id: Some("c1".to_string()),
},
],
tools: Some(vec![AiTool {
name: "t1".to_string(),
description: "d1".to_string(),
parameters: json!({}),
}]),
temperature: None,
response_format: None,
context_tag: None,
};
let gemini_req = translate_ai_request(request)?;
assert_eq!(gemini_req.contents.len(), 3);
assert_eq!(gemini_req.contents[0].role, "user");
assert_eq!(gemini_req.contents[1].role, "model");
assert_eq!(gemini_req.contents[2].role, "function");
if let Part::FunctionCall {
thought_signature, ..
} = &gemini_req.contents[1].parts[0]
{
assert_eq!(thought_signature.as_deref(), Some("s1"));
} else {
panic!("Expected FunctionCall in middle of chain");
}
assert!(gemini_req.tools.is_some());
assert_eq!(
gemini_req.tools.as_ref().unwrap()[0].function_declarations[0].name,
"t1"
);
Ok(())
}
#[test]
fn test_estimate_tokens_logic() {
let request = AiRequest {
system: None,
messages: vec![
AiMessage {
role: AiRole::User,
content: Some("Short message".to_string()),
thought: None,
thought_signature: None,
tool_calls: None,
tool_call_id: None,
},
AiMessage {
role: AiRole::Assistant,
content: None,
thought: None,
thought_signature: None,
tool_calls: Some(vec![ToolCall {
id: "c1".to_string(),
function_name: "my_function".to_string(),
arguments: json!({"key": "value"}),
thought_signature: None,
}]),
tool_call_id: None,
},
],
tools: Some(vec![AiTool {
name: "my_function".to_string(),
description: "Does something".to_string(),
parameters: json!({"type": "OBJECT"}),
}]),
temperature: None,
response_format: None,
context_tag: None,
};
let tokens = estimate_tokens_generic(&request);
assert!(tokens > 10);
assert!(tokens < 200);
}
#[test]
fn test_normalize_schema() {
let schema = json!({
"type": "object",
"properties": {
"files": {
"type": "array",
"items": {
"type": "object",
"properties": {
"path": { "type": "string" },
"start_line": { "type": "integer" }
}
}
}
}
});
let normalized = normalize_schema(schema);
assert_eq!(normalized["type"], "OBJECT");
assert_eq!(normalized["properties"]["files"]["type"], "ARRAY");
assert_eq!(normalized["properties"]["files"]["items"]["type"], "OBJECT");
assert_eq!(
normalized["properties"]["files"]["items"]["properties"]["path"]["type"],
"STRING"
);
assert_eq!(
normalized["properties"]["files"]["items"]["properties"]["start_line"]["type"],
"INTEGER"
);
}
}