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