Skip to repository content1177 lines · 44.5 KB · rust
tenant.openagents/omega
No repository description is available.
OpenAgents Git authority 2026-07-28T02:36:18.359Z Public web read
NIP-34 coordinate
30617:7649603503856e5148d571eac2766b288a8ff1e9e35d380337a1d2b0015b4f92:omegaMaintainersHidden in public view
References2 branches · 1 tag
Read-only clone
git clone https://openagents.com/git/tenant.openagents/omega.gitBrowse files
lmstudio.rs
1use anyhow::{Result, anyhow};
2use credentials_provider::CredentialsProvider;
3use fs::Fs;
4use futures::Stream;
5use futures::{FutureExt, StreamExt, future::BoxFuture, stream::BoxStream};
6use gpui::{App, AsyncApp, Context, Entity, Subscription, Task, TaskExt};
7use http_client::{CustomHeaders, HttpClient};
8use language_model::{
9 ApiKeyState, AuthenticateError, EnvVar, IconOrSvg, LanguageModel, LanguageModelCompletionError,
10 LanguageModelCompletionEvent, LanguageModelToolChoice, LanguageModelToolResultContent,
11 LanguageModelToolUse, MessageContent, StopReason, TokenUsage, env_var,
12};
13use language_model::{
14 InlineDescription, LanguageModelId, LanguageModelName, LanguageModelProvider,
15 LanguageModelProviderId, LanguageModelProviderName, LanguageModelProviderState,
16 LanguageModelRequest, ProviderSettingsView, RateLimiter, Role, SubPageProviderSettings,
17};
18use lmstudio::{LMSTUDIO_API_URL, ModelType, get_models};
19
20pub use settings::LmStudioAvailableModel as AvailableModel;
21use settings::{Settings, SettingsStore, update_settings_file};
22use std::pin::Pin;
23use std::sync::LazyLock;
24use std::{
25 collections::{BTreeMap, HashMap},
26 sync::Arc,
27};
28use ui::{ButtonLike, ConfiguredApiCard, Divider, List, ListBulletItem, Tooltip, prelude::*};
29use ui_input::InputField;
30
31use crate::AllLanguageModelSettings;
32use language_model::util::parse_tool_arguments;
33
34const LMSTUDIO_DOWNLOAD_URL: &str = "https://lmstudio.ai/download";
35const LMSTUDIO_CATALOG_URL: &str = "https://lmstudio.ai/models";
36const LMSTUDIO_SITE: &str = "https://lmstudio.ai/";
37
38const PROVIDER_ID: LanguageModelProviderId = LanguageModelProviderId::new("lmstudio");
39const PROVIDER_NAME: LanguageModelProviderName = LanguageModelProviderName::new("LM Studio");
40
41const API_KEY_ENV_VAR_NAME: &str = "LMSTUDIO_API_KEY";
42static API_KEY_ENV_VAR: LazyLock<EnvVar> = env_var!(API_KEY_ENV_VAR_NAME);
43
44#[derive(Default, Debug, Clone, PartialEq)]
45pub struct LmStudioSettings {
46 pub api_url: String,
47 pub available_models: Vec<AvailableModel>,
48 pub custom_headers: CustomHeaders,
49}
50
51pub struct LmStudioLanguageModelProvider {
52 http_client: Arc<dyn HttpClient>,
53 state: Entity<State>,
54}
55
56pub struct State {
57 api_key_state: ApiKeyState,
58 credentials_provider: Arc<dyn CredentialsProvider>,
59 http_client: Arc<dyn HttpClient>,
60 available_models: Vec<lmstudio::Model>,
61 fetch_model_task: Option<Task<Result<()>>>,
62 _subscription: Subscription,
63}
64
65impl State {
66 fn is_authenticated(&self) -> bool {
67 !self.available_models.is_empty()
68 }
69
70 fn set_api_key(&mut self, api_key: Option<String>, cx: &mut Context<Self>) -> Task<Result<()>> {
71 let credentials_provider = self.credentials_provider.clone();
72 let api_url = LmStudioLanguageModelProvider::api_url(cx).into();
73 let task = self.api_key_state.store(
74 api_url,
75 api_key,
76 |this| &mut this.api_key_state,
77 credentials_provider,
78 cx,
79 );
80 self.restart_fetch_models_task(cx);
81 task
82 }
83
84 fn fetch_models(&mut self, cx: &mut Context<Self>) -> Task<Result<()>> {
85 let settings = &AllLanguageModelSettings::get_global(cx).lmstudio;
86 let http_client = self.http_client.clone();
87 let api_url = settings.api_url.clone();
88 let api_key = self.api_key_state.key(&api_url);
89 let extra_headers = settings.custom_headers.clone();
90
91 // As a proxy for the server being "authenticated", we'll check if its up by fetching the models
92 cx.spawn(async move |this, cx| {
93 let models = get_models(
94 http_client.as_ref(),
95 &api_url,
96 api_key.as_deref(),
97 None,
98 &extra_headers,
99 )
100 .await?;
101
102 let mut models: Vec<lmstudio::Model> = models
103 .into_iter()
104 .filter(|model| model.r#type != ModelType::Embeddings)
105 .map(|model| {
106 lmstudio::Model::new(
107 &model.id,
108 None,
109 model
110 .loaded_context_length
111 .or_else(|| model.max_context_length),
112 model.capabilities.supports_tool_calls(),
113 model.capabilities.supports_images() || model.r#type == ModelType::Vlm,
114 )
115 })
116 .collect();
117
118 models.sort_by(|a, b| a.name.cmp(&b.name));
119
120 this.update(cx, |this, cx| {
121 this.available_models = models;
122 cx.notify();
123 })
124 })
125 }
126
127 fn restart_fetch_models_task(&mut self, cx: &mut Context<Self>) {
128 let task = self.fetch_models(cx);
129 self.fetch_model_task.replace(task);
130 }
131
132 fn authenticate(&mut self, cx: &mut Context<Self>) -> Task<Result<(), AuthenticateError>> {
133 let credentials_provider = self.credentials_provider.clone();
134 let api_url = LmStudioLanguageModelProvider::api_url(cx).into();
135 let _task = self.api_key_state.load_if_needed(
136 api_url,
137 |this| &mut this.api_key_state,
138 credentials_provider,
139 cx,
140 );
141
142 if self.is_authenticated() {
143 return Task::ready(Ok(()));
144 }
145
146 let fetch_models_task = self.fetch_models(cx);
147 cx.spawn(async move |_this, _cx| {
148 match fetch_models_task.await {
149 Ok(()) => Ok(()),
150 Err(err) => {
151 // If any cause in the error chain is an std::io::Error with
152 // ErrorKind::ConnectionRefused, treat this as "credentials not found"
153 // (i.e. LM Studio not running).
154 let mut connection_refused = false;
155 for cause in err.chain() {
156 if let Some(io_err) = cause.downcast_ref::<std::io::Error>() {
157 if io_err.kind() == std::io::ErrorKind::ConnectionRefused {
158 connection_refused = true;
159 break;
160 }
161 }
162 }
163 if connection_refused {
164 Err(AuthenticateError::ConnectionRefused)
165 } else {
166 Err(AuthenticateError::Other(err))
167 }
168 }
169 }
170 })
171 }
172}
173
174impl LmStudioLanguageModelProvider {
175 pub fn new(
176 http_client: Arc<dyn HttpClient>,
177 credentials_provider: Arc<dyn CredentialsProvider>,
178 cx: &mut App,
179 ) -> Self {
180 let this = Self {
181 http_client: http_client.clone(),
182 state: cx.new(|cx| {
183 let subscription = cx.observe_global::<SettingsStore>({
184 let mut settings = AllLanguageModelSettings::get_global(cx).lmstudio.clone();
185 move |this: &mut State, cx| {
186 let new_settings =
187 AllLanguageModelSettings::get_global(cx).lmstudio.clone();
188 if settings != new_settings {
189 let credentials_provider = this.credentials_provider.clone();
190 let api_url = Self::api_url(cx).into();
191 this.api_key_state.handle_url_change(
192 api_url,
193 |this| &mut this.api_key_state,
194 credentials_provider,
195 cx,
196 );
197 settings = new_settings;
198 this.restart_fetch_models_task(cx);
199 cx.notify();
200 }
201 }
202 });
203
204 State {
205 api_key_state: ApiKeyState::new(
206 Self::api_url(cx).into(),
207 (*API_KEY_ENV_VAR).clone(),
208 ),
209 credentials_provider,
210 http_client,
211 available_models: Default::default(),
212 fetch_model_task: None,
213 _subscription: subscription,
214 }
215 }),
216 };
217 this.state
218 .update(cx, |state, cx| state.restart_fetch_models_task(cx));
219 this
220 }
221
222 fn api_url(cx: &App) -> String {
223 AllLanguageModelSettings::get_global(cx)
224 .lmstudio
225 .api_url
226 .clone()
227 }
228
229 fn has_custom_url(cx: &App) -> bool {
230 Self::api_url(cx) != LMSTUDIO_API_URL
231 }
232}
233
234impl LanguageModelProviderState for LmStudioLanguageModelProvider {
235 type ObservableEntity = State;
236
237 fn observable_entity(&self) -> Option<Entity<Self::ObservableEntity>> {
238 Some(self.state.clone())
239 }
240}
241
242impl LanguageModelProvider for LmStudioLanguageModelProvider {
243 fn id(&self) -> LanguageModelProviderId {
244 PROVIDER_ID
245 }
246
247 fn name(&self) -> LanguageModelProviderName {
248 PROVIDER_NAME
249 }
250
251 fn icon(&self) -> IconOrSvg {
252 IconOrSvg::Icon(IconName::AiLmStudio)
253 }
254
255 fn default_model(&self, _: &App) -> Option<Arc<dyn LanguageModel>> {
256 // We shouldn't try to select default model, because it might lead to a load call for an unloaded model.
257 // In a constrained environment where user might not have enough resources it'll be a bad UX to select something
258 // to load by default.
259 None
260 }
261
262 fn default_fast_model(&self, _: &App) -> Option<Arc<dyn LanguageModel>> {
263 // See explanation for default_model.
264 None
265 }
266
267 fn provided_models(&self, cx: &App) -> Vec<Arc<dyn LanguageModel>> {
268 let mut models: BTreeMap<String, lmstudio::Model> = BTreeMap::default();
269
270 // Add models from the LM Studio API
271 for model in self.state.read(cx).available_models.iter() {
272 models.insert(model.name.clone(), model.clone());
273 }
274
275 // Override with available models from settings
276 for model in AllLanguageModelSettings::get_global(cx)
277 .lmstudio
278 .available_models
279 .iter()
280 {
281 models.insert(
282 model.name.clone(),
283 lmstudio::Model {
284 name: model.name.clone(),
285 display_name: model.display_name.clone(),
286 max_tokens: model.max_tokens,
287 supports_tool_calls: model.supports_tool_calls,
288 supports_images: model.supports_images,
289 },
290 );
291 }
292
293 models
294 .into_values()
295 .map(|model| {
296 Arc::new(LmStudioLanguageModel {
297 id: LanguageModelId::from(model.name.clone()),
298 model,
299 http_client: self.http_client.clone(),
300 request_limiter: RateLimiter::new(4),
301 state: self.state.clone(),
302 }) as Arc<dyn LanguageModel>
303 })
304 .collect()
305 }
306
307 fn is_authenticated(&self, cx: &App) -> bool {
308 self.state.read(cx).is_authenticated()
309 }
310
311 fn authenticate(&self, cx: &mut App) -> Task<Result<(), AuthenticateError>> {
312 self.state.update(cx, |state, cx| state.authenticate(cx))
313 }
314
315 fn settings_view(&self, _cx: &mut App) -> Option<ProviderSettingsView> {
316 let state = self.state.clone();
317 Some(ProviderSettingsView::SubPage(
318 SubPageProviderSettings::new(move |window, cx| {
319 cx.new(|cx| ConfigurationView::new(state.clone(), window, cx))
320 .into()
321 })
322 .description(InlineDescription::Text(
323 "Run local LLMs like Llama, Phi, and Qwen with LM Studio.".into(),
324 )),
325 ))
326 }
327}
328
329pub struct LmStudioLanguageModel {
330 id: LanguageModelId,
331 model: lmstudio::Model,
332 http_client: Arc<dyn HttpClient>,
333 request_limiter: RateLimiter,
334 state: Entity<State>,
335}
336
337impl LmStudioLanguageModel {
338 fn to_lmstudio_request(
339 &self,
340 request: LanguageModelRequest,
341 ) -> Result<lmstudio::ChatCompletionRequest> {
342 if request.contains_custom_tool_input() {
343 anyhow::bail!("LM Studio does not support custom tools");
344 }
345
346 let mut messages = Vec::new();
347
348 for message in request.messages {
349 for content in message.content {
350 match content {
351 MessageContent::Text(text) => add_message_content_part(
352 lmstudio::MessagePart::Text { text },
353 message.role,
354 &mut messages,
355 ),
356 MessageContent::Thinking { .. } => {}
357 MessageContent::RedactedThinking(_) => {}
358 MessageContent::Compaction(_) => {}
359 MessageContent::Image(image) => {
360 add_message_content_part(
361 lmstudio::MessagePart::Image {
362 image_url: lmstudio::ImageUrl {
363 url: image.to_base64_url(),
364 detail: None,
365 },
366 },
367 message.role,
368 &mut messages,
369 );
370 }
371 MessageContent::ToolUse(tool_use) => {
372 let input = tool_use.input.as_json().ok_or_else(|| {
373 anyhow!("LM Studio does not support custom tool calls")
374 })?;
375 let tool_call = lmstudio::ToolCall {
376 id: tool_use.id.to_string(),
377 content: lmstudio::ToolCallContent::Function {
378 function: lmstudio::FunctionContent {
379 name: tool_use.name.to_string(),
380 arguments: serde_json::to_string(input).unwrap_or_default(),
381 },
382 },
383 };
384
385 if let Some(lmstudio::ChatMessage::Assistant { tool_calls, .. }) =
386 messages.last_mut()
387 {
388 tool_calls.push(tool_call);
389 } else {
390 messages.push(lmstudio::ChatMessage::Assistant {
391 content: None,
392 tool_calls: vec![tool_call],
393 });
394 }
395 }
396 MessageContent::ToolResult(tool_result) => {
397 let content: Vec<lmstudio::MessagePart> = tool_result
398 .content
399 .iter()
400 .map(|part| match part {
401 LanguageModelToolResultContent::Text(text) => {
402 lmstudio::MessagePart::Text {
403 text: text.to_string(),
404 }
405 }
406 LanguageModelToolResultContent::Image(image) => {
407 lmstudio::MessagePart::Image {
408 image_url: lmstudio::ImageUrl {
409 url: image.to_base64_url(),
410 detail: None,
411 },
412 }
413 }
414 })
415 .collect();
416
417 messages.push(lmstudio::ChatMessage::Tool {
418 content: content.into(),
419 tool_call_id: tool_result.tool_use_id.to_string(),
420 });
421 }
422 }
423 }
424 }
425
426 Ok(lmstudio::ChatCompletionRequest {
427 model: self.model.name.clone(),
428 messages,
429 stream: true,
430 stream_options: Some(lmstudio::StreamOptions {
431 include_usage: true,
432 }),
433 max_tokens: Some(-1),
434 stop: Some(request.stop),
435 // In LM Studio you can configure specific settings you'd like to use for your model.
436 // For example Qwen3 is recommended to be used with 0.7 temperature.
437 // It would be a bad UX to silently override these settings from Zed, so we pass no temperature as a default.
438 temperature: request.temperature.or(None),
439 tools: request
440 .tools
441 .into_iter()
442 .map(|tool| {
443 let input_schema = match tool.input {
444 language_model::LanguageModelRequestToolInput::Function {
445 input_schema,
446 ..
447 } => input_schema,
448 language_model::LanguageModelRequestToolInput::Custom { .. } => {
449 return Err(anyhow::anyhow!("LM Studio does not support custom tools"));
450 }
451 };
452 Ok(lmstudio::ToolDefinition::Function {
453 function: lmstudio::FunctionDefinition {
454 name: tool.name,
455 description: Some(tool.description),
456 parameters: Some(input_schema),
457 },
458 })
459 })
460 .collect::<Result<_>>()?,
461 tool_choice: request.tool_choice.map(|choice| match choice {
462 LanguageModelToolChoice::Auto => lmstudio::ToolChoice::Auto,
463 LanguageModelToolChoice::Any => lmstudio::ToolChoice::Required,
464 LanguageModelToolChoice::None => lmstudio::ToolChoice::None,
465 }),
466 })
467 }
468
469 fn stream_completion(
470 &self,
471 request: lmstudio::ChatCompletionRequest,
472 cx: &AsyncApp,
473 ) -> BoxFuture<
474 'static,
475 Result<futures::stream::BoxStream<'static, Result<lmstudio::ResponseStreamEvent>>>,
476 > {
477 let http_client = self.http_client.clone();
478 let (api_key, api_url, extra_headers) = self.state.read_with(cx, |state, cx| {
479 let api_url = LmStudioLanguageModelProvider::api_url(cx);
480 let extra_headers = AllLanguageModelSettings::get_global(cx)
481 .lmstudio
482 .custom_headers
483 .clone();
484 (state.api_key_state.key(&api_url), api_url, extra_headers)
485 });
486
487 let future = self.request_limiter.stream(async move {
488 let stream = lmstudio::stream_chat_completion(
489 http_client.as_ref(),
490 &api_url,
491 api_key.as_deref(),
492 request,
493 &extra_headers,
494 )
495 .await?;
496 Ok(stream)
497 });
498
499 async move { Ok(future.await?.boxed()) }.boxed()
500 }
501}
502
503impl LanguageModel for LmStudioLanguageModel {
504 fn id(&self) -> LanguageModelId {
505 self.id.clone()
506 }
507
508 fn name(&self) -> LanguageModelName {
509 LanguageModelName::from(self.model.display_name().to_string())
510 }
511
512 fn provider_id(&self) -> LanguageModelProviderId {
513 PROVIDER_ID
514 }
515
516 fn provider_name(&self) -> LanguageModelProviderName {
517 PROVIDER_NAME
518 }
519
520 fn supports_tools(&self) -> bool {
521 self.model.supports_tool_calls()
522 }
523
524 fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool {
525 self.supports_tools()
526 && match choice {
527 LanguageModelToolChoice::Auto => true,
528 LanguageModelToolChoice::Any => true,
529 LanguageModelToolChoice::None => true,
530 }
531 }
532
533 fn supports_images(&self) -> bool {
534 self.model.supports_images
535 }
536
537 fn telemetry_id(&self) -> String {
538 format!("lmstudio/{}", self.model.id())
539 }
540
541 fn max_token_count(&self) -> u64 {
542 self.model.max_token_count()
543 }
544
545 fn stream_completion(
546 &self,
547 request: LanguageModelRequest,
548 cx: &AsyncApp,
549 ) -> BoxFuture<
550 'static,
551 Result<
552 BoxStream<'static, Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>,
553 LanguageModelCompletionError,
554 >,
555 > {
556 let request = match self.to_lmstudio_request(request) {
557 Ok(request) => request,
558 Err(error) => return async move { Err(error.into()) }.boxed(),
559 };
560 let completions = self.stream_completion(request, cx);
561 async move {
562 let mapper = LmStudioEventMapper::new();
563 Ok(mapper.map_stream(completions.await?).boxed())
564 }
565 .boxed()
566 }
567}
568
569struct LmStudioEventMapper {
570 tool_calls_by_index: HashMap<usize, RawToolCall>,
571}
572
573impl LmStudioEventMapper {
574 fn new() -> Self {
575 Self {
576 tool_calls_by_index: HashMap::default(),
577 }
578 }
579
580 pub fn map_stream(
581 mut self,
582 events: Pin<Box<dyn Send + Stream<Item = Result<lmstudio::ResponseStreamEvent>>>>,
583 ) -> impl Stream<Item = Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>
584 {
585 events.flat_map(move |event| {
586 futures::stream::iter(match event {
587 Ok(event) => self.map_event(event),
588 Err(error) => vec![Err(LanguageModelCompletionError::from(error))],
589 })
590 })
591 }
592
593 pub fn map_event(
594 &mut self,
595 event: lmstudio::ResponseStreamEvent,
596 ) -> Vec<Result<LanguageModelCompletionEvent, LanguageModelCompletionError>> {
597 let mut events = Vec::new();
598
599 if let Some(usage) = event.usage {
600 events.push(Ok(LanguageModelCompletionEvent::UsageUpdate(TokenUsage {
601 input_tokens: usage.prompt_tokens,
602 output_tokens: usage.completion_tokens,
603 cache_creation_input_tokens: 0,
604 cache_read_input_tokens: 0,
605 })));
606 }
607
608 // The final usage summary chunk from OpenAI-compatible servers has an empty choices array.
609 // Return accumulated events instead of treating it as an error.
610 let Some(choice) = event.choices.into_iter().next() else {
611 return events;
612 };
613
614 if let Some(content) = choice.delta.content {
615 events.push(Ok(LanguageModelCompletionEvent::Text(content)));
616 }
617
618 if let Some(reasoning_content) = choice.delta.reasoning_content {
619 events.push(Ok(LanguageModelCompletionEvent::Thinking {
620 text: reasoning_content,
621 signature: None,
622 }));
623 }
624
625 if let Some(tool_calls) = choice.delta.tool_calls {
626 for tool_call in tool_calls {
627 let entry = self.tool_calls_by_index.entry(tool_call.index).or_default();
628
629 if let Some(tool_id) = tool_call.id {
630 entry.id = tool_id;
631 }
632
633 if let Some(function) = tool_call.function {
634 if let Some(name) = function.name {
635 // At the time of writing this code LM Studio (0.3.15) is incompatible with the OpenAI API:
636 // 1. It sends function name in the first chunk
637 // 2. It sends empty string in the function name field in all subsequent chunks for arguments
638 // According to https://platform.openai.com/docs/guides/function-calling?api-mode=responses#streaming
639 // function name field should be sent only inside the first chunk.
640 if !name.is_empty() {
641 entry.name = name;
642 }
643 }
644
645 if let Some(arguments) = function.arguments {
646 entry.arguments.push_str(&arguments);
647 }
648 }
649 }
650 }
651
652 match choice.finish_reason.as_deref() {
653 Some("stop") => {
654 events.push(Ok(LanguageModelCompletionEvent::Stop(StopReason::EndTurn)));
655 }
656 Some("tool_calls") => {
657 events.extend(self.tool_calls_by_index.drain().map(|(_, tool_call)| {
658 match parse_tool_arguments(&tool_call.arguments) {
659 Ok(input) => Ok(LanguageModelCompletionEvent::ToolUse(
660 LanguageModelToolUse {
661 id: tool_call.id.into(),
662 name: tool_call.name.into(),
663 is_input_complete: true,
664 input: language_model::LanguageModelToolUseInput::Json(input),
665 raw_input: tool_call.arguments,
666 thought_signature: None,
667 },
668 )),
669 Err(error) => Ok(LanguageModelCompletionEvent::ToolUseJsonParseError {
670 id: tool_call.id.into(),
671 tool_name: tool_call.name.into(),
672 raw_input: tool_call.arguments.into(),
673 json_parse_error: error.to_string(),
674 }),
675 }
676 }));
677
678 events.push(Ok(LanguageModelCompletionEvent::Stop(StopReason::ToolUse)));
679 }
680 Some(stop_reason) => {
681 log::error!("Unexpected LMStudio stop_reason: {stop_reason:?}",);
682 events.push(Ok(LanguageModelCompletionEvent::Stop(StopReason::EndTurn)));
683 }
684 None => {}
685 }
686
687 events
688 }
689}
690
691#[derive(Default)]
692struct RawToolCall {
693 id: String,
694 name: String,
695 arguments: String,
696}
697
698#[cfg(test)]
699mod tests {
700 use super::*;
701 use lmstudio::{ChoiceDelta, ResponseMessageDelta, ResponseStreamEvent, Usage};
702
703 fn make_event(choices: Vec<ChoiceDelta>, usage: Option<Usage>) -> ResponseStreamEvent {
704 ResponseStreamEvent {
705 created: 0,
706 model: "test-model".to_string(),
707 object: "chat.completion.chunk".to_string(),
708 choices,
709 usage,
710 }
711 }
712
713 fn make_content_choice(content: &str) -> ChoiceDelta {
714 ChoiceDelta {
715 index: 0,
716 delta: ResponseMessageDelta {
717 role: None,
718 content: Some(content.to_string()),
719 reasoning_content: None,
720 tool_calls: None,
721 },
722 finish_reason: None,
723 }
724 }
725
726 fn make_stop_choice() -> ChoiceDelta {
727 ChoiceDelta {
728 index: 0,
729 delta: ResponseMessageDelta {
730 role: None,
731 content: None,
732 reasoning_content: None,
733 tool_calls: None,
734 },
735 finish_reason: Some("stop".to_string()),
736 }
737 }
738
739 // OpenAI-compatible servers send a final chunk with usage data and an empty
740 // choices array. Before this fix, the mapper returned an error for empty
741 // choices, discarding usage entirely.
742 #[test]
743 fn test_usage_in_final_empty_choices_chunk() {
744 let mut mapper = LmStudioEventMapper::new();
745 let event = make_event(
746 vec![],
747 Some(Usage {
748 prompt_tokens: 10,
749 completion_tokens: 20,
750 total_tokens: 30,
751 }),
752 );
753
754 let results: Vec<_> = mapper
755 .map_event(event)
756 .into_iter()
757 .map(|r| r.unwrap())
758 .collect();
759
760 assert_eq!(
761 results,
762 vec![LanguageModelCompletionEvent::UsageUpdate(TokenUsage {
763 input_tokens: 10,
764 output_tokens: 20,
765 cache_creation_input_tokens: 0,
766 cache_read_input_tokens: 0,
767 })]
768 );
769 }
770
771 #[test]
772 fn test_empty_choices_without_usage_returns_empty() {
773 let mut mapper = LmStudioEventMapper::new();
774 let event = make_event(vec![], None);
775
776 let results = mapper.map_event(event);
777
778 assert!(results.is_empty());
779 }
780
781 // Usage data can also arrive in a regular chunk that also contains content.
782 // Both events must be emitted, with UsageUpdate first.
783 #[test]
784 fn test_usage_emitted_alongside_content() {
785 let mut mapper = LmStudioEventMapper::new();
786 let event = make_event(
787 vec![make_content_choice("Hello!")],
788 Some(Usage {
789 prompt_tokens: 5,
790 completion_tokens: 3,
791 total_tokens: 8,
792 }),
793 );
794
795 let results: Vec<_> = mapper
796 .map_event(event)
797 .into_iter()
798 .map(|r| r.unwrap())
799 .collect();
800
801 assert_eq!(
802 results[0],
803 LanguageModelCompletionEvent::UsageUpdate(TokenUsage {
804 input_tokens: 5,
805 output_tokens: 3,
806 cache_creation_input_tokens: 0,
807 cache_read_input_tokens: 0,
808 })
809 );
810 assert_eq!(
811 results[1],
812 LanguageModelCompletionEvent::Text("Hello!".to_string())
813 );
814 }
815
816 #[test]
817 fn test_stop_event_emitted_on_finish_reason() {
818 let mut mapper = LmStudioEventMapper::new();
819 let event = make_event(vec![make_stop_choice()], None);
820
821 let results: Vec<_> = mapper
822 .map_event(event)
823 .into_iter()
824 .map(|r| r.unwrap())
825 .collect();
826
827 assert_eq!(
828 results,
829 vec![LanguageModelCompletionEvent::Stop(StopReason::EndTurn)]
830 );
831 }
832}
833
834fn add_message_content_part(
835 new_part: lmstudio::MessagePart,
836 role: Role,
837 messages: &mut Vec<lmstudio::ChatMessage>,
838) {
839 match (role, messages.last_mut()) {
840 (Role::User, Some(lmstudio::ChatMessage::User { content }))
841 | (
842 Role::Assistant,
843 Some(lmstudio::ChatMessage::Assistant {
844 content: Some(content),
845 ..
846 }),
847 )
848 | (Role::System, Some(lmstudio::ChatMessage::System { content })) => {
849 content.push_part(new_part);
850 }
851 _ => {
852 messages.push(match role {
853 Role::User => lmstudio::ChatMessage::User {
854 content: lmstudio::MessageContent::from(vec![new_part]),
855 },
856 Role::Assistant => lmstudio::ChatMessage::Assistant {
857 content: Some(lmstudio::MessageContent::from(vec![new_part])),
858 tool_calls: Vec::new(),
859 },
860 Role::System => lmstudio::ChatMessage::System {
861 content: lmstudio::MessageContent::from(vec![new_part]),
862 },
863 });
864 }
865 }
866}
867
868struct ConfigurationView {
869 state: Entity<State>,
870 api_key_editor: Entity<InputField>,
871 api_url_editor: Entity<InputField>,
872}
873
874impl ConfigurationView {
875 pub fn new(state: Entity<State>, _window: &mut Window, cx: &mut Context<Self>) -> Self {
876 let api_key_editor = cx.new(|cx| InputField::new(_window, cx, "sk-...").label("API key"));
877
878 let api_url_editor = cx.new(|cx| {
879 let input = InputField::new(_window, cx, LMSTUDIO_API_URL).label("API URL");
880 input.set_text(&LmStudioLanguageModelProvider::api_url(cx), _window, cx);
881 input
882 });
883
884 cx.observe(&state, |_, _, cx| {
885 cx.notify();
886 })
887 .detach();
888
889 Self {
890 state,
891 api_key_editor,
892 api_url_editor,
893 }
894 }
895
896 fn retry_connection(&mut self, _window: &mut Window, cx: &mut Context<Self>) {
897 let has_api_url = LmStudioLanguageModelProvider::has_custom_url(cx);
898 let has_api_key = self
899 .state
900 .read_with(cx, |state, _| state.api_key_state.has_key());
901 if !has_api_url {
902 self.save_api_url(cx);
903 }
904 if !has_api_key {
905 self.save_api_key(&Default::default(), _window, cx);
906 }
907
908 self.state.update(cx, |state, cx| {
909 state.restart_fetch_models_task(cx);
910 });
911 }
912
913 fn save_api_key(&mut self, _: &menu::Confirm, _window: &mut Window, cx: &mut Context<Self>) {
914 let api_key = self.api_key_editor.read(cx).text(cx).trim().to_string();
915 if api_key.is_empty() {
916 return;
917 }
918
919 self.api_key_editor
920 .update(cx, |input, cx| input.set_text("", _window, cx));
921
922 let state = self.state.clone();
923 cx.spawn_in(_window, async move |_, cx| {
924 state
925 .update(cx, |state, cx| state.set_api_key(Some(api_key), cx))
926 .await
927 })
928 .detach_and_log_err(cx);
929 }
930
931 fn reset_api_key(&mut self, _window: &mut Window, cx: &mut Context<Self>) {
932 self.api_key_editor
933 .update(cx, |input, cx| input.set_text("", _window, cx));
934
935 let state = self.state.clone();
936 cx.spawn_in(_window, async move |_, cx| {
937 state
938 .update(cx, |state, cx| state.set_api_key(None, cx))
939 .await
940 })
941 .detach_and_log_err(cx);
942
943 cx.notify();
944 }
945
946 fn save_api_url(&self, cx: &mut Context<Self>) {
947 let api_url = self.api_url_editor.read(cx).text(cx).trim().to_string();
948 let current_url = LmStudioLanguageModelProvider::api_url(cx);
949 if !api_url.is_empty() && &api_url != ¤t_url {
950 self.state
951 .update(cx, |state, cx| state.set_api_key(None, cx))
952 .detach_and_log_err(cx);
953
954 let fs = <dyn Fs>::global(cx);
955 update_settings_file(fs, cx, move |settings, _| {
956 settings
957 .language_models
958 .get_or_insert_default()
959 .lmstudio
960 .get_or_insert_default()
961 .api_url = Some(api_url);
962 });
963 }
964 }
965
966 fn reset_api_url(&mut self, _window: &mut Window, cx: &mut Context<Self>) {
967 self.api_url_editor
968 .update(cx, |input, cx| input.set_text("", _window, cx));
969
970 // Clear API key when URL changes since keys are URL-specific
971 self.state
972 .update(cx, |state, cx| state.set_api_key(None, cx))
973 .detach_and_log_err(cx);
974
975 let fs = <dyn Fs>::global(cx);
976 update_settings_file(fs, cx, |settings, _cx| {
977 if let Some(settings) = settings
978 .language_models
979 .as_mut()
980 .and_then(|models| models.lmstudio.as_mut())
981 {
982 settings.api_url = Some(LMSTUDIO_API_URL.into());
983 }
984 });
985 cx.notify();
986 }
987
988 fn render_api_url_editor(&self, cx: &Context<Self>) -> impl IntoElement {
989 let api_url = LmStudioLanguageModelProvider::api_url(cx);
990 let custom_api_url_set = api_url != LMSTUDIO_API_URL;
991
992 if custom_api_url_set {
993 ConfiguredApiCard::new("reset-api-url", api_url)
994 .on_click(cx.listener(|this, _, _window, cx| this.reset_api_url(_window, cx)))
995 .into_any_element()
996 } else {
997 v_flex()
998 .on_action(cx.listener(|this, _: &menu::Confirm, _window, cx| {
999 this.save_api_url(cx);
1000 cx.notify();
1001 }))
1002 .child(self.api_url_editor.clone())
1003 .into_any_element()
1004 }
1005 }
1006
1007 fn render_api_key_editor(&self, cx: &Context<Self>) -> impl IntoElement {
1008 let state = self.state.read(cx);
1009 let env_var_set = state.api_key_state.is_from_env_var();
1010 let configured_card_label = if env_var_set {
1011 format!("API key set in {API_KEY_ENV_VAR_NAME} environment variable.")
1012 } else {
1013 "API key configured".to_string()
1014 };
1015
1016 let api_key_control = if !state.api_key_state.has_key() {
1017 self.api_key_editor.clone().into_any_element()
1018 } else {
1019 ConfiguredApiCard::new("lmstudio-reset-key", configured_card_label)
1020 .disabled(env_var_set)
1021 .on_click(cx.listener(|this, _, _window, cx| this.reset_api_key(_window, cx)))
1022 .when(env_var_set, |this| {
1023 this.tooltip_label(format!(
1024 "To reset your API key, unset the {API_KEY_ENV_VAR_NAME} environment variable."
1025 ))
1026 })
1027 .into_any_element()
1028 };
1029
1030 v_flex()
1031 .on_action(cx.listener(Self::save_api_key))
1032 .child(api_key_control)
1033 .gap_1p5()
1034 .mb_2()
1035 .child(
1036 Label::new(format!(
1037 "You can also set the {API_KEY_ENV_VAR_NAME} environment variable and restart Omega."
1038 ))
1039 .size(LabelSize::Small)
1040 .color(Color::Muted),
1041 )
1042 }
1043}
1044
1045impl Render for ConfigurationView {
1046 fn render(&mut self, _window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
1047 let is_authenticated = self.state.read(cx).is_authenticated();
1048
1049 v_flex()
1050 .gap_2()
1051 .child(
1052 v_flex()
1053 .gap_1()
1054 .child(Headline::new("LM Studio").size(HeadlineSize::Small))
1055 .child(
1056 Label::new("Run local LLMs like Llama, Phi, and Qwen.").color(Color::Muted),
1057 )
1058 .child(
1059 List::new()
1060 .child(ListBulletItem::new(
1061 "LM Studio needs to be running with at least one model downloaded.",
1062 ).label_color(Color::Muted))
1063 .child(
1064 ListBulletItem::new("")
1065 .child(Label::new("To get your first model, try running").color(Color::Muted))
1066 .child(Label::new("lms get qwen2.5-coder-7b").inline_code(cx).color(Color::Muted).ml_1()),
1067 ),
1068 )
1069 .child(Label::new(
1070 "Alternatively, you can connect to an LM Studio server by specifying its \
1071 URL and API key (may not be required):",
1072 ).color(Color::Muted)),
1073 )
1074 .child(self.render_api_url_editor(cx))
1075 .child(self.render_api_key_editor(cx))
1076 .child(Divider::horizontal())
1077 .child(
1078 h_flex()
1079 .pt_2()
1080 .w_full()
1081 .justify_between()
1082 .gap_1()
1083 .child(
1084 h_flex()
1085 .w_full()
1086 .gap_1()
1087 .map(|this| {
1088 if is_authenticated {
1089 this.child(
1090 Button::new("lmstudio-site", "LM Studio")
1091 .style(ButtonStyle::OutlinedGhost)
1092 .size(ButtonSize::Medium)
1093 .end_icon(
1094 Icon::new(IconName::ArrowUpRight)
1095 .size(IconSize::Small)
1096 .color(Color::Muted),
1097 )
1098 .on_click(move |_, _window, cx| {
1099 cx.open_url(LMSTUDIO_SITE)
1100 })
1101 .into_any_element(),
1102 )
1103 } else {
1104 this.child(
1105 Button::new(
1106 "download_lmstudio_button",
1107 "Download LM Studio",
1108 )
1109 .style(ButtonStyle::OutlinedGhost)
1110 .size(ButtonSize::Medium)
1111 .end_icon(
1112 Icon::new(IconName::ArrowUpRight)
1113 .size(IconSize::Small)
1114 .color(Color::Muted),
1115 )
1116 .on_click(move |_, _window, cx| {
1117 cx.open_url(LMSTUDIO_DOWNLOAD_URL)
1118 })
1119 .into_any_element(),
1120 )
1121 }
1122 })
1123 .child(
1124 Button::new("view-models", "Model Catalog")
1125 .style(ButtonStyle::OutlinedGhost)
1126 .size(ButtonSize::Medium)
1127 .end_icon(
1128 Icon::new(IconName::ArrowUpRight)
1129 .size(IconSize::Small)
1130 .color(Color::Muted),
1131 )
1132 .on_click(move |_, _window, cx| {
1133 cx.open_url(LMSTUDIO_CATALOG_URL)
1134 }),
1135 ),
1136 )
1137 .map(|this| {
1138 if is_authenticated {
1139 this.child(
1140 ButtonLike::new("connected")
1141 .size(ButtonSize::Medium)
1142 .child(
1143 h_flex()
1144 .gap_1()
1145 .child(Icon::new(IconName::Check).color(Color::Success))
1146 .child(Label::new("Connected"))
1147 )
1148 .child(
1149 IconButton::new("refresh-models", IconName::RotateCcw)
1150 .tooltip(Tooltip::text("Refresh Models"))
1151 .icon_size(IconSize::Small)
1152 .on_click(cx.listener(|this, _, _window, cx| {
1153 this.state.update(cx, |state, _| {
1154 state.available_models.clear();
1155 });
1156 this.retry_connection(_window, cx);
1157 })),
1158 ),
1159 )
1160 } else {
1161 this.child(
1162 Button::new("retry_lmstudio_models", "Connect")
1163 .style(ButtonStyle::Outlined)
1164 .size(ButtonSize::Medium)
1165 .start_icon(
1166 Icon::new(IconName::PlayFilled).size(IconSize::XSmall),
1167 )
1168 .on_click(cx.listener(move |this, _, _window, cx| {
1169 this.retry_connection(_window, cx)
1170 })),
1171 )
1172 }
1173 }),
1174 )
1175 }
1176}
1177