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

assistant: Limit model access (#15820)

Browse source 
Release Notes:

- N/A

---------

Co-authored-by: Thorsten <thorsten@zed.dev>
This commit is contained in:
Bennet Bo Fenner 2024-08-06 12:19:19 +02:00 • committed by GitHub
parent 56abd68d0e
commit d6e5265e84
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 47 additions and 37 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -7,7 +7,7 @@ use crate::{
use anyhow::{anyhow, Context as _, Result}; use anyhow::{anyhow, Context as _, Result};
use client::{Client, PerformCompletionParams, UserStore, EXPIRED_LLM_TOKEN_HEADER_NAME}; use client::{Client, PerformCompletionParams, UserStore, EXPIRED_LLM_TOKEN_HEADER_NAME};
use collections::BTreeMap; use collections::BTreeMap;
use feature_flags::{FeatureFlag, FeatureFlagAppExt}; use feature_flags::{FeatureFlag, FeatureFlagAppExt, LanguageModels};
use futures::{future::BoxFuture, stream::BoxStream, AsyncBufReadExt, FutureExt, StreamExt}; use futures::{future::BoxFuture, stream::BoxStream, AsyncBufReadExt, FutureExt, StreamExt};
use gpui::{AnyView, AppContext, AsyncAppContext, Model, ModelContext, Subscription, Task}; use gpui::{AnyView, AppContext, AsyncAppContext, Model, ModelContext, Subscription, Task};
use http_client::{AsyncBody, HttpClient, Method, Response}; use http_client::{AsyncBody, HttpClient, Method, Response};
@ -142,46 +142,56 @@ impl LanguageModelProvider for CloudLanguageModelProvider {
fn provided_models(&self, cx: &AppContext) -> Vec<Arc<dyn LanguageModel>> { fn provided_models(&self, cx: &AppContext) -> Vec<Arc<dyn LanguageModel>> {
let mut models = BTreeMap::default(); let mut models = BTreeMap::default();
for model in anthropic::Model::iter() { let is_user = !cx.has_flag::<LanguageModels>();
if !matches!(model, anthropic::Model::Custom { .. }) { if is_user {
models.insert(model.id().to_string(), CloudModel::Anthropic(model)); models.insert(
anthropic::Model::Claude3_5Sonnet.id().to_string(),
CloudModel::Anthropic(anthropic::Model::Claude3_5Sonnet),
);
} else {
for model in anthropic::Model::iter() {
if !matches!(model, anthropic::Model::Custom { .. }) {
models.insert(model.id().to_string(), CloudModel::Anthropic(model));
}
} }
} for model in open_ai::Model::iter() {
for model in open_ai::Model::iter() { if !matches!(model, open_ai::Model::Custom { .. }) {
if !matches!(model, open_ai::Model::Custom { .. }) { models.insert(model.id().to_string(), CloudModel::OpenAi(model));
models.insert(model.id().to_string(), CloudModel::OpenAi(model)); }
} }
} for model in google_ai::Model::iter() {
for model in google_ai::Model::iter() { if !matches!(model, google_ai::Model::Custom { .. }) {
if !matches!(model, google_ai::Model::Custom { .. }) { models.insert(model.id().to_string(), CloudModel::Google(model));
models.insert(model.id().to_string(), CloudModel::Google(model)); }
}
for model in ZedModel::iter() {
models.insert(model.id().to_string(), CloudModel::Zed(model));
} }
}
for model in ZedModel::iter() {
models.insert(model.id().to_string(), CloudModel::Zed(model));
}
// Override with available models from settings // Override with available models from settings
for model in &AllLanguageModelSettings::get_global(cx) for model in &AllLanguageModelSettings::get_global(cx)
.zed_dot_dev .zed_dot_dev
.available_models .available_models
{ {
let model = match model.provider { let model = match model.provider {
AvailableProvider::Anthropic => CloudModel::Anthropic(anthropic::Model::Custom { AvailableProvider::Anthropic => {
name: model.name.clone(), CloudModel::Anthropic(anthropic::Model::Custom {
max_tokens: model.max_tokens, name: model.name.clone(),
tool_override: model.tool_override.clone(), max_tokens: model.max_tokens,
}), tool_override: model.tool_override.clone(),
AvailableProvider::OpenAi => CloudModel::OpenAi(open_ai::Model::Custom { })
name: model.name.clone(), }
max_tokens: model.max_tokens, AvailableProvider::OpenAi => CloudModel::OpenAi(open_ai::Model::Custom {
}), name: model.name.clone(),
AvailableProvider::Google => CloudModel::Google(google_ai::Model::Custom { max_tokens: model.max_tokens,
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.insert(model.id().to_string(), model.clone());
}
} }
models models