Skip to main content

mempal_embed/
lib.rs

1#![warn(clippy::all)]
2
3use std::path::PathBuf;
4
5use thiserror::Error;
6
7pub mod api;
8#[cfg(feature = "model2vec")]
9pub mod model2vec;
10#[cfg(feature = "onnx")]
11pub mod onnx;
12
13pub type Result<T> = std::result::Result<T, EmbedError>;
14
15#[derive(Debug, Error)]
16pub enum EmbedError {
17    #[error("failed to create model directory {path}")]
18    CreateModelDir {
19        path: PathBuf,
20        #[source]
21        source: std::io::Error,
22    },
23    #[error("failed to check whether {path} exists")]
24    CheckPathExists {
25        path: PathBuf,
26        #[source]
27        source: std::io::Error,
28    },
29    #[error("failed to download {url}")]
30    Download {
31        url: String,
32        #[source]
33        source: reqwest::Error,
34    },
35    #[error("download returned error status for {url}")]
36    DownloadStatus {
37        url: String,
38        #[source]
39        source: reqwest::Error,
40    },
41    #[error("failed to read download body from {url}")]
42    ReadDownloadBody {
43        url: String,
44        #[source]
45        source: reqwest::Error,
46    },
47    #[error("failed to write {path}")]
48    WriteFile {
49        path: PathBuf,
50        #[source]
51        source: std::io::Error,
52    },
53    #[error("failed to rename {from} to {to}")]
54    RenameFile {
55        from: PathBuf,
56        to: PathBuf,
57        #[source]
58        source: std::io::Error,
59    },
60    #[error("failed to initialize ONNX session builder: {0}")]
61    SessionBuilder(String),
62    #[error("failed to load ONNX model from {path}: {message}")]
63    LoadModel { path: PathBuf, message: String },
64    #[error("tokenizer error: {0}")]
65    Tokenizer(String),
66    #[error("embedding runtime error: {0}")]
67    Runtime(String),
68    #[error("embedding worker panicked")]
69    WorkerPanic(#[source] tokio::task::JoinError),
70    #[error("failed to call embedding endpoint {endpoint}")]
71    HttpRequest {
72        endpoint: String,
73        #[source]
74        source: reqwest::Error,
75    },
76    #[error("embedding endpoint returned error status {endpoint}")]
77    HttpStatus {
78        endpoint: String,
79        #[source]
80        source: reqwest::Error,
81    },
82    #[error("failed to decode embedding response from {endpoint}")]
83    DecodeResponse {
84        endpoint: String,
85        #[source]
86        source: reqwest::Error,
87    },
88    #[error("invalid embedding response: {0}")]
89    InvalidResponse(String),
90    #[error("embedding endpoint returned no vectors")]
91    EmptyVectors,
92    #[error(
93        "embedding endpoint returned vectors with unexpected dimensions; expected {expected}, got {actual}"
94    )]
95    InvalidDimensions { expected: usize, actual: usize },
96    #[error("unsupported embed backend: {0}")]
97    UnsupportedBackend(String),
98}
99
100#[async_trait::async_trait]
101pub trait Embedder: Send + Sync {
102    async fn embed(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>>;
103    fn dimensions(&self) -> usize;
104    fn name(&self) -> &str;
105}
106
107#[async_trait::async_trait]
108pub trait EmbedderFactory: Send + Sync {
109    async fn build(&self) -> Result<Box<dyn Embedder>>;
110}