//! Gemini API types, request builders, and response parsers. use serde::{Deserialize, Serialize}; use serde_json::{Value, json}; use super::super::super::TokenUsage; use super::super::RuntimeContext; // ── Gemini API types ───────────────────────────────────────────────── #[derive(Debug, Serialize, Deserialize)] pub(super) struct GeminiFunctionDeclaration { pub name: String, pub description: String, #[serde(skip_serializing_if = "Option::is_none")] pub parameters: Option, } // ── Request builders ───────────────────────────────────────────────── /// Build the system instruction content from the RuntimeContext. pub(super) fn build_system_instruction(ctx: &RuntimeContext) -> Value { // Use system_prompt from args if provided via --append-system-prompt, // otherwise use a sensible default. let system_text = ctx .args .iter() .position(|a| a == "--append-system-prompt") .and_then(|i| ctx.args.get(i + 1)) .cloned() .unwrap_or_else(|| { format!( "You are an AI coding agent working on story {}. \ You have access to tools via function calling. \ Use them to complete the task. \ Work in the directory: {}", ctx.story_id, ctx.cwd ) }); json!({ "parts": [{ "text": system_text }] }) } /// Build the full `generateContent` request body. pub(super) fn build_generate_content_request( system_instruction: &Value, contents: &[Value], gemini_tools: &[GeminiFunctionDeclaration], ) -> Value { let mut body = json!({ "system_instruction": system_instruction, "contents": contents, "generationConfig": { "temperature": 0.2, "maxOutputTokens": 65536, } }); if !gemini_tools.is_empty() { body["tools"] = json!([{ "functionDeclarations": gemini_tools }]); } body } // ── Response parsing ───────────────────────────────────────────────── /// Parse token usage metadata from a Gemini API response. pub(super) fn parse_usage_metadata(response: &Value) -> Option { let metadata = response.get("usageMetadata")?; Some(TokenUsage { input_tokens: metadata .get("promptTokenCount") .and_then(|v| v.as_u64()) .unwrap_or(0), output_tokens: metadata .get("candidatesTokenCount") .and_then(|v| v.as_u64()) .unwrap_or(0), // Gemini doesn't have cache token fields, but we keep the struct uniform. cache_creation_input_tokens: 0, cache_read_input_tokens: 0, // Google AI API doesn't report cost; leave at 0. total_cost_usd: 0.0, }) } // ── Tests ──────────────────────────────────────────────────────────── #[cfg(test)] mod tests { use super::*; use crate::http::context::AppContext; use std::sync::Arc; fn test_app_ctx() -> Arc { let tmp = tempfile::tempdir().unwrap(); Arc::new(AppContext::new_test(tmp.path().to_path_buf())) } #[test] fn build_system_instruction_uses_args() { let ctx = RuntimeContext { story_id: "42_story_test".to_string(), agent_name: "coder-1".to_string(), command: "gemini-2.5-pro".to_string(), args: vec![ "--append-system-prompt".to_string(), "Custom system prompt".to_string(), ], prompt: "Do the thing".to_string(), cwd: "/tmp/wt".to_string(), inactivity_timeout_secs: 300, app_ctx: Some(test_app_ctx()), session_id_to_resume: None, fresh_prompt: None, project_root: std::path::PathBuf::from("/tmp/project"), model: None, }; let instruction = build_system_instruction(&ctx); assert_eq!(instruction["parts"][0]["text"], "Custom system prompt"); } #[test] fn build_system_instruction_default() { let ctx = RuntimeContext { story_id: "42_story_test".to_string(), agent_name: "coder-1".to_string(), command: "gemini-2.5-pro".to_string(), args: vec![], prompt: "Do the thing".to_string(), cwd: "/tmp/wt".to_string(), inactivity_timeout_secs: 300, app_ctx: Some(test_app_ctx()), session_id_to_resume: None, fresh_prompt: None, project_root: std::path::PathBuf::from("/tmp/project"), model: None, }; let instruction = build_system_instruction(&ctx); let text = instruction["parts"][0]["text"].as_str().unwrap(); assert!(text.contains("42_story_test")); assert!(text.contains("/tmp/wt")); } #[test] fn build_generate_content_request_includes_tools() { let system = json!({"parts": [{"text": "system"}]}); let contents = vec![json!({"role": "user", "parts": [{"text": "hello"}]})]; let tools = vec![GeminiFunctionDeclaration { name: "my_tool".to_string(), description: "A tool".to_string(), parameters: Some(json!({"type": "object", "properties": {"x": {"type": "string"}}})), }]; let body = build_generate_content_request(&system, &contents, &tools); assert!(body["tools"][0]["functionDeclarations"].is_array()); assert_eq!( body["tools"][0]["functionDeclarations"][0]["name"], "my_tool" ); } #[test] fn build_generate_content_request_no_tools() { let system = json!({"parts": [{"text": "system"}]}); let contents = vec![json!({"role": "user", "parts": [{"text": "hello"}]})]; let tools: Vec = vec![]; let body = build_generate_content_request(&system, &contents, &tools); assert!(body.get("tools").is_none()); } #[test] fn parse_usage_metadata_valid() { let response = json!({ "usageMetadata": { "promptTokenCount": 100, "candidatesTokenCount": 50, "totalTokenCount": 150 } }); let usage = parse_usage_metadata(&response).unwrap(); assert_eq!(usage.input_tokens, 100); assert_eq!(usage.output_tokens, 50); assert_eq!(usage.cache_creation_input_tokens, 0); assert_eq!(usage.total_cost_usd, 0.0); } #[test] fn parse_usage_metadata_missing() { let response = json!({"candidates": []}); assert!(parse_usage_metadata(&response).is_none()); } }