chore(samples): update samples for async/sync refactor (#221)

* chore(samples): update samples for async/sync refactor

* clean

* fix

* clean

Author: averikitsch

samples/index_tuning_sample/create_vector_embeddings.py

@@ -106,6 +106,8 @@ async def create_vector_store_table(documents):
     ids = [str(uuid.uuid4()) for i in range(len(documents))]
     await vector_store.aadd_documents(documents, ids)
     print("Vector table created.")
+    await engine.close()
+    await engine._connector.close()
 
 
 async def main():

samples/index_tuning_sample/index_search.py

@@ -81,7 +81,10 @@ async def query_vector_with_timing(vector_store, query):
 
 async def hnsw_search(vector_store, knn_docs):
     hnsw_index = HNSWIndex(
-        name="hnsw", distance_strategy=DISTANCE_STRATEGY, m=36, ef_construction=96
+        name="hnsw",
+        distance_strategy=DISTANCE_STRATEGY,
+        m=36,
+        ef_construction=96,
     )
     await vector_store.aapply_vector_index(hnsw_index)
     assert await vector_store.is_valid_index(hnsw_index.name)
@@ -156,6 +159,8 @@ async def main():
     print(
         f"IVFFLAT average recall: {ivfflat_average_recall}    IVFFLAT latency: {ivfflat_average_latency}"
     )
+    await vector_store._engine.close()
+    await vector_store._engine._connector.close()
 
 
 if __name__ == "__main__":

samples/langchain_on_vertexai/clean_up.py

@@ -24,6 +24,7 @@
     TABLE_NAME,
     USER,
 )
+from sqlalchemy import text
 from vertexai.preview import reasoning_engines  # type: ignore
 
 from langchain_google_cloud_sql_pg import PostgresEngine
@@ -41,10 +42,12 @@ async def delete_tables():
         password=PASSWORD,
     )
 
-    await engine._aexecute_outside_tx(f"DROP TABLE IF EXISTS {TABLE_NAME}")
-    await engine._aexecute_outside_tx(f"DROP TABLE IF EXISTS {CHAT_TABLE_NAME}")
+    async with engine._pool.connect() as conn:
+        await conn.execute(text("COMMIT"))
+        await conn.execute(text(f"DROP TABLE IF EXISTS {TABLE_NAME}"))
+        await conn.execute(text(f"DROP TABLE IF EXISTS {CHAT_TABLE_NAME}"))
+    await engine.close()
     await engine._connector.close_async()
-    await engine._engine.dispose()
 
 
 def delete_engines():

samples/langchain_on_vertexai/create_embeddings.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import asyncio
-import os
 import uuid
 
 from config import (
@@ -28,6 +27,7 @@
 from google.cloud import resourcemanager_v3  # type: ignore
 from langchain_community.document_loaders.csv_loader import CSVLoader
 from langchain_google_vertexai import VertexAIEmbeddings
+from sqlalchemy import text
 
 from langchain_google_cloud_sql_pg import PostgresEngine, PostgresVectorStore
 
@@ -41,10 +41,11 @@ async def create_databases():
         user=USER,
         password=PASSWORD,
     )
-    await engine._aexecute_outside_tx(f'DROP DATABASE IF EXISTS "{DATABASE}"')
-    await engine._aexecute_outside_tx(f'CREATE DATABASE "{DATABASE}"')
-    await engine._connector.close_async()
-    await engine._engine.dispose()
+    async with engine._pool.connect() as conn:
+        await conn.execute(text("COMMIT"))
+        await conn.execute(text(f'DROP DATABASE IF EXISTS "{DATABASE}"'))
+        await conn.execute(text(f'CREATE DATABASE "{DATABASE}"'))
+    await engine.close()
 
 
 async def create_vectorstore():
@@ -69,7 +70,13 @@ async def create_vectorstore():
     )
     project_number = res.name.split("/")[1]
     IAM_USER = f"service-{project_number}@gcp-sa-aiplatform-re.iam"
-    await engine._aexecute(f'GRANT SELECT ON {TABLE_NAME} TO "{IAM_USER}";')
+
+    async def grant_select(engine):
+        async with engine._pool.connect() as conn:
+            await conn.execute(text(f'GRANT SELECT ON {TABLE_NAME} TO "{IAM_USER}";'))
+            await conn.commit()
+
+    await engine._run_as_async(grant_select(engine))
 
     metadata = [
         "show_id",
@@ -94,8 +101,6 @@ async def create_vectorstore():
 
     ids = [str(uuid.uuid4()) for i in range(len(docs))]
     await vector_store.aadd_documents(docs, ids=ids)
-    await engine._connector.close_async()
-    await engine._engine.dispose()
 
 
 async def main():

samples/langchain_on_vertexai/requirements.txt

@@ -1,5 +1,5 @@
-google-cloud-aiplatform[reasoningengine,langchain]==1.65.0
+google-cloud-aiplatform[reasoningengine,langchain]==1.68.0
 google-cloud-resource-manager==1.12.5
 langchain-community==0.2.16
-langchain-google-cloud-sql-pg==0.9.0
+langchain-google-cloud-sql-pg==0.10.0
 langchain-google-vertexai==1.0.10