vanna 0.0.32__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
 
@@ -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/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.32.dist-info → vanna-0.0.33.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vanna
-Version: 0.0.32
+Version: 0.0.33
 Summary: Generate SQL queries from natural language
 Author-email: Zain Hoda <zain@vanna.ai>
 Requires-Python: >=3.7

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

@@ -1,15 +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=XGUxkiZVSHg9U7JLF1xrxSPmrMOHrafuvQGKTn3yBDM,31187
+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
@@ -17,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.32.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
-vanna-0.0.32.dist-info/METADATA,sha256=cbzkVUYIywEmF-DHMg90iEzLU-jxZU05RJCwpjtdtiU,8780
-vanna-0.0.32.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,,