File size: 1,142 Bytes
3351f47 8a4f160 3351f47 8a4f160 3351f47 8a4f160 3351f47 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 | 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)
|