Fix token counting requests in Gemini (#29643)

Michael Sloan created 7 months ago

Release Notes:

- N/A

Change summary

crates/google_ai/src/google_ai.rs             | 6 ++----
crates/language_models/src/provider/google.rs | 4 +++-
2 files changed, 5 insertions(+), 5 deletions(-)

Detailed changes

crates/google_ai/src/google_ai.rs

@@ -76,12 +76,10 @@ pub async fn count_tokens(
     client: &dyn HttpClient,
     api_url: &str,
     api_key: &str,
+    model_id: &str,
     request: CountTokensRequest,
 ) -> Result<CountTokensResponse> {
-    let uri = format!(
-        "{}/v1beta/models/gemini-pro:countTokens?key={}",
-        api_url, api_key
-    );
+    let uri = format!("{api_url}/v1beta/models/{model_id}:countTokens?key={api_key}",);
     let request = serde_json::to_string(&request)?;
 
     let request_builder = HttpRequest::builder()

crates/language_models/src/provider/google.rs

@@ -327,7 +327,8 @@ impl LanguageModel for GoogleLanguageModel {
         request: LanguageModelRequest,
         cx: &App,
     ) -> BoxFuture<'static, Result<usize>> {
-        let request = into_google(request, self.model.id().to_string());
+        let model_id = self.model.id().to_string();
+        let request = into_google(request, model_id.clone());
         let http_client = self.http_client.clone();
         let api_key = self.state.read(cx).api_key.clone();
 
@@ -340,6 +341,7 @@ impl LanguageModel for GoogleLanguageModel {
                 http_client.as_ref(),
                 &api_url,
                 &api_key,
+                &model_id,
                 google_ai::CountTokensRequest {
                     contents: request.contents,
                 },