vanna 0.7.8__py3-none-any.whl → 2.0.0__py3-none-any.whl

vanna/{base → legacy/base}/base.py

@@ -136,7 +136,7 @@ class VannaBase(ABC):
         llm_response = self.submit_prompt(prompt, **kwargs)
         self.log(title="LLM Response", message=llm_response)
 
-        if 'intermediate_sql' in llm_response:
+        if "intermediate_sql" in llm_response:
             if not allow_llm_to_see_data:
                 return "The LLM is not allowed to see the data in your database. Your question requires database introspection to generate the necessary SQL. Please set allow_llm_to_see_data=True to enable this."
 
@@ -152,7 +152,11 @@ class VannaBase(ABC):
                         question=question,
                         question_sql_list=question_sql_list,
                         ddl_list=ddl_list,
-                        doc_list=doc_list+[f"The following is a pandas DataFrame with the results of the intermediate SQL query {intermediate_sql}: \n" + df.to_markdown()],
+                        doc_list=doc_list
+                        + [
+                            f"The following is a pandas DataFrame with the results of the intermediate SQL query {intermediate_sql}: \n"
+                            + df.to_markdown()
+                        ],
                         **kwargs,
                     )
                     self.log(title="Final SQL Prompt", message=prompt)
@@ -161,7 +165,6 @@ class VannaBase(ABC):
                 except Exception as e:
                     return f"Error running intermediate SQL: {e}"
 
-
         return self.extract_sql(llm_response)
 
     def extract_sql(self, llm_response: str) -> str:
@@ -181,30 +184,52 @@ class VannaBase(ABC):
             str: The extracted SQL query.
         """
 
-        # If the llm_response contains a CTE (with clause), extract the last sql between WITH and ;
-        sqls = re.findall(r"\bWITH\b .*?;", llm_response, re.DOTALL)
+        import re
+
+        """
+        Extracts the SQL query from the LLM response, handling various formats including:
+        - WITH clause
+        - SELECT statement
+        - CREATE TABLE AS SELECT
+        - Markdown code blocks
+        """
+
+        # Match CREATE TABLE ... AS SELECT
+        sqls = re.findall(
+            r"\bCREATE\s+TABLE\b.*?\bAS\b.*?;", llm_response, re.DOTALL | re.IGNORECASE
+        )
         if sqls:
             sql = sqls[-1]
             self.log(title="Extracted SQL", message=f"{sql}")
             return sql
 
-        # If the llm_response is not markdown formatted, extract last sql by finding select and ; in the response
-        sqls = re.findall(r"SELECT.*?;", llm_response, re.DOTALL)
+        # Match WITH clause (CTEs)
+        sqls = re.findall(r"\bWITH\b .*?;", llm_response, re.DOTALL | re.IGNORECASE)
         if sqls:
             sql = sqls[-1]
             self.log(title="Extracted SQL", message=f"{sql}")
             return sql
 
-        # If the llm_response contains a markdown code block, with or without the sql tag, extract the last sql from it
-        sqls = re.findall(r"```sql\n(.*)```", llm_response, re.DOTALL)
+        # Match SELECT ... ;
+        sqls = re.findall(r"\bSELECT\b .*?;", llm_response, re.DOTALL | re.IGNORECASE)
         if sqls:
             sql = sqls[-1]
             self.log(title="Extracted SQL", message=f"{sql}")
             return sql
 
-        sqls = re.findall(r"```(.*)```", llm_response, re.DOTALL)
+        # Match ```sql ... ``` blocks
+        sqls = re.findall(
+            r"```sql\s*\n(.*?)```", llm_response, re.DOTALL | re.IGNORECASE
+        )
         if sqls:
-            sql = sqls[-1]
+            sql = sqls[-1].strip()
+            self.log(title="Extracted SQL", message=f"{sql}")
+            return sql
+
+        # Match any ``` ... ``` code blocks
+        sqls = re.findall(r"```(.*?)```", llm_response, re.DOTALL | re.IGNORECASE)
+        if sqls:
+            sql = sqls[-1].strip()
             self.log(title="Extracted SQL", message=f"{sql}")
             return sql
 
@@ -229,7 +254,7 @@ class VannaBase(ABC):
         parsed = sqlparse.parse(sql)
 
         for statement in parsed:
-            if statement.get_type() == 'SELECT':
+            if statement.get_type() == "SELECT":
                 return True
 
         return False
@@ -251,12 +276,14 @@ class VannaBase(ABC):
             bool: True if a chart should be generated, False otherwise.
         """
 
