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, Deserialize, Serialize};
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
46/// Configuration for caching language model messages.
47#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)]
48pub struct LanguageModelCacheConfiguration {
49 pub max_cache_anchors: usize,
50 pub should_speculate: bool,
51 pub min_total_token: usize,
52}
53
54pub trait LanguageModel: Send + Sync {
55 fn id(&self) -> LanguageModelId;
56 fn name(&self) -> LanguageModelName;
57 /// If None, falls back to [LanguageModelProvider::icon]
58 fn icon(&self) -> Option<IconName> {
59 None
60 }
61 fn provider_id(&self) -> LanguageModelProviderId;
62 fn provider_name(&self) -> LanguageModelProviderName;
63 fn telemetry_id(&self) -> String;
64
65 /// Returns the availability of this language model.
66 fn availability(&self) -> LanguageModelAvailability {
67 LanguageModelAvailability::Public
68 }
69
70 fn max_token_count(&self) -> usize;
71 fn max_output_tokens(&self) -> Option<u32> {
72 None
73 }
74
75 fn count_tokens(
76 &self,
77 request: LanguageModelRequest,
78 cx: &AppContext,
79 ) -> BoxFuture<'static, Result<usize>>;
80
81 fn stream_completion(
82 &self,
83 request: LanguageModelRequest,
84 cx: &AsyncAppContext,
85 ) -> BoxFuture<'static, Result<BoxStream<'static, Result<String>>>>;
86
87 fn use_any_tool(
88 &self,
89 request: LanguageModelRequest,
90 name: String,
91 description: String,
92 schema: serde_json::Value,
93 cx: &AsyncAppContext,
94 ) -> BoxFuture<'static, Result<BoxStream<'static, Result<String>>>>;
95
96 fn cache_configuration(&self) -> Option<LanguageModelCacheConfiguration> {
97 None
98 }
99
100 #[cfg(any(test, feature = "test-support"))]
101 fn as_fake(&self) -> &provider::fake::FakeLanguageModel {
102 unimplemented!()
103 }
104}
105
106impl dyn LanguageModel {
107 pub fn use_tool<T: LanguageModelTool>(
108 &self,
109 request: LanguageModelRequest,
110 cx: &AsyncAppContext,
111 ) -> impl 'static + Future<Output = Result<T>> {
112 let schema = schemars::schema_for!(T);
113 let schema_json = serde_json::to_value(&schema).unwrap();
114 let stream = self.use_any_tool(request, T::name(), T::description(), schema_json, cx);
115 async move {
116 let stream = stream.await?;
117 let response = stream.try_collect::<String>().await?;
118 Ok(serde_json::from_str(&response)?)
119 }
120 }
121
122 pub fn use_tool_stream<T: LanguageModelTool>(
123 &self,
124 request: LanguageModelRequest,
125 cx: &AsyncAppContext,
126 ) -> BoxFuture<'static, Result<BoxStream<'static, Result<String>>>> {
127 let schema = schemars::schema_for!(T);
128 let schema_json = serde_json::to_value(&schema).unwrap();
129 self.use_any_tool(request, T::name(), T::description(), schema_json, cx)
130 }
131}
132
133pub trait LanguageModelTool: 'static + DeserializeOwned + JsonSchema {
134 fn name() -> String;
135 fn description() -> String;
136}
137
138pub trait LanguageModelProvider: 'static {
139 fn id(&self) -> LanguageModelProviderId;
140 fn name(&self) -> LanguageModelProviderName;
141 fn icon(&self) -> IconName {
142 IconName::ZedAssistant
143 }
144 fn provided_models(&self, cx: &AppContext) -> Vec<Arc<dyn LanguageModel>>;
145 fn load_model(&self, _model: Arc<dyn LanguageModel>, _cx: &AppContext) {}
146 fn is_authenticated(&self, cx: &AppContext) -> bool;
147 fn authenticate(&self, cx: &mut AppContext) -> Task<Result<()>>;
148 fn configuration_view(&self, cx: &mut WindowContext) -> AnyView;
149 fn must_accept_terms(&self, _cx: &AppContext) -> bool {
150 false
151 }
152 fn render_accept_terms(&self, _cx: &mut WindowContext) -> Option<AnyElement> {
153 None
154 }
155 fn reset_credentials(&self, cx: &mut AppContext) -> Task<Result<()>>;
156}
157
158pub trait LanguageModelProviderState: 'static {
159 type ObservableEntity;
160
161 fn observable_entity(&self) -> Option<gpui::Model<Self::ObservableEntity>>;
162
163 fn subscribe<T: 'static>(
164 &self,
165 cx: &mut gpui::ModelContext<T>,
166 callback: impl Fn(&mut T, &mut gpui::ModelContext<T>) + 'static,
167 ) -> Option<gpui::Subscription> {
168 let entity = self.observable_entity()?;
169 Some(cx.observe(&entity, move |this, _, cx| {
170 callback(this, cx);
171 }))
172 }
173}
174
175#[derive(Clone, Eq, PartialEq, Hash, Debug, Ord, PartialOrd)]
176pub struct LanguageModelId(pub SharedString);
177
178#[derive(Clone, Eq, PartialEq, Hash, Debug, Ord, PartialOrd)]
179pub struct LanguageModelName(pub SharedString);
180
181#[derive(Clone, Eq, PartialEq, Hash, Debug, Ord, PartialOrd)]
182pub struct LanguageModelProviderId(pub SharedString);
183
184#[derive(Clone, Eq, PartialEq, Hash, Debug, Ord, PartialOrd)]
185pub struct LanguageModelProviderName(pub SharedString);
186
187impl From<String> for LanguageModelId {
188 fn from(value: String) -> Self {
189 Self(SharedString::from(value))
190 }
191}
192
193impl From<String> for LanguageModelName {
194 fn from(value: String) -> Self {
195 Self(SharedString::from(value))
196 }
197}
198
199impl From<String> for LanguageModelProviderId {
200 fn from(value: String) -> Self {
201 Self(SharedString::from(value))
202 }
203}
204
205impl From<String> for LanguageModelProviderName {
206 fn from(value: String) -> Self {
207 Self(SharedString::from(value))
208 }
209}