From 8440ec03adf1babda4ff03ccfecaf7aa19b4d092 Mon Sep 17 00:00:00 2001
From: Michael Sloan <michael@zed.dev>
Date: Tue, 29 Apr 2025 23:32:51 -0600
Subject: [PATCH] Fixes after merge

---
 crates/google_ai/examples/google_ai_cli.rs | 25 ++++++++++++++++------
 1 file changed, 18 insertions(+), 7 deletions(-)

diff --git a/crates/google_ai/examples/google_ai_cli.rs b/crates/google_ai/examples/google_ai_cli.rs
index d32f6b02d7..5f8327ccb8 100644
--- a/crates/google_ai/examples/google_ai_cli.rs
+++ b/crates/google_ai/examples/google_ai_cli.rs
@@ -166,7 +166,7 @@ async fn main() -> Result<()> {
             };
 
             let request = GenerateContentRequest {
-                model: cli.model.clone(),
+                model: google_ai::ModelName { model_id: cli.model.clone() },
                 contents: vec![user_content],
                 system_instruction: system_instruction.as_ref().map(|instruction| {
                     SystemInstruction {
@@ -227,18 +227,29 @@ async fn main() -> Result<()> {
         } => {
             let prompt_text = get_prompt_text(prompt, prompt_file)?;
 
+            let user_content = Content {
+                role: Role::User,
+                parts: vec![Part::TextPart(TextPart { text: prompt_text })],
+            };
+
+            let generate_content_request = GenerateContentRequest {
+                model: google_ai::ModelName { model_id: cli.model.clone() },
+                contents: vec![user_content],
+                system_instruction: None,
+                generation_config: None,
+                safety_settings: None,
+                tools: None,
+                tool_config: None,
+            };
+
             let request = CountTokensRequest {
-                contents: vec![Content {
-                    role: Role::User,
-                    parts: vec![Part::TextPart(TextPart { text: prompt_text })],
-                }],
+                generate_content_request,
             };
 
             let response = count_tokens(
                 http_client.as_ref(),
                 &cli.api_url,
                 &cli.api_key,
-                &cli.model,
                 request,
             )
             .await?;
@@ -356,7 +367,7 @@ async fn main() -> Result<()> {
 
                 // Create request with history
                 let request = GenerateContentRequest {
-                    model: cli.model.clone(),
+                    model: google_ai::ModelName { model_id: cli.model.clone() },
                     contents: history
                         .iter()
                         .map(|content| Content {