1mod model;
  2pub mod provider;
  3mod rate_limiter;
  4mod registry;
  5mod request;
  6mod role;
  7pub mod settings;
  8
  9use anyhow::Result;
 10use client::{Client, UserStore};
 11use futures::{future::BoxFuture, stream::BoxStream};
 12use gpui::{
 13    AnyView, AppContext, AsyncAppContext, FocusHandle, Model, SharedString, Task, WindowContext,
 14};
 15pub use model::*;
 16use project::Fs;
 17use proto::Plan;
 18pub(crate) use rate_limiter::*;
 19pub use registry::*;
 20pub use request::*;
 21pub use role::*;
 22use schemars::JsonSchema;
 23use serde::de::DeserializeOwned;
 24use std::{future::Future, sync::Arc};
 25use ui::IconName;
 26
 27pub fn init(
 28    user_store: Model<UserStore>,
 29    client: Arc<Client>,
 30    fs: Arc<dyn Fs>,
 31    cx: &mut AppContext,
 32) {
 33    settings::init(fs, cx);
 34    registry::init(user_store, client, cx);
 35}
 36
 37/// The availability of a [`LanguageModel`].
 38#[derive(Debug, PartialEq, Eq, Clone, Copy)]
 39pub enum LanguageModelAvailability {
 40    /// The language model is available to the general public.
 41    Public,
 42    /// The language model is available to users on the indicated plan.
 43    RequiresPlan(Plan),
 44}
 45
 46pub trait LanguageModel: Send + Sync {
 47    fn id(&self) -> LanguageModelId;
 48    fn name(&self) -> LanguageModelName;
 49    fn provider_id(&self) -> LanguageModelProviderId;
 50    fn provider_name(&self) -> LanguageModelProviderName;
 51    fn telemetry_id(&self) -> String;
 52
 53    /// Returns the availability of this language model.
 54    fn availability(&self) -> LanguageModelAvailability {
 55        LanguageModelAvailability::Public
 56    }
 57
 58    fn max_token_count(&self) -> usize;
 59
 60    fn count_tokens(
 61        &self,
 62        request: LanguageModelRequest,
 63        cx: &AppContext,
 64    ) -> BoxFuture<'static, Result<usize>>;
 65
 66    fn stream_completion(
 67        &self,
 68        request: LanguageModelRequest,
 69        cx: &AsyncAppContext,
 70    ) -> BoxFuture<'static, Result<BoxStream<'static, Result<String>>>>;
 71
 72    fn use_any_tool(
 73        &self,
 74        request: LanguageModelRequest,
 75        name: String,
 76        description: String,
 77        schema: serde_json::Value,
 78        cx: &AsyncAppContext,
 79    ) -> BoxFuture<'static, Result<serde_json::Value>>;
 80}
 81
 82impl dyn LanguageModel {
 83    pub fn use_tool<T: LanguageModelTool>(
 84        &self,
 85        request: LanguageModelRequest,
 86        cx: &AsyncAppContext,
 87    ) -> impl 'static + Future<Output = Result<T>> {
 88        let schema = schemars::schema_for!(T);
 89        let schema_json = serde_json::to_value(&schema).unwrap();
 90        let request = self.use_any_tool(request, T::name(), T::description(), schema_json, cx);
 91        async move {
 92            let response = request.await?;
 93            Ok(serde_json::from_value(response)?)
 94        }
 95    }
 96}
 97
 98pub trait LanguageModelTool: 'static + DeserializeOwned + JsonSchema {
 99    fn name() -> String;
100    fn description() -> String;
101}
102
103pub trait LanguageModelProvider: 'static {
104    fn id(&self) -> LanguageModelProviderId;
105    fn name(&self) -> LanguageModelProviderName;
106    fn icon(&self) -> IconName {
107        IconName::ZedAssistant
108    }
109    fn provided_models(&self, cx: &AppContext) -> Vec<Arc<dyn LanguageModel>>;
110    fn load_model(&self, _model: Arc<dyn LanguageModel>, _cx: &AppContext) {}
111    fn is_authenticated(&self, cx: &AppContext) -> bool;
112    fn authenticate(&self, cx: &mut AppContext) -> Task<Result<()>>;
113    fn configuration_view(&self, cx: &mut WindowContext) -> (AnyView, Option<FocusHandle>);
114    fn reset_credentials(&self, cx: &mut AppContext) -> Task<Result<()>>;
115}
116
117pub trait LanguageModelProviderState: 'static {
118    type ObservableEntity;
119
120    fn observable_entity(&self) -> Option<gpui::Model<Self::ObservableEntity>>;
121
122    fn subscribe<T: 'static>(
123        &self,
124        cx: &mut gpui::ModelContext<T>,
125        callback: impl Fn(&mut T, &mut gpui::ModelContext<T>) + 'static,
126    ) -> Option<gpui::Subscription> {
127        let entity = self.observable_entity()?;
128        Some(cx.observe(&entity, move |this, _, cx| {
129            callback(this, cx);
130        }))
131    }
132}
133
134#[derive(Clone, Eq, PartialEq, Hash, Debug, Ord, PartialOrd)]
135pub struct LanguageModelId(pub SharedString);
136
137#[derive(Clone, Eq, PartialEq, Hash, Debug, Ord, PartialOrd)]
138pub struct LanguageModelName(pub SharedString);
139
140#[derive(Clone, Eq, PartialEq, Hash, Debug, Ord, PartialOrd)]
141pub struct LanguageModelProviderId(pub SharedString);
142
143#[derive(Clone, Eq, PartialEq, Hash, Debug, Ord, PartialOrd)]
144pub struct LanguageModelProviderName(pub SharedString);
145
146impl From<String> for LanguageModelId {
147    fn from(value: String) -> Self {
148        Self(SharedString::from(value))
149    }
150}
151
152impl From<String> for LanguageModelName {
153    fn from(value: String) -> Self {
154        Self(SharedString::from(value))
155    }
156}
157
158impl From<String> for LanguageModelProviderId {
159    fn from(value: String) -> Self {
160        Self(SharedString::from(value))
161    }
162}
163
164impl From<String> for LanguageModelProviderName {
165    fn from(value: String) -> Self {
166        Self(SharedString::from(value))
167    }
168}