Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add logic for closed beta LLM models (#16482)

Browse source 
Release Notes:

- N/A

---------

Co-authored-by: Marshall <marshall@zed.dev>
This commit is contained in:
Max Brunsfeld 2024-08-19 11:09:52 -07:00 • committed by GitHub
parent 41fc6d0885
commit b5bd8a5c5d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 104 additions and 47 deletions
Download patch file Download diff file Expand all files Collapse all files

94
crates/language_model/src/provider/cloud.rs
View file

@ -8,7 +8,7 @@ use anthropic::AnthropicError;
use anyhow::{anyhow, Result};
use client::{Client, PerformCompletionParams, UserStore, EXPIRED_LLM_TOKEN_HEADER_NAME};
use collections::BTreeMap;
use feature_flags::{FeatureFlagAppExt, ZedPro};
use feature_flags::{FeatureFlagAppExt, LlmClosedBeta, ZedPro};
use futures::{
future::BoxFuture, stream::BoxStream, AsyncBufReadExt, FutureExt, Stream, StreamExt,
TryStreamExt as _,
@ -26,7 +26,10 @@ use smol::{
io::{AsyncReadExt, BufReader},
lock::{RwLock, RwLockUpgradableReadGuard, RwLockWriteGuard},
};
use std::{future, sync::Arc};
use std::{
future,
sync::{Arc, LazyLock},
};
use strum::IntoEnumIterator;
use ui::prelude::*;
@ -37,6 +40,18 @@ use super::anthropic::count_anthropic_tokens;
pub const PROVIDER_ID: &str = "zed.dev";
pub const PROVIDER_NAME: &str = "Zed";
const ZED_CLOUD_PROVIDER_ADDITIONAL_MODELS_JSON: Option<&str> =
option_env!("ZED_CLOUD_PROVIDER_ADDITIONAL_MODELS_JSON");
fn zed_cloud_provider_additional_models() -> &'static [AvailableModel] {
static ADDITIONAL_MODELS: LazyLock<Vec<AvailableModel>> = LazyLock::new(|| {
ZED_CLOUD_PROVIDER_ADDITIONAL_MODELS_JSON
.map(|json| serde_json::from_str(json).unwrap())
.unwrap_or(Vec::new())
});
ADDITIONAL_MODELS.as_slice()
}
#[derive(Default, Clone, Debug, PartialEq)]
pub struct ZedDotDevSettings {
pub available_models: Vec<AvailableModel>,
@ -200,40 +215,6 @@ impl LanguageModelProvider for CloudLanguageModelProvider {
for model in ZedModel::iter() {
models.insert(model.id().to_string(), CloudModel::Zed(model));
}
// Override with available models from settings
for model in &AllLanguageModelSettings::get_global(cx)
.zed_dot_dev
.available_models
{
let model = match model.provider {
AvailableProvider::Anthropic => {
CloudModel::Anthropic(anthropic::Model::Custom {
name: model.name.clone(),
display_name: model.display_name.clone(),
max_tokens: model.max_tokens,
tool_override: model.tool_override.clone(),
cache_configuration: model.cache_configuration.as_ref().map(|config| {
anthropic::AnthropicModelCacheConfiguration {
max_cache_anchors: config.max_cache_anchors,
should_speculate: config.should_speculate,
min_total_token: config.min_total_token,
}
}),
max_output_tokens: model.max_output_tokens,
})
}
AvailableProvider::OpenAi => CloudModel::OpenAi(open_ai::Model::Custom {
name: model.name.clone(),
max_tokens: model.max_tokens,
}),
AvailableProvider::Google => CloudModel::Google(google_ai::Model::Custom {
name: model.name.clone(),
max_tokens: model.max_tokens,
}),
};
models.insert(model.id().to_string(), model.clone());
}
} else {
models.insert(
anthropic::Model::Claude3_5Sonnet.id().to_string(),
@ -241,6 +222,47 @@ impl LanguageModelProvider for CloudLanguageModelProvider {
);
}
let llm_closed_beta_models = if cx.has_flag::<LlmClosedBeta>() {
zed_cloud_provider_additional_models()
} else {
&[]
};
// Override with available models from settings
for model in AllLanguageModelSettings::get_global(cx)
.zed_dot_dev
.available_models
.iter()
.chain(llm_closed_beta_models)
.cloned()
{
let model = match model.provider {
AvailableProvider::Anthropic => CloudModel::Anthropic(anthropic::Model::Custom {
name: model.name.clone(),
display_name: model.display_name.clone(),
max_tokens: model.max_tokens,
tool_override: model.tool_override.clone(),
cache_configuration: model.cache_configuration.as_ref().map(|config| {
anthropic::AnthropicModelCacheConfiguration {
max_cache_anchors: config.max_cache_anchors,
should_speculate: config.should_speculate,
min_total_token: config.min_total_token,
}
}),
max_output_tokens: model.max_output_tokens,
}),
AvailableProvider::OpenAi => CloudModel::OpenAi(open_ai::Model::Custom {
name: model.name.clone(),
max_tokens: model.max_tokens,
}),
AvailableProvider::Google => CloudModel::Google(google_ai::Model::Custom {
name: model.name.clone(),
max_tokens: model.max_tokens,
}),
};
models.insert(model.id().to_string(), model.clone());
}
models
.into_values()
.map(|model| {