vanna 0.5.5__py3-none-any.whl → 0.6.1__py3-none-any.whl

vanna/advanced/__init__.py

@@ -0,0 +1,26 @@
+from abc import ABC, abstractmethod
+
+
+class VannaAdvanced(ABC):
+    def __init__(self, config=None):
+        self.config = config
+
+    @abstractmethod
+    def get_function(self, question: str, additional_data: dict = {}) -> dict:
+        pass
+
+    @abstractmethod
+    def create_function(self, question: str, sql: str, plotly_code: str, **kwargs) -> dict:
+        pass
+
+    @abstractmethod
+    def update_function(self, old_function_name: str, updated_function: dict) -> bool:
+        pass
+
+    @abstractmethod
+    def delete_function(self, function_name: str) -> bool:
+        pass
+
+    @abstractmethod
+    def get_all_functions(self) -> list:
+        pass

vanna/base/base.py

@@ -79,9 +79,10 @@ class VannaBase(ABC):
         self.static_documentation = ""
         self.dialect = self.config.get("dialect", "SQL")
         self.language = self.config.get("language", None)
+        self.max_tokens = self.config.get("max_tokens", 14000)
 
     def log(self, message: str, title: str = "Info"):
-        print(message)
+        print(f"{title}: {message}")
 
     def _response_language(self) -> str:
         if self.language is None:
@@ -559,14 +560,14 @@ class VannaBase(ABC):
             "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=14000
+            initial_prompt, ddl_list, max_tokens=self.max_tokens
         )
 
         if self.static_documentation != "":
             doc_list.append(self.static_documentation)
 
         initial_prompt = self.add_documentation_to_prompt(
-            initial_prompt, doc_list, max_tokens=14000
+            initial_prompt, doc_list, max_tokens=self.max_tokens
         )
 
         initial_prompt += (
@@ -603,15 +604,15 @@ class VannaBase(ABC):
         initial_prompt = f"The user initially asked the question: '{question}': \n\n"
 
         initial_prompt = self.add_ddl_to_prompt(
-            initial_prompt, ddl_list, max_tokens=14000
+            initial_prompt, ddl_list, max_tokens=self.max_tokens
         )
 
         initial_prompt = self.add_documentation_to_prompt(
-            initial_prompt, doc_list, max_tokens=14000
+            initial_prompt, doc_list, max_tokens=self.max_tokens
         )
 
         initial_prompt = self.add_sql_to_prompt(
-            initial_prompt, question_sql_list, max_tokens=14000
+            initial_prompt, question_sql_list, max_tokens=self.max_tokens
         )
 
         message_log = [self.system_message(initial_prompt)]
@@ -991,6 +992,7 @@ class VannaBase(ABC):
         def run_sql_mysql(sql: str) -> Union[pd.DataFrame, None]:
             if conn:
                 try:
+                    conn.ping(reconnect=True)
                     cs = conn.cursor()
                     cs.execute(sql)
                     results = cs.fetchall()
@@ -1022,11 +1024,11 @@ class VannaBase(ABC):
     ):
 
         try:
-            from clickhouse_driver import connect
+            import clickhouse_connect
         except ImportError:
             raise DependencyError(
                 "You need to install required dependencies to execute this method,"
-                " run command: \npip install clickhouse-driver"
+                " run command: \npip install clickhouse_connect"
             )
 
         if not host:
@@ -1062,12 +1064,13 @@ class VannaBase(ABC):
         conn = None
 
         try:
-            conn = connect(host=host,
-                                   user=user,
-                                   password=password,
-                                   database=dbname,
-                                   port=port,
-                                  )
+            conn = clickhouse_connect.get_client(
+                host=host,
+                port=port,
+                username=user,
+                password=password,
+                database=dbname,
+            )
             print(conn)
         except Exception as e:
             raise ValidationError(e)
@@ -1075,19 +1078,16 @@ class VannaBase(ABC):
         def run_sql_clickhouse(sql: str) -> Union[pd.DataFrame, None]:
             if conn:
                 try:
-                  cs = conn.cursor()
-                  cs.execute(sql)
-                  results = cs.fetchall()
+                    result = conn.query(sql)
+                    results = result.result_rows
 
-                  # Create a pandas dataframe from the results
-                  df = pd.DataFrame(
-                    results, columns=[desc[0] for desc in cs.description]
-                  )
-                  return df
+                    # Create a pandas dataframe from the results
+                    df = pd.DataFrame(results, columns=result.column_names)
+                    return df
 
                 except Exception as e:
                     raise e
