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