fix init and store training data of vanna

tryvanna.py

@@ -1,13 +1,7 @@
 import os
 
 from vanna.remote import VannaDefault
-vn = VannaDefault(model='gpt-3.5-turbo', api_key=os.getenv("VANNA_API_KEY"))
-
-# # vn.connect_to_postgres(os.getenv("DATABASE_URL"))
-#  export DATABASE_NAME=verceldb
-# ashwin@MacBook-Air-6 DDD % export DATABASE_USER=default
-# ashwin@MacBook-Air-6 DDD % export DATABASE_PASSWORD=lyzegA2r0ESO
-# ashwin@MacBook-Air-6 DDD % export DATABASE_HOST="ep-dawn-fire-a1i3ytre-pooler.ap-southeast-1.aws.neon.tech" 
+vn = VannaDefault(model='gpt-3.5-turbo', api_key=os.getenv("VANNA_API_KEY"), openai_api_key=os.getenv("OPENAI_API_KEY"))
 
 db_host = os.getenv("DATABASE_HOST")
 db_name = os.getenv("DATABASE_NAME")
@@ -15,9 +9,29 @@ db_user = os.getenv("DATABASE_USER")
 db_password = os.getenv("DATABASE_PASSWORD")
 db_port = 5432
 
+
 vn.connect_to_postgres(host=db_host, dbname=db_name, user=db_user, password=db_password, port=db_port)
 
-vn.ask('What are the top 10 artists by sales?')
+# Fetch information schema
+df_information_schema = vn.run_sql("SELECT * FROM INFORMATION_SCHEMA.COLUMNS")
+
+# Generate training plan
+plan = vn.get_training_plan_generic(df_information_schema)
+
+# Get training data
+training_data = vn.get_training_data()
+
+# Store training data in Vanna Hosted Vector DB
+vn.store_training_data(training_data)
+
+# Train Vanna with the information schema
+vn.train(df_information_schema)
+
+# Example query generation using OpenAI via Vanna.AI
+generated_query = vn.generate_sql("List all table names in the database")
+print("Generated SQL query:", generated_query)
+
+# Example natural language question
+answer = vn.ask('What is the data about?')
+print("Answer:", answer)
 
-from vanna.flask import VannaFlaskApp
-VannaFlaskApp(vn).run()