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