-        if len(df) > 1 and df.select_dtypes(include=['number']).shape[1] > 0:
+        if len(df) > 1 and df.select_dtypes(include=["number"]).shape[1] > 0:
             return True
 
         return False
 
-    def generate_rewritten_question(self, last_question: str, new_question: str, **kwargs) -> str:
+    def generate_rewritten_question(
+        self, last_question: str, new_question: str, **kwargs
+    ) -> str:
         """
         **Example:**
         ```python
@@ -277,8 +304,15 @@ class VannaBase(ABC):
             return new_question
 
         prompt = [
-            self.system_message("Your goal is to combine a sequence of questions into a singular question if they are related. If the second question does not relate to the first question and is fully self-contained, return the second question. Return just the new combined question with no additional explanations. The question should theoretically be answerable with a single SQL statement."),
-            self.user_message("First question: " + last_question + "\nSecond question: " + new_question),
+            self.system_message(
+                "Your goal is to combine a sequence of questions into a singular question if they are related. If the second question does not relate to the first question and is fully self-contained, return the second question. Return just the new combined question with no additional explanations. The question should theoretically be answerable with a single SQL statement."
+            ),
+            self.user_message(
+                "First question: "
+                + last_question
+                + "\nSecond question: "
+                + new_question
+            ),
         ]
 
         return self.submit_prompt(prompt=prompt, **kwargs)
@@ -309,8 +343,8 @@ class VannaBase(ABC):
                 f"You are a helpful data assistant. The user asked the question: '{question}'\n\nThe SQL query for this question was: {sql}\n\nThe following is a pandas DataFrame with the results of the query: \n{df.head(25).to_markdown()}\n\n"
             ),
             self.user_message(
-                f"Generate a list of {n_questions} followup questions that the user might ask about this data. Respond with a list of questions, one per line. Do not answer with any explanations -- just the questions. Remember that there should be an unambiguous SQL query that can be generated from the question. Prefer questions that are answerable outside of the context of this conversation. Prefer questions that are slight modifications of the SQL query that was generated that allow digging deeper into the data. Each question will be turned into a button that the user can click to generate a new SQL query so don't use 'example' type questions. Each question must have a one-to-one correspondence with an instantiated SQL query." +
-                self._response_language()
+                f"Generate a list of {n_questions} followup questions that the user might ask about this data. Respond with a list of questions, one per line. Do not answer with any explanations -- just the questions. Remember that there should be an unambiguous SQL query that can be generated from the question. Prefer questions that are answerable outside of the context of this conversation. Prefer questions that are slight modifications of the SQL query that was generated that allow digging deeper into the data. Each question will be turned into a button that the user can click to generate a new SQL query so don't use 'example' type questions. Each question must have a one-to-one correspondence with an instantiated SQL query."
+                + self._response_language()
             ),
         ]
 
@@ -354,8 +388,8 @@ class VannaBase(ABC):
                 f"You are a helpful data assistant. The user asked the question: '{question}'\n\nThe following is a pandas DataFrame with the results of the query: \n{df.to_markdown()}\n\n"
             ),
             self.user_message(
-                "Briefly summarize the data based on the question that was asked. Do not respond with any additional explanation beyond the summary." +
-                self._response_language()
+                "Briefly summarize the data based on the question that was asked. Do not respond with any additional explanation beyond the summary."
+                + self._response_language()
             ),
         ]
 
@@ -551,7 +585,7 @@ class VannaBase(ABC):
 
     def get_sql_prompt(
         self,
-        initial_prompt : str,
+        initial_prompt: str,
         question: str,
         question_sql_list: list,
         ddl_list: list,
@@ -583,8 +617,10 @@ class VannaBase(ABC):
         """
 
         if initial_prompt is None:
-            initial_prompt = f"You are a {self.dialect} expert. " + \
-            "Please help to generate a SQL query to answer the question. Your response should ONLY be based on the given context and follow the response guidelines and format instructions. "
+            initial_prompt = (
+                f"You are a {self.dialect} expert. "
+                + "Please help to generate a SQL query to answer the question. Your response should ONLY be based on the given context and follow the response guidelines and format instructions. "
+            )
 
         initial_prompt = self.add_ddl_to_prompt(
             initial_prompt, ddl_list, max_tokens=self.max_tokens
@@ -749,7 +785,7 @@ class VannaBase(ABC):
         database: str,
         role: Union[str, None] = None,
         warehouse: Union[str, None] = None,
-        **kwargs
+        **kwargs,
     ):
         try:
             snowflake = __import__("snowflake.connector")
