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