vanna 0.6.3__tar.gz → 0.6.4__tar.gz

{vanna-0.6.3 → vanna-0.6.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vanna
-Version: 0.6.3
+Version: 0.6.4
 Summary: Generate SQL queries from natural language
 Author-email: Zain Hoda <zain@vanna.ai>
 Requires-Python: >=3.9
@@ -16,6 +16,7 @@ Requires-Dist: sqlparse
 Requires-Dist: kaleido
 Requires-Dist: flask
 Requires-Dist: flask-sock
+Requires-Dist: flasgger
 Requires-Dist: sqlalchemy
 Requires-Dist: psycopg2-binary ; extra == "all"
 Requires-Dist: db-dtypes ; extra == "all"

{vanna-0.6.3 → vanna-0.6.4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "vanna"
-version = "0.6.3"
+version = "0.6.4"
 authors = [
   { name="Zain Hoda", email="zain@vanna.ai" },
 ]
@@ -18,7 +18,7 @@ classifiers = [
     "Operating System :: OS Independent",
 ]
 dependencies = [
-    "requests", "tabulate", "plotly", "pandas", "sqlparse", "kaleido", "flask", "flask-sock", "sqlalchemy"
+    "requests", "tabulate", "plotly", "pandas", "sqlparse", "kaleido", "flask", "flask-sock", "flasgger", "sqlalchemy"
 ]
 
 [project.urls]

{vanna-0.6.3 → vanna-0.6.4}/src/vanna/anthropic/anthropic_chat.py

@@ -8,15 +8,7 @@ from ..base import VannaBase
 class Anthropic_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 = anthropic.Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))
-            return
-
+      
         # default parameters - can be overrided using config
         self.temperature = 0.7
         self.max_tokens = 500
@@ -27,6 +19,14 @@ class Anthropic_Chat(VannaBase):
         if "max_tokens" in config:
             self.max_tokens = config["max_tokens"]
 
+        if client is not None:
+            self.client = client
+            return
+
+        if config is None and client is None:
+            self.client = anthropic.Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))
+            return
+      
         if "api_key" in config:
             self.client = anthropic.Anthropic(api_key=config["api_key"])
 

{vanna-0.6.3 → vanna-0.6.4}/src/vanna/base/base.py

@@ -577,6 +577,7 @@ class VannaBase(ABC):
             "3. If the provided context is insufficient, please explain why it can't be generated. \n"
             "4. Please use the most relevant table(s). \n"
             "5. If the question has been asked and answered before, please repeat the answer exactly as it was given before. \n"
+            f"6. Ensure that the output SQL is {self.dialect}-compliant and executable, and free of syntax errors. \n"
         )
 
         message_log = [self.system_message(initial_prompt)]
@@ -718,6 +719,7 @@ class VannaBase(ABC):
         database: str,
         role: Union[str, None] = None,
         warehouse: Union[str, None] = None,
+        **kwargs
     ):
         try:
             snowflake = __import__("snowflake.connector")
@@ -735,7 +737,7 @@ class VannaBase(ABC):
             else:
                 raise ImproperlyConfigured("Please set your Snowflake username.")
 
-        if password == "my-password":
+        if password == "mypassword":
             password_env = os.getenv("SNOWFLAKE_PASSWORD")
 
             if password_env is not None:
@@ -764,7 +766,8 @@ class VannaBase(ABC):
             password=password,
             account=account,
             database=database,
-            client_session_keep_alive=True
+            client_session_keep_alive=True,
+            **kwargs
         )
 
         def run_sql_snowflake(sql: str) -> pd.DataFrame:
@@ -790,13 +793,13 @@ class VannaBase(ABC):
         self.run_sql = run_sql_snowflake
         self.run_sql_is_set = True
 