@@ -797,7 +833,7 @@ class VannaBase(ABC):
             account=account,
             database=database,
             client_session_keep_alive=True,
-            **kwargs
+            **kwargs,
         )
 
         def run_sql_snowflake(sql: str) -> pd.DataFrame:
@@ -823,7 +859,7 @@ class VannaBase(ABC):
         self.run_sql = run_sql_snowflake
         self.run_sql_is_set = True
 
-    def connect_to_sqlite(self, url: str, check_same_thread: bool = False,  **kwargs):
+    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]
 
@@ -848,11 +884,7 @@ class VannaBase(ABC):
             url = path
 
         # Connect to the database
-        conn = sqlite3.connect(
-            url,
-            check_same_thread=check_same_thread,
-            **kwargs
-        )
+        conn = sqlite3.connect(url, check_same_thread=check_same_thread, **kwargs)
 
         def run_sql_sqlite(sql: str):
             return pd.read_sql_query(sql, conn)
@@ -868,9 +900,8 @@ class VannaBase(ABC):
         user: str = None,
         password: str = None,
         port: int = None,
-        **kwargs
+        **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]
         **Example:**
@@ -939,15 +970,20 @@ class VannaBase(ABC):
                 user=user,
                 password=password,
                 port=port,
-                **kwargs
+                **kwargs,
             )
         except psycopg2.Error as e:
             raise ValidationError(e)
 
         def connect_to_db():
-            return psycopg2.connect(host=host, dbname=dbname,
-                        user=user, password=password, port=port, **kwargs)
-
+            return psycopg2.connect(
+                host=host,
+                dbname=dbname,
+                user=user,
+                password=password,
+                port=port,
+                **kwargs,
+            )
 
         def run_sql_postgres(sql: str) -> Union[pd.DataFrame, None]:
             conn = None
@@ -980,14 +1016,13 @@ class VannaBase(ABC):
                     raise ValidationError(e)
 
             except Exception as e:
-                        conn.rollback()
-                        raise e
+                conn.rollback()
+                raise e
 
         self.dialect = "PostgreSQL"
         self.run_sql_is_set = True
         self.run_sql = run_sql_postgres
 
-
     def connect_to_mysql(
         self,
         host: str = None,
@@ -995,9 +1030,8 @@ class VannaBase(ABC):
         user: str = None,
         password: str = None,
         port: int = None,
-        **kwargs
+        **kwargs,
     ):
-
         try:
             import pymysql.cursors
         except ImportError:
@@ -1046,7 +1080,7 @@ class VannaBase(ABC):
                 database=dbname,
                 port=port,
                 cursorclass=pymysql.cursors.DictCursor,
-                **kwargs
+                **kwargs,
             )
         except pymysql.Error as e:
             raise ValidationError(e)
@@ -1083,9 +1117,8 @@ class VannaBase(ABC):
         user: str = None,
         password: str = None,
         port: int = None,
-        **kwargs
+        **kwargs,
     ):
-
         try:
             import clickhouse_connect
         except ImportError:
@@ -1133,7 +1166,7 @@ class VannaBase(ABC):
                 username=user,
                 password=password,
                 database=dbname,
-                **kwargs
+                **kwargs,
             )
             print(conn)
         except Exception as e:
@@ -1156,13 +1189,8 @@ class VannaBase(ABC):
         self.run_sql = run_sql_clickhouse
 
     def connect_to_oracle(
-        self,
-        user: str = None,
-        password: str = None,
-        dsn: str = None,
-        **kwargs
+        self, user: str = None, password: str = None, dsn: str = None, **kwargs
     ):
