vanna 0.2.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

vanna/mistral/mistral.py

@@ -1,33 +1,36 @@
 from mistralai.client import MistralClient
 from mistralai.models.chat_completion import ChatMessage
+
 from ..base import VannaBase
-import re
+
 
 class Mistral(VannaBase):
     def __init__(self, config=None):
         if config is None:
-            raise ValueError("For Mistral, config must be provided with an api_key and model")
+            raise ValueError(
+                "For Mistral, config must be provided with an api_key and model"
+            )
 
-        if 'api_key' not in config:
+        if "api_key" not in config:
             raise ValueError("config must contain a Mistral api_key")
-        
-        if 'model' not in config:
+
+        if "model" not in config:
             raise ValueError("config must contain a Mistral model")
 
-        api_key = config['api_key']
-        model = config['model']
+        api_key = config["api_key"]
+        model = config["model"]
         self.client = MistralClient(api_key=api_key)
         self.model = model
 
     def system_message(self, message: str) -> any:
         return ChatMessage(role="system", content=message)
-    
+
     def user_message(self, message: str) -> any:
         return ChatMessage(role="user", content=message)
-    
+
     def assistant_message(self, message: str) -> any:
         return ChatMessage(role="assistant", content=message)
-    
+
     def generate_sql(self, question: str, **kwargs) -> str:
         # Use the super generate_sql
         sql = super().generate_sql(question, **kwargs)
@@ -42,5 +45,5 @@ class Mistral(VannaBase):
             model=self.model,
             messages=prompt,
         )
-        
+
         return chat_response.choices[0].message.content

vanna/ollama/__init__.py

@@ -1,75 +0,0 @@
-from ..base import VannaBase
-import requests
-import json
-import re
-
-class Ollama(VannaBase):
-    def __init__(self, config=None):
-        if config is None or 'ollama_host' not in config:
-            self.host = "http://localhost:11434"
-        else:
-            self.host = config['ollama_host']
-
-        if config is None or 'model' not in config:
-            raise ValueError("config must contain a Ollama model")
-        else:
-            self.model = config['model']
-
-    def system_message(self, message: str) -> any:
-        return {"role": "system", "content": message}
-
-    def user_message(self, message: str) -> any:
-        return {"role": "user", "content": message}
-
-    def assistant_message(self, message: str) -> any:
-        return {"role": "assistant", "content": message}
-    
-    def extract_sql_query(self, text):
-        """
-        Extracts the first SQL statement after the word 'select', ignoring case,
-        matches until the first semicolon, three backticks, or the end of the string,
-        and removes three backticks if they exist in the extracted string.
-        
-        Args:
-        - text (str): The string to search within for an SQL statement.
-        
-        Returns:
-        - str: The first SQL statement found, with three backticks removed, or an empty string if no match is found.
-        """
-        # Regular expression to find 'select' (ignoring case) and capture until ';', '```', or end of string
-        pattern = re.compile(r'select.*?(?:;|```|$)', re.IGNORECASE | re.DOTALL)
-        
-        match = pattern.search(text)
-        if match:
-            # Remove three backticks from the matched string if they exist
-            return match.group(0).replace('```', '')
-        else:
-            return text
-
-    def generate_sql(self, question: str, **kwargs) -> str:
-        # Use the super generate_sql
-        sql = super().generate_sql(question, **kwargs)
-
-        # Replace "\_" with "_"
-        sql = sql.replace("\\_", "_")
-
-        sql = sql.replace("\\", "")
-
-        return self.extract_sql_query(sql)
-
-    def submit_prompt(self, prompt, **kwargs) -> str:
-        url = f"{self.host}/api/chat"
-        data = {
-            "model": self.model,
-            "stream": False,
-            "messages": prompt,
-        }
-
-        response = requests.post(url, json=data)
-
-        response_dict = response.json()
-
-        self.log(response.text)
-        
-        return response_dict['message']['content']
-

vanna/ollama/ollama.py

