Reify Embedding/Sha1 structs that can be (de)serialized from SQL

Co-Authored-By: Kyle Caverly <kyle@zed.dev>

crates/semantic_index/src/embedding.rs

@@ -8,6 +8,8 @@ use isahc::prelude::Configurable;
 use isahc::{AsyncBody, Response};
 use lazy_static::lazy_static;
 use parse_duration::parse;
+use rusqlite::types::{FromSql, FromSqlResult, ToSqlOutput, ValueRef};
+use rusqlite::ToSql;
 use serde::{Deserialize, Serialize};
 use std::env;
 use std::sync::Arc;
@@ -20,6 +22,62 @@ lazy_static! {
     static ref OPENAI_BPE_TOKENIZER: CoreBPE = cl100k_base().unwrap();
 }
 
+#[derive(Debug, PartialEq, Clone)]
+pub struct Embedding(Vec<f32>);
+
+impl From<Vec<f32>> for Embedding {
+    fn from(value: Vec<f32>) -> Self {
+        Embedding(value)
+    }
+}
+
+impl Embedding {
+    pub fn similarity(&self, other: &Self) -> f32 {
+        let len = self.0.len();
+        assert_eq!(len, other.0.len());
+
+        let mut result = 0.0;
+        unsafe {
+            matrixmultiply::sgemm(
+                1,
+                len,
+                1,
+                1.0,
+                self.0.as_ptr(),
+                len as isize,
+                1,
+                other.0.as_ptr(),
+                1,
+                len as isize,
+                0.0,
+                &mut result as *mut f32,
+                1,
+                1,
+            );
+        }
+        result
+    }
+}
+
+impl FromSql for Embedding {
+    fn column_result(value: ValueRef) -> FromSqlResult<Self> {
+        let bytes = value.as_blob()?;
+        let embedding: Result<Vec<f32>, Box<bincode::ErrorKind>> = bincode::deserialize(bytes);
+        if embedding.is_err() {
+            return Err(rusqlite::types::FromSqlError::Other(embedding.unwrap_err()));
+        }
+        Ok(Embedding(embedding.unwrap()))
+    }
+}
+
+impl ToSql for Embedding {
+    fn to_sql(&self) -> rusqlite::Result<ToSqlOutput> {
+        let bytes = bincode::serialize(&self.0)
+            .map_err(|err| rusqlite::Error::ToSqlConversionFailure(Box::new(err)))?;
+        Ok(ToSqlOutput::Owned(rusqlite::types::Value::Blob(bytes)))
+    }
+}
+
 #[derive(Clone)]
 pub struct OpenAIEmbeddings {
     pub client: Arc<dyn HttpClient>,
@@ -53,7 +111,7 @@ struct OpenAIEmbeddingUsage {
 
 #[async_trait]
 pub trait EmbeddingProvider: Sync + Send {
-    async fn embed_batch(&self, spans: Vec<String>) -> Result<Vec<Vec<f32>>>;
+    async fn embed_batch(&self, spans: Vec<String>) -> Result<Vec<Embedding>>;
     fn max_tokens_per_batch(&self) -> usize;
     fn truncate(&self, span: &str) -> (String, usize);
 }
@@ -62,10 +120,10 @@ pub struct DummyEmbeddings {}
 
 #[async_trait]
 impl EmbeddingProvider for DummyEmbeddings {
-    async fn embed_batch(&self, spans: Vec<String>) -> Result<Vec<Vec<f32>>> {
+    async fn embed_batch(&self, spans: Vec<String>) -> Result<Vec<Embedding>> {
         // 1024 is the OpenAI Embeddings size for ada models.
         // the model we will likely be starting with.
-        let dummy_vec = vec![0.32 as f32; 1536];
+        let dummy_vec = Embedding::from(vec![0.32 as f32; 1536]);
         return Ok(vec![dummy_vec; spans.len()]);
     }
 
@@ -137,7 +195,7 @@ impl EmbeddingProvider for OpenAIEmbeddings {
         (output, token_count)
     }
 
-    async fn embed_batch(&self, spans: Vec<String>) -> Result<Vec<Vec<f32>>> {
+    async fn embed_batch(&self, spans: Vec<String>) -> Result<Vec<Embedding>> {
         const BACKOFF_SECONDS: [usize; 4] = [3, 5, 15, 45];
         const MAX_RETRIES: usize = 4;
 
@@ -175,7 +233,7 @@ impl EmbeddingProvider for OpenAIEmbeddings {
                     return Ok(response
                         .data
                         .into_iter()
-                        .map(|embedding| embedding.embedding)
+                        .map(|embedding| Embedding::from(embedding.embedding))
                         .collect());
                 }
                 StatusCode::TOO_MANY_REQUESTS => {
@@ -218,3 +276,49 @@ impl EmbeddingProvider for OpenAIEmbeddings {
         Err(anyhow!("openai max retries"))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use rand::prelude::*;
+
+    #[gpui::test]
+    fn test_similarity(mut rng: StdRng) {
+        assert_eq!(
+            Embedding::from(vec![1., 0., 0., 0., 0.])
+                .similarity(&Embedding::from(vec![0., 1., 0., 0., 0.])),
+            0.
+        );
+        assert_eq!(
+            Embedding::from(vec![2., 0., 0., 0., 0.])
+                .similarity(&Embedding::from(vec![3., 1., 0., 0., 0.])),
+            6.
+        );
+
+        for _ in 0..100 {
+            let size = 1536;
+            let mut a = vec![0.; size];
+            let mut b = vec![0.; size];
+            for (a, b) in a.iter_mut().zip(b.iter_mut()) {
+                *a = rng.gen();
+                *b = rng.gen();
+            }
+            let a = Embedding::from(a);
+            let b = Embedding::from(b);
+
+            assert_eq!(
+                round_to_decimals(a.similarity(&b), 1),
+                round_to_decimals(reference_dot(&a.0, &b.0), 1)
+            );
+        }
+
+        fn round_to_decimals(n: f32, decimal_places: i32) -> f32 {
+            let factor = (10.0 as f32).powi(decimal_places);
+            (n * factor).round() / factor
+        }
+
+        fn reference_dot(a: &[f32], b: &[f32]) -> f32 {
+            a.iter().zip(b.iter()).map(|(a, b)| a * b).sum()
+        }
+    }
+}