-    def connect_to_sqlite(self, url: str):
+    def connect_to_sqlite(self, url: str, check_same_thread: bool = False,  **kwargs):
         """
         Connect to a SQLite database. This is just a helper function to set [`vn.run_sql`][vanna.base.base.VannaBase.run_sql]
 
         Args:
             url (str): The URL of the database to connect to.
-
+            check_same_thread (str): Allow the connection may be accessed in multiple threads.
         Returns:
             None
         """
@@ -815,7 +818,11 @@ class VannaBase(ABC):
             url = path
 
         # Connect to the database
-        conn = sqlite3.connect(url, check_same_thread=False)
+        conn = sqlite3.connect(
+            url,
+            check_same_thread=check_same_thread,
+            **kwargs
+        )
 
         def run_sql_sqlite(sql: str):
             return pd.read_sql_query(sql, conn)
@@ -831,6 +838,7 @@ class VannaBase(ABC):
         user: str = None,
         password: str = None,
         port: int = None,
+        **kwargs
     ):
         """
         Connect to postgres using the psycopg2 connector. This is just a helper function to set [`vn.run_sql`][vanna.base.base.VannaBase.run_sql]
@@ -900,6 +908,7 @@ class VannaBase(ABC):
                 user=user,
                 password=password,
                 port=port,
+                **kwargs
             )
         except psycopg2.Error as e:
             raise ValidationError(e)
@@ -931,12 +940,13 @@ class VannaBase(ABC):
 
 
     def connect_to_mysql(
-            self,
-            host: str = None,
-            dbname: str = None,
-            user: str = None,
-            password: str = None,
-            port: int = None,
+        self,
+        host: str = None,
+        dbname: str = None,
+        user: str = None,
+        password: str = None,
+        port: int = None,
+        **kwargs
     ):
 
         try:
@@ -980,12 +990,15 @@ class VannaBase(ABC):
         conn = None
 
         try:
-            conn = pymysql.connect(host=host,
-                                   user=user,
-                                   password=password,
-                                   database=dbname,
-                                   port=port,
-                                   cursorclass=pymysql.cursors.DictCursor)
+            conn = pymysql.connect(
+                host=host,
+                user=user,
+                password=password,
+                database=dbname,
+                port=port,
+                cursorclass=pymysql.cursors.DictCursor,
+                **kwargs
+            )
         except pymysql.Error as e:
             raise ValidationError(e)
 
@@ -1015,12 +1028,13 @@ class VannaBase(ABC):
         self.run_sql = run_sql_mysql
 
     def connect_to_clickhouse(
-            self,
-            host: str = None,
-            dbname: str = None,
-            user: str = None,
-            password: str = None,
-            port: int = None,
+        self,
+        host: str = None,
+        dbname: str = None,
+        user: str = None,
+        password: str = None,
+        port: int = None,
+        **kwargs
     ):
 
         try:
@@ -1070,6 +1084,7 @@ class VannaBase(ABC):
                 username=user,
                 password=password,
                 database=dbname,
+                **kwargs
             )
             print(conn)
         except Exception as e:
@@ -1087,15 +1102,16 @@ class VannaBase(ABC):
 
                 except Exception as e:
                     raise e
-        
+
         self.run_sql_is_set = True
         self.run_sql = run_sql_clickhouse
 
     def connect_to_oracle(
-    self,
-    user: str = None,
-    password: str = None,
-    dsn: str = None,
+        self,
+        user: str = None,
+        password: str = None,
+        dsn: str = None,
+        **kwargs
     ):
 
         """
@@ -1148,7 +1164,8 @@ class VannaBase(ABC):
                 user=user,
                 password=password,
                 dsn=dsn,
-                )
+                **kwargs
+            )
         except oracledb.Error as e:
             raise ValidationError(e)
 
@@ -1180,7 +1197,12 @@ class VannaBase(ABC):
         self.run_sql_is_set = True
         self.run_sql = run_sql_oracle
 
