vanna_model.py

import pandas as pd
import vanna as vn


# Create login and get API key
api_key = vn.get_api_key('saroj.dhiman@ociustechnologies.com') 
vn.set_api_key(api_key)

# Create models
vn.create_model(model="data_warehose", db_type="Snowflake")
vn.create_model(model="marketing_db", db_type="Postgres")
vn.create_model(model="marketing_db", db_type="BigQuery")

models = vn.get_models()

  # Set model
vn.set_model('chinook')

# Connect to database
url= "https://github.com/lerocha/chinook-database/raw/master/ChinookDatabase/DataSources/Chinook_Sqlite.sqlite"
vn.connect_to_sqlite(url=url)

# Check training data
training_data = vn.get_training_data()
training_data


vn.ask(question='What are the top 5 jazz artists by sales?')

training_data = vn.get_training_data()
training_data['question'].str.contains('What are the top 5 jazz artists by sales?').any()

vn.ask(question='Who is the oldest employee')