Skip to content

Commit 1a4af17

Browse files
fix(rig-1113): calculate_max_tokens assumes known model (anthropic) (#1216)
* fix(rig-1113): `calculate_max_tokens` assumes known model (anthropic) * refactor: amendments
1 parent d6452c5 commit 1a4af17

1 file changed

Lines changed: 26 additions & 10 deletions

File tree

rig/rig-core/src/providers/anthropic/completion.rs

Lines changed: 26 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -715,20 +715,36 @@ where
715715
/// set or if set too high, the request will fail. The following values are based on the models
716716
/// available at the time of writing.
717717
fn calculate_max_tokens(model: &str) -> Option<u64> {
718-
match model {
719-
CLAUDE_4_OPUS => Some(32_000),
720-
CLAUDE_4_SONNET | CLAUDE_3_7_SONNET => Some(64_000),
721-
CLAUDE_3_5_SONNET | CLAUDE_3_5_HAIKU => Some(8_192),
722-
_ => None,
718+
if model.starts_with("claude-opus-4") {
719+
Some(32000)
720+
} else if model.starts_with("claude-sonnet-4") || model.starts_with("claude-3-7-sonnet") {
721+
Some(64000)
722+
} else if model.starts_with("claude-3-5-sonnet") || model.starts_with("claude-3-5-haiku") {
723+
Some(8192)
724+
} else if model.starts_with("claude-3-opus")
725+
|| model.starts_with("claude-3-sonnet")
726+
|| model.starts_with("claude-3-haiku")
727+
{
728+
Some(4096)
729+
} else {
730+
None
723731
}
724732
}
725733

726734
fn calculate_max_tokens_custom(model: &str) -> u64 {
727-
match model {
728-
"claude-4-opus" => 32_000,
729-
"claude-4-sonnet" | "claude-3.7-sonnet" => 64_000,
730-
"claude-3.5-sonnet" | "claude-3.5-haiku" => 8_192,
731-
_ => 4_096,
735+
if model.starts_with("claude-opus-4") {
736+
32000
737+
} else if model.starts_with("claude-sonnet-4") || model.starts_with("claude-3-7-sonnet") {
738+
64000
739+
} else if model.starts_with("claude-3-5-sonnet") || model.starts_with("claude-3-5-haiku") {
740+
8192
741+
} else if model.starts_with("claude-3-opus")
742+
|| model.starts_with("claude-3-sonnet")
743+
|| model.starts_with("claude-3-haiku")
744+
{
745+
4096
746+
} else {
747+
2048
732748
}
733749
}
734750

0 commit comments

Comments
 (0)