Skip to main content

engram_embed/
lib.rs

1//! Embedding providers behind a single trait.
2//!
3//! v0 ships Gemini Embed 2 (cloud) + a deterministic stub for tests.
4//! Local fallbacks (candle/ort) land in Phase 4.
5
6pub mod error;
7pub mod gemini;
8pub mod stub;
9
10use async_trait::async_trait;
11pub use error::EmbedError;
12
13#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14pub enum TaskMode {
15    /// Embed a query for retrieval.
16    RetrievalQuery,
17    /// Embed a document being indexed.
18    RetrievalDocument,
19}
20
21#[async_trait]
22pub trait Embedder: Send + Sync {
23    fn name(&self) -> &'static str;
24    fn dimensions(&self) -> usize;
25    fn model(&self) -> String {
26        self.name().to_string()
27    }
28    fn prompt_format(&self) -> &'static str {
29        "legacy"
30    }
31
32    async fn embed_one(&self, text: &str, mode: TaskMode) -> Result<Vec<f32>, EmbedError>;
33
34    async fn embed_batch(
35        &self,
36        texts: &[&str],
37        mode: TaskMode,
38    ) -> Result<Vec<Vec<f32>>, EmbedError> {
39        let mut out = Vec::with_capacity(texts.len());
40        for t in texts {
41            out.push(self.embed_one(t, mode).await?);
42        }
43        Ok(out)
44    }
45}