-    def connect_to_bigquery(self, cred_file_path: str = None, project_id: str = None):
+    def connect_to_bigquery(
+        self,
+        cred_file_path: str = None,
+        project_id: str = None,
+        **kwargs
+    ):
         """
         Connect to gcs using the bigquery connector. This is just a helper function to set [`vn.run_sql`][vanna.base.base.VannaBase.run_sql]
         **Example:**
@@ -1242,7 +1264,11 @@ class VannaBase(ABC):
                 )
 
             try:
-                conn = bigquery.Client(project=project_id, credentials=credentials)
+                conn = bigquery.Client(
+                    project=project_id,
+                    credentials=credentials,
+                    **kwargs
+                )
             except:
                 raise ImproperlyConfigured(
                     "Could not connect to bigquery please correct credentials"
@@ -1265,7 +1291,7 @@ class VannaBase(ABC):
         self.run_sql_is_set = True
         self.run_sql = run_sql_bigquery
 
-    def connect_to_duckdb(self, url: str, init_sql: str = None):
+    def connect_to_duckdb(self, url: str, init_sql: str = None, **kwargs):
         """
         Connect to a DuckDB database. This is just a helper function to set [`vn.run_sql`][vanna.base.base.VannaBase.run_sql]
 
@@ -1303,7 +1329,7 @@ class VannaBase(ABC):
                         f.write(response.content)
 
         # Connect to the database
-        conn = duckdb.connect(path)
+        conn = duckdb.connect(path, **kwargs)
         if init_sql:
             conn.query(init_sql)
 
@@ -1314,7 +1340,7 @@ class VannaBase(ABC):
         self.run_sql = run_sql_duckdb
         self.run_sql_is_set = True
 
-    def connect_to_mssql(self, odbc_conn_str: str):
+    def connect_to_mssql(self, odbc_conn_str: str, **kwargs):
         """
         Connect to a Microsoft SQL Server database. This is just a helper function to set [`vn.run_sql`][vanna.base.base.VannaBase.run_sql]
 
@@ -1347,7 +1373,7 @@ class VannaBase(ABC):
 
         from sqlalchemy import create_engine
 
-        engine = create_engine(connection_url)
+        engine = create_engine(connection_url, **kwargs)
 
         def run_sql_mssql(sql: str):
             # Execute the SQL statement and return the result as a pandas DataFrame
@@ -1362,16 +1388,17 @@ class VannaBase(ABC):
         self.run_sql = run_sql_mssql
         self.run_sql_is_set = True
     def connect_to_presto(
-      self,
-      host: str,
-      catalog: str = 'hive',
-      schema: str = 'default',
-      user: str = None,
-      password: str = None,
-      port: int = None,
-      combined_pem_path: str = None,
-      protocol: str = 'https',
-      requests_kwargs: dict = None
+        self,
+        host: str,
+        catalog: str = 'hive',
+        schema: str = 'default',
+        user: str = None,
+        password: str = None,
+        port: int = None,
+        combined_pem_path: str = None,
+        protocol: str = 'https',
+        requests_kwargs: dict = None,
+        **kwargs
     ):
       """
         Connect to a Presto database using the specified parameters.
@@ -1444,7 +1471,8 @@ class VannaBase(ABC):
                                  schema=schema,
                                  port=port,
                                  protocol=protocol,
-                                 requests_kwargs=requests_kwargs)
+                                 requests_kwargs=requests_kwargs,
+                                 **kwargs)
       except presto.Error as e:
         raise ValidationError(e)
 
@@ -1477,13 +1505,14 @@ class VannaBase(ABC):
       self.run_sql = run_sql_presto
 
     def connect_to_hive(
-      self,
-      host: str = None,
-      dbname: str = 'default',
-      user: str = None,
-      password: str = None,
-      port: int = None,
-      auth: str = 'CUSTOM'
+        self,
+        host: str = None,
+        dbname: str = 'default',
+        user: str = None,
+        password: str = None,
+        port: int = None,
+        auth: str = 'CUSTOM',
+        **kwargs
     ):
       """
         Connect to a Hive database. This is just a helper function to set [`vn.run_sql`][vanna.base.base.VannaBase.run_sql]