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

16
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,6 +142,13 @@ 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();
let is_user = !cx.has_flag::<LanguageModels>();
if is_user {
models.insert(
anthropic::Model::Claude3_5Sonnet.id().to_string(),
CloudModel::Anthropic(anthropic::Model::Claude3_5Sonnet),
);
} else {
for model in anthropic::Model::iter() { for model in anthropic::Model::iter() {
if !matches!(model, anthropic::Model::Custom { .. }) { if !matches!(model, anthropic::Model::Custom { .. }) {
models.insert(model.id().to_string(), CloudModel::Anthropic(model)); models.insert(model.id().to_string(), CloudModel::Anthropic(model));
@ -167,11 +174,13 @@ impl LanguageModelProvider for CloudLanguageModelProvider {
.available_models .available_models
{ {
let model = match model.provider { let model = match model.provider {
AvailableProvider::Anthropic => CloudModel::Anthropic(anthropic::Model::Custom { AvailableProvider::Anthropic => {
CloudModel::Anthropic(anthropic::Model::Custom {
name: model.name.clone(), name: model.name.clone(),
max_tokens: model.max_tokens, max_tokens: model.max_tokens,
tool_override: model.tool_override.clone(), tool_override: model.tool_override.clone(),
}), })
}
AvailableProvider::OpenAi => CloudModel::OpenAi(open_ai::Model::Custom { AvailableProvider::OpenAi => CloudModel::OpenAi(open_ai::Model::Custom {
name: model.name.clone(), name: model.name.clone(),
max_tokens: model.max_tokens, max_tokens: model.max_tokens,
@ -183,6 +192,7 @@ impl LanguageModelProvider for CloudLanguageModelProvider {
}; };
models.insert(model.id().to_string(), model.clone()); models.insert(model.id().to_string(), model.clone());
} }
}
models models
.into_values() .into_values()