tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 6 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/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")];