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

assistant: Fix Google AI provider not respecting low_speed_timeout_in_seconds (#17423)

Browse source 
Release Notes:

- Fixed an issue when using Google Gemini models, where the setting
`low_speed_timeout_in_seconds` was not respected
This commit is contained in:
Bennet Bo Fenner 2024-09-05 18:16:30 +02:00 • committed by GitHub
parent a1c676128a
commit f413ea90bf
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 50 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

26
crates/language_model/src/provider/google.rs
View file

@ -257,10 +257,10 @@ impl LanguageModel for GoogleLanguageModel {
let request = request.into_google(self.model.id().to_string());
let http_client = self.http_client.clone();
let api_key = self.state.read(cx).api_key.clone();
let api_url = AllLanguageModelSettings::get_global(cx)
.google
.api_url
.clone();
let settings = &AllLanguageModelSettings::get_global(cx).google;
let api_url = settings.api_url.clone();
let low_speed_timeout = settings.low_speed_timeout;
async move {
let api_key = api_key.ok_or_else(|| anyhow!("missing api key"))?;
@ -271,6 +271,7 @@ impl LanguageModel for GoogleLanguageModel {
google_ai::CountTokensRequest {
contents: request.contents,
},
low_speed_timeout,
)
.await?;
Ok(response.total_tokens)
@ -289,17 +290,26 @@ impl LanguageModel for GoogleLanguageModel {
let request = request.into_google(self.model.id().to_string());
let http_client = self.http_client.clone();
let Ok((api_key, api_url)) = cx.read_model(&self.state, |state, cx| {
let Ok((api_key, api_url, low_speed_timeout)) = cx.read_model(&self.state, |state, cx| {
let settings = &AllLanguageModelSettings::get_global(cx).google;
(state.api_key.clone(), settings.api_url.clone())
(
state.api_key.clone(),
settings.api_url.clone(),
settings.low_speed_timeout,
)
}) else {
return futures::future::ready(Err(anyhow!("App state dropped"))).boxed();
};
let future = self.rate_limiter.stream(async move {
let api_key = api_key.ok_or_else(|| anyhow!("missing api key"))?;
let response =
stream_generate_content(http_client.as_ref(), &api_url, &api_key, request);
let response = stream_generate_content(
http_client.as_ref(),
&api_url,
&api_key,
request,
low_speed_timeout,
);
let events = response.await?;
Ok(google_ai::extract_text_from_events(events).boxed())
});