vanna 0.0.31__py3-none-any.whl → 0.0.33__py3-none-any.whl

vanna/base/base.py

@@ -1,8 +1,8 @@
 import json
 import os
+import re
 import sqlite3
 import traceback
-
 from abc import ABC, abstractmethod
 from typing import List, Tuple, Union
 from urllib.parse import urlparse
@@ -12,7 +12,6 @@ import plotly
 import plotly.express as px
 import plotly.graph_objects as go
 import requests
-import re
 
 from ..exceptions import DependencyError, ImproperlyConfigured, ValidationError
 from ..types import TrainingPlan, TrainingPlanItem
@@ -50,8 +49,8 @@ class VannaBase(ABC):
             **kwargs,
         )
         llm_response = self.submit_prompt(prompt, **kwargs)
-        
-        numbers_removed = re.sub(r'^\d+\.\s*', '', llm_response, flags=re.MULTILINE)
+
+        numbers_removed = re.sub(r"^\d+\.\s*", "", llm_response, flags=re.MULTILINE)
         return numbers_removed.split("\n")
 
     def generate_questions(self, **kwargs) -> list[str]:
@@ -65,7 +64,7 @@ class VannaBase(ABC):
         """
         question_sql = self.get_similar_question_sql(question="", **kwargs)
 
-        return [q['question'] for q in question_sql]
+        return [q["question"] for q in question_sql]
 
     # ----------------- Use Any Embeddings API ----------------- #
     @abstractmethod
@@ -94,7 +93,7 @@ class VannaBase(ABC):
         pass
 
     @abstractmethod
-    def add_documentation(self, doc: str, **kwargs) -> str:
+    def add_documentation(self, documentation: str, **kwargs) -> str:
         pass
 
     @abstractmethod
@@ -120,12 +119,12 @@ class VannaBase(ABC):
 
     @abstractmethod
     def get_followup_questions_prompt(
-        self, 
-        question: str, 
+        self,
+        question: str,
         question_sql_list: list,
         ddl_list: list,
-        doc_list: list, 
-        **kwargs
+        doc_list: list,
+        **kwargs,
     ):
         pass
 
@@ -248,7 +247,7 @@ class VannaBase(ABC):
             url = path
 
         # Connect to the database
-        conn = sqlite3.connect(url)
+        conn = sqlite3.connect(url, check_same_thread=False)
 
         def run_sql_sqlite(sql: str):
             return pd.read_sql_query(sql, conn)
@@ -829,9 +828,11 @@ class VannaBase(ABC):
             fig = ldict.get("fig", None)
         except Exception as e:
             # Inspect data types
-            numeric_cols = df.select_dtypes(include=['number']).columns.tolist()
-            categorical_cols = df.select_dtypes(include=['object', 'category']).columns.tolist()
-            
+            numeric_cols = df.select_dtypes(include=["number"]).columns.tolist()
+            categorical_cols = df.select_dtypes(
+                include=["object", "category"]
+            ).columns.tolist()
+
             # Decision-making for plot type
             if len(numeric_cols) >= 2:
                 # Use the first two numeric columns for a scatter plot

vanna/chromadb/chromadb_vector.py

@@ -3,9 +3,9 @@ import uuid
 from abc import abstractmethod
 
 import chromadb
+import pandas as pd
 from chromadb.config import Settings
 from chromadb.utils import embedding_functions
-import pandas as pd
 
 from ..base import VannaBase
 
@@ -47,7 +47,7 @@ class ChromaDB_VectorStore(VannaBase):
                 "sql": sql,
             }
         )
-        id = str(uuid.uuid4())+"-sql"
+        id = str(uuid.uuid4()) + "-sql"
         self.sql_collection.add(
             documents=question_sql_json,
             embeddings=self.generate_embedding(question_sql_json),
@@ -57,7 +57,7 @@ class ChromaDB_VectorStore(VannaBase):
         return id
 
     def add_ddl(self, ddl: str, **kwargs) -> str:
-        id = str(uuid.uuid4())+"-ddl"
+        id = str(uuid.uuid4()) + "-ddl"
         self.ddl_collection.add(
             documents=ddl,
             embeddings=self.generate_embedding(ddl),
@@ -65,11 +65,11 @@ class ChromaDB_VectorStore(VannaBase):
         )
         return id
 
-    def add_documentation(self, doc: str, **kwargs) -> str:
-        id = str(uuid.uuid4())+"-doc"
+    def add_documentation(self, documentation: str, **kwargs) -> str:
+        id = str(uuid.uuid4()) + "-doc"
         self.documentation_collection.add(
-            documents=doc,
-            embeddings=self.generate_embedding(doc),
+            documents=documentation,
+            embeddings=self.generate_embedding(documentation),
             ids=id,
         )
         return id
@@ -81,15 +81,17 @@ class ChromaDB_VectorStore(VannaBase):
 
         if sql_data is not None:
             # Extract the documents and ids
-            documents = [json.loads(doc) for doc in sql_data['documents']]
-            ids = sql_data['ids']
+            documents = [json.loads(doc) for doc in sql_data["documents"]]
+            ids = sql_data["ids"]
 
             # Create a DataFrame
-            df_sql = pd.DataFrame({
-                'id': ids,
-                'question': [doc['question'] for doc in documents],
-                'content': [doc['sql'] for doc in documents]
-            })
+            df_sql = pd.DataFrame(
+                {
+                    "id": ids,
+                    "question": [doc["question"] for doc in documents],
+                    "content": [doc["sql"] for doc in documents],
+                }
+            )
 
             df_sql["training_data_type"] = "sql"
 
@@ -99,15 +101,17 @@ class ChromaDB_VectorStore(VannaBase):
 
         if ddl_data is not None:
             # Extract the documents and ids
-            documents = [doc for doc in ddl_data['documents']]
-            ids = ddl_data['ids']
+            documents = [doc for doc in ddl_data["documents"]]
+            ids = ddl_data["ids"]
 
             # Create a DataFrame
-            df_ddl = pd.DataFrame({
-                'id': ids,
-                'question': [None for doc in documents],
-                'content': [doc for doc in documents]
-            })
+            df_ddl = pd.DataFrame(
+                {
+                    "id": ids,
+                    "question": [None for doc in documents],
+                    "content": [doc for doc in documents],
+                }
+            )
 
             df_ddl["training_data_type"] = "ddl"
 
@@ -117,15 +121,17 @@ class ChromaDB_VectorStore(VannaBase):
 
         if doc_data is not None:
             # Extract the documents and ids
-            documents = [doc for doc in doc_data['documents']]
-            ids = doc_data['ids']
+            documents = [doc for doc in doc_data["documents"]]
+            ids = doc_data["ids"]
 
             # Create a DataFrame
-            df_doc = pd.DataFrame({
-                'id': ids,
-                'question': [None for doc in documents],
-                'content': [doc for doc in documents]
-            })
+            df_doc = pd.DataFrame(
+                {
+                    "id": ids,
+                    "question": [None for doc in documents],
+                    "content": [doc for doc in documents],
+                }
+            )
 
             df_doc["training_data_type"] = "documentation"
 

vanna/flask.py

@@ -0,0 +1,331 @@
+import flask
+from flask import Flask, Response, jsonify, request
+import logging
+import requests
+from functools import wraps
+
+from abc import ABC, abstractmethod
+import uuid
+
+class Cache(ABC):
+    @abstractmethod
+    def generate_id(self, *args, **kwargs):
+        pass
+
+    @abstractmethod
+    def get(self, id, field):
+        pass
+
+    @abstractmethod
+    def get_all(self, field_list) -> list:
+        pass
+
+    @abstractmethod
+    def set(self, id, field, value):
+        pass
+
+    @abstractmethod
+    def delete(self, id):
+        pass
+
+
+class MemoryCache(Cache):
+    def __init__(self):
+        self.cache = {}
+
+    def generate_id(self, *args, **kwargs):
+        return str(uuid.uuid4())
+
+    def set(self, id, field, value):
+        if id not in self.cache:
+            self.cache[id] = {}
+
+        self.cache[id][field] = value
+
+    def get(self, id, field):
+        if id not in self.cache:
+            return None
+
+        if field not in self.cache[id]:
+            return None
+
+        return self.cache[id][field]
+
+    def get_all(self, field_list) -> list:
+        return [
+            {
+                "id": id,
+                **{
+                    field: self.get(id=id, field=field)
+                    for field in field_list
+                }
+            }
+            for id in self.cache
+        ]
+
+    def delete(self, id):
+        if id in self.cache:
+            del self.cache[id]
+
+class VannaFlaskApp:
+    flask_app = None
+
+    def requires_cache(self, fields):
+        def decorator(f):
+            @wraps(f)
+            def decorated(*args, **kwargs):
+                id = request.args.get('id')
+
+                if id is None:
+                    return jsonify({"type": "error", "error": "No id provided"})
+                
+                for field in fields:
+                    if self.cache.get(id=id, field=field) is None:
+                        return jsonify({"type": "error", "error": f"No {field} found"})
+                
+                field_values = {field: self.cache.get(id=id, field=field) for field in fields}
+                
+                # Add the id to the field_values
+                field_values['id'] = id
+
+                return f(*args, **field_values, **kwargs)
+            return decorated
+        return decorator
+
+    def __init__(self, vn, cache: Cache = MemoryCache()):
+        self.flask_app = Flask(__name__)
+        self.vn = vn
+        self.cache = cache
+
+        log = logging.getLogger('werkzeug')
+        log.setLevel(logging.ERROR)
+
+        @self.flask_app.route('/api/v0/generate_questions', methods=['GET'])
+        def generate_questions():
+            # If self has an _model attribute and model=='chinook'
+            if hasattr(self.vn, '_model') and self.vn._model == 'chinook':                
+                return jsonify({
+                    "type": "question_list", 
+                    "questions": ['What are the top 10 artists by sales?', 'What are the total sales per year by country?', 'Who is the top selling artist in each genre? Show the sales numbers.', 'How do the employees rank in terms of sales performance?', 'Which 5 cities have the most customers?'],
+                    "header": "Here are some questions you can ask:"
+                    })
+
+        @self.flask_app.route('/api/v0/generate_sql', methods=['GET'])
+        def generate_sql():
+            question = flask.request.args.get('question')
+
+            if question is None:
+                return jsonify({"type": "error", "error": "No question provided"})
+
+            id = self.cache.generate_id(question=question)
+            sql = vn.generate_sql(question=question)
+
+            self.cache.set(id=id, field='question', value=question)
+            self.cache.set(id=id, field='sql', value=sql)
+
+            return jsonify(
+                {
+                    "type": "sql", 
+                    "id": id,
+                    "text": sql,
+                })
+
+        @self.flask_app.route('/api/v0/run_sql', methods=['GET'])
+        @self.requires_cache(['sql'])
+        def run_sql(id: str, sql: str):
+            try:
+                if not vn.run_sql_is_set:
+                    return jsonify({"type": "error", "error": "Please connect to a database using vn.connect_to_... in order to run SQL queries."})
+
+                df = vn.run_sql(sql=sql)
+
+                cache.set(id=id, field='df', value=df)
+
+                return jsonify(
+                    {
+                        "type": "df", 
+                        "id": id,
+                        "df": df.head(10).to_json(orient='records'),
+                    })
+
+            except Exception as e:
+                return jsonify({"type": "error", "error": str(e)})
+
+        @self.flask_app.route('/api/v0/download_csv', methods=['GET'])
+        @self.requires_cache(['df'])
+        def download_csv(id: str, df):
+            csv = df.to_csv()
+
+            return Response(
+                csv,
+                mimetype="text/csv",
+                headers={"Content-disposition":
+                        f"attachment; filename={id}.csv"})
+
+        @self.flask_app.route('/api/v0/generate_plotly_figure', methods=['GET'])
+        @self.requires_cache(['df', 'question', 'sql'])
+        def generate_plotly_figure(id: str, df, question, sql):
+            try:
+                code = vn.generate_plotly_code(question=question, sql=sql, df_metadata=f"Running df.dtypes gives:\n {df.dtypes}")
+                fig = vn.get_plotly_figure(plotly_code=code, df=df, dark_mode=False)
+                fig_json = fig.to_json()
+
+                cache.set(id=id, field='fig_json', value=fig_json)
+
+                return jsonify(
+                    {
+                        "type": "plotly_figure", 
+                        "id": id,
+                        "fig": fig_json,
+                    })
+            except Exception as e:
+                # Print the stack trace
+                import traceback
+                traceback.print_exc()
+
+                return jsonify({"type": "error", "error": str(e)})
+
+        @self.flask_app.route('/api/v0/get_training_data', methods=['GET'])
+        def get_training_data():
+            df = vn.get_training_data()
+
+            return jsonify(
+            {
+                "type": "df", 
+                "id": "training_data",
+                "df": df.tail(25).to_json(orient='records'),
+            })
+
+        @self.flask_app.route('/api/v0/remove_training_data', methods=['POST'])
+        def remove_training_data():
+            # Get id from the JSON body
+            id = flask.request.json.get('id')
+
+            if id is None:
+                return jsonify({"type": "error", "error": "No id provided"})
+
+            if vn.remove_training_data(id=id):
+                return jsonify({"success": True})
+            else:
+                return jsonify({"type": "error", "error": "Couldn't remove training data"})
+
+        @self.flask_app.route('/api/v0/train', methods=['POST'])
+        def add_training_data():
+            question = flask.request.json.get('question')
+            sql = flask.request.json.get('sql')
+            ddl = flask.request.json.get('ddl')
+            documentation = flask.request.json.get('documentation')
+
+            try:
+                id = vn.train(question=question, sql=sql, ddl=ddl, documentation=documentation)
+
+                return jsonify({"id": id})
+            except Exception as e:
+                print("TRAINING ERROR", e)
+                return jsonify({"type": "error", "error": str(e)})
+
+        @self.flask_app.route('/api/v0/generate_followup_questions', methods=['GET'])
+        @self.requires_cache(['df', 'question'])
+        def generate_followup_questions(id: str, df, question):
+            followup_questions = []
+            # followup_questions = vn.generate_followup_questions(question=question, df=df)
+            # if followup_questions is not None and len(followup_questions) > 5:
+            #     followup_questions = followup_questions[:5]
+
+            cache.set(id=id, field='followup_questions', value=followup_questions)
+
+            return jsonify(
+                {
+                    "type": "question_list", 
+                    "id": id,
+                    "questions": followup_questions,
+                    "header": "Followup Questions can be enabled in a future version if you allow the LLM to 'see' your query results."
+                })
+
+        @self.flask_app.route('/api/v0/load_question', methods=['GET'])
+        @self.requires_cache(['question', 'sql', 'df', 'fig_json', 'followup_questions'])
+        def load_question(id: str, question, sql, df, fig_json, followup_questions):
+            try:
+                return jsonify(
+                    {
+                        "type": "question_cache", 
+                        "id": id,
+                        "question": question,
+                        "sql": sql,
+                        "df": df.head(10).to_json(orient='records'),
+                        "fig": fig_json,
+                        "followup_questions": followup_questions,
+                    })
+
+            except Exception as e:
+                return jsonify({"type": "error", "error": str(e)})
+
+        @self.flask_app.route('/api/v0/get_question_history', methods=['GET'])
+        def get_question_history():
+            return jsonify({"type": "question_history", "questions": cache.get_all(field_list=['question']) })
+
+
+        @self.flask_app.route('/api/v0/<path:catch_all>', methods=['GET', 'POST'])
+        def catch_all(catch_all):
+            return jsonify({"type": "error", "error": "The rest of the API is not ported yet."})
+
+        @self.flask_app.route('/assets/<path:filename>')
+        def proxy_assets(filename):
+            remote_url = f'https://vanna.ai/assets/{filename}'
+            response = requests.get(remote_url, stream=True)
+
+            # Check if the request to the remote URL was successful
+            if response.status_code == 200:
+                excluded_headers = ['content-encoding', 'content-length', 'transfer-encoding', 'connection']
+                headers = [(name, value) for (name, value) in response.raw.headers.items() if name.lower() not in excluded_headers]
+                return Response(response.content, response.status_code, headers)
+            else:
+                return 'Error fetching file from remote server', response.status_code
+
+        # Proxy the /vanna.svg file to the remote server
+        @self.flask_app.route('/vanna.svg')
+        def proxy_vanna_svg():
+            remote_url = f'https://vanna.ai/img/vanna.svg'
+            response = requests.get(remote_url, stream=True)
+
+            # Check if the request to the remote URL was successful
+            if response.status_code == 200:
+                excluded_headers = ['content-encoding', 'content-length', 'transfer-encoding', 'connection']
+                headers = [(name, value) for (name, value) in response.raw.headers.items() if name.lower() not in excluded_headers]
+                return Response(response.content, response.status_code, headers)
+            else:
+                return 'Error fetching file from remote server', response.status_code
+
+        @self.flask_app.route('/', defaults={'path': ''})
+        @self.flask_app.route('/<path:path>')
+        def hello(path: str):
+            return """
+<!doctype html>
+<html lang="en">
+  <head>
+    <meta charset="UTF-8" />
+    <link rel="icon" type="image/svg+xml" href="/vanna.svg" />
+    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
+    <link href="https://fonts.googleapis.com/css2?family=Roboto+Slab:wght@350&display=swap" rel="stylesheet">
+    <script src="https://cdn.plot.ly/plotly-latest.min.js" type="text/javascript"></script>
+    <title>Vanna.AI</title>
+    <script type="module" crossorigin src="/assets/index-d29524f4.js"></script>
+    <link rel="stylesheet" href="/assets/index-b1a5a2f1.css">
+  </head>
+  <body class="bg-white dark:bg-slate-900">
+    <div id="app"></div>
+  </body>
+</html>
+"""
+        
+    def run(self):
+        try:
+            from google.colab import output
+            output.serve_kernel_port_as_window(8084)
+            from google.colab.output import eval_js
+            print("Your app is running at:")
+            print(eval_js("google.colab.kernel.proxyPort(8084)"))
+        except:
+            print("Your app is running at:")
+            print("http://localhost:8084")
+        self.flask_app.run(host='0.0.0.0', port=8084, debug=False)

vanna/marqo/marqo.py

@@ -3,7 +3,6 @@ import uuid
 from abc import abstractmethod
 
 import marqo
-
 import pandas as pd
 
 from ..base import VannaBase
@@ -12,7 +11,7 @@ from ..base import VannaBase
 class Marqo_VectorStore(VannaBase):
     def __init__(self, config=None):
         VannaBase.__init__(self, config=config)
-        
+
         if config is not None and "marqo_url" in config:
             marqo_url = config["marqo_url"]
         else:
@@ -22,7 +21,7 @@ class Marqo_VectorStore(VannaBase):
             marqo_model = config["marqo_model"]
         else:
             marqo_model = "hf/all_datasets_v4_MiniLM-L6"
-        
+
         self.mq = marqo.Client(url=marqo_url)
 
         for index in ["vanna-sql", "vanna-ddl", "vanna-doc"]:
@@ -33,18 +32,17 @@ class Marqo_VectorStore(VannaBase):
                 print(f"Marqo index {index} already exists")
                 pass
 
-
     def generate_embedding(self, data: str, **kwargs) -> list[float]:
         # Marqo doesn't need to generate embeddings
-        pass        
+        pass
 
     def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
-        id = str(uuid.uuid4())+"-sql"
-        question_sql_dict ={
-                "question": question,
-                "sql": sql,
-                "_id": id,
-            }
+        id = str(uuid.uuid4()) + "-sql"
+        question_sql_dict = {
+            "question": question,
+            "sql": sql,
+            "_id": id,
+        }
 
         self.mq.index("vanna-sql").add_documents(
             [question_sql_dict],
@@ -54,11 +52,11 @@ class Marqo_VectorStore(VannaBase):
         return id
 
     def add_ddl(self, ddl: str, **kwargs) -> str:
-        id = str(uuid.uuid4())+"-ddl"
-        ddl_dict ={
-                "ddl": ddl,
-                "_id": id,
-            }
+        id = str(uuid.uuid4()) + "-ddl"
+        ddl_dict = {
+            "ddl": ddl,
+            "_id": id,
+        }
 
         self.mq.index("vanna-ddl").add_documents(
             [ddl_dict],
@@ -66,13 +64,13 @@ class Marqo_VectorStore(VannaBase):
         )
         return id
 
-    def add_documentation(self, doc: str, **kwargs) -> str:
-        id = str(uuid.uuid4())+"-doc"
-        doc_dict ={
-                "doc": doc,
-                "_id": id,
-            }
-        
+    def add_documentation(self, documentation: str, **kwargs) -> str:
+        id = str(uuid.uuid4()) + "-doc"
+        doc_dict = {
+            "doc": documentation,
+            "_id": id,
+        }
+
         self.mq.index("vanna-doc").add_documents(
             [doc_dict],
             tensor_fields=["doc"],
@@ -80,31 +78,37 @@ class Marqo_VectorStore(VannaBase):
         return id
 
     def get_training_data(self, **kwargs) -> pd.DataFrame:
-        data = []  
-
-        for hit in self.mq.index('vanna-doc').search("", limit=1000)['hits']:
-            data.append({
-                "id": hit["_id"],
-                "training_data_type": "documentation",
-                "question": "",
-                "content": hit["doc"],
-            })
-
-        for hit in self.mq.index('vanna-ddl').search("", limit=1000)['hits']:
-            data.append({
-                "id": hit["_id"],
-                "training_data_type": "ddl",
-                "question": "",
-                "content": hit["ddl"],
-            })
-
-        for hit in self.mq.index('vanna-sql').search("", limit=1000)['hits']:
-            data.append({
-                "id": hit["_id"],
-                "training_data_type": "sql",
-                "question": hit["question"],
-                "content": hit["sql"],
-            })
+        data = []
+
+        for hit in self.mq.index("vanna-doc").search("", limit=1000)["hits"]:
+            data.append(
+                {
+                    "id": hit["_id"],
+                    "training_data_type": "documentation",
+                    "question": "",
+                    "content": hit["doc"],
+                }
+            )
+
+        for hit in self.mq.index("vanna-ddl").search("", limit=1000)["hits"]:
+            data.append(
+                {
+                    "id": hit["_id"],
+                    "training_data_type": "ddl",
+                    "question": "",
+                    "content": hit["ddl"],
+                }
+            )
+
+        for hit in self.mq.index("vanna-sql").search("", limit=1000)["hits"]:
+            data.append(
+                {
+                    "id": hit["_id"],
+                    "training_data_type": "sql",
+                    "question": hit["question"],
+                    "content": hit["sql"],
+                }
+            )
 
         df = pd.DataFrame(data)
 
@@ -127,24 +131,24 @@ class Marqo_VectorStore(VannaBase):
     @staticmethod
     def _extract_documents(data) -> list:
         # Check if 'hits' key is in the dictionary and if it's a list
-        if 'hits' in data and isinstance(data['hits'], list):
+        if "hits" in data and isinstance(data["hits"], list):
             # Iterate over each item in 'hits'
 
-            if len(data['hits']) == 0:
+            if len(data["hits"]) == 0:
                 return []
 
             # If there is a "doc" key, return the value of that key
-            if "doc" in data['hits'][0]:
-                return [hit["doc"] for hit in data['hits']]
-            
+            if "doc" in data["hits"][0]:
+                return [hit["doc"] for hit in data["hits"]]
+
             # If there is a "ddl" key, return the value of that key
-            if "ddl" in data['hits'][0]:
-                return [hit["ddl"] for hit in data['hits']]
-            
+            if "ddl" in data["hits"][0]:
+                return [hit["ddl"] for hit in data["hits"]]
+
             # Otherwise return the entire hit
             return [
-                {key: value for key, value in hit.items() if not key.startswith('_')}
-                for hit in data['hits']
+                {key: value for key, value in hit.items() if not key.startswith("_")}
+                for hit in data["hits"]
             ]
         else:
             # Return an empty list if 'hits' is not found or not a list
@@ -152,15 +156,15 @@ class Marqo_VectorStore(VannaBase):
 
     def get_similar_question_sql(self, question: str, **kwargs) -> list:
         return Marqo_VectorStore._extract_documents(
-            self.mq.index('vanna-sql').search(question)
+            self.mq.index("vanna-sql").search(question)
         )
 
     def get_related_ddl(self, question: str, **kwargs) -> list:
         return Marqo_VectorStore._extract_documents(
-            self.mq.index('vanna-ddl').search(question)
+            self.mq.index("vanna-ddl").search(question)
         )
 
     def get_related_documentation(self, question: str, **kwargs) -> list:
         return Marqo_VectorStore._extract_documents(
-            self.mq.index('vanna-doc').search(question)
+            self.mq.index("vanna-doc").search(question)
         )

vanna/remote.py

@@ -197,7 +197,7 @@ class VannaDefault(VannaBase):
 
         return status.id
 
-    def add_documentation(self, doc: str, **kwargs) -> str:
+    def add_documentation(self, documentation: str, **kwargs) -> str:
         """
         Adds documentation to the model's training data
 
