import os

from vanna.remote import VannaDefault
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")
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)

# 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)