tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

Use Canonical Models to set context window sizes (#6723)

Browse Source
This commit is contained in:
David Katz
2026-02-17 11:43:10 -05:00
committed by GitHub
parent 576590d4c8
commit 3959805198
54 changed files with 465 additions and 581 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-acp/src/server.rs
+13 -3
View File
@@ -738,7 +738,8 @@ impl GooseAcpAgent {
let config_path = self.config_dir.join(CONFIG_YAML_NAME);
let config = Config::new(&config_path, "goose")?;
let model_id = config.get_goose_model()?;
goose::model::ModelConfig::new(&model_id)?
let provider_name = config.get_goose_provider()?;
goose::model::ModelConfig::new(&model_id)?.with_canonical_limits(&provider_name)
}
};
let provider = (self.provider_factory)(model_config, Vec::new()).await?;
@@ -956,9 +957,18 @@ impl GooseAcpAgent {
session_id: &str,
model_id: &str,
) -> Result<SetSessionModelResponse, sacp::Error> {
let model_config = goose::model::ModelConfig::new(model_id).map_err(|e| {
sacp::Error::invalid_params().data(format!("Invalid model config: {}", e))
let config_path = self.config_dir.join(CONFIG_YAML_NAME);
let config = Config::new(&config_path, "goose").map_err(|e| {
sacp::Error::internal_error().data(format!("Failed to read config: {}", e))
})?;
let provider_name = config.get_goose_provider().map_err(|_| {
sacp::Error::internal_error().data("No provider configured".to_string())
})?;
let model_config = goose::model::ModelConfig::new(model_id)
.map_err(|e| {
sacp::Error::invalid_params().data(format!("Invalid model config: {}", e))
})?
.with_canonical_limits(&provider_name);
let provider = (self.provider_factory)(model_config, Vec::new())
.await
.map_err(|e| {
crates/goose-acp/src/server_factory.rs
+1 -3
View File
@@ -38,9 +38,7 @@ impl AcpServer {
Box::pin(async move {
let config_path = config_dir.join(goose::config::base::CONFIG_YAML_NAME);
let config = goose::config::Config::new(&config_path, "goose")?;
let provider_name = config
.get_goose_provider()
.map_err(|_| anyhow::anyhow!("No provider configured"))?;
let provider_name = config.get_goose_provider()?;
goose::providers::create(&provider_name, model_config, extensions).await
})
});
crates/goose-acp/tests/common_tests/mod.rs
+1 -1
View File
@@ -25,7 +25,7 @@ pub async fn run_config_mcp<C: Connection>() {
let mcp = McpFixture::new(Some(expected_session_id.clone())).await;
let config_yaml = format!(
"GOOSE_MODEL: {TEST_MODEL}\nextensions:\n mcp-fixture:\n enabled: true\n type: streamable_http\n name: mcp-fixture\n description: MCP fixture\n uri: \"{}\"\n",
"GOOSE_MODEL: {TEST_MODEL}\nGOOSE_PROVIDER: openai\nextensions:\n mcp-fixture:\n enabled: true\n type: streamable_http\n name: mcp-fixture\n description: MCP fixture\n uri: \"{}\"\n",
mcp.url
);
fs::write(temp_dir.path().join(CONFIG_YAML_NAME), config_yaml).unwrap();
crates/goose-acp/tests/fixtures/mod.rs
Vendored
+5 -1
View File
@@ -199,7 +199,11 @@ pub async fn spawn_acp_server_in_process(
fs::create_dir_all(data_root).unwrap();
let config_path = data_root.join(goose::config::base::CONFIG_YAML_NAME);
if !config_path.exists() {
fs::write(&config_path, format!("GOOSE_MODEL: {TEST_MODEL}\n")).unwrap();
fs::write(
&config_path,
format!("GOOSE_MODEL: {TEST_MODEL}\nGOOSE_PROVIDER: openai\n"),
)
.unwrap();
}
let provider_factory = provider_factory.unwrap_or_else(|| {
let base_url = openai_base_url.to_string();
crates/goose-cli/src/commands/configure.rs
+6 -5
View File
@@ -327,7 +327,7 @@ async fn handle_oauth_configuration(provider_name: &str, key_name: &str) -> anyh
));
// Create a temporary provider instance to handle OAuth
let temp_model = ModelConfig::new("temp")?;
let temp_model = ModelConfig::new("temp")?.with_canonical_limits(provider_name);
match create(provider_name, temp_model, Vec::new()).await {
Ok(provider) => match provider.configure_oauth().await {
Ok(_) => {
@@ -682,7 +682,8 @@ pub async fn configure_provider_dialog() -> anyhow::Result<bool> {
let spin = spinner();
spin.start("Attempting to fetch supported models...");
let models_res = {
let temp_model_config = ModelConfig::new(&provider_meta.default_model)?;
let temp_model_config =
ModelConfig::new(&provider_meta.default_model)?.with_canonical_limits(provider_name);
let temp_provider = create(provider_name, temp_model_config, Vec::new()).await?;
retry_operation(&RetryConfig::default(), || async {
temp_provider.fetch_recommended_models().await
@@ -1442,7 +1443,7 @@ pub async fn configure_tool_permissions_dialog() -> anyhow::Result<()> {
let model: String = config
.get_goose_model()
.expect("No model configured. Please set model first");
let model_config = ModelConfig::new(&model)?;
let model_config = ModelConfig::new(&model)?.with_canonical_limits(&provider_name);
let agent = Agent::new();
@@ -1662,7 +1663,7 @@ pub async fn handle_openrouter_auth() -> anyhow::Result<()> {
println!("\nTesting configuration...");
let configured_model: String = config.get_goose_model()?;
let model_config = match goose::model::ModelConfig::new(&configured_model) {
Ok(config) => config,
Ok(config) => config.with_canonical_limits("openrouter"),
Err(e) => {
eprintln!("⚠️ Invalid model configuration: {}", e);
eprintln!("Your settings have been saved. Please check your model configuration.");
@@ -1742,7 +1743,7 @@ pub async fn handle_tetrate_auth() -> anyhow::Result<()> {
println!("\nTesting configuration...");
let configured_model: String = config.get_goose_model()?;
let model_config = match goose::model::ModelConfig::new(&configured_model) {
Ok(config) => config,
Ok(config) => config.with_canonical_limits("tetrate"),
Err(e) => {
eprintln!("⚠️ Invalid model configuration: {}", e);
eprintln!("Your settings have been saved. Please check your model configuration.");
crates/goose-cli/src/commands/term.rs
+7 -1
View File
@@ -291,7 +291,13 @@ pub async fn handle_term_info() -> Result<()> {
let context_limit = config
.get_goose_model()
.ok()
.and_then(|model_name| goose::model::ModelConfig::new(&model_name).ok())
.and_then(|model_name| {
config.get_goose_provider().ok().and_then(|provider_name| {
goose::model::ModelConfig::new(&model_name)
.ok()
.map(|c| c.with_canonical_limits(&provider_name))
})
})
.map(|mc| mc.context_limit())
.unwrap_or(128_000);
crates/goose-cli/src/commands/web.rs
+1 -1
View File
@@ -168,7 +168,7 @@ fn get_provider_and_model() -> (String, String) {
}
async fn create_agent(provider_name: &str, model: &str) -> Result<Agent> {
let model_config = goose::model::ModelConfig::new(model)?;
let model_config = goose::model::ModelConfig::new(model)?.with_canonical_limits(provider_name);
let agent = Agent::new();
crates/goose-cli/src/scenario_tests/scenario_runner.rs
+1 -1
View File
@@ -188,7 +188,7 @@ where
let inner_provider = create(
&factory_name,
ModelConfig::new(config.model_name)?,
ModelConfig::new(config.model_name)?.with_canonical_limits(&factory_name),
Vec::new(),
)
.await?;
crates/goose-cli/src/session/builder.rs
+1
View File
@@ -386,6 +386,7 @@ fn resolve_provider_and_model(
output::render_error(&format!("Failed to create model configuration: {}", e));
process::exit(1);
})
.with_canonical_limits(&provider_name)
.with_temperature(temperature)
};
crates/goose-cli/src/session/mod.rs
+1 -1
View File
@@ -1847,7 +1847,7 @@ async fn get_reasoner() -> Result<Arc<dyn Provider>, anyhow::Error> {
};
let model_config =
ModelConfig::new_with_context_env(model, Some("GOOSE_PLANNER_CONTEXT_LIMIT"))?;
ModelConfig::new_with_context_env(model, &provider, Some("GOOSE_PLANNER_CONTEXT_LIMIT"))?;
let extensions = goose::config::extensions::get_enabled_extensions_with_config(config);
let reasoner = create(&provider, model_config, extensions).await?;
crates/goose-server/src/openapi.rs
+4 -4
View File
@@ -357,7 +357,7 @@ derive_utoipa!(Icon as IconSchema);
super::routes::config_management::check_provider,
super::routes::config_management::set_config_provider,
super::routes::config_management::configure_provider_oauth,
super::routes::config_management::get_pricing,
super::routes::config_management::get_canonical_model_info,
super::routes::prompts::get_prompts,
super::routes::prompts::get_prompt,
super::routes::prompts::save_prompt,
@@ -443,9 +443,9 @@ derive_utoipa!(Icon as IconSchema);
super::routes::config_management::UpdateCustomProviderRequest,
super::routes::config_management::CheckProviderRequest,
super::routes::config_management::SetProviderRequest,
super::routes::config_management::PricingQuery,
super::routes::config_management::PricingResponse,
super::routes::config_management::PricingData,
super::routes::config_management::ModelInfoQuery,
super::routes::config_management::ModelInfoResponse,
super::routes::config_management::ModelInfoData,
super::routes::prompts::PromptsListResponse,
super::routes::prompts::PromptContentResponse,
super::routes::prompts::SavePromptRequest,
crates/goose-server/src/routes/agent.rs
+1
View File
@@ -549,6 +549,7 @@ async fn update_agent_provider(
format!("Invalid model config: {}", e),
)
})?
.with_canonical_limits(&payload.provider)
.with_context_limit(payload.context_limit)
.with_request_params(payload.request_params);
crates/goose-server/src/routes/config_management.rs
+43 -80
View File
@@ -227,13 +227,6 @@ fn is_valid_provider_name(provider_name: &str) -> bool {
pub async fn read_config(
Json(query): Json<ConfigKeyQuery>,
) -> Result<Json<ConfigValueResponse>, ErrorResponse> {
if query.key == "model-limits" {
let limits = ModelConfig::get_all_model_limits();
return Ok(Json(ConfigValueResponse::Value(serde_json::to_value(
limits,
)?)));
}
let config = Config::global();
let response_value = match config.get(&query.key, query.is_secret) {
@@ -386,7 +379,7 @@ pub async fn get_provider_models(
)));
}
let model_config = ModelConfig::new(&metadata.default_model)?;
let model_config = ModelConfig::new(&metadata.default_model)?.with_canonical_limits(&name);
let provider = goose::providers::create(&name, model_config, Vec::new()).await?;
let models_result = provider.fetch_recommended_models().await;
@@ -426,66 +419,60 @@ pub async fn get_slash_commands() -> Result<Json<SlashCommandsResponse>, ErrorRe
}
#[derive(Serialize, ToSchema)]
pub struct PricingData {
pub struct ModelInfoData {
pub provider: String,
pub model: String,
pub input_token_cost: f64,
pub output_token_cost: f64,
pub context_limit: usize,
pub max_output_tokens: Option<usize>,
pub input_token_cost: Option<f64>,
pub output_token_cost: Option<f64>,
pub cache_read_token_cost: Option<f64>,
pub cache_write_token_cost: Option<f64>,
pub currency: String,
pub context_length: Option<u32>,
}
#[derive(Serialize, ToSchema)]
pub struct PricingResponse {
pub pricing: Vec<PricingData>,
pub struct ModelInfoResponse {
pub model_info: Option<ModelInfoData>,
pub source: String,
}
#[derive(Deserialize, ToSchema)]
pub struct PricingQuery {
pub struct ModelInfoQuery {
pub provider: String,
pub model: String,
}
#[utoipa::path(
post,
path = "/config/pricing",
request_body = PricingQuery,
path = "/config/canonical-model-info",
request_body = ModelInfoQuery,
responses(
(status = 200, description = "Model pricing data retrieved successfully", body = PricingResponse)
(status = 200, description = "Model information retrieved successfully", body = ModelInfoResponse)
)
)]
pub async fn get_pricing(
Json(query): Json<PricingQuery>,
) -> Result<Json<PricingResponse>, ErrorResponse> {
let canonical_model =
maybe_get_canonical_model(&query.provider, &query.model).ok_or_else(|| {
ErrorResponse::not_found(format!(
"Model '{}/{}' not found",
query.provider, query.model
))
})?;
pub async fn get_canonical_model_info(
Json(query): Json<ModelInfoQuery>,
) -> Json<ModelInfoResponse> {
let canonical_model = maybe_get_canonical_model(&query.provider, &query.model);
let mut pricing_data = Vec::new();
let model_info = canonical_model.map(|canonical_model| ModelInfoData {
provider: query.provider.clone(),
model: query.model.clone(),
context_limit: canonical_model.limit.context,
max_output_tokens: canonical_model.limit.output,
// Costs are per million tokens - client handles division for display
input_token_cost: canonical_model.cost.input,
output_token_cost: canonical_model.cost.output,
cache_read_token_cost: canonical_model.cost.cache_read,
cache_write_token_cost: canonical_model.cost.cache_write,
currency: "$".to_string(),
});
if let (Some(input_cost), Some(output_cost)) =
(canonical_model.cost.input, canonical_model.cost.output)
{
pricing_data.push(PricingData {
provider: query.provider.clone(),
model: query.model.clone(),
// Canonical model costs are per million tokens, convert to per-token
input_token_cost: input_cost / 1_000_000.0,
output_token_cost: output_cost / 1_000_000.0,
currency: "$".to_string(),
context_length: Some(canonical_model.limit.context as u32),
});
}
Ok(Json(PricingResponse {
pricing: pricing_data,
Json(ModelInfoResponse {
model_info,
source: "canonical".to_string(),
}))
})
}
#[utoipa::path(
@@ -807,9 +794,11 @@ pub async fn configure_provider_oauth(
)));
}
let temp_model = ModelConfig::new("temp").map_err(|e| {
ErrorResponse::bad_request(format!("Failed to create temporary model config: {}", e))
})?;
let temp_model = ModelConfig::new("temp")
.map_err(|e| {
ErrorResponse::bad_request(format!("Failed to create temporary model config: {}", e))
})?
.with_canonical_limits(&provider_name);
// OAuth configuration does not use extensions.
let provider = create(&provider_name, temp_model, Vec::new())
@@ -849,7 +838,10 @@ pub fn routes(state: Arc<AppState>) -> Router {
.route("/config/providers/{name}/models", get(get_provider_models))
.route("/config/detect-provider", post(detect_provider))
.route("/config/slash_commands", get(get_slash_commands))
.route("/config/pricing", post(get_pricing))
.route(
"/config/canonical-model-info",
post(get_canonical_model_info),
)
.route("/config/init", post(init_config))
.route("/config/backup", post(backup_config))
.route("/config/recover", post(recover_config))
@@ -872,33 +864,4 @@ pub fn routes(state: Arc<AppState>) -> Router {
}
#[cfg(test)]
mod tests {
use http::HeaderMap;
use super::*;
#[tokio::test]
async fn test_read_model_limits() {
let mut headers = HeaderMap::new();
headers.insert("X-Secret-Key", "test".parse().unwrap());
let result = read_config(Json(ConfigKeyQuery {
key: "model-limits".to_string(),
is_secret: false,
}))
.await;
assert!(result.is_ok());
let response = match result.unwrap().0 {
ConfigValueResponse::Value(value) => value,
ConfigValueResponse::MaskedValue(_) => panic!("unexpected secret"),
};
let limits: Vec<goose::model::ModelLimitConfig> = serde_json::from_value(response).unwrap();
assert!(!limits.is_empty());
let gpt4_limit = limits.iter().find(|l| l.pattern == "gpt-4o");
assert!(gpt4_limit.is_some());
assert_eq!(gpt4_limit.unwrap().context_limit, 128_000);
}
}
mod tests {}
crates/goose/src/agents/agent.rs
+1
View File
@@ -1629,6 +1629,7 @@ impl Agent {
.ok_or_else(|| anyhow!("Could not configure agent: missing model"))?;
crate::model::ModelConfig::new(&model_name)
.map_err(|e| anyhow!("Could not configure agent: invalid model {}", e))?
.with_canonical_limits(&provider_name)
}
};
crates/goose/src/agents/platform_extensions/summon.rs
+4 -5
View File
@@ -1417,11 +1417,10 @@ impl SummonClient {
.or_else(|| session.provider_name.clone())
.ok_or_else(|| anyhow::anyhow!("No provider configured"))?;
let mut model_config = session
.model_config
.clone()
.map(Ok)
.unwrap_or_else(|| crate::model::ModelConfig::new("default"))?;
let mut model_config = session.model_config.clone().map(Ok).unwrap_or_else(|| {
crate::model::ModelConfig::new("default")
.map(|c| c.with_canonical_limits(&provider_name))
})?;
if let Some(model) = &params.model {
model_config.model_name = model.clone();
crates/goose/src/context_mgmt/mod.rs
+1 -1
View File
@@ -549,7 +549,7 @@ mod tests {
max_tokens: None,
toolshim: false,
toolshim_model: None,
fast_model: None,
fast_model_config: None,
request_params: None,
},
max_tool_responses: None,
crates/goose/src/model.rs
+58 -149
View File
@@ -44,65 +44,6 @@ pub enum ConfigError {
InvalidRange(String, String),
}
static MODEL_SPECIFIC_LIMITS: Lazy<Vec<(&'static str, usize)>> = Lazy::new(|| {
vec![
// openai
("gpt-5.2-codex", 400_000), // auto-compacting context
("gpt-5.2", 400_000), // auto-compacting context
("gpt-5.1-codex-max", 256_000),
("gpt-5.1-codex-mini", 256_000),
("gpt-4-turbo", 128_000),
("gpt-4.1", 1_000_000),
("gpt-4-1", 1_000_000),
("gpt-4o", 128_000),
("o4-mini", 200_000),
("o3-mini", 200_000),
("o3", 200_000),
// anthropic - all 200k
("claude", 200_000),
// google
("gemini-1.5-flash", 1_048_576),
("gemini-1", 128_000),
("gemini-2", 1_048_576),
("gemini-3-pro-image", 65_536),
("gemini-3-pro", 1_048_576),
("gemini-3-flash", 1_048_576),
("gemma-3-27b", 128_000),
("gemma-3-12b", 128_000),
("gemma-3-4b", 128_000),
("gemma-3-1b", 32_000),
("gemma3-27b", 128_000),
("gemma3-12b", 128_000),
("gemma3-4b", 128_000),
("gemma3-1b", 32_000),
("gemma-2-27b", 8_192),
("gemma-2-9b", 8_192),
("gemma-2-2b", 8_192),
("gemma2-", 8_192),
("gemma-7b", 8_192),
("gemma-2b", 8_192),
("gemma1", 8_192),
("gemma", 8_192),
// facebook
("llama-2-1b", 32_000),
("llama", 128_000),
// qwen
("qwen3-coder", 262_144),
("qwen2-7b", 128_000),
("qwen2-14b", 128_000),
("qwen2-32b", 131_072),
("qwen2-70b", 262_144),
("qwen2", 128_000),
("qwen3-32b", 131_072),
// xai
("grok-4", 256_000),
("grok-code-fast-1", 256_000),
("grok", 131_072),
// other
("kimi-k2", 131_072),
]
});
#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
pub struct ModelConfig {
pub model_name: String,
@@ -111,52 +52,49 @@ pub struct ModelConfig {
pub max_tokens: Option<i32>,
pub toolshim: bool,
pub toolshim_model: Option<String>,
pub fast_model: Option<String>,
#[serde(skip)]
pub fast_model_config: Option<Box<ModelConfig>>,
/// Provider-specific request parameters (e.g., anthropic_beta headers)
#[serde(default, skip_serializing_if = "Option::is_none")]
pub request_params: Option<HashMap<String, Value>>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelLimitConfig {
pub pattern: String,
pub context_limit: usize,
}
impl ModelConfig {
pub fn new(model_name: &str) -> Result<Self, ConfigError> {
Self::new_with_context_env(model_name.to_string(), None)
Self::new_base(model_name.to_string(), None)
}
pub fn new_with_context_env(
model_name: String,
provider_name: &str,
context_env_var: Option<&str>,
) -> Result<Self, ConfigError> {
let predefined = find_predefined_model(&model_name);
let config = Self::new_base(model_name, context_env_var)?;
Ok(config.with_canonical_limits(provider_name))
}
let context_limit = if let Some(ref pm) = predefined {
if let Some(env_var) = context_env_var {
if let Ok(val) = std::env::var(env_var) {
Some(Self::validate_context_limit(&val, env_var)?)
} else {
pm.context_limit
}
} else if let Ok(val) = std::env::var("GOOSE_CONTEXT_LIMIT") {
Some(Self::validate_context_limit(&val, "GOOSE_CONTEXT_LIMIT")?)
fn new_base(model_name: String, context_env_var: Option<&str>) -> Result<Self, ConfigError> {
let context_limit = if let Some(env_var) = context_env_var {
if let Ok(val) = std::env::var(env_var) {
Some(Self::validate_context_limit(&val, env_var)?)
} else {
pm.context_limit
None
}
} else if let Ok(val) = std::env::var("GOOSE_CONTEXT_LIMIT") {
Some(Self::validate_context_limit(&val, "GOOSE_CONTEXT_LIMIT")?)
} else {
Self::parse_context_limit(&model_name, None, context_env_var)?
None
};
let request_params = predefined.and_then(|pm| pm.request_params);
let temperature = Self::parse_temperature()?;
let max_tokens = Self::parse_max_tokens()?;
let temperature = Self::parse_temperature()?;
let toolshim = Self::parse_toolshim()?;
let toolshim_model = Self::parse_toolshim_model()?;
// Pick up request_params from predefined models (always applies)
let predefined = find_predefined_model(&model_name);
let request_params = predefined.and_then(|pm| pm.request_params);
Ok(Self {
model_name,
context_limit,
@@ -164,43 +102,34 @@ impl ModelConfig {
max_tokens,
toolshim,
toolshim_model,
fast_model: None,
fast_model_config: None,
request_params,
})
}
fn parse_context_limit(
model_name: &str,
fast_model: Option<&str>,
custom_env_var: Option<&str>,
) -> Result<Option<usize>, ConfigError> {
// First check if there's an explicit environment variable override
if let Some(env_var) = custom_env_var {
if let Ok(val) = std::env::var(env_var) {
return Self::validate_context_limit(&val, env_var).map(Some);
pub fn with_canonical_limits(mut self, provider_name: &str) -> Self {
if self.context_limit.is_none() || self.max_tokens.is_none() {
if let Some(canonical) = crate::providers::canonical::maybe_get_canonical_model(
provider_name,
&self.model_name,
) {
if self.context_limit.is_none() {
self.context_limit = Some(canonical.limit.context);
}
if self.max_tokens.is_none() {
self.max_tokens = canonical.limit.output.map(|o| o as i32);
}
}
}
if let Ok(val) = std::env::var("GOOSE_CONTEXT_LIMIT") {
return Self::validate_context_limit(&val, "GOOSE_CONTEXT_LIMIT").map(Some);
}
// Get the model's limit
let model_limit = Self::get_model_specific_limit(model_name);
// If there's a fast_model, get its limit and use the minimum
if let Some(fast_model_name) = fast_model {
let fast_model_limit = Self::get_model_specific_limit(fast_model_name);
// Return the minimum of both limits (if both exist)
match (model_limit, fast_model_limit) {
(Some(m), Some(f)) => Ok(Some(m.min(f))),
(Some(m), None) => Ok(Some(m)),
(None, Some(f)) => Ok(Some(f)),
(None, None) => Ok(None),
// Try filling remaining gaps from predefined models
if self.context_limit.is_none() {
if let Some(pm) = find_predefined_model(&self.model_name) {
self.context_limit = pm.context_limit;
}
} else {
Ok(model_limit)
}
self
}
fn validate_context_limit(val: &str, env_var: &str) -> Result<usize, ConfigError> {
@@ -291,23 +220,6 @@ impl ModelConfig {
}
}
fn get_model_specific_limit(model_name: &str) -> Option<usize> {
MODEL_SPECIFIC_LIMITS
.iter()
.find(|(pattern, _)| model_name.contains(pattern))
.map(|(_, limit)| *limit)
}
pub fn get_all_model_limits() -> Vec<ModelLimitConfig> {
MODEL_SPECIFIC_LIMITS
.iter()
.map(|(pattern, context_limit)| ModelLimitConfig {
pattern: pattern.to_string(),
context_limit: *context_limit,
})
.collect()
}
pub fn with_context_limit(mut self, limit: Option<usize>) -> Self {
if limit.is_some() {
self.context_limit = limit;
@@ -335,9 +247,15 @@ impl ModelConfig {
self
}
pub fn with_fast(mut self, fast_model: String) -> Self {
self.fast_model = Some(fast_model);
self
pub fn with_fast(
mut self,
fast_model_name: &str,
provider_name: &str,
) -> Result<Self, ConfigError> {
// Create a full ModelConfig for the fast model with proper canonical lookup
let fast_config = ModelConfig::new(fast_model_name)?.with_canonical_limits(provider_name);
self.fast_model_config = Some(Box::new(fast_config));
Ok(self)
}
pub fn with_request_params(mut self, params: Option<HashMap<String, Value>>) -> Self {
@@ -346,33 +264,24 @@ impl ModelConfig {
}
pub fn use_fast_model(&self) -> Self {
if let Some(fast_model) = &self.fast_model {
let mut config = self.clone();
config.model_name = fast_model.clone();
config
if let Some(fast_config) = &self.fast_model_config {
*fast_config.clone()
} else {
self.clone()
}
}
pub fn context_limit(&self) -> usize {
// If we have an explicit context limit set, use it
if let Some(limit) = self.context_limit {
return limit;
self.context_limit.unwrap_or(DEFAULT_CONTEXT_LIMIT)
}
pub fn max_output_tokens(&self) -> i32 {
if let Some(tokens) = self.max_tokens {
return tokens;
}
// Otherwise, get the model's default limit
let main_limit =
Self::get_model_specific_limit(&self.model_name).unwrap_or(DEFAULT_CONTEXT_LIMIT);
// If we have a fast_model, also check its limit and use the minimum
if let Some(fast_model) = &self.fast_model {
let fast_limit =
Self::get_model_specific_limit(fast_model).unwrap_or(DEFAULT_CONTEXT_LIMIT);
main_limit.min(fast_limit)
} else {
main_limit
}
// Priority 2: Global default
4_096
}
pub fn new_or_fail(model_name: &str) -> ModelConfig {
crates/goose/src/providers/anthropic.rs
+1 -1
View File
@@ -59,7 +59,7 @@ pub struct AnthropicProvider {
impl AnthropicProvider {
pub async fn from_env(model: ModelConfig) -> Result<Self> {
let model = model.with_fast(ANTHROPIC_DEFAULT_FAST_MODEL.to_string());
let model = model.with_fast(ANTHROPIC_DEFAULT_FAST_MODEL, ANTHROPIC_PROVIDER_NAME)?;
let config = crate::config::Config::global();
let api_key: String = config.get_secret("ANTHROPIC_API_KEY")?;
crates/goose/src/providers/auto_detect.rs
+1 -1
View File
@@ -21,7 +21,7 @@ pub async fn detect_provider_from_api_key(api_key: &str) -> Option<(String, Vec<
let result = match crate::providers::create(
provider_name,
ModelConfig::new_or_fail("default"),
ModelConfig::new_or_fail("default").with_canonical_limits(provider_name),
Vec::new(),
)
.await
crates/goose/src/providers/base.rs
+5 -3
View File
@@ -133,9 +133,11 @@ impl ProviderMetadata {
default_model: default_model.to_string(),
known_models: model_names
.iter()
.map(|&name| ModelInfo {
name: name.to_string(),
context_limit: ModelConfig::new_or_fail(name).context_limit(),
.map(|&model_name| ModelInfo {
name: model_name.to_string(),
context_limit: ModelConfig::new_or_fail(model_name)
.with_canonical_limits(name)
.context_limit(),
input_token_cost: None,
output_token_cost: None,
currency: None,
crates/goose/src/providers/canonical/name_builder.rs
+2 -2
View File
@@ -81,10 +81,10 @@ pub fn map_to_canonical_model(
if let Some(canonical) = registry.get(registry_provider, model) {
return Some(canonical.id.clone());
}
return None;
// If direct lookup failed, fall through to inference logic below
}
// For hosting/meta-providers do string matching magic to figure out the real provider and model
// For hosting/meta-providers (or unknown providers), do string matching magic to figure out the real provider and model
let model_stripped = strip_common_prefixes(model);
if let Some(swapped) = swap_claude_word_order(&model_stripped) {
crates/goose/src/providers/claude_code.rs
+3 -1
View File
@@ -1336,7 +1336,9 @@ mod tests {
fn make_provider() -> ClaudeCodeProvider {
ClaudeCodeProvider {
command: PathBuf::from("claude"),
model: ModelConfig::new(CLAUDE_CODE_DEFAULT_MODEL).unwrap(),
model: ModelConfig::new(CLAUDE_CODE_DEFAULT_MODEL)
.unwrap()
.with_canonical_limits(CLAUDE_CODE_PROVIDER_NAME),
name: "claude-code".to_string(),
mcp_config_file: None,
cli_process: tokio::sync::OnceCell::new(),
crates/goose/src/providers/databricks.rs
+4 -3
View File
@@ -150,7 +150,8 @@ impl DatabricksProvider {
fast_retry_config,
name: DATABRICKS_PROVIDER_NAME.to_string(),
};
provider.model = model.with_fast(DATABRICKS_DEFAULT_FAST_MODEL.to_string());
provider.model =
model.with_fast(DATABRICKS_DEFAULT_FAST_MODEL, DATABRICKS_PROVIDER_NAME)?;
Ok(provider)
}
@@ -300,9 +301,9 @@ impl Provider for DatabricksProvider {
// Use fast retry config if this is the fast model
let is_fast_model = self
.model
.fast_model
.fast_model_config
.as_ref()
.map(|fast| fast == &model_config.model_name)
.map(|fast| fast.model_name == model_config.model_name)
.unwrap_or(false);
let retry_config = if is_fast_model {
crates/goose/src/providers/formats/anthropic.rs
+1 -18
View File
@@ -395,30 +395,17 @@ pub fn create_request(
let tool_specs = format_tools(tools);
let system_spec = format_system(system);
// Check if we have any messages to send
if anthropic_messages.is_empty() {
return Err(anyhow!("No valid messages to send to Anthropic API"));
}
// https://platform.claude.com/docs/en/about-claude/models/overview
// 64k output tokens works for most claude models, but not old opus:
let max_tokens = model_config.max_tokens.unwrap_or_else(|| {
let name = &model_config.model_name;
if name.contains("claude-3-haiku") {
4096
} else if name.contains("claude-opus-4-0") || name.contains("claude-opus-4-1") {
32000
} else {
64000
}
});
let max_tokens = model_config.max_output_tokens();
let mut payload = json!({
"model": model_config.model_name,
"messages": anthropic_messages,
"max_tokens": max_tokens,
});
// Add system message if present
if !system.is_empty() {
payload
.as_object_mut()
@@ -426,7 +413,6 @@ pub fn create_request(
.insert("system".to_string(), json!(system_spec));
}
// Add tools if present
if !tool_specs.is_empty() {
payload
.as_object_mut()
@@ -434,7 +420,6 @@ pub fn create_request(
.insert("tools".to_string(), json!(tool_specs));
}
// Add temperature if specified and not using extended thinking model
if let Some(temp) = model_config.temperature {
payload
.as_object_mut()
@@ -442,10 +427,8 @@ pub fn create_request(
.insert("temperature".to_string(), json!(temp));
}
// Add thinking parameters when CLAUDE_THINKING_ENABLED is set
let is_thinking_enabled = std::env::var("CLAUDE_THINKING_ENABLED").is_ok();
if is_thinking_enabled {
// Minimum budget_tokens is 1024
let budget_tokens = std::env::var("CLAUDE_THINKING_BUDGET")
.unwrap_or_else(|_| "16000".to_string())
.parse()
crates/goose/src/providers/formats/databricks.rs
+4 -4
View File
@@ -1056,7 +1056,7 @@ mod tests {
max_tokens: Some(1024),
toolshim: false,
toolshim_model: None,
fast_model: None,
fast_model_config: None,
request_params: None,
};
let request = create_request(&model_config, "system", &[], &[], &ImageFormat::OpenAi)?;
@@ -1088,7 +1088,7 @@ mod tests {
max_tokens: Some(1024),
toolshim: false,
toolshim_model: None,
fast_model: None,
fast_model_config: None,
request_params: None,
};
let request = create_request(&model_config, "system", &[], &[], &ImageFormat::OpenAi)?;
@@ -1440,7 +1440,7 @@ mod tests {
max_tokens: Some(8192),
toolshim: false,
toolshim_model: None,
fast_model: None,
fast_model_config: None,
request_params: None,
};
@@ -1492,7 +1492,7 @@ mod tests {
max_tokens: Some(4096),
toolshim: false,
toolshim_model: None,
fast_model: None,
fast_model_config: None,
request_params: None,
};
crates/goose/src/providers/formats/openai.rs
+3 -3
View File
@@ -1496,7 +1496,7 @@ mod tests {
max_tokens: Some(1024),
toolshim: false,
toolshim_model: None,
fast_model: None,
fast_model_config: None,
request_params: None,
};
let request = create_request(
@@ -1536,7 +1536,7 @@ mod tests {
max_tokens: Some(1024),
toolshim: false,
toolshim_model: None,
fast_model: None,
fast_model_config: None,
request_params: None,
};
let request = create_request(
@@ -1577,7 +1577,7 @@ mod tests {
max_tokens: Some(1024),
toolshim: false,
toolshim_model: None,
fast_model: None,
fast_model_config: None,
request_params: None,
};
let request = create_request(
crates/goose/src/providers/formats/snowflake.rs
+4 -2
View File
@@ -568,7 +568,8 @@ data: {"id":"a9537c2c-2017-4906-9817-2456168d89fa","model":"claude-sonnet-4-2025
use crate::conversation::message::Message;
use crate::model::ModelConfig;
let model_config = ModelConfig::new_or_fail("claude-4-sonnet");
let model_config =
ModelConfig::new_or_fail("claude-4-sonnet").with_canonical_limits("snowflake");
let system = "You are a helpful assistant that can use tools to get information.";
let messages = vec![Message::user().with_text("What is the stock price of Nvidia?")];
@@ -677,7 +678,8 @@ data: {"id":"a9537c2c-2017-4906-9817-2456168d89fa","model":"claude-sonnet-4-2025
use crate::conversation::message::Message;
use crate::model::ModelConfig;
let model_config = ModelConfig::new_or_fail("claude-4-sonnet");
let model_config =
ModelConfig::new_or_fail("claude-4-sonnet").with_canonical_limits("snowflake");
let system = "Reply with only a description in four words or less";
let messages = vec![Message::user().with_text("Test message")];
let tools = vec![Tool::new(
crates/goose/src/providers/google.rs
+1 -1
View File
@@ -69,7 +69,7 @@ pub struct GoogleProvider {
impl GoogleProvider {
pub async fn from_env(model: ModelConfig) -> Result<Self> {
let model = model.with_fast(GOOGLE_DEFAULT_FAST_MODEL.to_string());
let model = model.with_fast(GOOGLE_DEFAULT_FAST_MODEL, GOOGLE_PROVIDER_NAME)?;
let config = crate::config::Config::global();
let api_key: String = config.get_secret("GOOGLE_API_KEY")?;
crates/goose/src/providers/init.rs
+15 -8
View File
@@ -141,7 +141,8 @@ pub async fn create_with_named_model(
model_name: &str,
extensions: Vec<ExtensionConfig>,
) -> Result<Arc<dyn Provider>> {
create(provider_name, ModelConfig::new(model_name)?, extensions).await
let config = ModelConfig::new(model_name)?.with_canonical_limits(provider_name);
create(provider_name, config, extensions).await
}
async fn create_lead_worker_from_env(
@@ -168,10 +169,11 @@ async fn create_lead_worker_from_env(
let lead_model_config = ModelConfig::new_with_context_env(
lead_model_name.to_string(),
&lead_provider_name,
Some("GOOSE_LEAD_CONTEXT_LIMIT"),
)?;
let worker_model_config = create_worker_model_config(default_model)?;
let worker_model_config = create_worker_model_config(default_model, default_provider_name)?;
let registry = get_registry().await;
@@ -207,8 +209,12 @@ async fn create_lead_worker_from_env(
)))
}
fn create_worker_model_config(default_model: &ModelConfig) -> Result<ModelConfig> {
fn create_worker_model_config(
default_model: &ModelConfig,
provider_name: &str,
) -> Result<ModelConfig> {
let mut worker_config = ModelConfig::new_or_fail(&default_model.model_name)
.with_canonical_limits(provider_name)
.with_context_limit(default_model.context_limit)
.with_temperature(default_model.temperature)
.with_max_tokens(default_model.max_tokens)
@@ -253,7 +259,7 @@ mod tests {
let provider = create(
"openai",
ModelConfig::new_or_fail("gpt-4o-mini"),
ModelConfig::new_or_fail("gpt-4o-mini").with_canonical_limits("openai"),
Vec::new(),
)
.await
@@ -282,7 +288,7 @@ mod tests {
let provider = create(
"openai",
ModelConfig::new_or_fail("gpt-4o-mini"),
ModelConfig::new_or_fail("gpt-4o-mini").with_canonical_limits("openai"),
Vec::new(),
)
.await
@@ -304,10 +310,11 @@ mod tests {
("GOOSE_CONTEXT_LIMIT", global_limit),
]);
let default_model =
ModelConfig::new_or_fail("gpt-3.5-turbo").with_context_limit(Some(16_000));
let default_model = ModelConfig::new_or_fail("gpt-3.5-turbo")
.with_canonical_limits("openai")
.with_context_limit(Some(16_000));
let result = create_worker_model_config(&default_model).unwrap();
let result = create_worker_model_config(&default_model, "openai").unwrap();
assert_eq!(result.context_limit, Some(expected_limit));
}
crates/goose/src/providers/openai.rs
+1 -1
View File
@@ -72,7 +72,7 @@ pub struct OpenAiProvider {
impl OpenAiProvider {
pub async fn from_env(model: ModelConfig) -> Result<Self> {
let model = model.with_fast(OPEN_AI_DEFAULT_FAST_MODEL.to_string());
let model = model.with_fast(OPEN_AI_DEFAULT_FAST_MODEL, OPEN_AI_PROVIDER_NAME)?;
let config = crate::config::Config::global();
let host: String = config
crates/goose/src/providers/openrouter.rs
+1 -1
View File
@@ -51,7 +51,7 @@ pub struct OpenRouterProvider {
impl OpenRouterProvider {
pub async fn from_env(model: ModelConfig) -> Result<Self> {
let model = model.with_fast(OPENROUTER_DEFAULT_FAST_MODEL.to_string());
let model = model.with_fast(OPENROUTER_DEFAULT_FAST_MODEL, OPENROUTER_PROVIDER_NAME)?;
let config = crate::config::Config::global();
let api_key: String = config.get_secret("OPENROUTER_API_KEY")?;
crates/goose/src/providers/provider_registry.rs
+3 -1
View File
@@ -25,7 +25,9 @@ impl ProviderEntry {
extensions: Vec<ExtensionConfig>,
) -> Result<Arc<dyn Provider>> {
let default_model = &self.metadata.default_model;
let model_config = ModelConfig::new(default_model.as_str())?;
let provider_name = &self.metadata.name;
let model_config =
ModelConfig::new(default_model.as_str())?.with_canonical_limits(provider_name);
(self.constructor)(model_config, extensions).await
}
}
crates/goose/src/providers/provider_test.rs
+1
View File
@@ -10,6 +10,7 @@ pub async fn test_provider_configuration(
toolshim_model: Option<String>,
) -> Result<()> {
let model_config = ModelConfig::new(model)?
.with_canonical_limits(provider_name)
.with_max_tokens(Some(50))
.with_toolshim(toolshim_enabled)
.with_toolshim_model(toolshim_model);
crates/goose/src/providers/toolshim.rs
+2 -1
View File
@@ -154,7 +154,8 @@ impl OllamaInterpreter {
messages.push(user_message);
let model_config = ModelConfig::new(model)
.map_err(|e| ProviderError::RequestFailed(format!("Model config error: {e}")))?;
.map_err(|e| ProviderError::RequestFailed(format!("Model config error: {e}")))?
.with_canonical_limits("ollama");
let mut payload = create_request(
&model_config,
crates/goose/src/providers/venice.rs
+1 -4
View File
@@ -85,7 +85,7 @@ pub struct VeniceProvider {
}
impl VeniceProvider {
pub async fn from_env(mut model: ModelConfig) -> Result<Self> {
pub async fn from_env(model: ModelConfig) -> Result<Self> {
let config = crate::config::Config::global();
let api_key: String = config.get_secret("VENICE_API_KEY")?;
let host: String = config
@@ -98,9 +98,6 @@ impl VeniceProvider {
.get_param("VENICE_MODELS_PATH")
.unwrap_or_else(|_| VENICE_DEFAULT_MODELS_PATH.to_string());
// Ensure we only keep the bare model id internally
model.model_name = strip_flags(&model.model_name).to_string();
let auth = AuthMethod::BearerToken(api_key);
let api_client = ApiClient::new(host, auth)?;
crates/goose/src/scheduler.rs
+2 -1
View File
@@ -739,7 +739,8 @@ async fn execute_job(
let config = Config::global();
let provider_name = config.get_goose_provider()?;
let model_name = config.get_goose_model()?;
let model_config = crate::model::ModelConfig::new(&model_name)?;
let model_config =
crate::model::ModelConfig::new(&model_name)?.with_canonical_limits(&provider_name);
let session = agent
.config
crates/goose/tests/providers.rs
+1 -1
View File
@@ -283,7 +283,7 @@ impl ProviderTester {
.model_switch_name
.as_deref()
.expect("model_switch_name required for test_model_switch");
let alt_config = goose::model::ModelConfig::new(alt)?;
let alt_config = goose::model::ModelConfig::new(alt)?.with_canonical_limits(&self.name);
let message = Message::user().with_text("Just say hello!");
let (response, _) = self
crates/goose/tests/tetrate_streaming.rs
+4 -2
View File
@@ -15,7 +15,8 @@ mod tetrate_streaming_tests {
async fn create_test_provider() -> Result<TetrateProvider> {
// Create a test provider with the default model
let model_config = ModelConfig::new("claude-3-5-sonnet-latest")?;
let model_config =
ModelConfig::new("claude-3-5-sonnet-latest")?.with_canonical_limits("tetrate");
TetrateProvider::from_env(model_config).await
}
@@ -237,7 +238,8 @@ mod tetrate_streaming_tests {
// Test with invalid API key to ensure error handling works
std::env::set_var("TETRATE_API_KEY", "invalid-key-for-testing");
let model_config = ModelConfig::new("claude-3-5-sonnet-latest")?;
let model_config =
ModelConfig::new("claude-3-5-sonnet-latest")?.with_canonical_limits("tetrate");
let provider = TetrateProvider::from_env(model_config).await?;
let messages = vec![Message::user().with_text("Hello")];