Embeddings support in the SDK (#1475)

levkk · web-flow · commit 03ca54e8deea · 2024-05-22T10:16:24.000-07:00
diff --git a/pgml-sdks/pgml/Cargo.lock b/pgml-sdks/pgml/Cargo.lock
diff --git a/pgml-sdks/pgml/Cargo.toml b/pgml-sdks/pgml/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "pgml"
-version = "1.0.3"
+version = "1.0.4"
 edition = "2021"
 authors = ["PosgresML <team@postgresml.org>"]
 homepage = "https://postgresml.org/"
diff --git a/pgml-sdks/pgml/python/tests/requirements.txt b/pgml-sdks/pgml/python/tests/requirements.txt
@@ -0,0 +1,2 @@
+pytest
+pytest-asyncio
diff --git a/pgml-sdks/pgml/python/tests/test.py b/pgml-sdks/pgml/python/tests/test.py
@@ -72,6 +72,18 @@ def test_can_create_builtins():
     builtins = pgml.Builtins()
     assert builtins is not None
 
+@pytest.mark.asyncio
+async def test_can_embed_with_builtins():
+    builtins = pgml.Builtins()
+    result = await builtins.embed("intfloat/e5-small-v2", "test")
+    assert result is not None
+
+@pytest.mark.asyncio
+async def test_can_embed_batch_with_builtins():
+    builtins = pgml.Builtins()
+    result = await builtins.embed_batch("intfloat/e5-small-v2", ["test"])
+    assert result is not None
+
 
 ###################################################
 ## Test searches ##################################
diff --git a/pgml-sdks/pgml/src/builtins.rs b/pgml-sdks/pgml/src/builtins.rs
@@ -1,3 +1,4 @@
+use anyhow::Context;
 use rust_bridge::{alias, alias_methods};
 use sqlx::Row;
 use tracing::instrument;
@@ -13,7 +14,7 @@ use crate::{get_or_initialize_pool, query_runner::QueryRunner, types::Json};
 #[cfg(feature = "python")]
 use crate::{query_runner::QueryRunnerPython, types::JsonPython};
 
-#[alias_methods(new, query, transform)]
+#[alias_methods(new, query, transform, embed, embed_batch)]
 impl Builtins {
     pub fn new(database_url: Option<String>) -> Self {
         Self { database_url }
@@ -87,6 +88,55 @@ impl Builtins {
         let results = results.first().unwrap().get::<serde_json::Value, _>(0);
         Ok(Json(results))
     }
+
+    /// Run the built-in `pgml.embed()` function.
+    ///
+    /// # Arguments
+    ///
+    /// * `model` - The model to use.
+    /// * `text` - The text to embed.
+    ///
+    pub async fn embed(&self, model: &str, text: &str) -> anyhow::Result<Json> {
+        let pool = get_or_initialize_pool(&self.database_url).await?;
+        let query = sqlx::query("SELECT embed FROM pgml.embed($1, $2)");
+        let result = query.bind(model).bind(text).fetch_one(&pool).await?;
+        let result = result.get::<Vec<f32>, _>(0);
+        let result = serde_json::to_value(result)?;
+        Ok(Json(result))
+    }
+
+    /// Run the built-in `pgml.embed()` function, but with handling for batch inputs and outputs.
+    ///
+    /// # Arguments
+    ///
+    /// * `model` - The model to use.
+    /// * `texts` - The texts to embed.
+    ///
+    pub async fn embed_batch(&self, model: &str, texts: Json) -> anyhow::Result<Json> {
+        let texts = texts
+            .0
+            .as_array()
+            .with_context(|| "embed_batch takes an array of strings")?
+            .into_iter()
+            .map(|v| {
+                v.as_str()
+                    .with_context(|| "only text embeddings are supported")
+                    .unwrap()
+                    .to_string()
+            })
+            .collect::<Vec<String>>();
+        let pool = get_or_initialize_pool(&self.database_url).await?;
+        let query = sqlx::query("SELECT embed AS embed_batch FROM pgml.embed($1, $2)");
+        let results = query
+            .bind(model)
+            .bind(texts)
+            .fetch_all(&pool)
+            .await?
+            .into_iter()
+            .map(|embeddings| embeddings.get::<Vec<f32>, _>(0))
+            .collect::<Vec<Vec<f32>>>();
+        Ok(Json(serde_json::to_value(results)?))
+    }
 }
 
 #[cfg(test)]
@@ -117,4 +167,28 @@ mod tests {
         assert!(results.as_array().is_some());
         Ok(())
     }
+
+    #[tokio::test]
+    async fn can_embed() -> anyhow::Result<()> {
+        internal_init_logger(None, None).ok();
+        let builtins = Builtins::new(None);
+        let results = builtins.embed("intfloat/e5-small-v2", "test").await?;
+        assert!(results.as_array().is_some());
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn can_embed_batch() -> anyhow::Result<()> {
+        internal_init_logger(None, None).ok();
+        let builtins = Builtins::new(None);
+        let results = builtins
+            .embed_batch(
+                "intfloat/e5-small-v2",
+                Json(serde_json::json!(["test", "test2",])),
+            )
+            .await?;
+        assert!(results.as_array().is_some());
+        assert_eq!(results.as_array().unwrap().len(), 2);
+        Ok(())
+    }
 }
