From 91ae0417f0dc210f2b477755ee832418e137447d Mon Sep 17 00:00:00 2001
From: Zain Hoda <7146154+zainhoda@users.noreply.github.com>
Date: Sat, 29 Jul 2023 09:57:49 -0400
Subject: [PATCH 1/3] filtering the training plan
---
src/vanna/__init__.py | 128 ++++++++++++++++++++++++++++++++++++++----
1 file changed, 118 insertions(+), 10 deletions(-)
diff --git a/src/vanna/__init__.py b/src/vanna/__init__.py
index 54ff337c..feb1355b 100644
--- a/src/vanna/__init__.py
+++ b/src/vanna/__init__.py
@@ -583,7 +583,24 @@ def remove_item(self, item: str):
-def get_training_plan() -> TrainingPlan:
+def __get_databases() -> List[str]:
+ try:
+ df_databases = run_sql("SELECT * FROM INFORMATION_SCHEMA.DATABASES")
+ except:
+ try:
+ df_databases = run_sql("SHOW DATABASES")
+ except:
+ return []
+
+ return df_databases['DATABASE_NAME'].unique().tolist()
+
+def __get_information_schema_tables(database: str) -> pd.DataFrame:
+ df_tables = run_sql(f'SELECT * FROM {database}.INFORMATION_SCHEMA.TABLES')
+
+ return df_tables
+
+
+def get_training_plan(filter_databases: Union[List[str], None] = None, filter_schemas: Union[List[str], None] = None, include_information_schema: bool = False, use_historical_queries: bool = True) -> TrainingPlan:
"""
**Example:**
```python
@@ -591,20 +608,111 @@ def get_training_plan() -> TrainingPlan:
vn.train(plan=plan)
```
+ """
- Get the training plan for the model.
+ plan = TrainingPlan([])
- Returns:
- TrainingPlan: The training plan for the model.
- """
- d = __rpc_call(method="get_training_plan", params=[])
+ if run_sql is None:
+ raise ValidationError("Please connect to a database first.")
- if 'result' not in d:
- raise ValidationError("Failed to get training plan")
+ databases = __get_databases()
+
+ for database in databases:
+ if filter_databases is not None and database not in filter_databases:
+ continue
+
+ try:
+ df_tables = __get_information_schema_tables(database=database)
+
+ print(f"Trying INFORMATION_SCHEMA.COLUMNS for {database}")
+ df_columns = run_sql(f"SELECT * FROM {database}.INFORMATION_SCHEMA.COLUMNS")
+
+ for schema in df_tables['TABLE_SCHEMA'].unique().tolist():
+ if filter_schemas is not None and schema not in filter_schemas:
+ continue
- training_plan = TrainingPlan(**d['result'])
+ if not include_information_schema and schema == "INFORMATION_SCHEMA":
+ continue
+
+ df_columns_filtered_to_schema = df_columns.query(f"TABLE_SCHEMA == '{schema}'")
+
+ try:
+ tables = df_columns_filtered_to_schema['TABLE_NAME'].unique().tolist()
+
+ for table in tables:
+ df_columns_filtered_to_table = df_columns_filtered_to_schema.query(f"TABLE_NAME == '{table}'")
+ doc = f"The following columns are in the {table} table in the {database} database:\n\n"
+ doc += df_columns_filtered_to_table[["TABLE_CATALOG", "TABLE_SCHEMA", "TABLE_NAME", "COLUMN_NAME", "DATA_TYPE", "COMMENT"]].to_markdown()
+
+ plan._plan.append(TrainingPlanItem(
+ item_type=TrainingPlanItem.ITEM_TYPE_IS,
+ item_group=f"{database}.{schema}",
+ item_name=table,
+ item_value=doc
+ ))
+
+ except Exception as e:
+ print(e)
+ pass
+ except Exception as e:
+ print(e)
+
+
+ # try:
+ # print("Trying SHOW TABLES")
+ # df_f = run_sql("SHOW TABLES")
+
+ # for schema in df_f.schema_name.unique():
+ # try:
+ # print(f"Trying GET_DDL for {schema}")
+ # ddl_df = run_sql(f"SELECT GET_DDL('schema', '{schema}')")
+
+ # plan._plan.append(TrainingPlanItem(
+ # item_type=TrainingPlanItem.ITEM_TYPE_DDL,
+ # item_group=schema,
+ # item_name="All Tables",
+ # item_value=ddl_df.iloc[0, 0]
+ # ))
+ # except:
+ # pass
+ # except:
+ # try:
+ # print("Trying INFORMATION_SCHEMA.TABLES")
+ # df = run_sql("SELECT * FROM INFORMATION_SCHEMA.TABLES")
+
+ # breakpoint()
+
+ # try:
+ # print("Trying SCHEMATA")
+ # df_schemata = run_sql("SELECT * FROM region-us.INFORMATION_SCHEMA.SCHEMATA")
+
+ # for schema in df_schemata.schema_name.unique():
+ # df = run_sql(f"SELECT * FROM {schema}.information_schema.tables")
+
+ # for table in df.table_name.unique():
+ # plan._plan.append(TrainingPlanItem(
+ # item_type=TrainingPlanItem.ITEM_TYPE_IS,
+ # item_group=schema,
+ # item_name=table,
+ # item_value=None
+ # ))
+
+ # try:
+ # ddl_df = run_sql(f"SELECT GET_DDL('schema', '{schema}')")
+
+ # plan._plan.append(TrainingPlanItem(
+ # item_type=TrainingPlanItem.ITEM_TYPE_DDL,
+ # item_group=schema,
+ # item_name=None,
+ # item_value=ddl_df.iloc[0, 0]
+ # ))
+ # except:
+ # pass
+ # except:
+ # pass
+
+ return plan
- return training_plan
def train(question: str = None, sql: str = None, ddl: str = None, documentation: str = None, json_file: str = None,
sql_file: str = None) -> bool:
From 6f9eb89adadf37f3ba66a93313388078898355e2 Mon Sep 17 00:00:00 2001
From: Zain Hoda <7146154+zainhoda@users.noreply.github.com>
Date: Sat, 29 Jul 2023 16:33:32 -0400
Subject: [PATCH 2/3] training plan
---
pyproject.toml | 2 +-
src/vanna/__init__.py | 57 +++++++++++++++++++++++++++++++++++++------
2 files changed, 51 insertions(+), 8 deletions(-)
diff --git a/pyproject.toml b/pyproject.toml
index 491fa5f6..47c8c94e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
[project]
name = "vanna"
-version = "0.0.16"
+version = "0.0.17"
authors = [
{ name="Zain Hoda", email="zain@vanna.ai" },
]
diff --git a/src/vanna/__init__.py b/src/vanna/__init__.py
index feb1355b..78459544 100644
--- a/src/vanna/__init__.py
+++ b/src/vanna/__init__.py
@@ -600,11 +600,15 @@ def __get_information_schema_tables(database: str) -> pd.DataFrame:
return df_tables
-def get_training_plan(filter_databases: Union[List[str], None] = None, filter_schemas: Union[List[str], None] = None, include_information_schema: bool = False, use_historical_queries: bool = True) -> TrainingPlan:
+def get_training_plan_experimental(filter_databases: Union[List[str], None] = None, filter_schemas: Union[List[str], None] = None, include_information_schema: bool = False, use_historical_queries: bool = True) -> TrainingPlan:
"""
+ **EXPERIMENTAL** : This method is experimental and may change in future versions.
+
+ Get a training plan based on the metadata in the database. Currently this only works for Snowflake.
+
**Example:**
```python
- plan = vn.get_training_plan()
+ plan = vn.get_training_plan_experimental(filter_databases=["employees"], filter_schemas=["public"])
vn.train(plan=plan)
```
@@ -615,6 +619,31 @@ def get_training_plan(filter_databases: Union[List[str], None] = None, filter_sc
if run_sql is None:
raise ValidationError("Please connect to a database first.")
+ if use_historical_queries:
+ try:
+ print("Trying query history")
+ df_history = run_sql(""" select * from table(information_schema.query_history(result_limit => 5000)) order by start_time""")
+
+ df_history_filtered = df_history.query('ROWS_PRODUCED > 1')
+ if filter_databases is not None:
+ mask = df_history_filtered['QUERY_TEXT'].str.lower().apply(lambda x: any(s in x for s in [s.lower() for s in filter_databases]))
+ df_history_filtered = df_history_filtered[mask]
+
+ if filter_schemas is not None:
+ mask = df_history_filtered['QUERY_TEXT'].str.lower().apply(lambda x: any(s in x for s in [s.lower() for s in filter_schemas]))
+ df_history_filtered = df_history_filtered[mask]
+
+ for query in df_history_filtered.sample(10)['QUERY_TEXT'].unique().tolist():
+ plan._plan.append(TrainingPlanItem(
+ item_type=TrainingPlanItem.ITEM_TYPE_SQL,
+ item_group="",
+ item_name=generate_question(query),
+ item_value=query
+ ))
+
+ except Exception as e:
+ print(e)
+
databases = __get_databases()
for database in databases:
@@ -657,7 +686,6 @@ def get_training_plan(filter_databases: Union[List[str], None] = None, filter_sc
except Exception as e:
print(e)
-
# try:
# print("Trying SHOW TABLES")
# df_f = run_sql("SHOW TABLES")
@@ -715,7 +743,7 @@ def get_training_plan(filter_databases: Union[List[str], None] = None, filter_sc
def train(question: str = None, sql: str = None, ddl: str = None, documentation: str = None, json_file: str = None,
- sql_file: str = None) -> bool:
+ sql_file: str = None, plan: TrainingPlan = None) -> bool:
"""
**Example:**
```python
@@ -728,6 +756,7 @@ def train(question: str = None, sql: str = None, ddl: str = None, documentation:
If you call it with the ddl argument, it's equivalent to [`add_ddl()`][vanna.add_ddl].
If you call it with the documentation argument, it's equivalent to [`add_documentation()`][vanna.add_documentation].
It can also accept a JSON file path or SQL file path to train on a batch of questions and SQL queries or a list of SQL queries respectively.
+ Additionally, you can pass a [`TrainingPlan`][vanna.TrainingPlan] object. Get a training plan with [`vn.get_training_plan()`][vanna.get_training_plan].
Args:
question (str): The question to train on.
@@ -736,6 +765,7 @@ def train(question: str = None, sql: str = None, ddl: str = None, documentation:
json_file (str): The JSON file path.
ddl (str): The DDL statement.
documentation (str): The documentation to train on.
+ plan (TrainingPlan): The training plan to train on.
"""
if question and not sql:
@@ -764,7 +794,7 @@ def train(question: str = None, sql: str = None, ddl: str = None, documentation:
print("Adding Questions And SQLs using file:", json_file)
for question in data:
if not add_sql(question=question['question'], sql=question['answer']):
- logger.warning(f"Not able to add sql for question: {question['question']} from {json_file}")
+ print(f"Not able to add sql for question: {question['question']} from {json_file}")
return False
return True
@@ -784,11 +814,24 @@ def train(question: str = None, sql: str = None, ddl: str = None, documentation:
if add_sql(question=question, sql=statement):
print("SQL added!")
return True
- logger.warning("Not able to add sql.")
+ print("Not able to add sql.")
return False
return False
- # Here we're going to attempt auto training
+ if plan:
+ for item in plan._plan:
+ if item.item_type == TrainingPlanItem.ITEM_TYPE_DDL:
+ if not add_ddl(item.item_value):
+ print(f"Not able to add ddl for {item.item_group}")
+ return False
+ elif item.item_type == TrainingPlanItem.ITEM_TYPE_IS:
+ if not add_documentation(item.item_value):
+ print(f"Not able to add documentation for {item.item_group}.{item.item_name}")
+ return False
+ elif item.item_type == TrainingPlanItem.ITEM_TYPE_SQL:
+ if not add_sql(question=item.item_name, sql=item.item_value):
+ print(f"Not able to add sql for {item.item_group}.{item.item_name}")
+ return False
def flag_sql_for_review(question: str, sql: Union[str, None] = None, error_msg: Union[str, None] = None) -> bool:
From 252fff3b282442ebc1472e63ded679604eedd823 Mon Sep 17 00:00:00 2001
From: Zain Hoda <7146154+zainhoda@users.noreply.github.com>
Date: Sat, 29 Jul 2023 17:58:06 -0400
Subject: [PATCH 3/3] notebooks
---
docs/notebooks/vn-ask.md | 179 ++++++--
docs/notebooks/vn-ask_files/vn-ask_10_2.png | Bin 0 -> 125336 bytes
docs/notebooks/vn-ask_files/vn-ask_11_2.png | Bin 0 -> 78743 bytes
docs/notebooks/vn-ask_files/vn-ask_12_2.png | Bin 0 -> 142409 bytes
docs/notebooks/vn-train.md | 266 ++++++++++++
mkdocs.yml | 6 +-
notebooks/vn-train.ipynb | 452 +++++++++++++++++++-
src/vanna/__init__.py | 2 +-
8 files changed, 855 insertions(+), 50 deletions(-)
create mode 100644 docs/notebooks/vn-ask_files/vn-ask_10_2.png
create mode 100644 docs/notebooks/vn-ask_files/vn-ask_11_2.png
create mode 100644 docs/notebooks/vn-ask_files/vn-ask_12_2.png
diff --git a/docs/notebooks/vn-ask.md b/docs/notebooks/vn-ask.md
index 1249c16d..f6f4bd90 100644
--- a/docs/notebooks/vn-ask.md
+++ b/docs/notebooks/vn-ask.md
@@ -1,10 +1,10 @@
-
+
-This notebook will help you unleash the full potential of AI-powered data analysis at your organization. We'll go through how to "bulk train" Vanna and generate SQL, tables, charts, and explanations, all with minimal code and effort. For more about Vanna, see our [intro blog post](https://medium.com/vanna-ai/intro-to-vanna-a-python-based-ai-sql-co-pilot-218c25b19c6a).
+The following notebook goes through the process of asking questions from your data using Vanna AI. Here we use a demo model that is pre-trained on the [TPC-H dataset](https://docs.snowflake.com/en/user-guide/sample-data-tpch.html) that is available in Snowflake.
-[](https://colab.research.google.com/github/vanna-ai/vanna-py/blob/main/notebooks/vn-full.ipynb)
+[](https://colab.research.google.com/github/vanna-ai/vanna-py/blob/main/notebooks/vn-ask.ipynb)
-[](https://github.com/vanna-ai/vanna-py/blob/main/notebooks/vn-full.ipynb)
+[](https://github.com/vanna-ai/vanna-py/blob/main/notebooks/vn-ask.ipynb)
# Install Vanna
First we install Vanna from [PyPI](https://pypi.org/project/vanna/) and import it.
@@ -20,7 +20,6 @@ Here, we'll also install the Snowflake connector. If you're using a different da
```python
import vanna as vn
import snowflake.connector
-import pandas as pd
```
# Login
@@ -37,7 +36,7 @@ You need to choose a globally unique model name. Try using your company name or
```python
-vn.set_model('my-model') # Enter your dataset name here. This is a globally unique identifier for your dataset.
+vn.set_model('tpc') # Enter your model name here. This is a globally unique identifier for your model.
```
# Set Database Connection
@@ -64,20 +63,105 @@ vn.ask("What are the top 10 customers by sales?")
GROUP BY customer_name
ORDER BY total_sales desc limit 10;
-
- AI-generated follow-up questions:
- What are the countries of the top 10 customers by sales?
- How many orders did each of the top 10 customers place?
- What is the average sales amount per customer in the top 10?
- Can you provide a breakdown of the sales by country for the top 10 customers?
- Who are the top 10 customers in terms of returned parts gross value?
- What are the total sales for each customer in the top 3?
- Can you provide a breakdown of the sales by region for the top customers?
- How many customers are there in each country?
- What is the total revenue for the top 5 countries?
- Can you provide a breakdown of the sales by customer for the top 5 countries?
+
+
+
+
+
+ |
+ CUSTOMER_NAME |
+ TOTAL_SALES |
+
+
+
+
+ | 0 |
+ Customer#000143500 |
+ 6757566.0218 |
+
+
+ | 1 |
+ Customer#000095257 |
+ 6294115.3340 |
+
+
+ | 2 |
+ Customer#000087115 |
+ 6184649.5176 |
+
+
+ | 3 |
+ Customer#000131113 |
+ 6080943.8305 |
+
+
+ | 4 |
+ Customer#000134380 |
+ 6075141.9635 |
+
+
+ | 5 |
+ Customer#000103834 |
+ 6059770.3232 |
+
+
+ | 6 |
+ Customer#000069682 |
+ 6057779.0348 |
+
+
+ | 7 |
+ Customer#000102022 |
+ 6039653.6335 |
+
+
+ | 8 |
+ Customer#000098587 |
+ 6027021.5855 |
+
+
+ | 9 |
+ Customer#000064660 |
+ 5905659.6159 |
+
+
+
+