diff --git a/samples/langchain_on_vertexai/clean_up.py b/samples/langchain_on_vertexai/clean_up.py index 45e57ae5..da1f0823 100644 --- a/samples/langchain_on_vertexai/clean_up.py +++ b/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() diff --git a/samples/langchain_on_vertexai/create_embeddings.py b/samples/langchain_on_vertexai/create_embeddings.py index 105a86df..2321624d 100644 --- a/samples/langchain_on_vertexai/create_embeddings.py +++ b/samples/langchain_on_vertexai/create_embeddings.py @@ -32,8 +32,8 @@ 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, @@ -41,15 +41,18 @@ async def create_databases(): 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() diff --git a/samples/langchain_on_vertexai/prebuilt_langchain_agent_template.py b/samples/langchain_on_vertexai/prebuilt_langchain_agent_template.py index 472b9da9..19bf67c7 100644 --- a/samples/langchain_on_vertexai/prebuilt_langchain_agent_template.py +++ b/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() diff --git a/samples/langchain_on_vertexai/retriever_agent_with_history_template.py b/samples/langchain_on_vertexai/retriever_agent_with_history_template.py index 7d8a520e..62ae6a63 100644 --- a/samples/langchain_on_vertexai/retriever_agent_with_history_template.py +++ b/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() diff --git a/samples/langchain_on_vertexai/retriever_chain_template.py b/samples/langchain_on_vertexai/retriever_chain_template.py index d05780c3..1601cd7a 100644 --- a/samples/langchain_on_vertexai/retriever_chain_template.py +++ b/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()