vanna 0.6.3__py3-none-any.whl → 0.6.4__py3-none-any.whl

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/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]

vanna/flask/__init__.py

@@ -8,11 +8,13 @@ from functools import wraps
 
 import flask
 import requests
+from flasgger import Swagger
 from flask import Flask, Response, jsonify, request, send_from_directory
 from flask_sock import Sock
 
 from .assets import css_content, html_content, js_content
 from .auth import AuthInterface, NoAuth
+from ..base import VannaBase
 
 
 class Cache(ABC):
@@ -88,7 +90,8 @@ class MemoryCache(Cache):
         if id in self.cache:
             del self.cache[id]
 
-class VannaFlaskApp:
+
+class VannaFlaskAPI:
     flask_app = None
 
     def requires_cache(self, required_fields, optional_fields=[]):
@@ -135,30 +138,17 @@ class VannaFlaskApp:
 
         return decorated
 
-    def __init__(self, vn, cache: Cache = MemoryCache(),
-                    auth: AuthInterface = NoAuth(),
-                    debug=True,
-                    allow_llm_to_see_data=False,
-                    logo="https://img.vanna.ai/vanna-flask.svg",
-                    title="Welcome to Vanna.AI",
-                    subtitle="Your AI-powered copilot for SQL queries.",
-                    show_training_data=True,
-                    suggested_questions=True,
-                    sql=True,
-                    table=True,
-                    csv_download=True,
-                    chart=True,
-                    redraw_chart=True,
-                    auto_fix_sql=True,
-                    ask_results_correct=True,
-                    followup_questions=True,
-                    summarization=True,
-                    function_generation=True,
-                    index_html_path=None,
-                    assets_folder=None,
-                 ):
+    def __init__(
+        self,
+        vn: VannaBase,
+        cache: Cache = MemoryCache(),
+        auth: AuthInterface = NoAuth(),
+        debug=True,
+        allow_llm_to_see_data=False,
+        chart=True,
+    ):
         """
-        Expose a Flask app that can be used to interact with a Vanna instance.
+        Expose a Flask API that can be used to interact with a Vanna instance.
 
         Args:
             vn: The Vanna instance to interact with.
@@ -166,52 +156,30 @@ class VannaFlaskApp:
             auth: The authentication method to use. Defaults to NoAuth, which doesn't require authentication. You can also pass in a custom authentication method that implements the AuthInterface interface.
             debug: Show the debug console. Defaults to True.
             allow_llm_to_see_data: Whether to allow the LLM to see data. Defaults to False.
-            logo: The logo to display in the UI. Defaults to the Vanna logo.
-            title: The title to display in the UI. Defaults to "Welcome to Vanna.AI".
-            subtitle: The subtitle to display in the UI. Defaults to "Your AI-powered copilot for SQL queries.".
-            show_training_data: Whether to show the training data in the UI. Defaults to True.
-            suggested_questions: Whether to show suggested questions in the UI. Defaults to True.
-            sql: Whether to show the SQL input in the UI. Defaults to True.
-            table: Whether to show the table output in the UI. Defaults to True.
-            csv_download: Whether to allow downloading the table output as a CSV file. Defaults to True.
             chart: Whether to show the chart output in the UI. Defaults to True.
-            redraw_chart: Whether to allow redrawing the chart. Defaults to True.
-            auto_fix_sql: Whether to allow auto-fixing SQL errors. Defaults to True.
-            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
         """
+
         self.flask_app = Flask(__name__)
+
+        self.swagger = Swagger(
+          self.flask_app, template={"info": {"title": "Vanna API"}}
+        )
         self.sock = Sock(self.flask_app)
         self.ws_clients = []
         self.vn = vn
-        self.debug = debug
         self.auth = auth
         self.cache = cache
+        self.debug = debug
         self.allow_llm_to_see_data = allow_llm_to_see_data
-        self.logo = logo
-        self.title = title
-        self.subtitle = subtitle
-        self.show_training_data = show_training_data
-        self.suggested_questions = suggested_questions
-        self.sql = sql
-        self.table = table
-        self.csv_download = csv_download
         self.chart = chart
