Skip to main content

qql_embed/
embedder.rs

1use async_trait::async_trait;
2use qql_core::error::QqlError;
3
4use crate::sparse::{self, SparseVector};
5
6#[cfg(not(target_arch = "wasm32"))]
7pub trait EmbedderBound: Send + Sync {}
8#[cfg(not(target_arch = "wasm32"))]
9impl<T: Send + Sync> EmbedderBound for T {}
10
11#[cfg(target_arch = "wasm32")]
12pub trait EmbedderBound {}
13#[cfg(target_arch = "wasm32")]
14impl<T> EmbedderBound for T {}
15
16/// Host-agnostic embedding backend.
17///
18/// Dense calls should batch when possible (`embed_dense_batch` → one HTTP
19/// request or one ONNX batch). Sparse defaults to local BM25-style hashing.
20#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
21#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
22pub trait Embedder: EmbedderBound {
23    async fn embed_dense(&self, text: &str, model: &str) -> Result<Vec<f32>, QqlError>;
24    async fn embed_sparse(&self, text: &str) -> Result<SparseVector, QqlError>;
25
26    /// Dense output dimension when it is known without running inference.
27    /// Custom and remote embedders may return `None`.
28    fn dimension(&self) -> Option<usize> {
29        None
30    }
31
32    /// Whether this embedder can satisfy a requested model identifier.
33    /// Dynamic providers may return `true` for every model.
34    fn accepts_model(&self, _model: &str) -> bool {
35        true
36    }
37
38    /// Embed many texts in one shot. Default loops `embed_dense`; override for
39    /// real batching (OpenAI-compatible `input: [...]`, fastembed batch, etc.).
40    async fn embed_dense_batch(
41        &self,
42        texts: &[String],
43        model: &str,
44    ) -> Result<Vec<Vec<f32>>, QqlError> {
45        let mut results = Vec::with_capacity(texts.len());
46        for text in texts {
47            results.push(self.embed_dense(text, model).await?);
48        }
49        Ok(results)
50    }
51}
52
53/// Local sparse-only helper (no dense model).
54pub struct SparseEmbedder;
55
56impl SparseEmbedder {
57    pub fn embed_sparse(text: &str) -> SparseVector {
58        sparse::build_query_default(text)
59    }
60}