utils/tools.py

import pandas as pd


kb_df = pd.read_csv("./data/kb.csv")
def get_kb(db_name, knowledge=None):
    if not knowledge:
        result = kb_df[(kb_df['db_name']==db_name)]
    else:
        result = kb_df[(kb_df['db_name']==db_name) & (kb_df['knowledge'].str.contains(knowledge))]
    return result

schema_df = pd.read_csv("./data/db_schema.csv")
def get_schema(db_name, table_name):
    result = schema_df[(schema_df['db_name']==db_name) & (schema_df['table_name']==table_name)]
    result = result[['schema', 'sample_data']]
    return result

def get_tables(db_name):
    result = schema_df[(schema_df['db_name']==db_name)]
    result = result.drop_duplicates(subset=['table_name'])
    tables = result['table_name'].to_list()
    return tables

meaning_df = pd.read_csv("./data/column_meanings.csv")
def get_meaning(db_name, table_name):
    result = meaning_df[(meaning_df['db_name']==db_name) & (meaning_df['table_name']==table_name)]
    result = result[['column_name', 'meaning']]
    return result

get_kb('solar', 'PP')
get_schema('solar', 'alerts')
get_tables('solar')
get_meaning('solar', 'alerts')