huskies: merge 803

2026-04-28 19:05:14 +00:00
parent fa54451ba6
commit e4af2d5c08
4 changed files with 853 additions and 815 deletions
@@ -0,0 +1,190 @@
+//! 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<Value>,
+}
+
+// ── 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<TokenUsage> {
+    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::*;
+
+    #[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,
+            mcp_port: 3001,
+            session_id_to_resume: None,
+            fresh_prompt: 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,
+            mcp_port: 3001,
+            session_id_to_resume: None,
+            fresh_prompt: 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<GeminiFunctionDeclaration> = 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());
+    }
+}