vrr/goose
2
0
Fork 0
mirror of https://github.com/block/goose.git synced 2026-07-09 16:09:22 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 22

fix(litellm): use context limit from /model/info for custom models (#9303)

Browse source 
Signed-off-by: Douwe Osinga <douwe@squareup.com>
Co-authored-by: Douwe Osinga <douwe@squareup.com>
This commit is contained in:
Douwe Osinga 2026-05-22 14:10:31 -04:00 • committed by GitHub
parent f288290fc5
commit 612cd89d51
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 30 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

37
crates/goose/src/providers/litellm.rs
View file

@ -30,6 +30,8 @@ pub struct LiteLLMProvider {
model: ModelConfig,
#[serde(skip)]
name: String,
#[serde(skip)]
cached_model_info: tokio::sync::OnceCell<Vec<ModelInfo>>,
}
impl LiteLLMProvider {
@ -77,10 +79,18 @@ impl LiteLLMProvider {
base_path,
model,
name: LITELLM_PROVIDER_NAME.to_string(),
cached_model_info: tokio::sync::OnceCell::new(),
})
}
async fn fetch_models(&self) -> Result<Vec<ModelInfo>, ProviderError> {
async fn get_or_fetch_models(&self) -> Result<&[ModelInfo], ProviderError> {
self.cached_model_info
.get_or_try_init(|| self.fetch_models_from_api())
.await
.map(|v| v.as_slice())
}
async fn fetch_models_from_api(&self) -> Result<Vec<ModelInfo>, ProviderError> {
let response = self
.api_client
.request(None, "model/info")
@ -184,7 +194,22 @@ impl Provider for LiteLLMProvider {
}
fn get_model_config(&self) -> ModelConfig {
self.model.clone()
let mut config = self.model.clone();
// The cache is populated lazily by the first stream() call (via
// supports_cache_control). On turn 1 this will be None and we fall
// back to DEFAULT_CONTEXT_LIMIT, which is fine — the conversation is
// too small to trigger compaction. From turn 2 onward the real limit
// from /model/info is used.
if config.context_limit.is_none() {
if let Some(models) = self.cached_model_info.get() {
if let Some(info) = models.iter().find(|m| m.name == config.model_name) {
if info.context_limit > 0 {
config.context_limit = Some(info.context_limit);
}
}
}
}
config
}
async fn stream(
@ -237,7 +262,7 @@ impl Provider for LiteLLMProvider {
}
async fn supports_cache_control(&self) -> bool {
if let Ok(models) = self.fetch_models().await {
if let Ok(models) = self.get_or_fetch_models().await {
if let Some(model_info) = models.iter().find(|m| m.name == self.model.model_name) {
return model_info.supports_cache_control.unwrap_or(false);
}
@ -247,10 +272,8 @@ impl Provider for LiteLLMProvider {
}
async fn fetch_supported_models(&self) -> Result<Vec<String>, ProviderError> {
let models = self.fetch_models().await.map_err(|e| {
ProviderError::RequestFailed(format!("Failed to fetch models from LiteLLM: {}", e))
})?;
Ok(models.into_iter().map(|m| m.name).collect())
let models = self.get_or_fetch_models().await?;
Ok(models.iter().map(|m| m.name.clone()).collect())
}
}