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