mempal_runtime/embed/
mod.rs1#![warn(clippy::all)]
2
3use async_trait::async_trait;
4
5use crate::core::config::Config;
6
7#[cfg(feature = "model2vec")]
8pub use mempal_embed::model2vec;
9#[cfg(feature = "onnx")]
10pub use mempal_embed::onnx;
11pub use mempal_embed::{EmbedError, Embedder, EmbedderFactory, Result, api};
12
13pub const DEFAULT_API_DIMENSIONS: usize = 384;
15
16#[derive(Clone)]
17pub struct ConfiguredEmbedderFactory {
18 config: Config,
19}
20
21impl ConfiguredEmbedderFactory {
22 pub fn new(config: Config) -> Self {
23 Self { config }
24 }
25}
26
27#[async_trait]
28impl EmbedderFactory for ConfiguredEmbedderFactory {
29 async fn build(&self) -> Result<Box<dyn Embedder>> {
30 match self.config.embed.backend.as_str() {
31 #[cfg(feature = "model2vec")]
32 "model2vec" => {
33 let model_id = self
34 .config
35 .embed
36 .model
37 .as_deref()
38 .unwrap_or("minishlab/potion-multilingual-128M");
39 Ok(Box::new(model2vec::Model2VecEmbedder::new(model_id)?))
40 }
41 #[cfg(not(feature = "model2vec"))]
42 "model2vec" => Err(EmbedError::UnsupportedBackend("model2vec".to_string())),
43 #[cfg(feature = "onnx")]
44 "onnx" => Ok(Box::new(onnx::OnnxEmbedder::new_or_download().await?)),
45 #[cfg(not(feature = "onnx"))]
46 "onnx" => Err(EmbedError::UnsupportedBackend("onnx".to_string())),
47 "api" => Ok(Box::new(api::ApiEmbedder::new(
48 self.config
49 .embed
50 .api_endpoint
51 .clone()
52 .unwrap_or_else(|| "http://localhost:11434/api/embeddings".to_string()),
53 self.config.embed.api_model.clone(),
54 self.config
55 .embed
56 .dimensions
57 .unwrap_or(DEFAULT_API_DIMENSIONS),
58 ))),
59 other => Err(EmbedError::UnsupportedBackend(other.to_string())),
60 }
61 }
62}