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
 80impl dyn LanguageModel {
 81    pub fn use_tool<T: LanguageModelTool>(
 82        &self,
 83        request: LanguageModelRequest,
 84        cx: &AsyncAppContext,
 85    ) -> impl 'static + Future<Output = Result<T>> {
 86        let schema = schemars::schema_for!(T);
 87        let schema_json = serde_json::to_value(&schema).unwrap();
 88        let request = self.use_any_tool(request, T::name(), T::description(), schema_json, cx);
 89        async move {
 90            let response = request.await?;
 91            Ok(serde_json::from_value(response)?)
 92        }
 93    }
 94}
 95
 96pub trait LanguageModelTool: 'static + DeserializeOwned + JsonSchema {
 97    fn name() -> String;
 98    fn description() -> String;
 99}
100
101pub trait LanguageModelProvider: 'static {
102    fn id(&self) -> LanguageModelProviderId;
103    fn name(&self) -> LanguageModelProviderName;
104    fn icon(&self) -> IconName {
105        IconName::ZedAssistant
106    }
107    fn provided_models(&self, cx: &AppContext) -> Vec<Arc<dyn LanguageModel>>;
108    fn load_model(&self, _model: Arc<dyn LanguageModel>, _cx: &AppContext) {}
109    fn is_authenticated(&self, cx: &AppContext) -> bool;
110    fn authenticate(&self, cx: &mut AppContext) -> Task<Result<()>>;
111    fn configuration_view(&self, cx: &mut WindowContext) -> AnyView;
112    fn reset_credentials(&self, cx: &mut AppContext) -> Task<Result<()>>;
113}
114
115pub trait LanguageModelProviderState: 'static {
116    type ObservableEntity;
117
118    fn observable_entity(&self) -> Option<gpui::Model<Self::ObservableEntity>>;
119
120    fn subscribe<T: 'static>(
121        &self,
122        cx: &mut gpui::ModelContext<T>,
123        callback: impl Fn(&mut T, &mut gpui::ModelContext<T>) + 'static,
124    ) -> Option<gpui::Subscription> {
125        let entity = self.observable_entity()?;
126        Some(cx.observe(&entity, move |this, _, cx| {
127            callback(this, cx);
128        }))
129    }
130}
131
132#[derive(Clone, Eq, PartialEq, Hash, Debug, Ord, PartialOrd)]
133pub struct LanguageModelId(pub SharedString);
134
135#[derive(Clone, Eq, PartialEq, Hash, Debug, Ord, PartialOrd)]
136pub struct LanguageModelName(pub SharedString);
137
138#[derive(Clone, Eq, PartialEq, Hash, Debug, Ord, PartialOrd)]
139pub struct LanguageModelProviderId(pub SharedString);
140
141#[derive(Clone, Eq, PartialEq, Hash, Debug, Ord, PartialOrd)]
142pub struct LanguageModelProviderName(pub SharedString);
143
144impl From<String> for LanguageModelId {
145    fn from(value: String) -> Self {
146        Self(SharedString::from(value))
147    }
148}
149
150impl From<String> for LanguageModelName {
151    fn from(value: String) -> Self {
152        Self(SharedString::from(value))
153    }
154}
155
156impl From<String> for LanguageModelProviderId {
157    fn from(value: String) -> Self {
158        Self(SharedString::from(value))
159    }
160}
161
162impl From<String> for LanguageModelProviderName {
163    fn from(value: String) -> Self {
164        Self(SharedString::from(value))
165    }
166}