-
+        
         self.run_sql_is_set = True
         self.run_sql = run_sql_clickhouse
 
@@ -1597,6 +1597,7 @@ class VannaBase(ABC):
         print_results: bool = True,
         auto_train: bool = True,
         visualize: bool = True,  # if False, will not generate plotly code
+        allow_llm_to_see_data: bool = False,
     ) -> Union[
         Tuple[
             Union[str, None],
@@ -1627,7 +1628,7 @@ class VannaBase(ABC):
             question = input("Enter a question: ")
 
         try:
-            sql = self.generate_sql(question=question)
+            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

vanna/flask/__init__.py

@@ -1,5 +1,6 @@
 import json
 import logging
+import os
 import sys
 import uuid
 from abc import ABC, abstractmethod
@@ -7,7 +8,7 @@ from functools import wraps
 
 import flask
 import requests
-from flask import Flask, Response, jsonify, request
+from flask import Flask, Response, jsonify, request, send_from_directory
 from flask_sock import Sock
 
 from .assets import css_content, html_content, js_content
@@ -151,7 +152,10 @@ class VannaFlaskApp:
                     auto_fix_sql=True,
                     ask_results_correct=True,
                     followup_questions=True,
-                    summarization=True
+                    summarization=True,
+                    function_generation=True,
+                    index_html_path=None,
+                    assets_folder=None,
                  ):
         """
         Expose a Flask app that can be used to interact with a Vanna instance.
@@ -176,6 +180,8 @@ class VannaFlaskApp:
             ask_results_correct: Whether to ask the user if the results are correct. Defaults to True.
             followup_questions: Whether to show followup questions. Defaults to True.
             summarization: Whether to show summarization. Defaults to True.
+            index_html_path: Path to the index.html. Defaults to None, which will use the default index.html
+            assets_folder: The location where you'd like to serve the static assets from. Defaults to None, which will use hardcoded Python variables.
 
         Returns:
             None
@@ -202,6 +208,9 @@ class VannaFlaskApp:
         self.ask_results_correct = ask_results_correct
         self.followup_questions = followup_questions
         self.summarization = summarization
+        self.function_generation = function_generation and hasattr(vn, "get_function")
+        self.index_html_path = index_html_path
+        self.assets_folder = assets_folder
 
         log = logging.getLogger("werkzeug")
         log.setLevel(logging.ERROR)
@@ -247,6 +256,7 @@ class VannaFlaskApp:
                 "ask_results_correct": self.ask_results_correct,
                 "followup_questions": self.followup_questions,
                 "summarization": self.summarization,
+                "function_generation": self.function_generation,
             }
 
             config = self.auth.override_config_for_user(user, config)
@@ -345,6 +355,56 @@ class VannaFlaskApp:
                     }
                 )
 
+        @self.flask_app.route("/api/v0/get_function", methods=["GET"])
+        @self.requires_auth
+        def get_function(user: any):
+            question = flask.request.args.get("question")
+
+            if question is None:
+                return jsonify({"type": "error", "error": "No question provided"})
+
+            if not hasattr(vn, "get_function"):
+                return jsonify({"type": "error", "error": "This setup does not support function generation."})
+
+            id = self.cache.generate_id(question=question)
+            function = vn.get_function(question=question)
+
+            if function is None:
+                return jsonify({"type": "error", "error": "No function found"})
+
+            if 'instantiated_sql' not in function:
+                self.vn.log(f"No instantiated SQL found for {question} in {function}")
+                return jsonify({"type": "error", "error": "No instantiated SQL found"})
+
+            self.cache.set(id=id, field="question", value=question)
+            self.cache.set(id=id, field="sql", value=function['instantiated_sql'])
+
+            if 'instantiated_post_processing_code' in function and function['instantiated_post_processing_code'] is not None and len(function['instantiated_post_processing_code']) > 0:
+                self.cache.set(id=id, field="plotly_code", value=function['instantiated_post_processing_code'])
+
+            return jsonify(
+                {
+                    "type": "function",
+                    "id": id,
+                    "function": function,
+                }
+            )
+
+        @self.flask_app.route("/api/v0/get_all_functions", methods=["GET"])
+        @self.requires_auth
+        def get_all_functions(user: any):
+            if not hasattr(vn, "get_all_functions"):
+                return jsonify({"type": "error", "error": "This setup does not support function generation."})
+
+            functions = vn.get_all_functions()
+
+            return jsonify(
+                {
+                    "type": "functions",
+                    "functions": functions,
+                }
+            )
+
         @self.flask_app.route("/api/v0/run_sql", methods=["GET"])
         @self.requires_auth
         @self.requires_cache(["sql"])
@@ -438,11 +498,18 @@ class VannaFlaskApp:
                 question = f"{question}. When generating the chart, use these special instructions: {chart_instructions}"
 
             try:
-                code = vn.generate_plotly_code(
-                    question=question,
-                    sql=sql,
-                    df_metadata=f"Running df.dtypes gives:\n {df.dtypes}",
-                )
+                # If chart_instructions is not set then attempt to retrieve the code from the cache
+                if chart_instructions is None or len(chart_instructions) == 0:
+                    code = self.cache.get(id=id, field="plotly_code")
+
+                if code is None:
+                    code = vn.generate_plotly_code(
+                        question=question,
+                        sql=sql,
+                        df_metadata=f"Running df.dtypes gives:\n {df.dtypes}",
+                    )
+                    self.cache.set(id=id, field="plotly_code", value=code)
+
                 fig = vn.get_plotly_figure(plotly_code=code, df=df, dark_mode=False)
                 fig_json = fig.to_json()
 
@@ -518,6 +585,46 @@ class VannaFlaskApp:
                 print("TRAINING ERROR", e)
                 return jsonify({"type": "error", "error": str(e)})
 
+        @self.flask_app.route("/api/v0/create_function", methods=["GET"])
+        @self.requires_auth
+        @self.requires_cache(["question", "sql"])
+        def create_function(user: any, id: str, question: str, sql: str):
+            plotly_code = self.cache.get(id=id, field="plotly_code")
+
+            if plotly_code is None:
+                plotly_code = ""
+
+            function_data = self.vn.create_function(question=question, sql=sql, plotly_code=plotly_code)
+
+            return jsonify(
+                {
+                    "type": "function_template",
+                    "id": id,
+                    "function_template": function_data,
+                }
+            )
+
+        @self.flask_app.route("/api/v0/update_function", methods=["POST"])
+        @self.requires_auth
+        def update_function(user: any):
+            old_function_name = flask.request.json.get("old_function_name")
+            updated_function = flask.request.json.get("updated_function")
+
+            print("old_function_name", old_function_name)
+            print("updated_function", updated_function)
+
+            updated = vn.update_function(old_function_name=old_function_name, updated_function=updated_function)
+
+            return jsonify({"success": updated})
+
+        @self.flask_app.route("/api/v0/delete_function", methods=["POST"])
+        @self.requires_auth
+        def delete_function(user: any):
+            function_name = flask.request.json.get("function_name")
+
+            return jsonify({"success": vn.delete_function(function_name=function_name)})
+
+
         @self.flask_app.route("/api/v0/generate_followup_questions", methods=["GET"])
         @self.requires_auth
         @self.requires_cache(["df", "question", "sql"])
@@ -578,8 +685,8 @@ class VannaFlaskApp:
         @self.flask_app.route("/api/v0/load_question", methods=["GET"])
         @self.requires_auth
         @self.requires_cache(
-            ["question", "sql", "df", "fig_json"],
-            optional_fields=["summary"]
+            ["question", "sql", "df"],
+            optional_fields=["summary", "fig_json"]
         )
         def load_question(user: any, id: str, question, sql, df, fig_json, summary):
             try:
@@ -616,6 +723,9 @@ class VannaFlaskApp:
 
         @self.flask_app.route("/assets/<path:filename>")
         def proxy_assets(filename):
+            if self.assets_folder:
+                return send_from_directory(self.assets_folder, filename)
+
             if ".css" in filename:
                 return Response(css_content, mimetype="text/css")
 
@@ -663,6 +773,10 @@ class VannaFlaskApp:
         @self.flask_app.route("/", defaults={"path": ""})
         @self.flask_app.route("/<path:path>")
         def hello(path: str):
+            if self.index_html_path:
+                directory = os.path.dirname(self.index_html_path)
+                filename = os.path.basename(self.index_html_path)
+                return send_from_directory(directory=directory, path=filename)
             return html_content
 
     def run(self, *args, **kwargs):
@@ -692,4 +806,4 @@ class VannaFlaskApp:
                 print("Your app is running at:")
                 print("http://localhost:8084")
 
-            self.flask_app.run(host="0.0.0.0", port=8084, debug=self.debug)
+            self.flask_app.run(host="0.0.0.0", port=8084, debug=self.debug, use_reloader=False)