-
         """
         Connect to an Oracle db using oracledb package. This is just a helper function to set [`vn.run_sql`][vanna.base.base.VannaBase.run_sql]
         **Example:**
@@ -1182,7 +1210,6 @@ class VannaBase(ABC):
         try:
             import oracledb
         except ImportError:
-
             raise DependencyError(
                 "You need to install required dependencies to execute this method,"
                 " run command: \npip install oracledb"
@@ -1192,7 +1219,9 @@ class VannaBase(ABC):
             dsn = os.getenv("DSN")
 
         if not dsn:
-            raise ImproperlyConfigured("Please set your Oracle dsn which should include host:port/sid")
+            raise ImproperlyConfigured(
+                "Please set your Oracle dsn which should include host:port/sid"
+            )
 
         if not user:
             user = os.getenv("USER")
@@ -1209,12 +1238,7 @@ class VannaBase(ABC):
         conn = None
 
         try:
-            conn = oracledb.connect(
-                user=user,
-                password=password,
-                dsn=dsn,
-                **kwargs
-            )
+            conn = oracledb.connect(user=user, password=password, dsn=dsn, **kwargs)
         except oracledb.Error as e:
             raise ValidationError(e)
 
@@ -1222,7 +1246,9 @@ class VannaBase(ABC):
             if conn:
                 try:
                     sql = sql.rstrip()
-                    if sql.endswith(';'): #fix for a known problem with Oracle db where an extra ; will cause an error.
+                    if sql.endswith(
+                        ";"
+                    ):  # fix for a known problem with Oracle db where an extra ; will cause an error.
                         sql = sql[:-1]
 
                     cs = conn.cursor()
@@ -1247,10 +1273,7 @@ class VannaBase(ABC):
         self.run_sql = run_sql_oracle
 
     def connect_to_bigquery(
-        self,
-        cred_file_path: str = None,
-        project_id: str = None,
-        **kwargs
+        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]
@@ -1299,7 +1322,7 @@ class VannaBase(ABC):
         if not cred_file_path:
             try:
                 conn = bigquery.Client(project=project_id)
-            except:
+            except Exception:
                 print("Could not found any google cloud implicit credentials")
         else:
             # Validate file path and pemissions
@@ -1314,11 +1337,9 @@ class VannaBase(ABC):
 
             try:
                 conn = bigquery.Client(
-                    project=project_id,
-                    credentials=credentials,
-                    **kwargs
+                    project=project_id, credentials=credentials, **kwargs
                 )
-            except:
+            except Exception:
                 raise ImproperlyConfigured(
                     "Could not connect to bigquery please correct credentials"
                 )
@@ -1430,20 +1451,21 @@ class VannaBase(ABC):
         self.dialect = "T-SQL / Microsoft SQL Server"
         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',
+        catalog: str = "hive",
+        schema: str = "default",
         user: str = None,
         password: str = None,
         port: int = None,
         combined_pem_path: str = None,
-        protocol: str = 'https',
+        protocol: str = "https",
         requests_kwargs: dict = None,
-        **kwargs
+        **kwargs,
     ):
-      """
+        """
         Connect to a Presto database using the specified parameters.
 
         Args:
@@ -1463,101 +1485,103 @@ class VannaBase(ABC):
 
         Returns:
             None
-      """
-      try:
-        from pyhive import presto
-      except ImportError:
-        raise DependencyError(
-          "You need to install required dependencies to execute this method,"
-          " run command: \npip install pyhive"
-        )
+        """
+        try:
+            from pyhive import presto
+        except ImportError:
+            raise DependencyError(
+                "You need to install required dependencies to execute this method,"
+                " run command: \npip install pyhive"
+            )
 
-      if not host:
-        host = os.getenv("PRESTO_HOST")
-
-      if not host:
-        raise ImproperlyConfigured("Please set your presto host")
-
-      if not catalog:
-        catalog = os.getenv("PRESTO_CATALOG")
-
-      if not catalog:
-        raise ImproperlyConfigured("Please set your presto catalog")
-
-      if not user:
-        user = os.getenv("PRESTO_USER")
-
-      if not user:
-        raise ImproperlyConfigured("Please set your presto user")
-
-      if not password:
-        password = os.getenv("PRESTO_PASSWORD")
-
-      if not port:
-        port = os.getenv("PRESTO_PORT")
-
-      if not port:
-        raise ImproperlyConfigured("Please set your presto port")
-
-      conn = None
-
-      try:
-        if requests_kwargs is None and combined_pem_path is not None:
-          # use the combined pem file to verify the SSL connection
-          requests_kwargs = {
-            'verify': combined_pem_path,  # 使用转换后得到的 PEM 文件进行 SSL 验证
-          }
-        conn = presto.Connection(host=host,
-                                 username=user,
-                                 password=password,
-                                 catalog=catalog,
-                                 schema=schema,
-                                 port=port,
-                                 protocol=protocol,
-                                 requests_kwargs=requests_kwargs,
-                                 **kwargs)
-      except presto.Error as e:
-        raise ValidationError(e)
-
-      def run_sql_presto(sql: str) -> Union[pd.DataFrame, None]:
-        if conn:
-          try:
-            sql = sql.rstrip()
-            # fix for a known problem with presto db where an extra ; will cause an error.
-            if sql.endswith(';'):
-                sql = sql[:-1]
-            cs = conn.cursor()
-            cs.execute(sql)
-            results = cs.fetchall()
+        if not host:
+            host = os.getenv("PRESTO_HOST")
 
-            # Create a pandas dataframe from the results
-            df = pd.DataFrame(
-              results, columns=[desc[0] for desc in cs.description]
-            )
-            return df
+        if not host:
+            raise ImproperlyConfigured("Please set your presto host")
 
-          except presto.Error as e:
-            print(e)
+        if not catalog:
+            catalog = os.getenv("PRESTO_CATALOG")
+
+        if not catalog:
+            raise ImproperlyConfigured("Please set your presto catalog")
+
+        if not user:
+            user = os.getenv("PRESTO_USER")
+
+        if not user:
+            raise ImproperlyConfigured("Please set your presto user")
+
+        if not password:
+            password = os.getenv("PRESTO_PASSWORD")
+
+        if not port:
+            port = os.getenv("PRESTO_PORT")
+
+        if not port:
+            raise ImproperlyConfigured("Please set your presto port")
+
+        conn = None
+
+        try:
+            if requests_kwargs is None and combined_pem_path is not None:
+                # use the combined pem file to verify the SSL connection
+                requests_kwargs = {
+                    "verify": combined_pem_path,  # 使用转换后得到的 PEM 文件进行 SSL 验证
+                }
+            conn = presto.Connection(
+                host=host,
+                username=user,
+                password=password,
+                catalog=catalog,
+                schema=schema,
+                port=port,
+                protocol=protocol,
+                requests_kwargs=requests_kwargs,
+                **kwargs,
+            )
+        except presto.Error as e:
             raise ValidationError(e)
 
-          except Exception as e:
-            print(e)
-            raise e
+        def run_sql_presto(sql: str) -> Union[pd.DataFrame, None]:
+            if conn:
+                try:
+                    sql = sql.rstrip()
+                    # fix for a known problem with presto db where an extra ; will cause an error.
+                    if sql.endswith(";"):
+                        sql = sql[:-1]
+                    cs = conn.cursor()
+                    cs.execute(sql)
+                    results = cs.fetchall()
+
+                    # Create a pandas dataframe from the results
+                    df = pd.DataFrame(
+                        results, columns=[desc[0] for desc in cs.description]
+                    )
+                    return df
 
-      self.run_sql_is_set = True
-      self.run_sql = run_sql_presto
+                except presto.Error as e:
+                    print(e)
+                    raise ValidationError(e)
+
+                except Exception as e:
+                    print(e)
+                    raise e
+
+        self.run_sql_is_set = True
+        self.run_sql = run_sql_presto
 
     def connect_to_hive(
         self,
         host: str = None,
-        dbname: str = 'default',
+        dbname: str = "default",
         user: str = None,
         password: str = None,
         port: int = None,
-        auth: str = 'CUSTOM',
-        **kwargs
+        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]
         Connect to a Hive database. This is just a helper function to set [`vn.run_sql`][vanna.base.base.VannaBase.run_sql]
 
@@ -1571,78 +1595,80 @@ class VannaBase(ABC):
 
         Returns:
             None
-      """
-
-      try:
-        from pyhive import hive
-      except ImportError:
-        raise DependencyError(
-          "You need to install required dependencies to execute this method,"
-          " run command: \npip install pyhive"
-        )
-
-      if not host:
-        host = os.getenv("HIVE_HOST")
+        """
 
-      if not host:
-        raise ImproperlyConfigured("Please set your hive host")
+        try:
+            from pyhive import hive
+        except ImportError:
+            raise DependencyError(
+                "You need to install required dependencies to execute this method,"
+                " run command: \npip install pyhive"
+            )
 
-      if not dbname:
-        dbname = os.getenv("HIVE_DATABASE")
+        if not host:
+            host = os.getenv("HIVE_HOST")
 
-      if not dbname:
-        raise ImproperlyConfigured("Please set your hive database")
+        if not host:
+            raise ImproperlyConfigured("Please set your hive host")
 
-      if not user:
-        user = os.getenv("HIVE_USER")
+        if not dbname:
+            dbname = os.getenv("HIVE_DATABASE")
 
-      if not user:
-        raise ImproperlyConfigured("Please set your hive user")
+        if not dbname:
+            raise ImproperlyConfigured("Please set your hive database")
 
-      if not password:
-        password = os.getenv("HIVE_PASSWORD")
+        if not user:
+            user = os.getenv("HIVE_USER")
 
-      if not port:
-        port = os.getenv("HIVE_PORT")
+        if not user:
+            raise ImproperlyConfigured("Please set your hive user")
 
-      if not port:
-        raise ImproperlyConfigured("Please set your hive port")
+        if not password:
+            password = os.getenv("HIVE_PASSWORD")
 
-      conn = None
+        if not port:
+            port = os.getenv("HIVE_PORT")
 
-      try:
-        conn = hive.Connection(host=host,
-                               username=user,
-                               password=password,
-                               database=dbname,
-                               port=port,
-                               auth=auth)
-      except hive.Error as e:
-        raise ValidationError(e)
+        if not port:
+            raise ImproperlyConfigured("Please set your hive port")
 
-      def run_sql_hive(sql: str) -> Union[pd.DataFrame, None]:
-        if conn:
-          try:
-            cs = conn.cursor()
-            cs.execute(sql)
-            results = cs.fetchall()
+        conn = None
 
-            # Create a pandas dataframe from the results
-            df = pd.DataFrame(
-              results, columns=[desc[0] for desc in cs.description]
+        try:
+            conn = hive.Connection(
+                host=host,
+                username=user,
+                password=password,
+                database=dbname,
+                port=port,
+                auth=auth,
             )
-            return df
-
-          except hive.Error as e:
-            print(e)
+        except hive.Error as e:
             raise ValidationError(e)
 
-          except Exception as e:
-            print(e)
-            raise e
+        def run_sql_hive(sql: str) -> Union[pd.DataFrame, None]:
+            if conn:
+                try:
+                    cs = conn.cursor()
+                    cs.execute(sql)
+                    results = cs.fetchall()
 
-      self.run_sql_is_set = True
-      self.run_sql = run_sql_hive
+                    # Create a pandas dataframe from the results
+                    df = pd.DataFrame(
+                        results, columns=[desc[0] for desc in cs.description]
+                    )
+                    return df
+
+                except hive.Error as e:
+                    print(e)
+                    raise ValidationError(e)
+
+                except Exception as e:
+                    print(e)
+                    raise e
+
+        self.run_sql_is_set = True
+        self.run_sql = run_sql_hive
 
     def run_sql(self, sql: str, **kwargs) -> pd.DataFrame:
         """