@@ -0,0 +1,77 @@
+import json
+import re
+
+import requests
+
+from ..base import VannaBase
+
+
+class Ollama(VannaBase):
+    def __init__(self, config=None):
+        if config is None or "ollama_host" not in config:
+            self.host = "http://localhost:11434"
+        else:
+            self.host = config["ollama_host"]
+
+        if config is None or "model" not in config:
+            raise ValueError("config must contain a Ollama model")
+        else:
+            self.model = config["model"]
+
+    def system_message(self, message: str) -> any:
+        return {"role": "system", "content": message}
+
+    def user_message(self, message: str) -> any:
+        return {"role": "user", "content": message}
+
+    def assistant_message(self, message: str) -> any:
+        return {"role": "assistant", "content": message}
+
+    def extract_sql_query(self, text):
+        """
+        Extracts the first SQL statement after the word 'select', ignoring case,
+        matches until the first semicolon, three backticks, or the end of the string,
+        and removes three backticks if they exist in the extracted string.
+
+        Args:
+        - text (str): The string to search within for an SQL statement.
+
+        Returns:
+        - str: The first SQL statement found, with three backticks removed, or an empty string if no match is found.
+        """
+        # Regular expression to find 'select' (ignoring case) and capture until ';', '```', or end of string
+        pattern = re.compile(r"select.*?(?:;|```|$)", re.IGNORECASE | re.DOTALL)
+
+        match = pattern.search(text)
+        if match:
+            # Remove three backticks from the matched string if they exist
+            return match.group(0).replace("```", "")
+        else:
+            return text
+
+    def generate_sql(self, question: str, **kwargs) -> str:
+        # Use the super generate_sql
+        sql = super().generate_sql(question, **kwargs)
+
+        # Replace "\_" with "_"
+        sql = sql.replace("\\_", "_")
+
+        sql = sql.replace("\\", "")
+
+        return self.extract_sql_query(sql)
+
+    def submit_prompt(self, prompt, **kwargs) -> str:
+        url = f"{self.host}/api/chat"
+        data = {
+            "model": self.model,
+            "stream": False,
+            "messages": prompt,
+        }
+
+        response = requests.post(url, json=data)
+
+        response_dict = response.json()
+
+        self.log(response.text)
+
+        return response_dict["message"]["content"]

vanna/openai/openai_chat.py

@@ -9,14 +9,6 @@ class OpenAI_Chat(VannaBase):
     def __init__(self, client=None, config=None):
         VannaBase.__init__(self, config=config)
 
-        if client is not None:
-            self.client = client
-            return
-
-        if config is None and client is None:
-            self.client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
-            return
-
         # default parameters - can be overrided using config
         self.temperature = 0.7
         self.max_tokens = 500
@@ -42,6 +34,14 @@ class OpenAI_Chat(VannaBase):
                 "Passing api_version is now deprecated. Please pass an OpenAI client instead."
             )
 
+        if client is not None:
+            self.client = client
+            return
+
+        if config is None and client is None:
+            self.client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
+            return
+      
         if "api_key" in config:
             self.client = OpenAI(api_key=config["api_key"])
 

