1mod model;
2mod rate_limiter;
3mod registry;
4mod request;
5mod role;
6mod telemetry;
7
8#[cfg(any(test, feature = "test-support"))]
9pub mod fake_provider;
10
11use anyhow::Result;
12use client::Client;
13use futures::FutureExt;
14use futures::{future::BoxFuture, stream::BoxStream, StreamExt, TryStreamExt as _};
15use gpui::{AnyElement, AnyView, App, AsyncApp, SharedString, Task, Window};
16use icons::IconName;
17use proto::Plan;
18use schemars::JsonSchema;
19use serde::{de::DeserializeOwned, Deserialize, Serialize};
20use std::fmt;
21use std::ops::{Add, Sub};
22use std::{future::Future, sync::Arc};
23use thiserror::Error;
24use util::serde::is_default;
25
26pub use crate::model::*;
27pub use crate::rate_limiter::*;
28pub use crate::registry::*;
29pub use crate::request::*;
30pub use crate::role::*;
31pub use crate::telemetry::*;
32
33pub const ZED_CLOUD_PROVIDER_ID: &str = "zed.dev";
34
35pub fn init(client: Arc<Client>, cx: &mut App) {
36 registry::init(cx);
37 RefreshLlmTokenListener::register(client.clone(), cx);
38}
39
40/// The availability of a [`LanguageModel`].
41#[derive(Debug, PartialEq, Eq, Clone, Copy)]
42pub enum LanguageModelAvailability {
43 /// The language model is available to the general public.
44 Public,
45 /// The language model is available to users on the indicated plan.
46 RequiresPlan(Plan),
47}
48
49/// Configuration for caching language model messages.
50#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)]
51pub struct LanguageModelCacheConfiguration {
52 pub max_cache_anchors: usize,
53 pub should_speculate: bool,
54 pub min_total_token: usize,
55}
56
57/// A completion event from a language model.
58#[derive(Debug, PartialEq, Clone, Serialize, Deserialize)]
59pub enum LanguageModelCompletionEvent {
60 Stop(StopReason),
61 Text(String),
62 Thinking(String),
63 ToolUse(LanguageModelToolUse),
64 StartMessage { message_id: String },
65 UsageUpdate(TokenUsage),
66}
67
68#[derive(Debug, PartialEq, Clone, Copy, Serialize, Deserialize)]
69#[serde(rename_all = "snake_case")]
70pub enum StopReason {
71 EndTurn,
72 MaxTokens,
73 ToolUse,
74}
75
76#[derive(Debug, PartialEq, Clone, Serialize, Deserialize, Default)]
77pub struct TokenUsage {
78 #[serde(default, skip_serializing_if = "is_default")]
79 pub input_tokens: u32,
80 #[serde(default, skip_serializing_if = "is_default")]
81 pub output_tokens: u32,
82 #[serde(default, skip_serializing_if = "is_default")]
83 pub cache_creation_input_tokens: u32,
84 #[serde(default, skip_serializing_if = "is_default")]
85 pub cache_read_input_tokens: u32,
86}
87
88impl Add<TokenUsage> for TokenUsage {
89 type Output = Self;
90
91 fn add(self, other: Self) -> Self {
92 Self {
93 input_tokens: self.input_tokens + other.input_tokens,
94 output_tokens: self.output_tokens + other.output_tokens,
95 cache_creation_input_tokens: self.cache_creation_input_tokens
96 + other.cache_creation_input_tokens,
97 cache_read_input_tokens: self.cache_read_input_tokens + other.cache_read_input_tokens,
98 }
99 }
100}
101
102impl Sub<TokenUsage> for TokenUsage {
103 type Output = Self;
104
105 fn sub(self, other: Self) -> Self {
106 Self {
107 input_tokens: self.input_tokens - other.input_tokens,
108 output_tokens: self.output_tokens - other.output_tokens,
109 cache_creation_input_tokens: self.cache_creation_input_tokens
110 - other.cache_creation_input_tokens,
111 cache_read_input_tokens: self.cache_read_input_tokens - other.cache_read_input_tokens,
112 }
113 }
114}
115
116#[derive(Debug, PartialEq, Eq, Hash, Clone, Serialize, Deserialize)]
117pub struct LanguageModelToolUseId(Arc<str>);
118
119impl fmt::Display for LanguageModelToolUseId {
120 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
121 write!(f, "{}", self.0)
122 }
123}
124
125impl<T> From<T> for LanguageModelToolUseId
126where
127 T: Into<Arc<str>>,
128{
129 fn from(value: T) -> Self {
130 Self(value.into())
131 }
132}
133
134#[derive(Debug, PartialEq, Eq, Hash, Clone, Serialize, Deserialize)]
135pub struct LanguageModelToolUse {
136 pub id: LanguageModelToolUseId,
137 pub name: Arc<str>,
138 pub input: serde_json::Value,
139}
140
141pub struct LanguageModelTextStream {
142 pub message_id: Option<String>,
143 pub stream: BoxStream<'static, Result<String>>,
144}
145
146impl Default for LanguageModelTextStream {
147 fn default() -> Self {
148 Self {
149 message_id: None,
150 stream: Box::pin(futures::stream::empty()),
151 }
152 }
153}
154
155pub trait LanguageModel: Send + Sync {
156 fn id(&self) -> LanguageModelId;
157 fn name(&self) -> LanguageModelName;
158 /// If None, falls back to [LanguageModelProvider::icon]
159 fn icon(&self) -> Option<IconName> {
160 None
161 }
162 fn provider_id(&self) -> LanguageModelProviderId;
163 fn provider_name(&self) -> LanguageModelProviderName;
164 fn telemetry_id(&self) -> String;
165
166 fn api_key(&self, _cx: &App) -> Option<String> {
167 None
168 }
169
170 /// Returns the availability of this language model.
171 fn availability(&self) -> LanguageModelAvailability {
172 LanguageModelAvailability::Public
173 }
174
175 fn max_token_count(&self) -> usize;
176 fn max_output_tokens(&self) -> Option<u32> {
177 None
178 }
179
180 fn count_tokens(
181 &self,
182 request: LanguageModelRequest,
183 cx: &App,
184 ) -> BoxFuture<'static, Result<usize>>;
185
186 fn stream_completion(
187 &self,
188 request: LanguageModelRequest,
189 cx: &AsyncApp,
190 ) -> BoxFuture<'static, Result<BoxStream<'static, Result<LanguageModelCompletionEvent>>>>;
191
192 fn stream_completion_text(
193 &self,
194 request: LanguageModelRequest,
195 cx: &AsyncApp,
196 ) -> BoxFuture<'static, Result<LanguageModelTextStream>> {
197 let events = self.stream_completion(request, cx);
198
199 async move {
200 let mut events = events.await?.fuse();
201 let mut message_id = None;
202 let mut first_item_text = None;
203
204 if let Some(first_event) = events.next().await {
205 match first_event {
206 Ok(LanguageModelCompletionEvent::StartMessage { message_id: id }) => {
207 message_id = Some(id.clone());
208 }
209 Ok(LanguageModelCompletionEvent::Text(text)) => {
210 first_item_text = Some(text);
211 }
212 _ => (),
213 }
214 }
215
216 let stream = futures::stream::iter(first_item_text.map(Ok))
217 .chain(events.filter_map(|result| async move {
218 match result {
219 Ok(LanguageModelCompletionEvent::StartMessage { .. }) => None,
220 Ok(LanguageModelCompletionEvent::Text(text)) => Some(Ok(text)),
221 Ok(LanguageModelCompletionEvent::Thinking(_)) => None,
222 Ok(LanguageModelCompletionEvent::Stop(_)) => None,
223 Ok(LanguageModelCompletionEvent::ToolUse(_)) => None,
224 Ok(LanguageModelCompletionEvent::UsageUpdate(_)) => None,
225 Err(err) => Some(Err(err)),
226 }
227 }))
228 .boxed();
229
230 Ok(LanguageModelTextStream { message_id, stream })
231 }
232 .boxed()
233 }
234
235 fn use_any_tool(
236 &self,
237 request: LanguageModelRequest,
238 name: String,
239 description: String,
240 schema: serde_json::Value,
241 cx: &AsyncApp,
242 ) -> BoxFuture<'static, Result<BoxStream<'static, Result<String>>>>;
243
244 fn cache_configuration(&self) -> Option<LanguageModelCacheConfiguration> {
245 None
246 }
247
248 #[cfg(any(test, feature = "test-support"))]
249 fn as_fake(&self) -> &fake_provider::FakeLanguageModel {
250 unimplemented!()
251 }
252}
253
254impl dyn LanguageModel {
255 pub fn use_tool<T: LanguageModelTool>(
256 &self,
257 request: LanguageModelRequest,
258 cx: &AsyncApp,
259 ) -> impl 'static + Future<Output = Result<T>> {
260 let schema = schemars::schema_for!(T);
261 let schema_json = serde_json::to_value(&schema).unwrap();
262 let stream = self.use_any_tool(request, T::name(), T::description(), schema_json, cx);
263 async move {
264 let stream = stream.await?;
265 let response = stream.try_collect::<String>().await?;
266 Ok(serde_json::from_str(&response)?)
267 }
268 }
269
270 pub fn use_tool_stream<T: LanguageModelTool>(
271 &self,
272 request: LanguageModelRequest,
273 cx: &AsyncApp,
274 ) -> BoxFuture<'static, Result<BoxStream<'static, Result<String>>>> {
275 let schema = schemars::schema_for!(T);
276 let schema_json = serde_json::to_value(&schema).unwrap();
277 self.use_any_tool(request, T::name(), T::description(), schema_json, cx)
278 }
279}
280
281pub trait LanguageModelTool: 'static + DeserializeOwned + JsonSchema {
282 fn name() -> String;
283 fn description() -> String;
284}
285
286/// An error that occurred when trying to authenticate the language model provider.
287#[derive(Debug, Error)]
288pub enum AuthenticateError {
289 #[error("credentials not found")]
290 CredentialsNotFound,
291 #[error(transparent)]
292 Other(#[from] anyhow::Error),
293}
294
295pub trait LanguageModelProvider: 'static {
296 fn id(&self) -> LanguageModelProviderId;
297 fn name(&self) -> LanguageModelProviderName;
298 fn icon(&self) -> IconName {
299 IconName::ZedAssistant
300 }
301 fn default_model(&self, cx: &App) -> Option<Arc<dyn LanguageModel>>;
302 fn provided_models(&self, cx: &App) -> Vec<Arc<dyn LanguageModel>>;
303 fn load_model(&self, _model: Arc<dyn LanguageModel>, _cx: &App) {}
304 fn is_authenticated(&self, cx: &App) -> bool;
305 fn authenticate(&self, cx: &mut App) -> Task<Result<(), AuthenticateError>>;
306 fn configuration_view(&self, window: &mut Window, cx: &mut App) -> AnyView;
307 fn must_accept_terms(&self, _cx: &App) -> bool {
308 false
309 }
310 fn render_accept_terms(
311 &self,
312 _view: LanguageModelProviderTosView,
313 _cx: &mut App,
314 ) -> Option<AnyElement> {
315 None
316 }
317 fn reset_credentials(&self, cx: &mut App) -> Task<Result<()>>;
318}
319
320#[derive(PartialEq, Eq)]
321pub enum LanguageModelProviderTosView {
322 ThreadEmptyState,
323 PromptEditorPopup,
324 Configuration,
325}
326
327pub trait LanguageModelProviderState: 'static {
328 type ObservableEntity;
329
330 fn observable_entity(&self) -> Option<gpui::Entity<Self::ObservableEntity>>;
331
332 fn subscribe<T: 'static>(
333 &self,
334 cx: &mut gpui::Context<T>,
335 callback: impl Fn(&mut T, &mut gpui::Context<T>) + 'static,
336 ) -> Option<gpui::Subscription> {
337 let entity = self.observable_entity()?;
338 Some(cx.observe(&entity, move |this, _, cx| {
339 callback(this, cx);
340 }))
341 }
342}
343
344#[derive(Clone, Eq, PartialEq, Hash, Debug, Ord, PartialOrd)]
345pub struct LanguageModelId(pub SharedString);
346
347#[derive(Clone, Eq, PartialEq, Hash, Debug, Ord, PartialOrd)]
348pub struct LanguageModelName(pub SharedString);
349
350#[derive(Clone, Eq, PartialEq, Hash, Debug, Ord, PartialOrd)]
351pub struct LanguageModelProviderId(pub SharedString);
352
353#[derive(Clone, Eq, PartialEq, Hash, Debug, Ord, PartialOrd)]
354pub struct LanguageModelProviderName(pub SharedString);
355
356impl fmt::Display for LanguageModelProviderId {
357 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
358 write!(f, "{}", self.0)
359 }
360}
361
362impl From<String> for LanguageModelId {
363 fn from(value: String) -> Self {
364 Self(SharedString::from(value))
365 }
366}
367
368impl From<String> for LanguageModelName {
369 fn from(value: String) -> Self {
370 Self(SharedString::from(value))
371 }
372}
373
374impl From<String> for LanguageModelProviderId {
375 fn from(value: String) -> Self {
376 Self(SharedString::from(value))
377 }
378}
379
380impl From<String> for LanguageModelProviderName {
381 fn from(value: String) -> Self {
382 Self(SharedString::from(value))
383 }
384}