@@ -1700,22 +1726,23 @@ class VannaBase(ABC):
             question = input("Enter a question: ")
 
         try:
-            sql = self.generate_sql(question=question, allow_llm_to_see_data=allow_llm_to_see_data)
+            sql = self.generate_sql(
+                question=question, allow_llm_to_see_data=allow_llm_to_see_data
+            )
         except Exception as e:
             print(e)
             return None, None, None
 
         if print_results:
             try:
-                Code = __import__("IPython.display", fromList=["Code"]).Code
+                from IPython.display import Code, display
+
                 display(Code(sql))
             except Exception as e:
                 print(sql)
 
         if self.run_sql_is_set is False:
-            print(
-                "If you want to run the SQL query, connect to a database first."
-            )
+            print("If you want to run the SQL query, connect to a database first.")
 
             if print_results:
                 return None
@@ -1759,6 +1786,7 @@ class VannaBase(ABC):
                             fig.show()
                 except Exception as e:
                     # Print stack trace
+                    traceback.print_stack()
                     traceback.print_exc()
                     print("Couldn't run plotly code: ", e)
                     if print_results:
@@ -1874,12 +1902,8 @@ class VannaBase(ABC):
         table_column = df.columns[
             df.columns.str.lower().str.contains("table_name")
         ].to_list()[0]
-        columns = [database_column,
-                    schema_column,
-                    table_column]
-        candidates = ["column_name",
-                      "data_type",
-                      "comment"]
+        columns = [database_column, schema_column, table_column]
+        candidates = ["column_name", "data_type", "comment"]
         matches = df.columns.str.lower().str.contains("|".join(candidates), regex=True)
         columns += df.columns[matches].to_list()