{vanna-0.2.0.dist-info → vanna-0.3.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vanna
-Version: 0.2.0
+Version: 0.3.0
 Summary: Generate SQL queries from natural language
 Author-email: Zain Hoda <zain@vanna.ai>
 Requires-Python: >=3.9
@@ -18,18 +18,22 @@ Requires-Dist: flask
 Requires-Dist: sqlalchemy
 Requires-Dist: psycopg2-binary ; extra == "all"
 Requires-Dist: db-dtypes ; extra == "all"
+Requires-Dist: PyMySQL ; extra == "all"
 Requires-Dist: google-cloud-bigquery ; extra == "all"
 Requires-Dist: snowflake-connector-python ; extra == "all"
 Requires-Dist: duckdb ; extra == "all"
 Requires-Dist: openai ; extra == "all"
 Requires-Dist: mistralai ; extra == "all"
 Requires-Dist: chromadb ; extra == "all"
+Requires-Dist: anthropic ; extra == "all"
+Requires-Dist: anthropic ; extra == "anthropic"
 Requires-Dist: google-cloud-bigquery ; extra == "bigquery"
 Requires-Dist: chromadb ; extra == "chromadb"
 Requires-Dist: duckdb ; extra == "duckdb"
 Requires-Dist: google-generativeai ; extra == "gemini"
 Requires-Dist: marqo ; extra == "marqo"
 Requires-Dist: mistralai ; extra == "mistralai"
+Requires-Dist: PyMySQL ; extra == "mysql"
 Requires-Dist: openai ; extra == "openai"
 Requires-Dist: psycopg2-binary ; extra == "postgres"
 Requires-Dist: db-dtypes ; extra == "postgres"
@@ -38,12 +42,14 @@ 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: anthropic
 Provides-Extra: bigquery
 Provides-Extra: chromadb
 Provides-Extra: duckdb
 Provides-Extra: gemini
 Provides-Extra: marqo
 Provides-Extra: mistralai
+Provides-Extra: mysql
 Provides-Extra: openai
 Provides-Extra: postgres
 Provides-Extra: snowflake

{vanna-0.2.0.dist-info → vanna-0.3.0.dist-info}/RECORD

@@ -5,25 +5,27 @@ vanna/utils.py,sha256=Q0H4eugPYg9SVpEoTWgvmuoJZZxOVRhNzrP97E5lyak,1472
 vanna/ZhipuAI/ZhipuAI_Chat.py,sha256=hcx__0ZKHr5wtmIv0Ye2Xake8E-sOi5Qyc5nFIIdFvg,8957
 vanna/ZhipuAI/ZhipuAI_embeddings.py,sha256=lUqzJg9fOx7rVFhjdkFjXcDeVGV4aAB5Ss0oERsa8pE,2849
 vanna/ZhipuAI/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+vanna/anthropic/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+vanna/anthropic/anthropic_chat.py,sha256=Wk0o-NMW1uvR2fhSWxrR_2FqNh-dLprNG4uuVqpqAkY,2615
 vanna/base/__init__.py,sha256=Sl-HM1RRYzAZoSqmL1CZQmF3ZF-byYTCFQP3JZ2A5MU,28
-vanna/base/base.py,sha256=6tSgklJdaZwhK9vMhI7ZilZsB_bsngUWk7bW1vVSoOE,52175
+vanna/base/base.py,sha256=25tHTMsCnoek6X9C0hqpPyNM3H0mihcL4I8Mlq-74c0,54890
 vanna/chromadb/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vanna/chromadb/chromadb_vector.py,sha256=fa7uj_knzSfzsVLvpSunwwu1ZJNC3GbiNZ4Yy09v4l4,8372
 vanna/exceptions/__init__.py,sha256=N76unE7sjbGGBz6LmCrPQAugFWr9cUFv8ErJxBrCTts,717
-vanna/flask/__init__.py,sha256=i7eh55uB64Kio3nnrAK4eg45gAoY_b2uziLLlCmj1Dk,15186
-vanna/flask/assets.py,sha256=1DcZNu3Uo93hz7M8ldEYKDeEslROIL3IK44sSS8PZgM,163301
+vanna/flask/__init__.py,sha256=UgM0Ce5pGDdadWV6ZEAXj7RXDE1E420DW1wtR-juBMw,21212
+vanna/flask/assets.py,sha256=uluOgW41cHYmqPyRFHn8HYgaozmPm9nhDLZkHLuSBM4,180728
 vanna/marqo/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vanna/marqo/marqo.py,sha256=2OBuC5IZmGcFXN2Ah6GVPKHBYtkDXeSwhXsqUbxyU94,5285
 vanna/mistral/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vanna/mistral/mistral.py,sha256=6NarIVLfh9qxJIyUv_9pdRzbbMfC54gYsqg99Q-K7aA,1514
-vanna/ollama/__init__.py,sha256=GW1ek7zw_fpL2yFNgrnN5RNjV2PdkK8CmTmcJF9YpDU,2464
-vanna/ollama/ollama.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+vanna/mistral/mistral.py,sha256=DAEqAT9SzC91rfMM_S3SuzBZ34MrKHw9qAj6EP2MGVk,1508
+vanna/ollama/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+vanna/ollama/ollama.py,sha256=U02yy6DCAR-Ar7YXF98aBg3Kc-c9CiSWnMIGv4gs8zE,2430
 vanna/openai/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vanna/openai/openai_chat.py,sha256=2A6YEibwnMgdc6r1mvInqi_32xDDkdGi-bx73UEUaJM,3846
+vanna/openai/openai_chat.py,sha256=Y3-Fhz9c6D-5vrMR1zGibavAPNPAu-hMTwqGfKhAg3Q,3852
 vanna/openai/openai_embeddings.py,sha256=g4pNh9LVcYP9wOoO8ecaccDFWmCUYMInebfHucAa2Gc,1260
 vanna/types/__init__.py,sha256=Qhn_YscKtJh7mFPCyCDLa2K8a4ORLMGVnPpTbv9uB2U,4957
 vanna/vannadb/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vanna/vannadb/vannadb_vector.py,sha256=f4kddaJgTpZync7wnQi09QdODUuMtiHsK7WfKBUAmSo,5644
-vanna-0.2.0.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
-vanna-0.2.0.dist-info/METADATA,sha256=UzXxOVxudUIRanqrhsqp_ssL-lpvLqulSrn_PMyjnSY,9741
-vanna-0.2.0.dist-info/RECORD,,
+vanna-0.3.0.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
+vanna-0.3.0.dist-info/METADATA,sha256=lg_gRR1TR9RgdxDh3UmthfyudyEdcWdxHpwXEMrz7AU,9961
+vanna-0.3.0.dist-info/RECORD,,