-        self.redraw_chart = redraw_chart
-        self.auto_fix_sql = auto_fix_sql
-        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
-
+        self.config = {
+          "debug": debug,
+          "allow_llm_to_see_data": allow_llm_to_see_data,
+          "chart": chart,
+        }
         log = logging.getLogger("werkzeug")
         log.setLevel(logging.ERROR)
 
@@ -225,42 +193,27 @@ class VannaFlaskApp:
 
             self.vn.log = log
 
-        @self.flask_app.route("/auth/login", methods=["POST"])
-        def login():
-            return self.auth.login_handler(flask.request)
-
-        @self.flask_app.route("/auth/callback", methods=["GET"])
-        def callback():
-            return self.auth.callback_handler(flask.request)
-
-        @self.flask_app.route("/auth/logout", methods=["GET"])
-        def logout():
-            return self.auth.logout_handler(flask.request)
-
         @self.flask_app.route("/api/v0/get_config", methods=["GET"])
         @self.requires_auth
         def get_config(user: any):
-            config = {
-                "debug": self.debug,
-                "logo": self.logo,
-                "title": self.title,
-                "subtitle": self.subtitle,
-                "show_training_data": self.show_training_data,
-                "suggested_questions": self.suggested_questions,
-                "sql": self.sql,
-                "table": self.table,
-                "csv_download": self.csv_download,
-                "chart": self.chart,
-                "redraw_chart": self.redraw_chart,
-                "auto_fix_sql": self.auto_fix_sql,
-                "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)
-
+            """
+            Get the configuration for a user
+            ---
+            parameters:
+              - name: user
+                in: query
+            responses:
+              200:
+                schema:
+                  type: object
+                  properties:
+                    type:
+                      type: string
+                      default: config
+                    config:
+                      type: object
+            """
+            config = self.auth.override_config_for_user(user, self.config)
             return jsonify(
                 {
                     "type": "config",
@@ -271,6 +224,28 @@ class VannaFlaskApp:
         @self.flask_app.route("/api/v0/generate_questions", methods=["GET"])
         @self.requires_auth
         def generate_questions(user: any):
+            """
+            Generate questions
+            ---
+            parameters:
+              - name: user
+                in: query
+            responses:
+              200:
+                schema:
+                  type: object
+                  properties:
+                    type:
+                      type: string
+                      default: question_list
+                    questions:
+                      type: array
+                      items:
+                        type: string
+                    header:
+                      type: string
+                      default: Here are some questions you can ask
+            """
             # If self has an _model attribute and model=='chinook'
             if hasattr(self.vn, "_model") and self.vn._model == "chinook":
                 return jsonify(
@@ -327,6 +302,29 @@ class VannaFlaskApp:
         @self.flask_app.route("/api/v0/generate_sql", methods=["GET"])
         @self.requires_auth
         def generate_sql(user: any):
+            """
+            Generate SQL from a question
+            ---
+            parameters:
+              - name: user
+                in: query
+              - name: question
+                in: query
+                type: string
+                required: true
+            responses:
+              200:
+                schema:
+                  type: object
+                  properties:
+                    type:
+                      type: string
+                      default: sql
+                    id:
+                      type: string
+                    text:
+                      type: string
+            """
             question = flask.request.args.get("question")
 
             if question is None:
@@ -358,6 +356,29 @@ class VannaFlaskApp:
         @self.flask_app.route("/api/v0/get_function", methods=["GET"])
         @self.requires_auth
         def get_function(user: any):
+            """
+            Get a function from a question
+            ---
+            parameters:
+              - name: user
+                in: query
+              - name: question
+                in: query
+                type: string
+                required: true
+            responses:
+              200:
+                schema:
+                  type: object
+                  properties:
+                    type:
+                      type: string
+                      default: function
+                    id:
+                      type: object
+                    function:
+                      type: string
+            """
             question = flask.request.args.get("question")
 
             if question is None:
@@ -393,6 +414,23 @@ class VannaFlaskApp:
         @self.flask_app.route("/api/v0/get_all_functions", methods=["GET"])
         @self.requires_auth
         def get_all_functions(user: any):
+            """
+            Get all the functions
+            ---
+            parameters:
+              - name: user
+                in: query
+            responses:
+              200:
+                schema:
+                  type: object
+                  properties:
+                    type:
+                      type: string
+                      default: functions
+                    functions:
+                      type: array
+            """
             if not hasattr(vn, "get_all_functions"):
                 return jsonify({"type": "error", "error": "This setup does not support function generation."})
 
@@ -409,6 +447,31 @@ class VannaFlaskApp:
         @self.requires_auth
         @self.requires_cache(["sql"])
         def run_sql(user: any, id: str, sql: str):
+            """
+            Run SQL
+            ---
+            parameters:
+              - name: user
+                in: query
+              - name: id
+                in: query|body
+                type: string
+                required: true
+            responses:
+              200:
+                schema:
+                  type: object
+                  properties:
+                    type:
+                      type: string
+                      default: df
+                    id:
+                      type: string
+                    df:
+                      type: object
+                    should_generate_chart:
+                      type: boolean
+            """
             try:
                 if not vn.run_sql_is_set:
                     return jsonify(
@@ -437,7 +500,34 @@ class VannaFlaskApp:
         @self.flask_app.route("/api/v0/fix_sql", methods=["POST"])
         @self.requires_auth
         @self.requires_cache(["question", "sql"])
-        def fix_sql(user: any, id: str, question:str, sql: str):
+        def fix_sql(user: any, id: str, question: str, sql: str):
+            """
+            Fix SQL
+            ---
+            parameters:
+              - name: user
+                in: query
+              - name: id
+                in: query|body
+                type: string
+                required: true
+              - name: error
+                in: body
+                type: string
+                required: true
+            responses:
+              200:
+                schema:
+                  type: object
+                  properties:
+                    type:
+                      type: string
+                      default: sql
+                    id:
+                      type: string
+                    text:
+                      type: string
+            """
             error = flask.request.json.get("error")
 
             if error is None:
@@ -462,6 +552,33 @@ class VannaFlaskApp:
         @self.requires_auth
         @self.requires_cache([])
         def update_sql(user: any, id: str):
+            """
+            Update SQL
+            ---
+            parameters:
+              - name: user
+                in: query
+              - name: id
+                in: query|body
+                type: string
+                required: true
+              - name: sql
+                in: body
+                type: string
+                required: true
+            responses:
+              200:
+                schema:
+                  type: object
+                  properties:
+                    type:
+                      type: string
+                      default: sql
+                    id:
+                      type: string
+                    text:
+                      type: string
+            """
             sql = flask.request.json.get('sql')
 
             if sql is None:
@@ -480,6 +597,20 @@ class VannaFlaskApp:
         @self.requires_auth
         @self.requires_cache(["df"])
         def download_csv(user: any, id: str, df):
+            """
+            Download CSV
+            ---
+            parameters:
+              - name: user
+                in: query
+              - name: id
+                in: query|body
+                type: string
+                required: true
+            responses:
+              200:
+                description: download CSV
+            """
             csv = df.to_csv()
 
             return Response(
@@ -492,6 +623,32 @@ class VannaFlaskApp:
         @self.requires_auth
         @self.requires_cache(["df", "question", "sql"])
         def generate_plotly_figure(user: any, id: str, df, question, sql):
+            """
+            Generate plotly figure
+            ---
+            parameters:
+              - name: user
+                in: query
+              - name: id
+                in: query|body
+                type: string
+                required: true
+              - name: chart_instructions
+                in: body
+                type: string
+            responses:
+              200:
+                schema:
+                  type: object
+                  properties:
+                    type:
+                      type: string
+                      default: plotly_figure
+                    id:
+                      type: string
+                    fig:
+                      type: object
+            """
             chart_instructions = flask.request.args.get('chart_instructions')
 
             try:
@@ -530,6 +687,26 @@ class VannaFlaskApp:
         @self.flask_app.route("/api/v0/get_training_data", methods=["GET"])
         @self.requires_auth
         def get_training_data(user: any):
+            """
+            Get all training data
+            ---
+            parameters:
+              - name: user
+                in: query
+            responses:
+              200:
+                schema:
+                  type: object
+                  properties:
+                    type:
+                      type: string
+                      default: df
+                    id:
+                      type: string
+                      default: training_data
+                    df:
+                      type: object
+            """
             df = vn.get_training_data()
 
             if df is None or len(df) == 0:
@@ -551,6 +728,24 @@ class VannaFlaskApp:
         @self.flask_app.route("/api/v0/remove_training_data", methods=["POST"])
         @self.requires_auth
         def remove_training_data(user: any):
+            """
+            Remove training data
+            ---
+            parameters:
+              - name: user
+                in: query
+              - name: id
+                in: body
+                type: string
+                required: true
+            responses:
+              200:
+                schema:
+                  type: object
+                  properties:
+                    success:
+                      type: boolean
+            """
             # Get id from the JSON body
             id = flask.request.json.get("id")
 
@@ -567,6 +762,32 @@ class VannaFlaskApp:
         @self.flask_app.route("/api/v0/train", methods=["POST"])
         @self.requires_auth
         def add_training_data(user: any):
+            """
+            Add training data
+            ---
+            parameters:
+              - name: user
+                in: query
+              - name: question
+                in: body
+                type: string
+              - name: sql
+                in: body
+                type: string
+              - name: ddl
+                in: body
+                type: string
+              - name: documentation
+                in: body
+                type: string
+            responses:
+              200:
+                schema:
+                  type: object
+                  properties:
+                    id:
+                      type: string
+            """
             question = flask.request.json.get("question")
             sql = flask.request.json.get("sql")
             ddl = flask.request.json.get("ddl")
@@ -586,6 +807,29 @@ class VannaFlaskApp:
         @self.requires_auth
         @self.requires_cache(["question", "sql"])
         def create_function(user: any, id: str, question: str, sql: str):
+            """
+            Create function
+            ---
+            parameters:
+              - name: user
+                in: query
+              - name: id
+                in: query|body
+                type: string
+                required: true
+            responses:
+              200:
+                schema:
+                  type: object
+                  properties:
+                    type:
+                      type: string
+                      default: function_template
+                    id:
+                      type: string
+                    function_template:
+                      type: object
+            """
             plotly_code = self.cache.get(id=id, field="plotly_code")
 
             if plotly_code is None:
@@ -604,6 +848,28 @@ class VannaFlaskApp:
         @self.flask_app.route("/api/v0/update_function", methods=["POST"])
         @self.requires_auth
         def update_function(user: any):
+            """
+            Update function
+            ---
+            parameters:
+              - name: user
+                in: query
+              - name: old_function_name
+                in: body
+                type: string
+                required: true
+              - name: updated_function
+                in: body
+                type: object
+                required: true
+            responses:
+              200:
+                schema:
+                  type: object
+                  properties:
+                    success:
+                      type: boolean
+            """
             old_function_name = flask.request.json.get("old_function_name")
             updated_function = flask.request.json.get("updated_function")
 
@@ -617,15 +883,57 @@ class VannaFlaskApp:
         @self.flask_app.route("/api/v0/delete_function", methods=["POST"])
         @self.requires_auth
         def delete_function(user: any):
+            """
+            Delete function
+            ---
+            parameters:
+              - name: user
+                in: query
+              - name: function_name
+                in: body
+                type: string
+                required: true
+            responses:
+              200:
+                schema:
+                  type: object
+                  properties:
+                    success:
+                      type: boolean
+            """
             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"])
         def generate_followup_questions(user: any, id: str, df, question, sql):
+            """
+            Generate followup questions
+            ---
+            parameters:
+              - name: user
+                in: query
+              - name: id
+                in: query|body
+                type: string
+                required: true
+            responses:
+              200:
+                schema:
+                  type: object
+                  properties:
+                    type:
+                      type: string
+                      default: question_list
+                    questions:
+                      type: array
+                      items:
+                        type: string
+                    header:
+                      type: string
+            """
             if self.allow_llm_to_see_data:
                 followup_questions = vn.generate_followup_questions(
                     question=question, sql=sql, df=df
@@ -658,6 +966,29 @@ class VannaFlaskApp:
         @self.requires_auth
         @self.requires_cache(["df", "question"])
         def generate_summary(user: any, id: str, df, question):
+            """
+            Generate summary
+            ---
+            parameters:
+              - name: user
+                in: query
+              - name: id
+                in: query|body
+                type: string
+                required: true
+            responses:
+              200:
+                schema:
+                  type: object
+                  properties:
+                    type:
+                      type: string
+                      default: text
+                    id:
+                      type: string
+                    text:
+                      type: string
+            """
             if self.allow_llm_to_see_data:
                 summary = vn.generate_summary(question=question, df=df)
 
@@ -686,6 +1017,37 @@ class VannaFlaskApp:
             optional_fields=["summary", "fig_json"]
         )
         def load_question(user: any, id: str, question, sql, df, fig_json, summary):
+            """
+            Load question
+            ---
+            parameters:
+              - name: user
+                in: query
+              - name: id
+                in: query|body
+                type: string
+                required: true
+            responses:
+              200:
+                schema:
+                  type: object
+                  properties:
+                    type:
+                      type: string
+                      default: question_cache
+                    id:
+                      type: string
+                    question:
+                      type: string
+                    sql:
+                      type: string
+                    df:
+                      type: object
+                    fig:
+                      type: object
+                    summary:
+                      type: string
+            """
             try:
                 return jsonify(
                     {
@@ -705,6 +1067,25 @@ class VannaFlaskApp:
         @self.flask_app.route("/api/v0/get_question_history", methods=["GET"])
         @self.requires_auth
         def get_question_history(user: any):
+            """
+            Get question history
+            ---
+            parameters:
+              - name: user
+                in: query
+            responses:
+              200:
+                schema:
+                  type: object
+                  properties:
+                    type:
+                      type: string
+                      default: question_history
+                    questions:
+                      type: array
+                      items:
+                        type: string
+            """
             return jsonify(
                 {
                     "type": "question_history",
@@ -718,6 +1099,136 @@ class VannaFlaskApp:
                 {"type": "error", "error": "The rest of the API is not ported yet."}
             )
 
+        if self.debug:
+            @self.sock.route("/api/v0/log")
+            def sock_log(ws):
+                self.ws_clients.append(ws)
+
+                try:
+                    while True:
+                        message = ws.receive()  # This example just reads and ignores to keep the socket open
+                finally:
+                    self.ws_clients.remove(ws)
+
+    def run(self, *args, **kwargs):
+        """
+        Run the Flask app.
+
+        Args:
+            *args: Arguments to pass to Flask's run method.
+            **kwargs: Keyword arguments to pass to Flask's run method.
+
+        Returns:
+            None
+        """
+        if args or kwargs:
+            self.flask_app.run(*args, **kwargs)
+
+        else:
+            try:
+                from google.colab import output
+
+                output.serve_kernel_port_as_window(8084)
+                from google.colab.output import eval_js
+
+                print("Your app is running at:")
+                print(eval_js("google.colab.kernel.proxyPort(8084)"))
+            except:
+                print("Your app is running at:")
+                print("http://localhost:8084")
+
+            self.flask_app.run(host="0.0.0.0", port=8084, debug=self.debug, use_reloader=False)
+
+
+class VannaFlaskApp(VannaFlaskAPI):
+    def __init__(
+        self,
+        vn: VannaBase,
+        cache: Cache = MemoryCache(),
+        auth: AuthInterface = NoAuth(),
+        debug=True,
+        allow_llm_to_see_data=False,
+        logo="https://img.vanna.ai/vanna-flask.svg",
+        title="Welcome to Vanna.AI",
+        subtitle="Your AI-powered copilot for SQL queries.",
+        show_training_data=True,
+        suggested_questions=True,
+        sql=True,
+        table=True,
+        csv_download=True,
+        chart=True,
+        redraw_chart=True,
+        auto_fix_sql=True,
+        ask_results_correct=True,
+        followup_questions=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.
+
+        Args:
+            vn: The Vanna instance to interact with.
+            cache: The cache to use. Defaults to MemoryCache, which uses an in-memory cache. You can also pass in a custom cache that implements the Cache interface.
+            auth: The authentication method to use. Defaults to NoAuth, which doesn't require authentication. You can also pass in a custom authentication method that implements the AuthInterface interface.
+            debug: Show the debug console. Defaults to True.
+            allow_llm_to_see_data: Whether to allow the LLM to see data. Defaults to False.
+            logo: The logo to display in the UI. Defaults to the Vanna logo.
+            title: The title to display in the UI. Defaults to "Welcome to Vanna.AI".
+            subtitle: The subtitle to display in the UI. Defaults to "Your AI-powered copilot for SQL queries.".
+            show_training_data: Whether to show the training data in the UI. Defaults to True.
+            suggested_questions: Whether to show suggested questions in the UI. Defaults to True.
+            sql: Whether to show the SQL input in the UI. Defaults to True.
+            table: Whether to show the table output in the UI. Defaults to True.
+            csv_download: Whether to allow downloading the table output as a CSV file. Defaults to True.
+            chart: Whether to show the chart output in the UI. Defaults to True.
+            redraw_chart: Whether to allow redrawing the chart. Defaults to True.
+            auto_fix_sql: Whether to allow auto-fixing SQL errors. Defaults to True.
+            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
+        """
+        super().__init__(vn, cache, auth, debug, allow_llm_to_see_data, chart)
+
+        self.config["logo"] = logo
+        self.config["title"] = title
+        self.config["subtitle"] = subtitle
+        self.config["show_training_data"] = show_training_data
+        self.config["suggested_questions"] = suggested_questions
+        self.config["sql"] = sql
+        self.config["table"] = table
+        self.config["csv_download"] = csv_download
+        self.config["chart"] = chart
+        self.config["redraw_chart"] = redraw_chart
+        self.config["auto_fix_sql"] = auto_fix_sql
+        self.config["ask_results_correct"] = ask_results_correct
+        self.config["followup_questions"] = followup_questions
+        self.config["summarization"] = summarization
+        self.config["function_generation"] = function_generation
+
+        self.index_html_path = index_html_path
+        self.assets_folder = assets_folder
+
+        @self.flask_app.route("/auth/login", methods=["POST"])
+        def login():
+            return self.auth.login_handler(flask.request)
+
+        @self.flask_app.route("/auth/callback", methods=["GET"])
+        def callback():
+            return self.auth.callback_handler(flask.request)
+
+        @self.flask_app.route("/auth/logout", methods=["GET"])
+        def logout():
+            return self.auth.logout_handler(flask.request)
+
+
         @self.flask_app.route("/assets/<path:filename>")
         def proxy_assets(filename):
             if self.assets_folder:
@@ -755,18 +1266,6 @@ class VannaFlaskApp:
             else:
                 return "Error fetching file from remote server", response.status_code
 
-        if self.debug:
-            @self.sock.route("/api/v0/log")
-            def sock_log(ws):
-                self.ws_clients.append(ws)
-
-                try:
-                    while True:
-                        message = ws.receive()  # This example just reads and ignores to keep the socket open
-                finally:
-                    self.ws_clients.remove(ws)
-
-
         @self.flask_app.route("/", defaults={"path": ""})
         @self.flask_app.route("/<path:path>")
         def hello(path: str):
@@ -775,32 +1274,3 @@ class VannaFlaskApp:
                 filename = os.path.basename(self.index_html_path)
                 return send_from_directory(directory=directory, path=filename)
             return html_content
-
-    def run(self, *args, **kwargs):
-        """
-        Run the Flask app.
-
-        Args:
-            *args: Arguments to pass to Flask's run method.
-            **kwargs: Keyword arguments to pass to Flask's run method.
-
-        Returns:
-            None
-        """
-        if args or kwargs:
-            self.flask_app.run(*args, **kwargs)
-
-        else:
-            try:
-                from google.colab import output
-
-                output.serve_kernel_port_as_window(8084)
-                from google.colab.output import eval_js
-
-                print("Your app is running at:")
-                print(eval_js("google.colab.kernel.proxyPort(8084)"))
-            except:
-                print("Your app is running at:")
-                print("http://localhost:8084")
-
-            self.flask_app.run(host="0.0.0.0", port=8084, debug=self.debug, use_reloader=False)

vanna/openai/openai_chat.py

@@ -11,14 +11,10 @@ class OpenAI_Chat(VannaBase):
 
         # default parameters - can be overrided using config
         self.temperature = 0.7
-        self.max_tokens = 500
 
         if "temperature" in config:
             self.temperature = config["temperature"]
 
-        if "max_tokens" in config:
-            self.max_tokens = config["max_tokens"]
-
         if "api_type" in config:
             raise Exception(
                 "Passing api_type is now deprecated. Please pass an OpenAI client instead."
@@ -75,7 +71,6 @@ class OpenAI_Chat(VannaBase):
             response = self.client.chat.completions.create(
                 model=model,
                 messages=prompt,
-                max_tokens=self.max_tokens,
                 stop=None,
                 temperature=self.temperature,
             )
@@ -87,7 +82,6 @@ class OpenAI_Chat(VannaBase):
             response = self.client.chat.completions.create(
                 engine=engine,
                 messages=prompt,
-                max_tokens=self.max_tokens,
                 stop=None,
                 temperature=self.temperature,
             )
@@ -98,7 +92,6 @@ class OpenAI_Chat(VannaBase):
             response = self.client.chat.completions.create(
                 engine=self.config["engine"],
                 messages=prompt,
-                max_tokens=self.max_tokens,
                 stop=None,
                 temperature=self.temperature,
             )
@@ -109,7 +102,6 @@ class OpenAI_Chat(VannaBase):
             response = self.client.chat.completions.create(
                 model=self.config["model"],
                 messages=prompt,
-                max_tokens=self.max_tokens,
                 stop=None,
                 temperature=self.temperature,
             )
@@ -123,7 +115,6 @@ class OpenAI_Chat(VannaBase):
             response = self.client.chat.completions.create(
                 model=model,
                 messages=prompt,
-                max_tokens=self.max_tokens,
                 stop=None,
                 temperature=self.temperature,
             )

{vanna-0.6.3.dist-info → vanna-0.6.4.dist-info}/METADATA

@@ -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.dist-info → vanna-0.6.4.dist-info}/RECORD

@@ -7,15 +7,15 @@ vanna/ZhipuAI/ZhipuAI_embeddings.py,sha256=lUqzJg9fOx7rVFhjdkFjXcDeVGV4aAB5Ss0oE
 vanna/ZhipuAI/__init__.py,sha256=NlsijtcZp5Tj9jkOe9fNcOQND_QsGgu7otODsCLBPr0,116
 vanna/advanced/__init__.py,sha256=oDj9g1JbrbCfp4WWdlr_bhgdMqNleyHgr6VXX6DcEbo,658
 vanna/anthropic/__init__.py,sha256=85s_2mAyyPxc0T_0JEvYeAkEKWJwkwqoyUwSC5dw9Gk,43
-vanna/anthropic/anthropic_chat.py,sha256=Wk0o-NMW1uvR2fhSWxrR_2FqNh-dLprNG4uuVqpqAkY,2615
+vanna/anthropic/anthropic_chat.py,sha256=7X3x8SYwDY28aGyBnt0YNRMG8YY1p_t-foMfKGj8_Oo,2627
 vanna/base/__init__.py,sha256=Sl-HM1RRYzAZoSqmL1CZQmF3ZF-byYTCFQP3JZ2A5MU,28
-vanna/base/base.py,sha256=Cz5fW0Odg5GaRg2svvl6YPeBULKXKMNODAb3Zc1Q8bU,70993
+vanna/base/base.py,sha256=3Du70NrXQMn_LOif2YFPRRVKo4wH5-f6eZcLlXEX0X8,71705
 vanna/bedrock/__init__.py,sha256=hRT2bgJbHEqViLdL-t9hfjSfFdIOkPU2ADBt-B1En-8,46
 vanna/bedrock/bedrock_converse.py,sha256=Nx5kYm-diAfYmsWAnTP5xnv7V84Og69-AP9b3seIe0E,2869
 vanna/chromadb/__init__.py,sha256=-iL0nW_g4uM8nWKMuWnNePfN4nb9uk8P3WzGvezOqRg,50
 vanna/chromadb/chromadb_vector.py,sha256=eKyPck99Y6Jt-BNWojvxLG-zvAERzLSm-3zY-bKXvaA,8792
 vanna/exceptions/__init__.py,sha256=dJ65xxxZh1lqBeg6nz6Tq_r34jLVmjvBvPO9Q6hFaQ8,685
-vanna/flask/__init__.py,sha256=urPrHUqM1mpx96VHiQWVXCy3NQwDh6OsSkm4V4wqccY,30211
+vanna/flask/__init__.py,sha256=r1ucQupb6wuTcjVVKpkdrg6R38FZe6KQoKw9AtcghDQ,42889
 vanna/flask/assets.py,sha256=_UoUr57sS0QL2BuTxAOe9k4yy8T7-fp2NpbRSVtW3IM,451769
 vanna/flask/auth.py,sha256=UpKxh7W5cd43W0LGch0VqhncKwB78L6dtOQkl1JY5T0,1246
 vanna/google/__init__.py,sha256=M-dCxCZcKL4bTQyMLj6r6VRs65YNX9Tl2aoPCuqGm-8,41
@@ -35,7 +35,7 @@ vanna/mock/vectordb.py,sha256=h45znfYMUnttE2BBC8v6TKeMaA58pFJL-5B3OGeRNFI,2681
 vanna/ollama/__init__.py,sha256=4xyu8aHPdnEHg5a-QAMwr5o0ns5wevsp_zkI-ndMO2k,27
 vanna/ollama/ollama.py,sha256=rXa7cfvdlO1E5SLysXIl3IZpIaA2r0RBvV5jX2-upiE,3794
 vanna/openai/__init__.py,sha256=tGkeQ7wTIPsando7QhoSHehtoQVdYLwFbKNlSmCmNeQ,86
-vanna/openai/openai_chat.py,sha256=lm-hUsQxu6Q1t06A2csC037zI4VkMk0wFbQ-_Lj74Wg,4764
+vanna/openai/openai_chat.py,sha256=KU6ynOQ5v7vwrQQ13phXoUXeQUrH6_vmhfiPvWddTrQ,4427
 vanna/openai/openai_embeddings.py,sha256=g4pNh9LVcYP9wOoO8ecaccDFWmCUYMInebfHucAa2Gc,1260
 vanna/opensearch/__init__.py,sha256=0unDevWOTs7o8S79TOHUKF1mSiuQbBUVm-7k9jV5WW4,54
 vanna/opensearch/opensearch_vector.py,sha256=VhIcrSyNzWR9ZrqrJnyGFOyuQZs3swfbhr8QyVGI0eI,12226
@@ -50,6 +50,6 @@ vanna/vllm/__init__.py,sha256=aNlUkF9tbURdeXAJ8ytuaaF1gYwcG3ny1MfNl_cwQYg,23
 vanna/vllm/vllm.py,sha256=oM_aA-1Chyl7T_Qc_yRKlL6oSX1etsijY9zQdjeMGMQ,2827
 vanna/weaviate/__init__.py,sha256=HL6PAl7ePBAkeG8uln-BmM7IUtWohyTPvDfcPzSGSCg,46
 vanna/weaviate/weaviate_vector.py,sha256=GEiu4Vd9w-7j10aB-zTxJ8gefqe_F-LUUGvttFs1vlg,7539
-vanna-0.6.3.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
-vanna-0.6.3.dist-info/METADATA,sha256=-CFXJd2aKy2tq2Gxipp-52YCPHJBlI0wIOPzcOV5s80,11865
-vanna-0.6.3.dist-info/RECORD,,
+vanna-0.6.4.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
+vanna-0.6.4.dist-info/METADATA,sha256=LqIi4Hg1y_aTEH79PX48nnY1TM-u6ese9K8Os9Cqkg0,11889
+vanna-0.6.4.dist-info/RECORD,,