1use anyhow::{Result, anyhow};
2use futures::{FutureExt, StreamExt, future::BoxFuture, stream::BoxStream};
3use gpui::{AnyView, App, AsyncApp, Context, Subscription, Task};
4use http_client::HttpClient;
5use language_model::{
6 AuthenticateError, LanguageModelCompletionError, LanguageModelCompletionEvent,
7};
8use language_model::{
9 LanguageModel, LanguageModelId, LanguageModelName, LanguageModelProvider,
10 LanguageModelProviderId, LanguageModelProviderName, LanguageModelProviderState,
11 LanguageModelRequest, RateLimiter, Role,
12};
13use lmstudio::{
14 ChatCompletionRequest, ChatMessage, ModelType, get_models, preload_model,
15 stream_chat_completion,
16};
17use schemars::JsonSchema;
18use serde::{Deserialize, Serialize};
19use settings::{Settings, SettingsStore};
20use std::{collections::BTreeMap, sync::Arc};
21use ui::{ButtonLike, Indicator, List, prelude::*};
22use util::ResultExt;
23
24use crate::AllLanguageModelSettings;
25use crate::ui::InstructionListItem;
26
27const LMSTUDIO_DOWNLOAD_URL: &str = "https://lmstudio.ai/download";
28const LMSTUDIO_CATALOG_URL: &str = "https://lmstudio.ai/models";
29const LMSTUDIO_SITE: &str = "https://lmstudio.ai/";
30
31const PROVIDER_ID: &str = "lmstudio";
32const PROVIDER_NAME: &str = "LM Studio";
33
34#[derive(Default, Debug, Clone, PartialEq)]
35pub struct LmStudioSettings {
36 pub api_url: String,
37 pub available_models: Vec<AvailableModel>,
38}
39
40#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)]
41pub struct AvailableModel {
42 /// The model name in the LM Studio API. e.g. qwen2.5-coder-7b, phi-4, etc
43 pub name: String,
44 /// The model's name in Zed's UI, such as in the model selector dropdown menu in the assistant panel.
45 pub display_name: Option<String>,
46 /// The model's context window size.
47 pub max_tokens: usize,
48}
49
50pub struct LmStudioLanguageModelProvider {
51 http_client: Arc<dyn HttpClient>,
52 state: gpui::Entity<State>,
53}
54
55pub struct State {
56 http_client: Arc<dyn HttpClient>,
57 available_models: Vec<lmstudio::Model>,
58 fetch_model_task: Option<Task<Result<()>>>,
59 _subscription: Subscription,
60}
61
62impl State {
63 fn is_authenticated(&self) -> bool {
64 !self.available_models.is_empty()
65 }
66
67 fn fetch_models(&mut self, cx: &mut Context<Self>) -> Task<Result<()>> {
68 let settings = &AllLanguageModelSettings::get_global(cx).lmstudio;
69 let http_client = self.http_client.clone();
70 let api_url = settings.api_url.clone();
71
72 // As a proxy for the server being "authenticated", we'll check if its up by fetching the models
73 cx.spawn(async move |this, cx| {
74 let models = get_models(http_client.as_ref(), &api_url, None).await?;
75
76 let mut models: Vec<lmstudio::Model> = models
77 .into_iter()
78 .filter(|model| model.r#type != ModelType::Embeddings)
79 .map(|model| lmstudio::Model::new(&model.id, None, None))
80 .collect();
81
82 models.sort_by(|a, b| a.name.cmp(&b.name));
83
84 this.update(cx, |this, cx| {
85 this.available_models = models;
86 cx.notify();
87 })
88 })
89 }
90
91 fn restart_fetch_models_task(&mut self, cx: &mut Context<Self>) {
92 let task = self.fetch_models(cx);
93 self.fetch_model_task.replace(task);
94 }
95
96 fn authenticate(&mut self, cx: &mut Context<Self>) -> Task<Result<(), AuthenticateError>> {
97 if self.is_authenticated() {
98 return Task::ready(Ok(()));
99 }
100
101 let fetch_models_task = self.fetch_models(cx);
102 cx.spawn(async move |_this, _cx| Ok(fetch_models_task.await?))
103 }
104}
105
106impl LmStudioLanguageModelProvider {
107 pub fn new(http_client: Arc<dyn HttpClient>, cx: &mut App) -> Self {
108 let this = Self {
109 http_client: http_client.clone(),
110 state: cx.new(|cx| {
111 let subscription = cx.observe_global::<SettingsStore>({
112 let mut settings = AllLanguageModelSettings::get_global(cx).lmstudio.clone();
113 move |this: &mut State, cx| {
114 let new_settings = &AllLanguageModelSettings::get_global(cx).lmstudio;
115 if &settings != new_settings {
116 settings = new_settings.clone();
117 this.restart_fetch_models_task(cx);
118 cx.notify();
119 }
120 }
121 });
122
123 State {
124 http_client,
125 available_models: Default::default(),
126 fetch_model_task: None,
127 _subscription: subscription,
128 }
129 }),
130 };
131 this.state
132 .update(cx, |state, cx| state.restart_fetch_models_task(cx));
133 this
134 }
135}
136
137impl LanguageModelProviderState for LmStudioLanguageModelProvider {
138 type ObservableEntity = State;
139
140 fn observable_entity(&self) -> Option<gpui::Entity<Self::ObservableEntity>> {
141 Some(self.state.clone())
142 }
143}
144
145impl LanguageModelProvider for LmStudioLanguageModelProvider {
146 fn id(&self) -> LanguageModelProviderId {
147 LanguageModelProviderId(PROVIDER_ID.into())
148 }
149
150 fn name(&self) -> LanguageModelProviderName {
151 LanguageModelProviderName(PROVIDER_NAME.into())
152 }
153
154 fn icon(&self) -> IconName {
155 IconName::AiLmStudio
156 }
157
158 fn default_model(&self, cx: &App) -> Option<Arc<dyn LanguageModel>> {
159 self.provided_models(cx).into_iter().next()
160 }
161
162 fn default_fast_model(&self, cx: &App) -> Option<Arc<dyn LanguageModel>> {
163 self.default_model(cx)
164 }
165
166 fn provided_models(&self, cx: &App) -> Vec<Arc<dyn LanguageModel>> {
167 let mut models: BTreeMap<String, lmstudio::Model> = BTreeMap::default();
168
169 // Add models from the LM Studio API
170 for model in self.state.read(cx).available_models.iter() {
171 models.insert(model.name.clone(), model.clone());
172 }
173
174 // Override with available models from settings
175 for model in AllLanguageModelSettings::get_global(cx)
176 .lmstudio
177 .available_models
178 .iter()
179 {
180 models.insert(
181 model.name.clone(),
182 lmstudio::Model {
183 name: model.name.clone(),
184 display_name: model.display_name.clone(),
185 max_tokens: model.max_tokens,
186 },
187 );
188 }
189
190 models
191 .into_values()
192 .map(|model| {
193 Arc::new(LmStudioLanguageModel {
194 id: LanguageModelId::from(model.name.clone()),
195 model: model.clone(),
196 http_client: self.http_client.clone(),
197 request_limiter: RateLimiter::new(4),
198 }) as Arc<dyn LanguageModel>
199 })
200 .collect()
201 }
202
203 fn load_model(&self, model: Arc<dyn LanguageModel>, cx: &App) {
204 let settings = &AllLanguageModelSettings::get_global(cx).lmstudio;
205 let http_client = self.http_client.clone();
206 let api_url = settings.api_url.clone();
207 let id = model.id().0.to_string();
208 cx.spawn(async move |_| preload_model(http_client, &api_url, &id).await)
209 .detach_and_log_err(cx);
210 }
211
212 fn is_authenticated(&self, cx: &App) -> bool {
213 self.state.read(cx).is_authenticated()
214 }
215
216 fn authenticate(&self, cx: &mut App) -> Task<Result<(), AuthenticateError>> {
217 self.state.update(cx, |state, cx| state.authenticate(cx))
218 }
219
220 fn configuration_view(&self, _window: &mut Window, cx: &mut App) -> AnyView {
221 let state = self.state.clone();
222 cx.new(|cx| ConfigurationView::new(state, cx)).into()
223 }
224
225 fn reset_credentials(&self, cx: &mut App) -> Task<Result<()>> {
226 self.state.update(cx, |state, cx| state.fetch_models(cx))
227 }
228}
229
230pub struct LmStudioLanguageModel {
231 id: LanguageModelId,
232 model: lmstudio::Model,
233 http_client: Arc<dyn HttpClient>,
234 request_limiter: RateLimiter,
235}
236
237impl LmStudioLanguageModel {
238 fn to_lmstudio_request(&self, request: LanguageModelRequest) -> ChatCompletionRequest {
239 ChatCompletionRequest {
240 model: self.model.name.clone(),
241 messages: request
242 .messages
243 .into_iter()
244 .map(|msg| match msg.role {
245 Role::User => ChatMessage::User {
246 content: msg.string_contents(),
247 },
248 Role::Assistant => ChatMessage::Assistant {
249 content: Some(msg.string_contents()),
250 tool_calls: None,
251 },
252 Role::System => ChatMessage::System {
253 content: msg.string_contents(),
254 },
255 })
256 .collect(),
257 stream: true,
258 max_tokens: Some(-1),
259 stop: Some(request.stop),
260 temperature: request.temperature.or(Some(0.0)),
261 tools: vec![],
262 }
263 }
264}
265
266impl LanguageModel for LmStudioLanguageModel {
267 fn id(&self) -> LanguageModelId {
268 self.id.clone()
269 }
270
271 fn name(&self) -> LanguageModelName {
272 LanguageModelName::from(self.model.display_name().to_string())
273 }
274
275 fn provider_id(&self) -> LanguageModelProviderId {
276 LanguageModelProviderId(PROVIDER_ID.into())
277 }
278
279 fn provider_name(&self) -> LanguageModelProviderName {
280 LanguageModelProviderName(PROVIDER_NAME.into())
281 }
282
283 fn supports_tools(&self) -> bool {
284 false
285 }
286
287 fn telemetry_id(&self) -> String {
288 format!("lmstudio/{}", self.model.id())
289 }
290
291 fn max_token_count(&self) -> usize {
292 self.model.max_token_count()
293 }
294
295 fn count_tokens(
296 &self,
297 request: LanguageModelRequest,
298 _cx: &App,
299 ) -> BoxFuture<'static, Result<usize>> {
300 // Endpoint for this is coming soon. In the meantime, hacky estimation
301 let token_count = request
302 .messages
303 .iter()
304 .map(|msg| msg.string_contents().split_whitespace().count())
305 .sum::<usize>();
306
307 let estimated_tokens = (token_count as f64 * 0.75) as usize;
308 async move { Ok(estimated_tokens) }.boxed()
309 }
310
311 fn stream_completion(
312 &self,
313 request: LanguageModelRequest,
314 cx: &AsyncApp,
315 ) -> BoxFuture<
316 'static,
317 Result<
318 BoxStream<'static, Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>,
319 >,
320 > {
321 let request = self.to_lmstudio_request(request);
322
323 let http_client = self.http_client.clone();
324 let Ok(api_url) = cx.update(|cx| {
325 let settings = &AllLanguageModelSettings::get_global(cx).lmstudio;
326 settings.api_url.clone()
327 }) else {
328 return futures::future::ready(Err(anyhow!("App state dropped"))).boxed();
329 };
330
331 let future = self.request_limiter.stream(async move {
332 let response = stream_chat_completion(http_client.as_ref(), &api_url, request).await?;
333 let stream = response
334 .filter_map(|response| async move {
335 match response {
336 Ok(fragment) => {
337 // Skip empty deltas
338 if fragment.choices[0].delta.is_object()
339 && fragment.choices[0].delta.as_object().unwrap().is_empty()
340 {
341 return None;
342 }
343
344 // Try to parse the delta as ChatMessage
345 if let Ok(chat_message) = serde_json::from_value::<ChatMessage>(
346 fragment.choices[0].delta.clone(),
347 ) {
348 let content = match chat_message {
349 ChatMessage::User { content } => content,
350 ChatMessage::Assistant { content, .. } => {
351 content.unwrap_or_default()
352 }
353 ChatMessage::System { content } => content,
354 };
355 if !content.is_empty() {
356 Some(Ok(content))
357 } else {
358 None
359 }
360 } else {
361 None
362 }
363 }
364 Err(error) => Some(Err(error)),
365 }
366 })
367 .boxed();
368 Ok(stream)
369 });
370
371 async move {
372 Ok(future
373 .await?
374 .map(|result| {
375 result
376 .map(LanguageModelCompletionEvent::Text)
377 .map_err(LanguageModelCompletionError::Other)
378 })
379 .boxed())
380 }
381 .boxed()
382 }
383}
384
385struct ConfigurationView {
386 state: gpui::Entity<State>,
387 loading_models_task: Option<Task<()>>,
388}
389
390impl ConfigurationView {
391 pub fn new(state: gpui::Entity<State>, cx: &mut Context<Self>) -> Self {
392 let loading_models_task = Some(cx.spawn({
393 let state = state.clone();
394 async move |this, cx| {
395 if let Some(task) = state
396 .update(cx, |state, cx| state.authenticate(cx))
397 .log_err()
398 {
399 task.await.log_err();
400 }
401 this.update(cx, |this, cx| {
402 this.loading_models_task = None;
403 cx.notify();
404 })
405 .log_err();
406 }
407 }));
408
409 Self {
410 state,
411 loading_models_task,
412 }
413 }
414
415 fn retry_connection(&self, cx: &mut App) {
416 self.state
417 .update(cx, |state, cx| state.fetch_models(cx))
418 .detach_and_log_err(cx);
419 }
420}
421
422impl Render for ConfigurationView {
423 fn render(&mut self, _window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
424 let is_authenticated = self.state.read(cx).is_authenticated();
425
426 let lmstudio_intro = "Run local LLMs like Llama, Phi, and Qwen.";
427
428 if self.loading_models_task.is_some() {
429 div().child(Label::new("Loading models...")).into_any()
430 } else {
431 v_flex()
432 .gap_2()
433 .child(
434 v_flex().gap_1().child(Label::new(lmstudio_intro)).child(
435 List::new()
436 .child(InstructionListItem::text_only(
437 "LM Studio needs to be running with at least one model downloaded.",
438 ))
439 .child(InstructionListItem::text_only(
440 "To get your first model, try running `lms get qwen2.5-coder-7b`",
441 )),
442 ),
443 )
444 .child(
445 h_flex()
446 .w_full()
447 .justify_between()
448 .gap_2()
449 .child(
450 h_flex()
451 .w_full()
452 .gap_2()
453 .map(|this| {
454 if is_authenticated {
455 this.child(
456 Button::new("lmstudio-site", "LM Studio")
457 .style(ButtonStyle::Subtle)
458 .icon(IconName::ArrowUpRight)
459 .icon_size(IconSize::XSmall)
460 .icon_color(Color::Muted)
461 .on_click(move |_, _window, cx| {
462 cx.open_url(LMSTUDIO_SITE)
463 })
464 .into_any_element(),
465 )
466 } else {
467 this.child(
468 Button::new(
469 "download_lmstudio_button",
470 "Download LM Studio",
471 )
472 .style(ButtonStyle::Subtle)
473 .icon(IconName::ArrowUpRight)
474 .icon_size(IconSize::XSmall)
475 .icon_color(Color::Muted)
476 .on_click(move |_, _window, cx| {
477 cx.open_url(LMSTUDIO_DOWNLOAD_URL)
478 })
479 .into_any_element(),
480 )
481 }
482 })
483 .child(
484 Button::new("view-models", "Model Catalog")
485 .style(ButtonStyle::Subtle)
486 .icon(IconName::ArrowUpRight)
487 .icon_size(IconSize::XSmall)
488 .icon_color(Color::Muted)
489 .on_click(move |_, _window, cx| {
490 cx.open_url(LMSTUDIO_CATALOG_URL)
491 }),
492 ),
493 )
494 .map(|this| {
495 if is_authenticated {
496 this.child(
497 ButtonLike::new("connected")
498 .disabled(true)
499 .cursor_style(gpui::CursorStyle::Arrow)
500 .child(
501 h_flex()
502 .gap_2()
503 .child(Indicator::dot().color(Color::Success))
504 .child(Label::new("Connected"))
505 .into_any_element(),
506 ),
507 )
508 } else {
509 this.child(
510 Button::new("retry_lmstudio_models", "Connect")
511 .icon_position(IconPosition::Start)
512 .icon_size(IconSize::XSmall)
513 .icon(IconName::Play)
514 .on_click(cx.listener(move |this, _, _window, cx| {
515 this.retry_connection(cx)
516 })),
517 )
518 }
519 }),
520 )
521 .into_any()
522 }
523 }
524}