vanna 0.0.27__tar.gz → 0.0.29__tar.gz

vanna-0.0.29/PKG-INFO

@@ -0,0 +1,245 @@
+Metadata-Version: 2.1
+Name: vanna
+Version: 0.0.29
+Summary: Generate SQL queries from natural language
+Author-email: Zain Hoda <zain@vanna.ai>
+Requires-Python: >=3.7
+Description-Content-Type: text/markdown
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Dist: requests
+Requires-Dist: tabulate
+Requires-Dist: plotly
+Requires-Dist: pandas
+Requires-Dist: sqlparse
+Requires-Dist: kaleido
+Requires-Dist: psycopg2 ; extra == "all"
+Requires-Dist: db-dtypes ; extra == "all"
+Requires-Dist: google-cloud-bigquery ; extra == "all"
+Requires-Dist: snowflake-connector-python ; extra == "all"
+Requires-Dist: google-cloud-bigquery ; extra == "bigquery"
+Requires-Dist: chromadb ; extra == "chromadb"
+Requires-Dist: openai ; extra == "openai"
+Requires-Dist: psycopg2 ; extra == "postgres"
+Requires-Dist: db-dtypes ; extra == "postgres"
+Requires-Dist: snowflake-connector-python ; extra == "snowflake"
+Requires-Dist: tox ; extra == "test"
+Project-URL: Bug Tracker, https://github.com/vanna-ai/vanna/issues
+Project-URL: Homepage, https://github.com/vanna-ai/vanna
+Provides-Extra: all
+Provides-Extra: bigquery
+Provides-Extra: chromadb
+Provides-Extra: openai
+Provides-Extra: postgres
+Provides-Extra: snowflake
+Provides-Extra: test
+
+![](https://img.vanna.ai/vanna-github.svg)
+
+| GitHub | PyPI | Documentation |
+| ------ | ---- | ------------- |
+| [![GitHub](https://img.shields.io/badge/GitHub-vanna-blue?logo=github)](https://github.com/vanna-ai/vanna) | [![PyPI](https://img.shields.io/pypi/v/vanna?logo=pypi)](https://pypi.org/project/vanna/) | [![Documentation](https://img.shields.io/badge/Documentation-vanna-blue?logo=read-the-docs)](https://vanna.ai/docs/) |
+
+# Vanna
+Vanna is an MIT-licensed open-source Python RAG (Retrieval-Augmented Generation) framework for SQL generation and related functionality.
+
+https://github.com/vanna-ai/vanna/assets/7146154/1901f47a-515d-4982-af50-f12761a3b2ce
+
+## How Vanna works
+Vanna works in two easy steps - train a RAG "model" on your data, and then ask questions which will return SQL queries that can be set up to automatically run on your database.
+
+1. **Train a RAG "model" on your data**.
+2. **Ask questions**.
+
+![](img/vanna-readme-diagram.png)
+
+If you don't know what RAG is, don't worry -- you don't need to know how this works under the hood to use it. You just need to know that you "train" a model, which stores some metadata and then use it to "ask" questions.
+
+See the [base class](src/vanna/base/base.py) for more details on how this works under the hood.
+
+## User Interfaces
+These are some of the user interfaces that we've built using Vanna. You can use these as-is or as a starting point for your own custom interface.
+
+- [Jupyter Notebook](https://github.com/vanna-ai/vanna/blob/main/notebooks/getting-started.ipynb)
+- [vanna-ai/vanna-streamlit](https://github.com/vanna-ai/vanna-streamlit)
+- [vanna-ai/vanna-flask](https://github.com/vanna-ai/vanna-flask)
+- [vanna-ai/vanna-slack](https://github.com/vanna-ai/vanna-slack)
+
+
+## Getting started
+See the [documentation](https://vanna.ai/docs/) for specifics on your desired database, LLM, etc.
+
+If you want to get a feel for how it works after training, you can try this [Colab notebook](https://colab.research.google.com/github/vanna-ai/vanna/blob/main/notebooks/getting-started.ipynb).
+
+
+### Install
+```bash
+pip install vanna
+```
+
+There are a number of optional packages that can be installed so see the [documentation](https://vanna.ai/docs/) for more details.
+
+### Import
+See the [documentation](https://vanna.ai/docs/) if you're customizing the LLM or vector database.
+
+```python
+import vanna as vn
+```
+
+
+## Training
+You may or may not need to run these `vn.train` commands depending on your use case. See the [documentation](https://vanna.ai/docs/) for more details.
+
+These statements are shown to give you a feel for how it works.
+
+### Train with DDL Statements
+DDL statements contain information about the table names, columns, data types, and relationships in your database.
+
+```python
+vn.train(ddl="""
+    CREATE TABLE IF NOT EXISTS my-table (
+        id INT PRIMARY KEY,
+        name VARCHAR(100),
+        age INT
+    )
+""")
+```
+
+### Train with Documentation
+Sometimes you may want to add documentation about your business terminology or definitions.
+
+```python
+vn.train(documentation="Our business defines XYZ as ...")
+```
+
+### Train with SQL
+You can also add SQL queries to your training data. This is useful if you have some queries already laying around. You can just copy and paste those from your editor to begin generating new SQL.
+
+```python
+vn.train(sql="SELECT name, age FROM my-table WHERE name = 'John Doe'")
+```
+
+
+## Asking questions
+```python
+vn.ask("What are the top 10 customers by sales?")
+```
+
+You'll get SQL
+```sql
+SELECT c.c_name as customer_name,
+        sum(l.l_extendedprice * (1 - l.l_discount)) as total_sales
+FROM   snowflake_sample_data.tpch_sf1.lineitem l join snowflake_sample_data.tpch_sf1.orders o
+        ON l.l_orderkey = o.o_orderkey join snowflake_sample_data.tpch_sf1.customer c
+        ON o.o_custkey = c.c_custkey
+GROUP BY customer_name
+ORDER BY total_sales desc limit 10;
+```
+
+If you've connected to a database, you'll get the table:
+<div>
+<table border="1" class="dataframe">
+  <thead>
+    <tr style="text-align: right;">
+      <th></th>
+      <th>CUSTOMER_NAME</th>
+      <th>TOTAL_SALES</th>
+    </tr>
+  </thead>
+  <tbody>
+    <tr>
+      <th>0</th>
+      <td>Customer#000143500</td>
+      <td>6757566.0218</td>
+    </tr>
+    <tr>
+      <th>1</th>
+      <td>Customer#000095257</td>
+      <td>6294115.3340</td>
+    </tr>
+    <tr>
+      <th>2</th>
+      <td>Customer#000087115</td>
+      <td>6184649.5176</td>
+    </tr>
+    <tr>
+      <th>3</th>
+      <td>Customer#000131113</td>
+      <td>6080943.8305</td>
+    </tr>
+    <tr>
+      <th>4</th>
+      <td>Customer#000134380</td>
+      <td>6075141.9635</td>
+    </tr>
+    <tr>
+      <th>5</th>
+      <td>Customer#000103834</td>
+      <td>6059770.3232</td>
+    </tr>
+    <tr>
+      <th>6</th>
+      <td>Customer#000069682</td>
+      <td>6057779.0348</td>
+    </tr>
+    <tr>
+      <th>7</th>
+      <td>Customer#000102022</td>
+      <td>6039653.6335</td>
+    </tr>
+    <tr>
+      <th>8</th>
+      <td>Customer#000098587</td>
+      <td>6027021.5855</td>
+    </tr>
+    <tr>
+      <th>9</th>
+      <td>Customer#000064660</td>
+      <td>5905659.6159</td>
+    </tr>
+  </tbody>
+</table>
+</div>
+
+You'll also get an automated Plotly chart:
+![](img/top-10-customers.png)
+
+## RAG vs. Fine-Tuning
+RAG
+- Portable across LLMs
+- Easy to remove training data if any of it becomes obsolete
+- Much cheaper to run than fine-tuning
+- More future-proof -- if a better LLM comes out, you can just swap it out
+
+Fine-Tuning
+- Good if you need to minimize tokens in the prompt
+- Slow to get started
+- Expensive to train and run (generally)
+
+## Why Vanna?
+
+1. **High accuracy on complex datasets.**
+    - Vanna’s capabilities are tied to the training data you give it
+    - More training data means better accuracy for large and complex datasets
+2. **Secure and private.**
+    - Your database contents are never sent to the LLM or the vector database
+    - SQL execution happens in your local environment
+3. **Self learning.**
+    - If using via Jupyter, you can choose to "auto-train" it on the queries that were successfully executed
+    - If using via other interfaces, you can have the interface prompt the user to provide feedback on the results
+    - Correct question to SQL pairs are stored for future reference and make the future results more accurate
+4. **Supports any SQL database.**
+    - The package allows you to connect to any SQL database that you can otherwise connect to with Python
+5. **Choose your front end.**
+    - Most people start in a Jupyter Notebook.
+    - Expose to your end users via Slackbot, web app, Streamlit app, or a custom front end.
+
+## Extending Vanna
+Vanna is designed to connect to any database, LLM, and vector database. There's a [VannaBase](src/vanna/base/base.py) abstract base class that defines some basic functionality. The package provides implementations for use with OpenAI and ChromaDB. You can easily extend Vanna to use your own LLM or vector database. See the [documentation](https://vanna.ai/docs/) for more details.
+
+## More resources
+ - [Full Documentation](https://vanna.ai/docs/)
+ - [Website](https://vanna.ai)
+ - [Discord group for support](https://discord.gg/qUZYKHremx)
+

vanna-0.0.29/README.md

@@ -0,0 +1,207 @@
+![](https://img.vanna.ai/vanna-github.svg)
+
+| GitHub | PyPI | Documentation |
+| ------ | ---- | ------------- |
+| [![GitHub](https://img.shields.io/badge/GitHub-vanna-blue?logo=github)](https://github.com/vanna-ai/vanna) | [![PyPI](https://img.shields.io/pypi/v/vanna?logo=pypi)](https://pypi.org/project/vanna/) | [![Documentation](https://img.shields.io/badge/Documentation-vanna-blue?logo=read-the-docs)](https://vanna.ai/docs/) |
+
+# Vanna
+Vanna is an MIT-licensed open-source Python RAG (Retrieval-Augmented Generation) framework for SQL generation and related functionality.
+
+https://github.com/vanna-ai/vanna/assets/7146154/1901f47a-515d-4982-af50-f12761a3b2ce
+
+## How Vanna works
+Vanna works in two easy steps - train a RAG "model" on your data, and then ask questions which will return SQL queries that can be set up to automatically run on your database.
+
+1. **Train a RAG "model" on your data**.
+2. **Ask questions**.
+
+![](img/vanna-readme-diagram.png)
+
+If you don't know what RAG is, don't worry -- you don't need to know how this works under the hood to use it. You just need to know that you "train" a model, which stores some metadata and then use it to "ask" questions.
+
+See the [base class](src/vanna/base/base.py) for more details on how this works under the hood.
+
+## User Interfaces
+These are some of the user interfaces that we've built using Vanna. You can use these as-is or as a starting point for your own custom interface.
+
+- [Jupyter Notebook](https://github.com/vanna-ai/vanna/blob/main/notebooks/getting-started.ipynb)
+- [vanna-ai/vanna-streamlit](https://github.com/vanna-ai/vanna-streamlit)
+- [vanna-ai/vanna-flask](https://github.com/vanna-ai/vanna-flask)
+- [vanna-ai/vanna-slack](https://github.com/vanna-ai/vanna-slack)
+
+
+## Getting started
+See the [documentation](https://vanna.ai/docs/) for specifics on your desired database, LLM, etc.
+
+If you want to get a feel for how it works after training, you can try this [Colab notebook](https://colab.research.google.com/github/vanna-ai/vanna/blob/main/notebooks/getting-started.ipynb).
+
+
+### Install
+```bash
+pip install vanna
+```
+
+There are a number of optional packages that can be installed so see the [documentation](https://vanna.ai/docs/) for more details.
+
+### Import
+See the [documentation](https://vanna.ai/docs/) if you're customizing the LLM or vector database.
+
+```python
+import vanna as vn
+```
+
+
+## Training
+You may or may not need to run these `vn.train` commands depending on your use case. See the [documentation](https://vanna.ai/docs/) for more details.
+
+These statements are shown to give you a feel for how it works.
+
+### Train with DDL Statements
+DDL statements contain information about the table names, columns, data types, and relationships in your database.
+
+```python
+vn.train(ddl="""
+    CREATE TABLE IF NOT EXISTS my-table (
+        id INT PRIMARY KEY,
+        name VARCHAR(100),
+        age INT
+    )
+""")
+```
+
+### Train with Documentation
+Sometimes you may want to add documentation about your business terminology or definitions.
+
+```python
+vn.train(documentation="Our business defines XYZ as ...")
+```
+
+### Train with SQL
+You can also add SQL queries to your training data. This is useful if you have some queries already laying around. You can just copy and paste those from your editor to begin generating new SQL.
+
+```python
+vn.train(sql="SELECT name, age FROM my-table WHERE name = 'John Doe'")
+```
+
+
+## Asking questions
+```python
+vn.ask("What are the top 10 customers by sales?")
+```
+
+You'll get SQL
+```sql
+SELECT c.c_name as customer_name,
+        sum(l.l_extendedprice * (1 - l.l_discount)) as total_sales
+FROM   snowflake_sample_data.tpch_sf1.lineitem l join snowflake_sample_data.tpch_sf1.orders o
+        ON l.l_orderkey = o.o_orderkey join snowflake_sample_data.tpch_sf1.customer c
+        ON o.o_custkey = c.c_custkey
+GROUP BY customer_name
+ORDER BY total_sales desc limit 10;
+```
+
+If you've connected to a database, you'll get the table:
+<div>
+<table border="1" class="dataframe">
+  <thead>
+    <tr style="text-align: right;">
+      <th></th>
+      <th>CUSTOMER_NAME</th>
+      <th>TOTAL_SALES</th>
+    </tr>
+  </thead>
+  <tbody>
+    <tr>
+      <th>0</th>
+      <td>Customer#000143500</td>
+      <td>6757566.0218</td>
+    </tr>
+    <tr>
+      <th>1</th>
+      <td>Customer#000095257</td>
+      <td>6294115.3340</td>
+    </tr>
+    <tr>
+      <th>2</th>
+      <td>Customer#000087115</td>
+      <td>6184649.5176</td>
+    </tr>
+    <tr>
+      <th>3</th>
+      <td>Customer#000131113</td>
+      <td>6080943.8305</td>
+    </tr>
+    <tr>
+      <th>4</th>
+      <td>Customer#000134380</td>
+      <td>6075141.9635</td>
+    </tr>
+    <tr>
+      <th>5</th>
+      <td>Customer#000103834</td>
+      <td>6059770.3232</td>
+    </tr>
+    <tr>
+      <th>6</th>
+      <td>Customer#000069682</td>
+      <td>6057779.0348</td>
+    </tr>
+    <tr>
+      <th>7</th>
+      <td>Customer#000102022</td>
+      <td>6039653.6335</td>
+    </tr>
+    <tr>
+      <th>8</th>
+      <td>Customer#000098587</td>
+      <td>6027021.5855</td>
+    </tr>
+    <tr>
+      <th>9</th>
+      <td>Customer#000064660</td>
+      <td>5905659.6159</td>
+    </tr>
+  </tbody>
+</table>
+</div>
+
+You'll also get an automated Plotly chart:
+![](img/top-10-customers.png)
+
+## RAG vs. Fine-Tuning
+RAG
+- Portable across LLMs
+- Easy to remove training data if any of it becomes obsolete
+- Much cheaper to run than fine-tuning
+- More future-proof -- if a better LLM comes out, you can just swap it out
+
+Fine-Tuning
+- Good if you need to minimize tokens in the prompt
+- Slow to get started
+- Expensive to train and run (generally)
+
+## Why Vanna?
+
+1. **High accuracy on complex datasets.**
+    - Vanna’s capabilities are tied to the training data you give it
+    - More training data means better accuracy for large and complex datasets
+2. **Secure and private.**
+    - Your database contents are never sent to the LLM or the vector database
+    - SQL execution happens in your local environment
+3. **Self learning.**
+    - If using via Jupyter, you can choose to "auto-train" it on the queries that were successfully executed
+    - If using via other interfaces, you can have the interface prompt the user to provide feedback on the results
+    - Correct question to SQL pairs are stored for future reference and make the future results more accurate
+4. **Supports any SQL database.**
+    - The package allows you to connect to any SQL database that you can otherwise connect to with Python
+5. **Choose your front end.**
+    - Most people start in a Jupyter Notebook.
+    - Expose to your end users via Slackbot, web app, Streamlit app, or a custom front end.
+
+## Extending Vanna
+Vanna is designed to connect to any database, LLM, and vector database. There's a [VannaBase](src/vanna/base/base.py) abstract base class that defines some basic functionality. The package provides implementations for use with OpenAI and ChromaDB. You can easily extend Vanna to use your own LLM or vector database. See the [documentation](https://vanna.ai/docs/) for more details.
+
+## More resources
+ - [Full Documentation](https://vanna.ai/docs/)
+ - [Website](https://vanna.ai)
+ - [Discord group for support](https://discord.gg/qUZYKHremx)

{vanna-0.0.27 → vanna-0.0.29}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "vanna"
-version = "0.0.27"
+version = "0.0.29"
 authors = [
   { name="Zain Hoda", email="zain@vanna.ai" },
 ]

{vanna-0.0.27 → vanna-0.0.29}/src/vanna/__init__.py

@@ -1065,7 +1065,7 @@ def train(
     if sql:
         if question is None:
             question = generate_question(sql)
-            print("Question generated with sql:", Question, "\nAdding SQL...")
+            print("Question generated with sql:", question, "\nAdding SQL...")
         return add_sql(question=question, sql=sql)
 
     if ddl:

{vanna-0.0.27 → vanna-0.0.29}/src/vanna/base/base.py

@@ -612,6 +612,65 @@ class VannaBase(ABC):
 
         return df_tables
 
+    def get_training_plan_generic(self, df) -> TrainingPlan:
+        # For each of the following, we look at the df columns to see if there's a match:
+        database_column = df.columns[
+            df.columns.str.lower().str.contains("database")
+            | df.columns.str.lower().str.contains("table_catalog")
+        ].to_list()[0]
+        schema_column = df.columns[
+            df.columns.str.lower().str.contains("table_schema")
+        ].to_list()[0]
+        table_column = df.columns[
+            df.columns.str.lower().str.contains("table_name")
+        ].to_list()[0]
+        column_column = df.columns[
+            df.columns.str.lower().str.contains("column_name")
+        ].to_list()[0]
+        data_type_column = df.columns[
+            df.columns.str.lower().str.contains("data_type")
+        ].to_list()[0]
+
+        plan = TrainingPlan([])
+
+        for database in df[database_column].unique().tolist():
+            for schema in (
+                df.query(f'{database_column} == "{database}"')[schema_column]
+                .unique()
+                .tolist()
+            ):
+                for table in (
+                    df.query(
+                        f'{database_column} == "{database}" and {schema_column} == "{schema}"'
+                    )[table_column]
+                    .unique()
+                    .tolist()
+                ):
+                    df_columns_filtered_to_table = df.query(
+                        f'{database_column} == "{database}" and {schema_column} == "{schema}" and {table_column} == "{table}"'
+                    )
+                    doc = f"The following columns are in the {table} table in the {database} database:\n\n"
+                    doc += df_columns_filtered_to_table[
+                        [
+                            database_column,
+                            schema_column,
+                            table_column,
+                            column_column,
+                            data_type_column,
+                        ]
+                    ].to_markdown()
+
+                    plan._plan.append(
+                        TrainingPlanItem(
+                            item_type=TrainingPlanItem.ITEM_TYPE_IS,
+                            item_group=f"{database}.{schema}",
+                            item_name=table,
+                            item_value=doc,
+                        )
+                    )
+
+        return plan
+
     def get_training_plan_snowflake(
         self,
         filter_databases: Union[List[str], None] = None,

{vanna-0.0.27 → vanna-0.0.29}/src/vanna/openai/openai_chat.py

@@ -204,7 +204,7 @@ class OpenAI_Chat(VannaBase):
             print(
                 f"Using engine {self.config['engine']} for {num_tokens} tokens (approx)"
             )
-            response = openai.ChatCompletion.create(
+            response = openai.chat.completions.create(
                 engine=self.config["engine"],
                 messages=prompt,
                 max_tokens=500,
@@ -215,7 +215,7 @@ class OpenAI_Chat(VannaBase):
             print(
                 f"Using model {self.config['model']} for {num_tokens} tokens (approx)"
             )
-            response = openai.ChatCompletion.create(
+            response = openai.chat.completions.create(
                 model=self.config["model"],
                 messages=prompt,
                 max_tokens=500,
@@ -229,7 +229,7 @@ class OpenAI_Chat(VannaBase):
                 model = "gpt-3.5-turbo"
 
             print(f"Using model {model} for {num_tokens} tokens (approx)")
-            response = openai.ChatCompletion.create(
+            response = openai.chat.completions.create(
                 model=model, messages=prompt, max_tokens=500, stop=None, temperature=0.7
             )
 

{vanna-0.0.27 → vanna-0.0.29}/src/vanna/openai/openai_embeddings.py

@@ -1,6 +1,6 @@
 from abc import abstractmethod
 
-import openai
+from openai import OpenAI
 
 from ..base import VannaBase
 
@@ -9,29 +9,31 @@ class OpenAI_Embeddings(VannaBase):
     def __init__(self, config=None):
         VannaBase.__init__(self, config=config)
 
+        self.client = OpenAI()
+
         if config is None:
             return
 
         if "api_type" in config:
-            openai.api_type = config["api_type"]
+            self.client.api_type = config["api_type"]
 
         if "api_base" in config:
-            openai.api_base = config["api_base"]
+            self.client.api_base = config["api_base"]
 
         if "api_version" in config:
-            openai.api_version = config["api_version"]
+            self.client.api_version = config["api_version"]
 
         if "api_key" in config:
-            openai.api_key = config["api_key"]
+            self.client.api_key = config["api_key"]
 
     def generate_embedding(self, data: str, **kwargs) -> list[float]:
         if self.config is not None and "engine" in self.config:
-            embedding = openai.Embedding.create(
+            embedding = self.client.embeddings.create(
                 engine=self.config["engine"],
                 input=data,
             )
         else:
-            embedding = openai.Embedding.create(
+            embedding = self.client.embeddings.create(
                 model="text-embedding-ada-002",
                 input=data,
             )

{vanna-0.0.27 → vanna-0.0.29}/src/vanna/remote.py

@@ -197,7 +197,7 @@ class VannaDefault(VannaBase):
 
         return status.id
 
-    def add_documentation(self, documentation: str, **kwargs) -> str:
+    def add_documentation(self, doc: str, **kwargs) -> str:
         """
         Adds documentation to the model's training data
 

vanna-0.0.29/src/vanna/vannadb/vannadb_vector.py

@@ -0,0 +1,187 @@
+from ..base import VannaBase
+from ..types import (
+    QuestionSQLPair, 
+    StatusWithId,
+    StringData,
+    DataFrameJSON,
+    Status,
+    TrainingData,
+    Question,
+)
+from io import StringIO
+import pandas as pd
+import requests
+import json
+import dataclasses
+
+class VannaDB_VectorStore(VannaBase):
+    def __init__(self, vanna_model: str, vanna_api_key: str, config=None):
+        VannaBase.__init__(self, config=config)
+
+        self._model = vanna_model
+        self._api_key = vanna_api_key
+
+        self._endpoint = (
+            "https://ask.vanna.ai/rpc"
+            if config is None or "endpoint" not in config
+            else config["endpoint"]
+        )
+        self._unauthenticated_endpoint = (
+            "https://ask.vanna.ai/unauthenticated_rpc"
+            if config is None or "unauthenticated_endpoint" not in config
+            else config["unauthenticated_endpoint"]
+        )
+        self.related_training_data = {}
+
+    def _unauthenticated_rpc_call(self, method, params):
+        headers = {
+            "Content-Type": "application/json",
+        }
+        data = {
+            "method": method,
+            "params": [self._dataclass_to_dict(obj) for obj in params],
+        }
+
+        response = requests.post(
+            self._unauthenticated_endpoint, headers=headers, data=json.dumps(data)
+        )
+        return response.json()
+
+    def _rpc_call(self, method, params):
+        if method != "list_orgs":
+            headers = {
+                "Content-Type": "application/json",
+                "Vanna-Key": self._api_key,
+                "Vanna-Org": self._model,
+            }
+        else:
+            headers = {
+                "Content-Type": "application/json",
+                "Vanna-Key": self._api_key,
+                "Vanna-Org": "demo-tpc-h",
+            }
+
+        data = {
+            "method": method,
+            "params": [self._dataclass_to_dict(obj) for obj in params],
+        }
+
+        response = requests.post(self._endpoint, headers=headers, data=json.dumps(data))
+        return response.json()
+
+    def _dataclass_to_dict(self, obj):
+        return dataclasses.asdict(obj)
+
+    def generate_embedding(self, data: str, **kwargs) -> list[float]:
+        # This is done server-side
+        pass
+
+    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
+        if "tag" in kwargs:
+            tag = kwargs["tag"]
+        else:
+            tag = "Manually Trained"
+
+        params = [QuestionSQLPair(question=question, sql=sql, tag=tag)]
+
+        d = self._rpc_call(method="add_sql", params=params)
+
+        if "result" not in d:
+            raise Exception("Error adding question and SQL pair", d)
+
+        status = StatusWithId(**d["result"])
+
+        return status.id
+
+    def add_ddl(self, ddl: str, **kwargs) -> str:
+        params = [StringData(data=ddl)]
+
+        d = self._rpc_call(method="add_ddl", params=params)
+
+        if "result" not in d:
+            raise Exception("Error adding DDL", d)
+
+        status = StatusWithId(**d["result"])
+
+        return status.id
+
+    def add_documentation(self, doc: str, **kwargs) -> str:
+        params = [StringData(data=doc)]
+
+        d = self._rpc_call(method="add_documentation", params=params)
+
+        if "result" not in d:
+            raise Exception("Error adding documentation", d)
+
+        status = StatusWithId(**d["result"])
+
+        return status.id
+
+    def get_training_data(self, **kwargs) -> pd.DataFrame:
+        params = []
+
+        d = self._rpc_call(method="get_training_data", params=params)
+
+        if "result" not in d:
+            return None
+
+        # Load the result into a dataclass
+        training_data = DataFrameJSON(**d["result"])
+
+        df = pd.read_json(StringIO(training_data.data))
+
+        return df
+
+    def remove_training_data(self, id: str, **kwargs) -> bool:
+        params = [StringData(data=id)]
+
+        d = self._rpc_call(method="remove_training_data", params=params)
+
+        if "result" not in d:
+            raise Exception(f"Error removing training data")
+
+        status = Status(**d["result"])
+
+        if not status.success:
+            raise Exception(f"Error removing training data: {status.message}")
+
+        return status.success
+
+    def get_related_training_data_cached(self, question: str) -> TrainingData:
+        params = [Question(question=question)]
+
+        d = self._rpc_call(method="get_related_training_data", params=params)
+
+        if "result" not in d:
+            return None
+
+        # Load the result into a dataclass
+        training_data = TrainingData(**d["result"])
+
+        self.related_training_data[question] = training_data
+
+        return training_data
+
+    def get_similar_question_sql(self, question: str, **kwargs) -> list:
+        if question in self.related_training_data:
+            training_data = self.related_training_data[question]
+        else:
+            training_data = self.get_related_training_data_cached(question)
+        
+        return training_data.questions
+
+    def get_related_ddl(self, question: str, **kwargs) -> list:
+        if question in self.related_training_data:
+            training_data = self.related_training_data[question]
+        else:
+            training_data = self.get_related_training_data_cached(question)
+        
+        return training_data.ddl
+
+    def get_related_documentation(self, question: str, **kwargs) -> list:
+        if question in self.related_training_data:
+            training_data = self.related_training_data[question]
+        else:
+            training_data = self.get_related_training_data_cached(question)
+        
+        return training_data.documentation

vanna-0.0.27/PKG-INFO

@@ -1,213 +0,0 @@
-Metadata-Version: 2.1
-Name: vanna
-Version: 0.0.27
-Summary: Generate SQL queries from natural language
-Author-email: Zain Hoda <zain@vanna.ai>
-Requires-Python: >=3.7
-Description-Content-Type: text/markdown
-Classifier: Programming Language :: Python :: 3
-Classifier: License :: OSI Approved :: MIT License
-Classifier: Operating System :: OS Independent
-Requires-Dist: requests
-Requires-Dist: tabulate
-Requires-Dist: plotly
-Requires-Dist: pandas
-Requires-Dist: sqlparse
-Requires-Dist: kaleido
-Requires-Dist: psycopg2 ; extra == "all"
-Requires-Dist: db-dtypes ; extra == "all"
-Requires-Dist: google-cloud-bigquery ; extra == "all"
-Requires-Dist: snowflake-connector-python ; extra == "all"
-Requires-Dist: google-cloud-bigquery ; extra == "bigquery"
-Requires-Dist: chromadb ; extra == "chromadb"
-Requires-Dist: openai ; extra == "openai"
-Requires-Dist: psycopg2 ; extra == "postgres"
-Requires-Dist: db-dtypes ; extra == "postgres"
-Requires-Dist: snowflake-connector-python ; extra == "snowflake"
-Requires-Dist: tox ; extra == "test"
-Project-URL: Bug Tracker, https://github.com/vanna-ai/vanna/issues
-Project-URL: Homepage, https://github.com/vanna-ai/vanna
-Provides-Extra: all
-Provides-Extra: bigquery
-Provides-Extra: chromadb
-Provides-Extra: openai
-Provides-Extra: postgres
-Provides-Extra: snowflake
-Provides-Extra: test
-
-![](https://img.vanna.ai/vanna-github.svg)
-
-| GitHub | PyPI | Documentation |
-| ------ | ---- | ------------- |
-| [![GitHub](https://img.shields.io/badge/GitHub-vanna-blue?logo=github)](https://github.com/vanna-ai/vanna) | [![PyPI](https://img.shields.io/pypi/v/vanna?logo=pypi)](https://pypi.org/project/vanna/) | [![Documentation](https://img.shields.io/badge/Documentation-vanna-blue?logo=read-the-docs)](https://vanna.ai/docs/) |
-
-# Vanna.AI - Personalized AI SQL Agent
-
-https://github.com/vanna-ai/vanna/assets/7146154/1901f47a-515d-4982-af50-f12761a3b2ce
-
-## How Vanna works
-Vanna works in two easy steps - train a model on your data, and then ask questions.
-
-1. **Train a model on your data**.
-2. **Ask questions**.
-
-When you ask a question, we utilize a custom model for your dataset to generate SQL, as seen below. Your model performance and accuracy depends on the quality and quantity of training data you use to train your model.
-<img width="1725" alt="how-vanna-works" src="https://github.com/vanna-ai/vanna/assets/7146154/5e2e2179-ed7a-4df4-92a2-1c017923a675">
-
-## Getting started
-You can start by [automatically training Vanna (currently works for Snowflake)](https://vanna.ai/docs/vn-train.html) or add manual training data.
-
-### Install Vanna
-```
-pip install vanna
-```
-
-Depending on the database you're using, you can also install the associated database drivers
-```
-pip install 'vanna[snowflake]'
-```
-
-### Import Vanna
-```python
-import vanna as vn
-```
-
-### Train with DDL Statements
-If you prefer to manually train, you do not need to connect to a database. You can use the train function with other parmaeters like ddl
-
-
-```python
-vn.train(ddl="""
-    CREATE TABLE IF NOT EXISTS my-table (
-        id INT PRIMARY KEY,
-        name VARCHAR(100),
-        age INT
-    )
-""")
-```
-
-### Train with Documentation
-Sometimes you may want to add documentation about your business terminology or definitions.
-
-```python
-vn.train(documentation="Our business defines OTIF score as the percentage of orders that are delivered on time and in full")
-```
-
-### Train with SQL
-You can also add SQL queries to your training data. This is useful if you have some queries already laying around. You can just copy and paste those from your editor to begin generating new SQL.
-
-```python
-vn.train(sql="SELECT * FROM my-table WHERE name = 'John Doe'")
-```
-
-
-## Asking questions
-```python
-vn.ask("What are the top 10 customers by sales?")
-```
-
-    SELECT c.c_name as customer_name,
-           sum(l.l_extendedprice * (1 - l.l_discount)) as total_sales
-    FROM   snowflake_sample_data.tpch_sf1.lineitem l join snowflake_sample_data.tpch_sf1.orders o
-            ON l.l_orderkey = o.o_orderkey join snowflake_sample_data.tpch_sf1.customer c
-            ON o.o_custkey = c.c_custkey
-    GROUP BY customer_name
-    ORDER BY total_sales desc limit 10;
-
-
-
-<div>
-<table border="1" class="dataframe">
-  <thead>
-    <tr style="text-align: right;">
-      <th></th>
-      <th>CUSTOMER_NAME</th>
-      <th>TOTAL_SALES</th>
-    </tr>
-  </thead>
-  <tbody>
-    <tr>
-      <th>0</th>
-      <td>Customer#000143500</td>
-      <td>6757566.0218</td>
-    </tr>
-    <tr>
-      <th>1</th>
-      <td>Customer#000095257</td>
-      <td>6294115.3340</td>
-    </tr>
-    <tr>
-      <th>2</th>
-      <td>Customer#000087115</td>
-      <td>6184649.5176</td>
-    </tr>
-    <tr>
-      <th>3</th>
-      <td>Customer#000131113</td>
-      <td>6080943.8305</td>
-    </tr>
-    <tr>
-      <th>4</th>
-      <td>Customer#000134380</td>
-      <td>6075141.9635</td>
-    </tr>
-    <tr>
-      <th>5</th>
-      <td>Customer#000103834</td>
-      <td>6059770.3232</td>
-    </tr>
-    <tr>
-      <th>6</th>
-      <td>Customer#000069682</td>
-      <td>6057779.0348</td>
-    </tr>
-    <tr>
-      <th>7</th>
-      <td>Customer#000102022</td>
-      <td>6039653.6335</td>
-    </tr>
-    <tr>
-      <th>8</th>
-      <td>Customer#000098587</td>
-      <td>6027021.5855</td>
-    </tr>
-    <tr>
-      <th>9</th>
-      <td>Customer#000064660</td>
-      <td>5905659.6159</td>
-    </tr>
-  </tbody>
-</table>
-</div>
-
-## Why Vanna?
-
-1. **High accuracy on complex datasets.**
-    - Vanna’s capabilities are tied to the training data you give it
-    - More training data means better accuracy for large and complex datasets
-2. **Secure and private.**
-    - Your database contents  are never sent to Vanna’s servers
-    - We only see the bare minimum - schemas & queries.
-3. **Isolated, custom model.**
-    - You train a custom model specific to your database and your schema.
-    - Nobody else can use your model or view your model’s training data unless you choose to add members to your model or make it public
-    - We use a combination of third-party foundational models (OpenAI, Google) and our own LLM.
-4. **Self learning.**
-    - As you use Vanna more, your model continuously improves as we augment your training data
-5. **Supports many databases.**
-    - We have out-of-the-box support Snowflake, BigQuery, Postgres
-    - You can easily make a connector for any [database](https://docs.vanna.ai/databases/)
-6. **Pretrained models.**
-    - If you’re a data provider you can publish your models for anyone to use
-    - As part of our roadmap, we are in the process of pre-training models for common datasets (Google Ads, Facebook ads, etc)
-7. **Choose your front end.**
-    - Start in a Jupyter Notebook.
-    - Expose to business users via Slackbot, web app, Streamlit app, or Excel plugin.
-    - Even integrate in your web app for customers.
-
-## More resources
- - [Full Documentation](https://vanna.ai/docs/)
- - [Website](https://vanna.ai)
- - [Slack channel for support](https://join.slack.com/t/vanna-ai/shared_invite/zt-1unu0ipog-iE33QCoimQiBDxf2o7h97w)
- - [LinkedIn](https://www.linkedin.com/company/vanna-ai/)
-

vanna-0.0.27/README.md

@@ -1,175 +0,0 @@
-![](https://img.vanna.ai/vanna-github.svg)
-
-| GitHub | PyPI | Documentation |
-| ------ | ---- | ------------- |
-| [![GitHub](https://img.shields.io/badge/GitHub-vanna-blue?logo=github)](https://github.com/vanna-ai/vanna) | [![PyPI](https://img.shields.io/pypi/v/vanna?logo=pypi)](https://pypi.org/project/vanna/) | [![Documentation](https://img.shields.io/badge/Documentation-vanna-blue?logo=read-the-docs)](https://vanna.ai/docs/) |
-
-# Vanna.AI - Personalized AI SQL Agent
-
-https://github.com/vanna-ai/vanna/assets/7146154/1901f47a-515d-4982-af50-f12761a3b2ce
-
-## How Vanna works
-Vanna works in two easy steps - train a model on your data, and then ask questions.
-
-1. **Train a model on your data**.
-2. **Ask questions**.
-
-When you ask a question, we utilize a custom model for your dataset to generate SQL, as seen below. Your model performance and accuracy depends on the quality and quantity of training data you use to train your model.
-<img width="1725" alt="how-vanna-works" src="https://github.com/vanna-ai/vanna/assets/7146154/5e2e2179-ed7a-4df4-92a2-1c017923a675">
-
-## Getting started
-You can start by [automatically training Vanna (currently works for Snowflake)](https://vanna.ai/docs/vn-train.html) or add manual training data.
-
-### Install Vanna
-```
-pip install vanna
-```
-
-Depending on the database you're using, you can also install the associated database drivers
-```
-pip install 'vanna[snowflake]'
-```
-
-### Import Vanna
-```python
-import vanna as vn
-```
-
-### Train with DDL Statements
-If you prefer to manually train, you do not need to connect to a database. You can use the train function with other parmaeters like ddl
-
-
-```python
-vn.train(ddl="""
-    CREATE TABLE IF NOT EXISTS my-table (
-        id INT PRIMARY KEY,
-        name VARCHAR(100),
-        age INT
-    )
-""")
-```
-
-### Train with Documentation
-Sometimes you may want to add documentation about your business terminology or definitions.
-
-```python
-vn.train(documentation="Our business defines OTIF score as the percentage of orders that are delivered on time and in full")
-```
-
-### Train with SQL
-You can also add SQL queries to your training data. This is useful if you have some queries already laying around. You can just copy and paste those from your editor to begin generating new SQL.
-
-```python
-vn.train(sql="SELECT * FROM my-table WHERE name = 'John Doe'")
-```
-
-
-## Asking questions
-```python
-vn.ask("What are the top 10 customers by sales?")
-```
-
-    SELECT c.c_name as customer_name,
-           sum(l.l_extendedprice * (1 - l.l_discount)) as total_sales
-    FROM   snowflake_sample_data.tpch_sf1.lineitem l join snowflake_sample_data.tpch_sf1.orders o
-            ON l.l_orderkey = o.o_orderkey join snowflake_sample_data.tpch_sf1.customer c
-            ON o.o_custkey = c.c_custkey
-    GROUP BY customer_name
-    ORDER BY total_sales desc limit 10;
-
-
-
-<div>
-<table border="1" class="dataframe">
-  <thead>
-    <tr style="text-align: right;">
-      <th></th>
-      <th>CUSTOMER_NAME</th>
-      <th>TOTAL_SALES</th>
-    </tr>
-  </thead>
-  <tbody>
-    <tr>
-      <th>0</th>
-      <td>Customer#000143500</td>
-      <td>6757566.0218</td>
-    </tr>
-    <tr>
-      <th>1</th>
-      <td>Customer#000095257</td>
-      <td>6294115.3340</td>
-    </tr>
-    <tr>
-      <th>2</th>
-      <td>Customer#000087115</td>
-      <td>6184649.5176</td>
-    </tr>
-    <tr>
-      <th>3</th>
-      <td>Customer#000131113</td>
-      <td>6080943.8305</td>
-    </tr>
-    <tr>
-      <th>4</th>
-      <td>Customer#000134380</td>
-      <td>6075141.9635</td>
-    </tr>
-    <tr>
-      <th>5</th>
-      <td>Customer#000103834</td>
-      <td>6059770.3232</td>
-    </tr>
-    <tr>
-      <th>6</th>
-      <td>Customer#000069682</td>
-      <td>6057779.0348</td>
-    </tr>
-    <tr>
-      <th>7</th>
-      <td>Customer#000102022</td>
-      <td>6039653.6335</td>
-    </tr>
-    <tr>
-      <th>8</th>
-      <td>Customer#000098587</td>
-      <td>6027021.5855</td>
-    </tr>
-    <tr>
-      <th>9</th>
-      <td>Customer#000064660</td>
-      <td>5905659.6159</td>
-    </tr>
-  </tbody>
-</table>
-</div>
-
-## Why Vanna?
-
-1. **High accuracy on complex datasets.**
-    - Vanna’s capabilities are tied to the training data you give it
-    - More training data means better accuracy for large and complex datasets
-2. **Secure and private.**
-    - Your database contents  are never sent to Vanna’s servers
-    - We only see the bare minimum - schemas & queries.
-3. **Isolated, custom model.**
-    - You train a custom model specific to your database and your schema.
-    - Nobody else can use your model or view your model’s training data unless you choose to add members to your model or make it public
-    - We use a combination of third-party foundational models (OpenAI, Google) and our own LLM.
-4. **Self learning.**
-    - As you use Vanna more, your model continuously improves as we augment your training data
-5. **Supports many databases.**
-    - We have out-of-the-box support Snowflake, BigQuery, Postgres
-    - You can easily make a connector for any [database](https://docs.vanna.ai/databases/)
-6. **Pretrained models.**
-    - If you’re a data provider you can publish your models for anyone to use
-    - As part of our roadmap, we are in the process of pre-training models for common datasets (Google Ads, Facebook ads, etc)
-7. **Choose your front end.**
-    - Start in a Jupyter Notebook.
-    - Expose to business users via Slackbot, web app, Streamlit app, or Excel plugin.
-    - Even integrate in your web app for customers.
-
-## More resources
- - [Full Documentation](https://vanna.ai/docs/)
- - [Website](https://vanna.ai)
- - [Slack channel for support](https://join.slack.com/t/vanna-ai/shared_invite/zt-1unu0ipog-iE33QCoimQiBDxf2o7h97w)
- - [LinkedIn](https://www.linkedin.com/company/vanna-ai/)