chore: Refactor async functions to synchronous in reasoning engine sample

samples/langchain_on_vertexai/clean_up.py

@@ -32,8 +32,8 @@
 TEST_NAME = os.getenv("DISPLAY_NAME")
 
 
-async def delete_tables():
-    engine = await PostgresEngine.afrom_instance(
+def delete_tables():
+    engine = PostgresEngine.from_instance(
         PROJECT_ID,
         REGION,
         INSTANCE,
@@ -42,12 +42,14 @@ async def delete_tables():
         password=PASSWORD,
     )
 
-    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()
+    async def _cleanup_logic():
+        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}"))
+
+    engine._run_as_sync(_cleanup_logic())
+    engine._run_as_sync(engine.close())
 
 
 def delete_engines():
@@ -56,9 +58,10 @@ def delete_engines():
         app.delete()
 
 
-async def main():
-    await delete_tables()
+def main():
+    delete_tables()
     delete_engines()
 
 
-asyncio.run(main())
+if __name__ == "__main__":
+    main()

samples/langchain_on_vertexai/create_embeddings.py

@@ -32,24 +32,27 @@
 from langchain_google_cloud_sql_pg import PostgresEngine, PostgresVectorStore
 
 
-async def create_databases():
-    engine = await PostgresEngine.afrom_instance(
+def create_databases():
+    engine = PostgresEngine.from_instance(
         PROJECT_ID,
         REGION,
         INSTANCE,
         database="postgres",
         user=USER,
         password=PASSWORD,
     )
-    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_logic():
+        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}"'))
+
+    engine._run_as_sync(_create_logic())
 
 
-async def create_vectorstore():
-    engine = await PostgresEngine.afrom_instance(
+def create_vectorstore():
+    engine = PostgresEngine.from_instance(
         PROJECT_ID,
         REGION,
         INSTANCE,
@@ -58,11 +61,11 @@ async def create_vectorstore():
         password=PASSWORD,
     )
 
-    await engine.ainit_vectorstore_table(
+    engine.init_vectorstore_table(
         table_name=TABLE_NAME, vector_size=768, overwrite_existing=True
     )
 
-    await engine.ainit_chat_history_table(table_name=CHAT_TABLE_NAME)
+    engine.init_chat_history_table(table_name=CHAT_TABLE_NAME)
 
     rm = resourcemanager_v3.ProjectsClient()
     res = rm.get_project(
@@ -76,7 +79,7 @@ async def grant_select(engine):
             await conn.execute(text(f'GRANT SELECT ON {TABLE_NAME} TO "{IAM_USER}";'))
             await conn.commit()
 
-    await engine._run_as_async(grant_select(engine))
+    engine._run_as_sync(grant_select(engine))
 
     metadata = [
         "show_id",
@@ -91,21 +94,22 @@ async def grant_select(engine):
     loader = CSVLoader(file_path="./movies.csv", metadata_columns=metadata)
     docs = loader.load()
 
-    vector_store = await PostgresVectorStore.create(
+    vector_store = PostgresVectorStore.create_sync(
         engine,
         table_name=TABLE_NAME,
         embedding_service=VertexAIEmbeddings(
-            model_name="textembedding-gecko@latest", project=PROJECT_ID
+            model_name="text-embedding-004", project=PROJECT_ID
         ),
     )
 
     ids = [str(uuid.uuid4()) for i in range(len(docs))]
-    await vector_store.aadd_documents(docs, ids=ids)
+    vector_store.add_documents(docs, ids=ids)
 
 
-async def main():
-    await create_databases()
-    await create_vectorstore()
+def main():
+    create_databases()
+    create_vectorstore()
 
 
-asyncio.run(main())
+if __name__ == "__main__":
+    main()

samples/langchain_on_vertexai/prebuilt_langchain_agent_template.py

@@ -65,7 +65,7 @@ def similarity_search(query: str) -> list[Document]:
         engine,
         table_name=TABLE_NAME,
         embedding_service=VertexAIEmbeddings(
-            model_name="textembedding-gecko@latest", project=PROJECT_ID
+            model_name="text-embedding-004", project=PROJECT_ID
         ),
     )
     retriever = vector_store.as_retriever()

samples/langchain_on_vertexai/retriever_agent_with_history_template.py

@@ -91,7 +91,7 @@ def set_up(self):
             engine,
             table_name=self.table,
             embedding_service=VertexAIEmbeddings(
-                model_name="textembedding-gecko@latest", project=self.project
+                model_name="text-embedding-004", project=self.project
             ),
         )
         retriever = vector_store.as_retriever()

samples/langchain_on_vertexai/retriever_chain_template.py

@@ -97,7 +97,7 @@ def set_up(self):
             engine,
             table_name=self.table,
             embedding_service=VertexAIEmbeddings(
-                model_name="textembedding-gecko@latest", project=self.project
+                model_name="text-embedding-004", project=self.project
             ),
         )
         retriever = vector_store.as_retriever()