language_model: Denote the availability of language models (#15660)

This PR updates the `LanguageModel` trait with a new method for denoting
the availability of a model.

Right now we have two variants:

- `Public` for models that have no additional restrictions (other than
their respective setup/authentication requirements)
- `RequiresPlan` for models that require a specific Zed plan

Release Notes:

- N/A
This commit is contained in:
Marshall Bowers 2024-08-01 18:26:27 -04:00 committed by GitHub
parent 906d9736d5
commit 5e011ab029
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 73 additions and 13 deletions

View file

@ -14,6 +14,7 @@ use gpui::{
};
pub use model::*;
use project::Fs;
use proto::Plan;
pub(crate) use rate_limiter::*;
pub use registry::*;
pub use request::*;
@ -32,6 +33,15 @@ pub fn init(
registry::init(user_store, client, cx);
}
/// The availability of a [`LanguageModel`].
#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum LanguageModelAvailability {
/// The language model is available to the general public.
Public,
/// The language model is available to users on the indicated plan.
RequiresPlan(Plan),
}
pub trait LanguageModel: Send + Sync {
fn id(&self) -> LanguageModelId;
fn name(&self) -> LanguageModelName;
@ -39,6 +49,11 @@ pub trait LanguageModel: Send + Sync {
fn provider_name(&self) -> LanguageModelProviderName;
fn telemetry_id(&self) -> String;
/// Returns the availability of this language model.
fn availability(&self) -> LanguageModelAvailability {
LanguageModelAvailability::Public
}
fn max_token_count(&self) -> usize;
fn count_tokens(