SDK - Make collection.upsert_documents batch embeddings (#1457)

SilasMarvin · web-flow · commit 88238176a1bb · 2024-05-14T09:33:15.000-07:00
diff --git a/pgml-sdks/pgml/src/queries.rs b/pgml-sdks/pgml/src/queries.rs
@@ -217,9 +217,9 @@ ON CONFLICT (chunk_id) DO UPDATE SET ts = EXCLUDED.ts;
 pub const GENERATE_EMBEDDINGS_FOR_CHUNK_IDS: &str = r#"
 INSERT INTO %s (chunk_id, embedding) 
 SELECT 
-  id, 
+  unnest(array_agg(id)), 
   pgml.embed(
-    text => chunk, 
+    inputs => array_agg(chunk), 
     transformer => $1, 
     kwargs => $2 
   ) 
