vanna 0.2.1__py3-none-any.whl → 0.3.1__py3-none-any.whl

vanna/ZhipuAI/ZhipuAI_Chat.py

@@ -1,11 +1,11 @@
 import re
 from typing import List
+
 import pandas as pd
 from zhipuai import ZhipuAI
+
 from ..base import VannaBase
-import re
-from typing import List
-import pandas as pd
+
 
 class ZhipuAI_Chat(VannaBase):
     def __init__(self, config=None):

vanna/ZhipuAI/__init__.py

@@ -0,0 +1,2 @@
+from .ZhipuAI_Chat import ZhipuAI_Chat
+from .ZhipuAI_embeddings import ZhipuAI_Embeddings, ZhipuAIEmbeddingFunction

vanna/anthropic/__init__.py

@@ -0,0 +1 @@
+from .anthropic_chat import Anthropic_Chat

vanna/anthropic/anthropic_chat.py

@@ -0,0 +1,78 @@
+import os
+
+import anthropic
+
+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
+
+        if "temperature" in config:
+            self.temperature = config["temperature"]
+
+        if "max_tokens" in config:
+            self.max_tokens = config["max_tokens"]
+
+        if "api_key" in config:
+            self.client = anthropic.Anthropic(api_key=config["api_key"])
+
+    def system_message(self, message: str) -> any:
+        return {"role": "system", "content": message}
+
+    def user_message(self, message: str) -> any:
+        return {"role": "user", "content": message}
+
+    def assistant_message(self, message: str) -> any:
+        return {"role": "assistant", "content": message}
+
+    def submit_prompt(self, prompt, **kwargs) -> str:
+        if prompt is None:
+            raise Exception("Prompt is None")
+
+        if len(prompt) == 0:
+            raise Exception("Prompt is empty")
+
+        # Count the number of tokens in the message log
+        # Use 4 as an approximation for the number of characters per token
+        num_tokens = 0
+        for message in prompt:
+            num_tokens += len(message["content"]) / 4
+
+        if self.config is not None and "model" in self.config:
+            print(
+                f"Using model {self.config['model']} for {num_tokens} tokens (approx)"
+            )
+            # claude required system message is a single filed
+            # https://docs.anthropic.com/claude/reference/messages_post
+            system_message = ''
+            no_system_prompt = []
+            for prompt_message in prompt:
+                role = prompt_message['role']
+                if role == 'system':
+                    system_message = prompt_message['content']
+                else:
+                    no_system_prompt.append({"role": role, "content": prompt_message['content']})
+
+            response = self.client.messages.create(
+                model=self.config["model"],
+                messages=no_system_prompt,
+                system=system_message,
+                max_tokens=self.max_tokens,
+                temperature=self.temperature,
+            )
+
+        return response.content[0].text

vanna/base/base.py

@@ -103,10 +103,15 @@ class VannaBase(ABC):
         Returns:
             str: The SQL query that answers the question.
         """
+        if self.config is not None:
+            initial_prompt = self.config.get("initial_prompt", None)
+        else:
+            initial_prompt = None
         question_sql_list = self.get_similar_question_sql(question, **kwargs)
         ddl_list = self.get_related_ddl(question, **kwargs)
         doc_list = self.get_related_documentation(question, **kwargs)
         prompt = self.get_sql_prompt(
+            initial_prompt=initial_prompt,
             question=question,
             question_sql_list=question_sql_list,
             ddl_list=ddl_list,
@@ -405,6 +410,7 @@ class VannaBase(ABC):
 
     def get_sql_prompt(
         self,
+        initial_prompt : str,
         question: str,
         question_sql_list: list,
         ddl_list: list,
@@ -434,7 +440,9 @@ class VannaBase(ABC):
         Returns:
             any: The prompt for the LLM to generate SQL.
         """
