Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix token counting requests in Gemini (#29643)

Browse source 
Release Notes:

- N/A
This commit is contained in:
Michael Sloan 2025-04-29 22:55:07 -06:00 • committed by GitHub
parent 8d77efa781
commit edf78e770d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 5 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

6
crates/google_ai/src/google_ai.rs
View file

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

4
crates/language_models/src/provider/google.rs
View file

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