vanna

import os from vanna.local import VannaDefault import pandas as pd # For local testing, VannaDefault uses DuckDB and OpenAI as defaults. # Ensure OPENAI_API_KEY is set in your environment. # pip install "vanna[duckdb]" "vanna[openai]" # Initialize VannaDefault with your OpenAI API key and preferred model vn = VannaDefault( model='gpt-4o', api_key=os.environ.get('OPENAI_API_KEY', 'YOUR_OPENAI_API_KEY'), # Replace with actual key or set env var path='my_vanna_db' # Path for local ChromaDB and DuckDB storage ) # Connect to a database (DuckDB is default for VannaDefault) vn.run_sql_query("CREATE TABLE sales (id INT, product TEXT, amount INT)") vn.run_sql_query("INSERT INTO sales (id, product, amount) VALUES (1, 'Apple', 100), (2, 'Banana', 150)") # Train Vanna with DDL, documentation, and example questions/SQL vn.train(ddl="CREATE TABLE employees (id INT, name TEXT, salary INT)") vn.train(question="What are the total sales for each product?", sql="SELECT product, SUM(amount) FROM sales GROUP BY product") vn.train(sql="SELECT * FROM sales", df=pd.DataFrame({'id': [1], 'product': ['Apple'], 'amount': [100]})) # Ask a question question = "Show me the sum of sales for each product" sql = vn.generate_sql(question=question) print(f"Generated SQL: {sql}") # Run the generated SQL if sql: results = vn.run_sql(sql=sql) print(f"Query Results:\n{results}") else: print("Could not generate SQL for the question.")