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 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::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 stream_completion(
319 &self,
320 request: LanguageModelRequest,
321 cx: &AsyncApp,
322 ) -> BoxFuture<
323 'static,
324 Result<
325 futures::stream::BoxStream<
326 'static,
327 Result<LanguageModelCompletionEvent, LanguageModelCompletionError>,
328 >,
329 LanguageModelCompletionError,
330 >,
331 > {
332 let request = crate::provider::open_ai::into_open_ai(
333 request,
334 self.model.id(),
335 self.model.supports_parallel_tool_calls(),
336 self.model.supports_prompt_cache_key(),
337 self.max_output_tokens(),
338 None,
339 false,
340 );
341 let completions = self.stream_completion(request, cx);
342 async move {
343 let mapper = crate::provider::open_ai::OpenAiEventMapper::new();
344 Ok(mapper.map_stream(completions.await?).boxed())
345 }
346 .boxed()
347 }
348}
349
350struct ConfigurationView {
351 api_key_editor: Entity<InputField>,
352 state: Entity<State>,
353 load_credentials_task: Option<Task<()>>,
354}
355
356impl ConfigurationView {
357 fn new(state: Entity<State>, window: &mut Window, cx: &mut Context<Self>) -> Self {
358 let api_key_editor = cx.new(|cx| {
359 InputField::new(
360 window,
361 cx,
362 "xai-0000000000000000000000000000000000000000000000000",
363 )
364 .label("API key")
365 });
366
367 cx.observe(&state, |_, _, cx| {
368 cx.notify();
369 })
370 .detach();
371
372 let load_credentials_task = Some(cx.spawn_in(window, {
373 let state = state.clone();
374 async move |this, cx| {
375 if let Some(task) = Some(state.update(cx, |state, cx| state.authenticate(cx))) {
376 // We don't log an error, because "not signed in" is also an error.
377 let _ = task.await;
378 }
379 this.update(cx, |this, cx| {
380 this.load_credentials_task = None;
381 cx.notify();
382 })
383 .log_err();
384 }
385 }));
386
387 Self {
388 api_key_editor,
389 state,
390 load_credentials_task,
391 }
392 }
393
394 fn save_api_key(&mut self, _: &menu::Confirm, window: &mut Window, cx: &mut Context<Self>) {
395 let api_key = self.api_key_editor.read(cx).text(cx).trim().to_string();
396 if api_key.is_empty() {
397 return;
398 }
399
400 // url changes can cause the editor to be displayed again
401 self.api_key_editor
402 .update(cx, |editor, cx| editor.set_text("", window, cx));
403
404 let state = self.state.clone();
405 cx.spawn_in(window, async move |_, cx| {
406 state
407 .update(cx, |state, cx| state.set_api_key(Some(api_key), cx))
408 .await
409 })
410 .detach_and_log_err(cx);
411 }
412
413 fn reset_api_key(&mut self, window: &mut Window, cx: &mut Context<Self>) {
414 self.api_key_editor
415 .update(cx, |input, cx| input.set_text("", window, cx));
416
417 let state = self.state.clone();
418 cx.spawn_in(window, async move |_, cx| {
419 state
420 .update(cx, |state, cx| state.set_api_key(None, cx))
421 .await
422 })
423 .detach_and_log_err(cx);
424 }
425
426 fn should_render_editor(&self, cx: &mut Context<Self>) -> bool {
427 !self.state.read(cx).is_authenticated()
428 }
429}
430
431impl Render for ConfigurationView {
432 fn render(&mut self, _: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
433 let env_var_set = self.state.read(cx).api_key_state.is_from_env_var();
434 let configured_card_label = if env_var_set {
435 format!("API key set in {API_KEY_ENV_VAR_NAME} environment variable")
436 } else {
437 let api_url = XAiLanguageModelProvider::api_url(cx);
438 if api_url == XAI_API_URL {
439 "API key configured".to_string()
440 } else {
441 format!("API key configured for {}", api_url)
442 }
443 };
444
445 let api_key_section = if self.should_render_editor(cx) {
446 v_flex()
447 .on_action(cx.listener(Self::save_api_key))
448 .child(Label::new("To use Zed's agent with xAI, you need to add an API key. Follow these steps:"))
449 .child(
450 List::new()
451 .child(
452 ListBulletItem::new("")
453 .child(Label::new("Create one by visiting"))
454 .child(ButtonLink::new("xAI console", "https://console.x.ai/team/default/api-keys"))
455 )
456 .child(
457 ListBulletItem::new("Paste your API key below and hit enter to start using the agent")
458 ),
459 )
460 .child(self.api_key_editor.clone())
461 .child(
462 Label::new(format!(
463 "You can also set the {API_KEY_ENV_VAR_NAME} environment variable and restart Zed."
464 ))
465 .size(LabelSize::Small)
466 .color(Color::Muted),
467 )
468 .child(
469 Label::new("Note that xAI is a custom OpenAI-compatible provider.")
470 .size(LabelSize::Small)
471 .color(Color::Muted),
472 )
473 .into_any_element()
474 } else {
475 ConfiguredApiCard::new(configured_card_label)
476 .disabled(env_var_set)
477 .when(env_var_set, |this| {
478 this.tooltip_label(format!("To reset your API key, unset the {API_KEY_ENV_VAR_NAME} environment variable."))
479 })
480 .on_click(cx.listener(|this, _, window, cx| this.reset_api_key(window, cx)))
481 .into_any_element()
482 };
483
484 if self.load_credentials_task.is_some() {
485 div().child(Label::new("Loading credentials…")).into_any()
486 } else {
487 v_flex().size_full().child(api_key_section).into_any()
488 }
489 }
490}