insert_many fix attempt #2

OoriData · Nov 24, 2023 · f7a75aa · f7a75aa
1 parent 0f5500e
commit f7a75aa
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 45 deletions.
diff --git a/pylib/embedding/pgvector.py b/pylib/embedding/pgvector.py
@@ -341,7 +341,9 @@ async def insert_many(
         await self.conn.executemany(
             INSERT_DOCS.format(table_name=self.table_name),
             [
-                (self._embedding_model.encode(content), content, title, page_numbers, tags)
+                (
+                    self._embedding_model.encode(content).tolist(), content, title, page_numbers, tags
+                )
                 for content, title, page_numbers, tags in content_list
             ]
         )

diff --git a/test/embedding/test_pgvector.py b/test/embedding/test_pgvector.py
@@ -85,49 +85,49 @@ async def test_PGv_embed_pacer():
 
     await vDB.drop_table()
 
-# @pytest.mark.asyncio
-# async def test_PGv_embed_many_pacer():
-#     dummy_model = SentenceTransformer('mock_transformer')
-#     dummy_model.encode.return_value = np.array([1, 2, 3])
-#     print(f'EMODEL: {dummy_model}')
-#     TABLE_NAME = 'embedding_test'
-#     try:
-#         vDB = await DocDB.from_conn_params(
-#             embedding_model=dummy_model,
-#             table_name=TABLE_NAME,
-#             db_name=DB_NAME,
-#             host=HOST,
-#             port=int(PORT),
-#             user=USER,
-#             password=PASSWORD)
-#     except ConnectionRefusedError:
-#         pytest.skip("No Postgres instance made available for test. Skipping.", allow_module_level=True)
+@pytest.mark.asyncio
+async def test_PGv_embed_many_pacer():
+    dummy_model = SentenceTransformer('mock_transformer')
+    dummy_model.encode.return_value = np.array([1, 2, 3])
+    print(f'EMODEL: {dummy_model}')
+    TABLE_NAME = 'embedding_test'
+    try:
+        vDB = await DocDB.from_conn_params(
+            embedding_model=dummy_model,
+            table_name=TABLE_NAME,
+            db_name=DB_NAME,
+            host=HOST,
+            port=int(PORT),
+            user=USER,
+            password=PASSWORD)
+    except ConnectionRefusedError:
+        pytest.skip("No Postgres instance made available for test. Skipping.", allow_module_level=True)
 
-#     assert vDB is not None, ConnectionError("Postgres docker instance not available for testing PG code")
+    assert vDB is not None, ConnectionError("Postgres docker instance not available for testing PG code")
 
-#     # Create tables
-#     await vDB.drop_table()
-#     assert await vDB.table_exists() is False, Exception("Table exists before creation")
-#     await vDB.create_table()
-#     assert await vDB.table_exists() is True, Exception("Table does not exist after creation")
-
-#     # Insert data using insert_many()
-#     documents = [
-#         {
-#             'content': text,
-#             'title': f'Pacer Copypasta line {index}',
-#             'page_numbers': [1, 2, 3],
-#             'tags': ['fitness', 'pacer', 'copypasta']
-#         }
-#         for index, text in enumerate(pacer_copypasta)
-#     ]
-#     await vDB.insert_many(documents)
-
-#     assert await vDB.count_items() == len(pacer_copypasta), Exception("Not all documents inserted")
-
-#     # Search table with perfect match
-#     search_string = '[beep] A single lap should be completed each time you hear this sound.'
-#     sim_search = await vDB.search(query_string=search_string, limit=3)
-#     assert sim_search is not None, Exception("No results returned from perfect search")
-
-#     await vDB.drop_table()
+    # Create tables
+    await vDB.drop_table()
+    assert await vDB.table_exists() is False, Exception("Table exists before creation")
+    await vDB.create_table()
+    assert await vDB.table_exists() is True, Exception("Table does not exist after creation")
+
+    # Insert data using insert_many()
+    documents = [
+        {
+            'content': text,
+            'title': f'Pacer Copypasta line {index}',
+            'page_numbers': [1, 2, 3],
+            'tags': ['fitness', 'pacer', 'copypasta']
+        }
+        for index, text in enumerate(pacer_copypasta)
+    ]
+    await vDB.insert_many(documents)
+
+    assert await vDB.count_items() == len(pacer_copypasta), Exception("Not all documents inserted")
+
+    # Search table with perfect match
+    search_string = '[beep] A single lap should be completed each time you hear this sound.'
+    sim_search = await vDB.search(query_string=search_string, limit=3)
+    assert sim_search is not None, Exception("No results returned from perfect search")
+
+    await vDB.drop_table()