-        initial_prompt = "The user provides a question and you provide SQL. You will only respond with SQL code and not with any explanations.\n\nRespond with only SQL code. Do not answer with any explanations -- just the code.\n"
+
+        if initial_prompt is None:
+            initial_prompt = "The user provides a question and you provide SQL. You will only respond with SQL code and not with any explanations.\n\nRespond with only SQL code. Do not answer with any explanations -- just the code.\n"
 
         initial_prompt = self.add_ddl_to_prompt(
             initial_prompt, ddl_list, max_tokens=14000
@@ -796,6 +804,91 @@ class VannaBase(ABC):
         self.run_sql_is_set = True
         self.run_sql = run_sql_postgres
 
+
+    def connect_to_mysql(
+            self,
+            host: str = None,
+            dbname: str = None,
+            user: str = None,
+            password: str = None,
+            port: int = None,
+    ):
+
+        try:
+            import pymysql.cursors
+        except ImportError:
+            raise DependencyError(
+                "You need to install required dependencies to execute this method,"
+                " run command: \npip install PyMySQL"
+            )
+
+        if not host:
+            host = os.getenv("HOST")
+
+        if not host:
+            raise ImproperlyConfigured("Please set your MySQL host")
+
+        if not dbname:
+            dbname = os.getenv("DATABASE")
+
+        if not dbname:
+            raise ImproperlyConfigured("Please set your MySQL database")
+
+        if not user:
+            user = os.getenv("USER")
+
+        if not user:
+            raise ImproperlyConfigured("Please set your MySQL user")
+
+        if not password:
+            password = os.getenv("PASSWORD")
+
+        if not password:
+            raise ImproperlyConfigured("Please set your MySQL password")
+
+        if not port:
+            port = os.getenv("PORT")
+
+        if not port:
+            raise ImproperlyConfigured("Please set your MySQL port")
+
+        conn = None
+
+        try:
+            conn = pymysql.connect(host=host,
+                                   user=user,
+                                   password=password,
+                                   database=dbname,
+                                   port=port,
+                                   cursorclass=pymysql.cursors.DictCursor)
+        except pymysql.Error as e:
+            raise ValidationError(e)
+
+        def run_sql_mysql(sql: str) -> Union[pd.DataFrame, None]:
+            if conn:
+                try:
+                    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
+
+                except pymysql.Error as e:
+                    conn.rollback()
+                    raise ValidationError(e)
+
+                except Exception as e:
+                    conn.rollback()
+                    raise e
+
+        self.run_sql_is_set = True
+        self.run_sql = run_sql_mysql
+
+
     def connect_to_bigquery(self, cred_file_path: str = None, project_id: str = None):
         """
         Connect to gcs using the bigquery connector. This is just a helper function to set [`vn.run_sql`][vanna.base.base.VannaBase.run_sql]

vanna/chromadb/__init__.py

@@ -0,0 +1 @@
+from .chromadb_vector import ChromaDB_VectorStore

vanna/flask/__init__.py

@@ -75,7 +75,9 @@ class VannaFlaskApp:
                 id = request.args.get("id")
 
                 if id is None:
-                    return jsonify({"type": "error", "error": "No id provided"})
+                    id = request.json.get("id")
+                    if id is None:
+                        return jsonify({"type": "error", "error": "No id provided"})
 
                 for field in fields:
                     if self.cache.get(id=id, field=field) is None:
@@ -94,15 +96,94 @@ class VannaFlaskApp:
 
         return decorator
 
-    def __init__(self, vn, cache: Cache = MemoryCache(), allow_llm_to_see_data=False):
+    def __init__(self, vn, cache: Cache = MemoryCache(),
+                    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
+                 ):
+        """
+        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.
+            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.
+
+        Returns:
+            None
+        """
         self.flask_app = Flask(__name__)
         self.vn = vn
         self.cache = cache
         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
 
         log = logging.getLogger("werkzeug")
         log.setLevel(logging.ERROR)
 
+        @self.flask_app.route("/api/v0/get_config", methods=["GET"])
+        def get_config():
+            return jsonify(
+                {
+                    "type": "config",
+                    "config": {
+                        "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,
+                    },
+                }
+            )
+
         @self.flask_app.route("/api/v0/generate_questions", methods=["GET"])
         def generate_questions():
             # If self has an _model attribute and model=='chinook'
@@ -199,12 +280,52 @@ class VannaFlaskApp:
                     {
                         "type": "df",
                         "id": id,
-                        "df": df.head(10).to_json(orient="records"),
+                        "df": df.head(10).to_json(orient='records', date_format='iso'),
                     }
                 )
 
             except Exception as e:
-                return jsonify({"type": "error", "error": str(e)})
+                return jsonify({"type": "sql_error", "error": str(e)})
+
+        @self.flask_app.route("/api/v0/fix_sql", methods=["POST"])
+        @self.requires_cache(["question", "sql"])
+        def fix_sql(id: str, question:str, sql: str):
+            error = flask.request.json.get("error")
+
+            if error is None:
+                return jsonify({"type": "error", "error": "No error provided"})
+
+            question = f"I have an error: {error}\n\nHere is the SQL I tried to run: {sql}\n\nThis is the question I was trying to answer: {question}\n\nCan you rewrite the SQL to fix the error?"
+
+            fixed_sql = vn.generate_sql(question=question)
+
+            self.cache.set(id=id, field="sql", value=fixed_sql)
+
+            return jsonify(
+                {
+                    "type": "sql",
+                    "id": id,
+                    "text": fixed_sql,
+                }
+            )
+
+
+        @self.flask_app.route('/api/v0/update_sql', methods=['POST'])
+        @self.requires_cache([])
+        def update_sql(id: str):
+            sql = flask.request.json.get('sql')
+
+            if sql is None:
+                return jsonify({"type": "error", "error": "No sql provided"})
+
+            cache.set(id=id, field='sql', value=sql)
+
+            return jsonify(
+                {
+                    "type": "sql",
+                    "id": id,
+                    "text": sql,
+                })
 
         @self.flask_app.route("/api/v0/download_csv", methods=["GET"])
         @self.requires_cache(["df"])
@@ -220,6 +341,11 @@ class VannaFlaskApp:
         @self.flask_app.route("/api/v0/generate_plotly_figure", methods=["GET"])
         @self.requires_cache(["df", "question", "sql"])
         def generate_plotly_figure(id: str, df, question, sql):
+            chart_instructions = flask.request.args.get('chart_instructions')
+
+            if chart_instructions is not None:
+                question = f"{question}. When generating the chart, use these special instructions: {chart_instructions}"
+
             try:
                 code = vn.generate_plotly_code(
                     question=question,
@@ -319,6 +445,7 @@ class VannaFlaskApp:
                     }
                 )
             else:
+                cache.set(id=id, field="followup_questions", value=[])
                 return jsonify(
                     {
                         "type": "question_list",
@@ -351,9 +478,9 @@ class VannaFlaskApp:
 
         @self.flask_app.route("/api/v0/load_question", methods=["GET"])
         @self.requires_cache(
-            ["question", "sql", "df", "fig_json", "followup_questions"]
+            ["question", "sql", "df", "fig_json"]
         )
-        def load_question(id: str, question, sql, df, fig_json, followup_questions):
+        def load_question(id: str, question, sql, df, fig_json):
             try:
                 return jsonify(
                     {
@@ -363,7 +490,6 @@ class VannaFlaskApp:
                         "sql": sql,
                         "df": df.head(10).to_json(orient="records"),
                         "fig": fig_json,
-                        "followup_questions": followup_questions,
                     }
                 )
 
@@ -424,16 +550,31 @@ class VannaFlaskApp:
         def hello(path: str):
             return html_content
 
-    def run(self):
-        try:
-            from google.colab import output
+    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
 
-            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")
 
-            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=False)
+            self.flask_app.run(host="0.0.0.0", port=8084, debug=False)