@@ -469,4 +469,4 @@ class VannaDefault(VannaBase):
         # Load the result into a dataclass
         question_string_list = QuestionStringList(**d["result"])
 
-        return question_string_list.questions        
+        return question_string_list.questions        

vanna/vannadb/vannadb_vector.py

@@ -1,18 +1,21 @@
+import dataclasses
+import json
+from io import StringIO
+
+import pandas as pd
+import requests
+
 from ..base import VannaBase
 from ..types import (
-    QuestionSQLPair, 
-    StatusWithId,
-    StringData,
     DataFrameJSON,
+    Question,
+    QuestionSQLPair,
     Status,
+    StatusWithId,
+    StringData,
     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):
@@ -105,8 +108,8 @@ class VannaDB_VectorStore(VannaBase):
 
         return status.id
 
-    def add_documentation(self, doc: str, **kwargs) -> str:
-        params = [StringData(data=doc)]
+    def add_documentation(self, documentation: str, **kwargs) -> str:
+        params = [StringData(data=documentation)]
 
         d = self._rpc_call(method="add_documentation", params=params)
 
@@ -167,7 +170,7 @@ class VannaDB_VectorStore(VannaBase):
             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:
@@ -175,7 +178,7 @@ class VannaDB_VectorStore(VannaBase):
             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:
@@ -183,5 +186,5 @@ class VannaDB_VectorStore(VannaBase):
             training_data = self.related_training_data[question]
         else:
             training_data = self.get_related_training_data_cached(question)
-        
-        return training_data.documentation
+
+        return training_data.documentation

{vanna-0.0.31.dist-info → vanna-0.0.33.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vanna
-Version: 0.0.31
+Version: 0.0.33
 Summary: Generate SQL queries from natural language
 Author-email: Zain Hoda <zain@vanna.ai>
 Requires-Python: >=3.7
@@ -41,7 +41,7 @@ Provides-Extra: postgres
 Provides-Extra: snowflake
 Provides-Extra: test
 
-![](https://img.vanna.ai/vanna-github.svg)
+
 
 | GitHub | PyPI | Documentation |
 | ------ | ---- | ------------- |
@@ -52,6 +52,8 @@ Vanna is an MIT-licensed open-source Python RAG (Retrieval-Augmented Generation)
 
 https://github.com/vanna-ai/vanna/assets/7146154/1901f47a-515d-4982-af50-f12761a3b2ce
 
+![vanna-quadrants](https://github.com/vanna-ai/vanna/assets/7146154/1c7c88ba-c144-4ecf-a028-cf5ba7344ca2)
+
 ## 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.
 

{vanna-0.0.31.dist-info → vanna-0.0.33.dist-info}/RECORD

@@ -1,14 +1,15 @@
 vanna/__init__.py,sha256=thjmOUgHCboSxIkzQRKw-JvZLLFbnuyM7G5YIzmmmPQ,61545
+vanna/flask.py,sha256=Kw7qjObb39J1BWX5PsjLRIbMJqntZI91K3tGlUxx5M0,12496
 vanna/local.py,sha256=U5s8ybCRQhBUizi8I69o3jqOpTeu_6KGYY6DMwZxjG4,313
-vanna/remote.py,sha256=xWlF48eQXuc03NZrDpMQgvrM6dbbfbEjX_FEmQf_b5c,13573
+vanna/remote.py,sha256=yr0QSJCAKzziiPOa-mfsjGg1pVa5-vLj9vYl2VDlAfU,13584
 vanna/utils.py,sha256=Q0H4eugPYg9SVpEoTWgvmuoJZZxOVRhNzrP97E5lyak,1472
 vanna/base/__init__.py,sha256=Sl-HM1RRYzAZoSqmL1CZQmF3ZF-byYTCFQP3JZ2A5MU,28
-vanna/base/base.py,sha256=4y0FFUNIb8Y948RFAzt6sYOs7iKqp3ORwMvtgrda00o,31162
+vanna/base/base.py,sha256=w3qYB-8LhcP0bvV0MCc4VlGgmVzQ4TY-N3Ufix8DZs8,31204
 vanna/chromadb/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vanna/chromadb/chromadb_vector.py,sha256=af1n7htIkSnpd7h9906mkKSK9BpvNNQa48_z4FS-_nE,5716
+vanna/chromadb/chromadb_vector.py,sha256=4YGgWQNIw4QJFwtBRIW53vieXH8rTBez-cs7EZwxsNI,5893
 vanna/exceptions/__init__.py,sha256=N76unE7sjbGGBz6LmCrPQAugFWr9cUFv8ErJxBrCTts,717
 vanna/marqo/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vanna/marqo/marqo.py,sha256=8L2W6XRu37BvnIpVnlyGbg_w2r2bceDyLuVCDAhDqs0,5206
+vanna/marqo/marqo.py,sha256=2OBuC5IZmGcFXN2Ah6GVPKHBYtkDXeSwhXsqUbxyU94,5285
 vanna/mistral/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vanna/mistral/mistral.py,sha256=A9dD8U-c12whGx8h_WOX15zUzaAJV-XLu_tpSiLamHo,8095
 vanna/openai/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -16,7 +17,7 @@ vanna/openai/openai_chat.py,sha256=U6wkXztJnQtABItUMDlBIDN6m3fqD6pMpa9gyQAQx8A,9
 vanna/openai/openai_embeddings.py,sha256=kPtOqrKQYJnXe6My3pO9BWg-L3KIR1sJVqE3YoW0roA,1139
 vanna/types/__init__.py,sha256=Qhn_YscKtJh7mFPCyCDLa2K8a4ORLMGVnPpTbv9uB2U,4957
 vanna/vannadb/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vanna/vannadb/vannadb_vector.py,sha256=zX_oT66LQSDeqO87I5xdKA87uQRQDl-ZrGOh8BYkUOU,5645
-vanna-0.0.31.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
-vanna-0.0.31.dist-info/METADATA,sha256=ILA-ZXW-UHaYqL6xKJrrJfu4WfoX45kawA8VN4X5itE,8715
-vanna-0.0.31.dist-info/RECORD,,
+vanna/vannadb/vannadb_vector.py,sha256=f4kddaJgTpZync7wnQi09QdODUuMtiHsK7WfKBUAmSo,5644
+vanna-0.0.33.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
+vanna-0.0.33.dist-info/METADATA,sha256=iLDk_AR5D179kN3392f67KeKoUyIDVNEWKEtLiomSTA,8780
+vanna-0.0.33.dist-info/RECORD,,