Add LM Studio support to the Assistant (#23097)

#### Release Notes:

- Added support for [LM Studio](https://lmstudio.ai/) to the Assistant.

#### Quick demo:


https://github.com/user-attachments/assets/af58fc13-1abc-4898-9747-3511016da86a

#### Future enhancements:
- wire up tool calling (new in [LM Studio
0.3.6](https://lmstudio.ai/blog/lmstudio-v0.3.6))

---------

Co-authored-by: Marshall Bowers <elliott.codes@gmail.com>
This commit is contained in:
Yagil Burowski 2025-01-14 15:41:58 -05:00 committed by GitHub
parent 4445679f3c
commit c038696aa8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
24 changed files with 1153 additions and 2 deletions

View file

@ -1,8 +1,10 @@
mod cloud;
mod lmstudio;
mod ollama;
mod open_ai;
pub use cloud::*;
pub use lmstudio::*;
pub use ollama::*;
pub use open_ai::*;
use sha2::{Digest, Sha256};

View file

@ -0,0 +1,70 @@
use anyhow::{Context as _, Result};
use futures::{future::BoxFuture, AsyncReadExt as _, FutureExt};
use http_client::HttpClient;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use crate::{Embedding, EmbeddingProvider, TextToEmbed};
pub enum LmStudioEmbeddingModel {
NomicEmbedText,
}
pub struct LmStudioEmbeddingProvider {
client: Arc<dyn HttpClient>,
model: LmStudioEmbeddingModel,
}
#[derive(Serialize)]
struct LmStudioEmbeddingRequest {
model: String,
prompt: String,
}
#[derive(Deserialize)]
struct LmStudioEmbeddingResponse {
embedding: Vec<f32>,
}
impl LmStudioEmbeddingProvider {
pub fn new(client: Arc<dyn HttpClient>, model: LmStudioEmbeddingModel) -> Self {
Self { client, model }
}
}
impl EmbeddingProvider for LmStudioEmbeddingProvider {
fn embed<'a>(&'a self, texts: &'a [TextToEmbed<'a>]) -> BoxFuture<'a, Result<Vec<Embedding>>> {
let model = match self.model {
LmStudioEmbeddingModel::NomicEmbedText => "nomic-embed-text",
};
futures::future::try_join_all(texts.iter().map(|to_embed| {
let request = LmStudioEmbeddingRequest {
model: model.to_string(),
prompt: to_embed.text.to_string(),
};
let request = serde_json::to_string(&request).unwrap();
async {
let response = self
.client
.post_json("http://localhost:1234/api/v0/embeddings", request.into())
.await?;
let mut body = String::new();
response.into_body().read_to_string(&mut body).await?;
let response: LmStudioEmbeddingResponse =
serde_json::from_str(&body).context("Unable to parse response")?;
Ok(Embedding::new(response.embedding))
}
}))
.boxed()
}
fn batch_size(&self) -> usize {
256
}
}