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:
parent
4445679f3c
commit
c038696aa8
24 changed files with 1153 additions and 2 deletions
|
@ -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};
|
||||
|
|
70
crates/semantic_index/src/embedding/lmstudio.rs
Normal file
70
crates/semantic_index/src/embedding/lmstudio.rs
Normal 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
|
||||
}
|
||||
}
|
Loading…
Add table
Add a link
Reference in a new issue