1use anyhow::Result;
2use collections::BTreeMap;
3use credentials_provider::CredentialsProvider;
4use futures::{FutureExt, StreamExt, future::BoxFuture};
5use gpui::{AnyView, App, AsyncApp, Context, Entity, Task, Window};
6use http_client::HttpClient;
7use language_model::{
8 ApiKeyState, AuthenticateError, EnvVar, IconOrSvg, LanguageModel, LanguageModelCompletionError,
9 LanguageModelCompletionEvent, LanguageModelId, LanguageModelName, LanguageModelProvider,
10 LanguageModelProviderId, LanguageModelProviderName, LanguageModelProviderState,
11 LanguageModelRequest, LanguageModelToolChoice, LanguageModelToolSchemaFormat, RateLimiter,
12 Role, env_var,
13};
14use open_ai::ResponseStreamEvent;
15pub use settings::XaiAvailableModel as AvailableModel;
16use settings::{Settings, SettingsStore};
17use std::sync::{Arc, LazyLock};
18use strum::IntoEnumIterator;
19use ui::{ButtonLink, ConfiguredApiCard, List, ListBulletItem, prelude::*};
20use ui_input::InputField;
21use util::ResultExt;
22use x_ai::{Model, XAI_API_URL};
23
24const PROVIDER_ID: LanguageModelProviderId = LanguageModelProviderId::new("x_ai");
25const PROVIDER_NAME: LanguageModelProviderName = LanguageModelProviderName::new("xAI");
26
27const API_KEY_ENV_VAR_NAME: &str = "XAI_API_KEY";
28static API_KEY_ENV_VAR: LazyLock<EnvVar> = env_var!(API_KEY_ENV_VAR_NAME);
29
30#[derive(Default, Clone, Debug, PartialEq)]
31pub struct XAiSettings {
32 pub api_url: String,
33 pub available_models: Vec<AvailableModel>,
34}
35
36pub struct XAiLanguageModelProvider {
37 http_client: Arc<dyn HttpClient>,
38 state: Entity<State>,
39}
40
41pub struct State {
42 api_key_state: ApiKeyState,
43 credentials_provider: Arc<dyn CredentialsProvider>,
44}
45
46impl State {
47 fn is_authenticated(&self) -> bool {
48 self.api_key_state.has_key()
49 }
50
51 fn set_api_key(&mut self, api_key: Option<String>, cx: &mut Context<Self>) -> Task<Result<()>> {
52 let credentials_provider = self.credentials_provider.clone();
53 let api_url = XAiLanguageModelProvider::api_url(cx);
54 self.api_key_state.store(
55 api_url,
56 api_key,
57 |this| &mut this.api_key_state,
58 credentials_provider,
59 cx,
60 )
61 }
62
63 fn authenticate(&mut self, cx: &mut Context<Self>) -> Task<Result<(), AuthenticateError>> {
64 let credentials_provider = self.credentials_provider.clone();
65 let api_url = XAiLanguageModelProvider::api_url(cx);
66 self.api_key_state.load_if_needed(
67 api_url,
68 |this| &mut this.api_key_state,
69 credentials_provider,
70 cx,
71 )
72 }
73}
74
75impl XAiLanguageModelProvider {
76 pub fn new(
77 http_client: Arc<dyn HttpClient>,
78 credentials_provider: Arc<dyn CredentialsProvider>,
79 cx: &mut App,
80 ) -> Self {
81 let state = cx.new(|cx| {
82 cx.observe_global::<SettingsStore>(|this: &mut State, cx| {
83 let credentials_provider = this.credentials_provider.clone();
84 let api_url = Self::api_url(cx);
85 this.api_key_state.handle_url_change(
86 api_url,
87 |this| &mut this.api_key_state,
88 credentials_provider,
89 cx,
90 );
91 cx.notify();
92 })
93 .detach();
94 State {
95 api_key_state: ApiKeyState::new(Self::api_url(cx), (*API_KEY_ENV_VAR).clone()),
96 credentials_provider,
97 }
98 });
99
100 Self { http_client, state }
101 }
102
103 fn create_language_model(&self, model: x_ai::Model) -> Arc<dyn LanguageModel> {
104 Arc::new(XAiLanguageModel {
105 id: LanguageModelId::from(model.id().to_string()),
106 model,
107 state: self.state.clone(),
108 http_client: self.http_client.clone(),
109 request_limiter: RateLimiter::new(4),
110 })
111 }
112
113 fn settings(cx: &App) -> &XAiSettings {
114 &crate::AllLanguageModelSettings::get_global(cx).x_ai
115 }
116
117 fn api_url(cx: &App) -> SharedString {
118 let api_url = &Self::settings(cx).api_url;
119 if api_url.is_empty() {
120 XAI_API_URL.into()
121 } else {
122 SharedString::new(api_url.as_str())
123 }
124 }
125}
126
127impl LanguageModelProviderState for XAiLanguageModelProvider {
128 type ObservableEntity = State;
129
130 fn observable_entity(&self) -> Option<Entity<Self::ObservableEntity>> {
131 Some(self.state.clone())
132 }
133}
134
135impl LanguageModelProvider for XAiLanguageModelProvider {
136 fn id(&self) -> LanguageModelProviderId {
137 PROVIDER_ID
138 }
139
140 fn name(&self) -> LanguageModelProviderName {
141 PROVIDER_NAME
142 }
143
144 fn icon(&self) -> IconOrSvg {
145 IconOrSvg::Icon(IconName::AiXAi)
146 }
147
148 fn default_model(&self, _cx: &App) -> Option<Arc<dyn LanguageModel>> {
149 Some(self.create_language_model(x_ai::Model::default()))
150 }
151
152 fn default_fast_model(&self, _cx: &App) -> Option<Arc<dyn LanguageModel>> {
153 Some(self.create_language_model(x_ai::Model::default_fast()))
154 }
155
156 fn provided_models(&self, cx: &App) -> Vec<Arc<dyn LanguageModel>> {
157 let mut models = BTreeMap::default();
158
159 for model in x_ai::Model::iter() {
160 if !matches!(model, x_ai::Model::Custom { .. }) {
161 models.insert(model.id().to_string(), model);
162 }
163 }
164
165 for model in &Self::settings(cx).available_models {
166 models.insert(
167 model.name.clone(),
168 x_ai::Model::Custom {
169 name: model.name.clone(),
170 display_name: model.display_name.clone(),
171 max_tokens: model.max_tokens,
172 max_output_tokens: model.max_output_tokens,
173 max_completion_tokens: model.max_completion_tokens,
174 supports_images: model.supports_images,
175 supports_tools: model.supports_tools,
176 parallel_tool_calls: model.parallel_tool_calls,
177 },
178 );
179 }
180
181 models
182 .into_values()
183 .map(|model| self.create_language_model(model))
184 .collect()
185 }
186
187 fn is_authenticated(&self, cx: &App) -> bool {
188 self.state.read(cx).is_authenticated()
189 }
190
191 fn authenticate(&self, cx: &mut App) -> Task<Result<(), AuthenticateError>> {
192 self.state.update(cx, |state, cx| state.authenticate(cx))
193 }
194
195 fn configuration_view(
196 &self,
197 _target_agent: language_model::ConfigurationViewTargetAgent,
198 window: &mut Window,
199 cx: &mut App,
200 ) -> AnyView {
201 cx.new(|cx| ConfigurationView::new(self.state.clone(), window, cx))
202 .into()
203 }
204
205 fn reset_credentials(&self, cx: &mut App) -> Task<Result<()>> {
206 self.state
207 .update(cx, |state, cx| state.set_api_key(None, cx))
208 }
209}
210
211pub struct XAiLanguageModel {
212 id: LanguageModelId,
213 model: x_ai::Model,
214 state: Entity<State>,
215 http_client: Arc<dyn HttpClient>,
216 request_limiter: RateLimiter,
217}
218
219impl XAiLanguageModel {
220 fn stream_completion(
221 &self,
222 request: open_ai::Request,
223 cx: &AsyncApp,
224 ) -> BoxFuture<
225 'static,
226 Result<
227 futures::stream::BoxStream<'static, Result<ResponseStreamEvent>>,
228 LanguageModelCompletionError,
229 >,
230 > {
231 let http_client = self.http_client.clone();
232
233 let (api_key, api_url) = self.state.read_with(cx, |state, cx| {
234 let api_url = XAiLanguageModelProvider::api_url(cx);
235 (state.api_key_state.key(&api_url), api_url)
236 });
237
238 let future = self.request_limiter.stream(async move {
239 let provider = PROVIDER_NAME;
240 let Some(api_key) = api_key else {
241 return Err(LanguageModelCompletionError::NoApiKey { provider });
242 };
243 let request = open_ai::stream_completion(
244 http_client.as_ref(),
245 provider.0.as_str(),
246 &api_url,
247 &api_key,
248 request,
249 );
250 let response = request.await?;
251 Ok(response)
252 });
253
254 async move { Ok(future.await?.boxed()) }.boxed()
255 }
256}
257
258impl LanguageModel for XAiLanguageModel {
259 fn id(&self) -> LanguageModelId {
260 self.id.clone()
261 }
262
263 fn name(&self) -> LanguageModelName {
264 LanguageModelName::from(self.model.display_name().to_string())
265 }
266
267 fn provider_id(&self) -> LanguageModelProviderId {
268 PROVIDER_ID
269 }
270
271 fn provider_name(&self) -> LanguageModelProviderName {
272 PROVIDER_NAME
273 }
274
275 fn supports_tools(&self) -> bool {
276 self.model.supports_tool()
277 }
278
279 fn supports_images(&self) -> bool {
280 self.model.supports_images()
281 }
282
283 fn supports_streaming_tools(&self) -> bool {
284 true
285 }
286
287 fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool {
288 match choice {
289 LanguageModelToolChoice::Auto
290 | LanguageModelToolChoice::Any
291 | LanguageModelToolChoice::None => true,
292 }
293 }
294 fn tool_input_format(&self) -> LanguageModelToolSchemaFormat {
295 if self.model.requires_json_schema_subset() {
296 LanguageModelToolSchemaFormat::JsonSchemaSubset
297 } else {
298 LanguageModelToolSchemaFormat::JsonSchema
299 }
300 }
301
302 fn telemetry_id(&self) -> String {
303 format!("x_ai/{}", self.model.id())
304 }
305
306 fn max_token_count(&self) -> u64 {
307 self.model.max_token_count()
308 }
309
310 fn max_output_tokens(&self) -> Option<u64> {
311 self.model.max_output_tokens()
312 }
313
314 fn supports_split_token_display(&self) -> bool {
315 true
316 }
317
318 fn count_tokens(
319 &self,
320 request: LanguageModelRequest,
321 cx: &App,
322 ) -> BoxFuture<'static, Result<u64>> {
323 count_xai_tokens(request, self.model.clone(), cx)
324 }
325
326 fn stream_completion(
327 &self,
328 request: LanguageModelRequest,
329 cx: &AsyncApp,
330 ) -> BoxFuture<
331 'static,
332 Result<
333 futures::stream::BoxStream<
334 'static,
335 Result<LanguageModelCompletionEvent, LanguageModelCompletionError>,
336 >,
337 LanguageModelCompletionError,
338 >,
339 > {
340 let request = crate::provider::open_ai::into_open_ai(
341 request,
342 self.model.id(),
343 self.model.supports_parallel_tool_calls(),
344 self.model.supports_prompt_cache_key(),
345 self.max_output_tokens(),
346 None,
347 );
348 let completions = self.stream_completion(request, cx);
349 async move {
350 let mapper = crate::provider::open_ai::OpenAiEventMapper::new();
351 Ok(mapper.map_stream(completions.await?).boxed())
352 }
353 .boxed()
354 }
355}
356
357pub fn count_xai_tokens(
358 request: LanguageModelRequest,
359 model: Model,
360 cx: &App,
361) -> BoxFuture<'static, Result<u64>> {
362 cx.background_spawn(async move {
363 let messages = request
364 .messages
365 .into_iter()
366 .map(|message| tiktoken_rs::ChatCompletionRequestMessage {
367 role: match message.role {
368 Role::User => "user".into(),
369 Role::Assistant => "assistant".into(),
370 Role::System => "system".into(),
371 },
372 content: Some(message.string_contents()),
373 name: None,
374 function_call: None,
375 })
376 .collect::<Vec<_>>();
377
378 let model_name = if model.max_token_count() >= 100_000 {
379 "gpt-4o"
380 } else {
381 "gpt-4"
382 };
383 tiktoken_rs::num_tokens_from_messages(model_name, &messages).map(|tokens| tokens as u64)
384 })
385 .boxed()
386}
387
388struct ConfigurationView {
389 api_key_editor: Entity<InputField>,
390 state: Entity<State>,
391 load_credentials_task: Option<Task<()>>,
392}
393
394impl ConfigurationView {
395 fn new(state: Entity<State>, window: &mut Window, cx: &mut Context<Self>) -> Self {
396 let api_key_editor = cx.new(|cx| {
397 InputField::new(
398 window,
399 cx,
400 "xai-0000000000000000000000000000000000000000000000000",
401 )
402 .label("API key")
403 });
404
405 cx.observe(&state, |_, _, cx| {
406 cx.notify();
407 })
408 .detach();
409
410 let load_credentials_task = Some(cx.spawn_in(window, {
411 let state = state.clone();
412 async move |this, cx| {
413 if let Some(task) = Some(state.update(cx, |state, cx| state.authenticate(cx))) {
414 // We don't log an error, because "not signed in" is also an error.
415 let _ = task.await;
416 }
417 this.update(cx, |this, cx| {
418 this.load_credentials_task = None;
419 cx.notify();
420 })
421 .log_err();
422 }
423 }));
424
425 Self {
426 api_key_editor,
427 state,
428 load_credentials_task,
429 }
430 }
431
432 fn save_api_key(&mut self, _: &menu::Confirm, window: &mut Window, cx: &mut Context<Self>) {
433 let api_key = self.api_key_editor.read(cx).text(cx).trim().to_string();
434 if api_key.is_empty() {
435 return;
436 }
437
438 // url changes can cause the editor to be displayed again
439 self.api_key_editor
440 .update(cx, |editor, cx| editor.set_text("", window, cx));
441
442 let state = self.state.clone();
443 cx.spawn_in(window, async move |_, cx| {
444 state
445 .update(cx, |state, cx| state.set_api_key(Some(api_key), cx))
446 .await
447 })
448 .detach_and_log_err(cx);
449 }
450
451 fn reset_api_key(&mut self, window: &mut Window, cx: &mut Context<Self>) {
452 self.api_key_editor
453 .update(cx, |input, cx| input.set_text("", window, cx));
454
455 let state = self.state.clone();
456 cx.spawn_in(window, async move |_, cx| {
457 state
458 .update(cx, |state, cx| state.set_api_key(None, cx))
459 .await
460 })
461 .detach_and_log_err(cx);
462 }
463
464 fn should_render_editor(&self, cx: &mut Context<Self>) -> bool {
465 !self.state.read(cx).is_authenticated()
466 }
467}
468
469impl Render for ConfigurationView {
470 fn render(&mut self, _: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
471 let env_var_set = self.state.read(cx).api_key_state.is_from_env_var();
472 let configured_card_label = if env_var_set {
473 format!("API key set in {API_KEY_ENV_VAR_NAME} environment variable")
474 } else {
475 let api_url = XAiLanguageModelProvider::api_url(cx);
476 if api_url == XAI_API_URL {
477 "API key configured".to_string()
478 } else {
479 format!("API key configured for {}", api_url)
480 }
481 };
482
483 let api_key_section = if self.should_render_editor(cx) {
484 v_flex()
485 .on_action(cx.listener(Self::save_api_key))
486 .child(Label::new("To use Zed's agent with xAI, you need to add an API key. Follow these steps:"))
487 .child(
488 List::new()
489 .child(
490 ListBulletItem::new("")
491 .child(Label::new("Create one by visiting"))
492 .child(ButtonLink::new("xAI console", "https://console.x.ai/team/default/api-keys"))
493 )
494 .child(
495 ListBulletItem::new("Paste your API key below and hit enter to start using the agent")
496 ),
497 )
498 .child(self.api_key_editor.clone())
499 .child(
500 Label::new(format!(
501 "You can also set the {API_KEY_ENV_VAR_NAME} environment variable and restart Zed."
502 ))
503 .size(LabelSize::Small)
504 .color(Color::Muted),
505 )
506 .child(
507 Label::new("Note that xAI is a custom OpenAI-compatible provider.")
508 .size(LabelSize::Small)
509 .color(Color::Muted),
510 )
511 .into_any_element()
512 } else {
513 ConfiguredApiCard::new(configured_card_label)
514 .disabled(env_var_set)
515 .when(env_var_set, |this| {
516 this.tooltip_label(format!("To reset your API key, unset the {API_KEY_ENV_VAR_NAME} environment variable."))
517 })
518 .on_click(cx.listener(|this, _, window, cx| this.reset_api_key(window, cx)))
519 .into_any_element()
520 };
521
522 if self.load_credentials_task.is_some() {
523 div().child(Label::new("Loading credentials…")).into_any()
524 } else {
525 v_flex().size_full().child(api_key_section).into_any()
526 }
527 }
528}