zed-industries/zed
1
1
Fork 0
mirror of https://github.com/zed-industries/zed.git synced 2024-11-08 07:35:01 +03:00
Code Issues Packages Projects Releases Wiki Activity

Hard code max token counts for supported models (#9675)

Browse Source
This commit is contained in:
Nathan Sobo 2024-03-21 20:30:33 -06:00 committed by GitHub
parent 441677c90a
commit 6d5787cfdc
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 19 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
crates/assistant/src/assistant.rs
View File

@ -97,13 +97,8 @@ impl LanguageModel {
pub fn max_token_count(&self) -> usize {
match self {
LanguageModel::OpenAi(model) => tiktoken_rs::model::get_context_size(model.id()),
LanguageModel::ZedDotDev(model) => match model {
ZedDotDevModel::GptThreePointFiveTurbo
| ZedDotDevModel::GptFour
| ZedDotDevModel::GptFourTurbo => tiktoken_rs::model::get_context_size(model.id()),
ZedDotDevModel::Custom(_) => 30720, // TODO: Base this on the selected model.
},
LanguageModel::OpenAi(model) => model.max_token_count(),
LanguageModel::ZedDotDev(model) => model.max_token_count(),
}
}

9
crates/assistant/src/assistant_settings.rs
View File

@ -109,6 +109,15 @@ impl ZedDotDevModel {
Self::Custom(id) => id.as_str(),
}
}
pub fn max_token_count(&self) -> usize {
match self {
Self::GptThreePointFiveTurbo => 2048,
Self::GptFour => 4096,
Self::GptFourTurbo => 128000,
Self::Custom(_) => 4096, // TODO: Make this configurable
}
}
}
#[derive(Copy, Clone, Default, Debug, Serialize, Deserialize, JsonSchema)]

8
crates/open_ai/src/open_ai.rs
View File

@ -72,6 +72,14 @@ impl Model {
Self::FourTurbo => "gpt-4-turbo",
}
}
pub fn max_token_count(&self) -> usize {
match self {
Model::ThreePointFiveTurbo => 4096,
Model::Four => 8192,
Model::FourTurbo => 128000,
}
}
}
#[derive(Debug, Serialize)]