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