Tests

pgml-sdks/pgml/python/tests/requirements.txt

@@ -0,0 +1,2 @@
+pytest
+pytest-asyncio

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 ##################################

pgml-sdks/pgml/src/builtins.rs

@@ -116,7 +116,7 @@ impl Builtins {
  let texts = texts
  .0
  .as_array()
- .with_context(|| "embed_batch takes an array of texts")?
+ .with_context(|| "embed_batch takes an array of strings")?
  .into_iter()
  .map(|v| {
  v.as_str()
@@ -167,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(())
+ }
 }