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