chore: Use psycopg2.extras.execute_values for batch insert in PGVector

langgenius · May 5, 2024 · 86cb30d · 86cb30d
1 parent a932466
commit 86cb30d
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/api/core/rag/datasource/vdb/pgvector/pgvector.py b/api/core/rag/datasource/vdb/pgvector/pgvector.py
@@ -3,6 +3,7 @@
 from contextlib import contextmanager
 from typing import Any
 
+import psycopg2.extras
 import psycopg2.pool
 from pydantic import BaseModel
 
@@ -80,8 +81,8 @@ def add_texts(self, documents: list[Document], embeddings: list[list[float]], **
                 )
             )
         with self._get_cursor() as cur:
-            cur.executemany(
-                f"INSERT INTO {self.table_name} (id, text, meta, embedding) VALUES (%s, %s, %s, %s)", values
+            psycopg2.extras.execute_values(
+                cur, f"INSERT INTO {self.table_name} (id, text, meta, embedding) VALUES %s", values
             )
         return pks