vanna 0.3.3__py3-none-any.whl → 0.4.0__py3-none-any.whl

vanna/ZhipuAI/ZhipuAI_Chat.py

@@ -40,7 +40,7 @@ class ZhipuAI_Chat(VannaBase):
         initial_prompt: str, ddl_list: List[str], max_tokens: int = 14000
     ) -> str:
         if len(ddl_list) > 0:
-            initial_prompt += f"\nYou may use the following DDL statements as a reference for what tables might be available. Use responses to past questions also to guide you:\n\n"
+            initial_prompt += "\nYou may use the following DDL statements as a reference for what tables might be available. Use responses to past questions also to guide you:\n\n"
 
             for ddl in ddl_list:
                 if (
@@ -57,7 +57,7 @@ class ZhipuAI_Chat(VannaBase):
         initial_prompt: str, documentation_List: List[str], max_tokens: int = 14000
     ) -> str:
         if len(documentation_List) > 0:
-            initial_prompt += f"\nYou may use the following documentation as a reference for what tables might be available. Use responses to past questions also to guide you:\n\n"
+            initial_prompt += "\nYou may use the following documentation as a reference for what tables might be available. Use responses to past questions also to guide you:\n\n"
 
             for documentation in documentation_List:
                 if (
@@ -74,7 +74,7 @@ class ZhipuAI_Chat(VannaBase):
         initial_prompt: str, sql_List: List[str], max_tokens: int = 14000
     ) -> str:
         if len(sql_List) > 0:
-            initial_prompt += f"\nYou may use the following SQL statements as a reference for what tables might be available. Use responses to past questions also to guide you:\n\n"
+            initial_prompt += "\nYou may use the following SQL statements as a reference for what tables might be available. Use responses to past questions also to guide you:\n\n"
 
             for question in sql_List:
                 if (

vanna/base/base.py

@@ -124,6 +124,18 @@ class VannaBase(ABC):
         return self.extract_sql(llm_response)
 
     def extract_sql(self, llm_response: str) -> str:
+        # If the llm_response is not markdown formatted, extract sql by finding select and ; in the response
+        sql = re.search(r"SELECT.*?;", llm_response, re.DOTALL)
+        if sql:
+            self.log(f"Output from LLM: {llm_response} \nExtracted SQL: {sql.group(0)}"
+            )
+            return sql.group(0)
+
+        # If the llm_response contains a CTE (with clause), extract the sql bewteen WITH and ;
+        sql = re.search(r"WITH.*?;", llm_response, re.DOTALL)
+        if sql:
+            self.log(f"Output from LLM: {llm_response} \nExtracted SQL: {sql.group(0)}")
+            return sql.group(0)
         # If the llm_response contains a markdown code block, with or without the sql tag, extract the sql from it
         sql = re.search(r"```sql\n(.*)```", llm_response, re.DOTALL)
         if sql:
@@ -363,7 +375,7 @@ class VannaBase(ABC):
         self, initial_prompt: str, ddl_list: list[str], max_tokens: int = 14000
     ) -> str:
         if len(ddl_list) > 0:
-            initial_prompt += f"\nYou may use the following DDL statements as a reference for what tables might be available. Use responses to past questions also to guide you:\n\n"
+            initial_prompt += "\nYou may use the following DDL statements as a reference for what tables might be available. Use responses to past questions also to guide you:\n\n"
 
             for ddl in ddl_list:
                 if (
@@ -382,7 +394,7 @@ class VannaBase(ABC):
         max_tokens: int = 14000,
     ) -> str:
         if len(documentation_list) > 0:
-            initial_prompt += f"\nYou may use the following documentation as a reference for what tables might be available. Use responses to past questions also to guide you:\n\n"
+            initial_prompt += "\nYou may use the following documentation as a reference for what tables might be available. Use responses to past questions also to guide you:\n\n"
 
             for documentation in documentation_list:
                 if (
@@ -398,7 +410,7 @@ class VannaBase(ABC):
         self, initial_prompt: str, sql_list: list[str], max_tokens: int = 14000
     ) -> str:
         if len(sql_list) > 0:
-            initial_prompt += f"\nYou may use the following SQL statements as a reference for what tables might be available. Use responses to past questions also to guide you:\n\n"
+            initial_prompt += "\nYou may use the following SQL statements as a reference for what tables might be available. Use responses to past questions also to guide you:\n\n"
 
             for question in sql_list:
                 if (
@@ -642,6 +654,7 @@ class VannaBase(ABC):
             password=password,
             account=account,
             database=database,
+            client_session_keep_alive=True
         )
 
         def run_sql_snowflake(sql: str) -> pd.DataFrame:
@@ -890,6 +903,94 @@ class VannaBase(ABC):
         self.run_sql_is_set = True
         self.run_sql = run_sql_mysql
 
+    def connect_to_oracle(
+    self,
+    user: str = None,
+    password: str = None,
+    dsn: str = None,
+    ):
+
+        """
+        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:**
+        ```python
+        vn.connect_to_oracle(
+        user="username",
+        password="password",
+        dns="host:port/sid",
+        )
+        ```
+        Args:
+            USER (str): Oracle db user name.
+            PASSWORD (str): Oracle db user password.
+            DSN (str): Oracle db host ip - host:port/sid.
+        """
+
+        try:
+            import oracledb
+        except ImportError:
+
+            raise DependencyError(
+                "You need to install required dependencies to execute this method,"
+                " run command: \npip install oracledb"
+            )
+
+        if not dsn:
+            dsn = os.getenv("DSN")
+
+        if not dsn:
+            raise ImproperlyConfigured("Please set your Oracle dsn which should include host:port/sid")
+
+        if not user:
+            user = os.getenv("USER")
+
+        if not user:
+            raise ImproperlyConfigured("Please set your Oracle db user")
+
+        if not password:
+            password = os.getenv("PASSWORD")
+
+        if not password:
+            raise ImproperlyConfigured("Please set your Oracle db password")
+
+        conn = None
+
+        try:
+            conn = oracledb.connect(
+                user=user,
+                password=password,
+                dsn=dsn,
+                )
+        except oracledb.Error as e:
+            raise ValidationError(e)
+
+        def run_sql_oracle(sql: str) -> Union[pd.DataFrame, None]:
+            if conn:
+                try:
+                    sql = sql.rstrip()
+                    if sql.endswith(';'): #fix for a known problem with Oracle db where an extra ; will cause an error.
+                        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
+
+                except oracledb.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_oracle
 
     def connect_to_bigquery(self, cred_file_path: str = None, project_id: str = None):
         """
@@ -1238,7 +1339,7 @@ class VannaBase(ABC):
         """
 
         if question and not sql:
-            raise ValidationError(f"Please also provide a SQL query")
+            raise ValidationError("Please also provide a SQL query")
 
         if documentation:
             print("Adding documentation....")

vanna/chromadb/chromadb_vector.py

@@ -1,5 +1,4 @@
 import json
-import uuid
 from typing import List
 
 import chromadb
@@ -16,17 +15,16 @@ default_ef = embedding_functions.DefaultEmbeddingFunction()
 class ChromaDB_VectorStore(VannaBase):
     def __init__(self, config=None):
         VannaBase.__init__(self, config=config)
+        if config is None:
+            config = {}
 
-        if config is not None:
-            path = config.get("path", ".")
-            self.embedding_function = config.get("embedding_function", default_ef)
-            curr_client = config.get("client", "persistent")
-            self.n_results = config.get("n_results", 10)
-        else:
-            path = "."
-            self.embedding_function = default_ef
-            curr_client = "persistent"  # defaults to persistent storage
-            self.n_results = 10  # defaults to 10 documents
+        path = config.get("path", ".")
+        self.embedding_function = config.get("embedding_function", default_ef)
+        curr_client = config.get("client", "persistent")
+        collection_metadata = config.get("collection_metadata", None)
+        self.n_results_sql = config.get("n_results_sql", config.get("n_results", 10))
+        self.n_results_documentation = config.get("n_results_documentation", config.get("n_results", 10))
+        self.n_results_ddl = config.get("n_results_ddl", config.get("n_results", 10))
 
         if curr_client == "persistent":
             self.chroma_client = chromadb.PersistentClient(
@@ -43,13 +41,19 @@ class ChromaDB_VectorStore(VannaBase):
             raise ValueError(f"Unsupported client was set in config: {curr_client}")
 
         self.documentation_collection = self.chroma_client.get_or_create_collection(
-            name="documentation", embedding_function=self.embedding_function
+            name="documentation",
+            embedding_function=self.embedding_function,
+            metadata=collection_metadata,
         )
         self.ddl_collection = self.chroma_client.get_or_create_collection(
-            name="ddl", embedding_function=self.embedding_function
+            name="ddl",
+            embedding_function=self.embedding_function,
+            metadata=collection_metadata,
         )
         self.sql_collection = self.chroma_client.get_or_create_collection(
-            name="sql", embedding_function=self.embedding_function
+            name="sql",
+            embedding_function=self.embedding_function,
+            metadata=collection_metadata,
         )
 
     def generate_embedding(self, data: str, **kwargs) -> List[float]:
@@ -232,7 +236,7 @@ class ChromaDB_VectorStore(VannaBase):
         return ChromaDB_VectorStore._extract_documents(
             self.sql_collection.query(
                 query_texts=[question],
-                n_results=self.n_results,
+                n_results=self.n_results_sql,
             )
         )
 
@@ -240,6 +244,7 @@ class ChromaDB_VectorStore(VannaBase):
         return ChromaDB_VectorStore._extract_documents(
             self.ddl_collection.query(
                 query_texts=[question],
+                n_results=self.n_results_ddl,
             )
         )
 
@@ -247,5 +252,6 @@ class ChromaDB_VectorStore(VannaBase):
         return ChromaDB_VectorStore._extract_documents(
             self.documentation_collection.query(
                 query_texts=[question],
+                n_results=self.n_results_documentation,
             )
         )

vanna/flask/__init__.py

@@ -8,27 +8,47 @@ import requests
 from flask import Flask, Response, jsonify, request
 
 from .assets import css_content, html_content, js_content
+from .auth import AuthInterface, NoAuth
 
 
 class Cache(ABC):
+    """
+    Define the interface for a cache that can be used to store data in a Flask app.
+    """
+
     @abstractmethod
     def generate_id(self, *args, **kwargs):
+        """
+        Generate a unique ID for the cache.
+        """
         pass
 
     @abstractmethod
     def get(self, id, field):
+        """
+        Get a value from the cache.
+        """
         pass
 
     @abstractmethod
     def get_all(self, field_list) -> list:
+        """
+        Get all values from the cache.
+        """
         pass
 
     @abstractmethod
     def set(self, id, field, value):
+        """
+        Set a value in the cache.
+        """
         pass
 
     @abstractmethod
     def delete(self, id):
+        """
+        Delete a value from the cache.
+        """
         pass
 
 
@@ -64,11 +84,10 @@ class MemoryCache(Cache):
         if id in self.cache:
             del self.cache[id]
 
-
 class VannaFlaskApp:
     flask_app = None
 
-    def requires_cache(self, fields):
+    def requires_cache(self, required_fields, optional_fields=[]):
         def decorator(f):
             @wraps(f)
             def decorated(*args, **kwargs):
@@ -79,14 +98,17 @@ class VannaFlaskApp:
                     if id is None:
                         return jsonify({"type": "error", "error": "No id provided"})
 
-                for field in fields:
+                for field in required_fields:
                     if self.cache.get(id=id, field=field) is None:
                         return jsonify({"type": "error", "error": f"No {field} found"})
 
                 field_values = {
-                    field: self.cache.get(id=id, field=field) for field in fields
+                    field: self.cache.get(id=id, field=field) for field in required_fields
                 }
 
+                for field in optional_fields:
+                    field_values[field] = self.cache.get(id=id, field=field)
+
                 # Add the id to the field_values
                 field_values["id"] = id
 
@@ -96,7 +118,21 @@ class VannaFlaskApp:
 
         return decorator
 
+    def requires_auth(self, f):
+        @wraps(f)
+        def decorated(*args, **kwargs):
+            user = self.auth.get_user(flask.request)
+
+            if not self.auth.is_logged_in(user):
+                return jsonify({"type": "not_logged_in", "html": self.auth.login_form()})
+
+            # Pass the user to the function
+            return f(*args, user=user, **kwargs)
+
+        return decorated
+
     def __init__(self, vn, cache: Cache = MemoryCache(),
+                    auth: AuthInterface = NoAuth(),
                     allow_llm_to_see_data=False,
                     logo="https://img.vanna.ai/vanna-flask.svg",
                     title="Welcome to Vanna.AI",
@@ -119,6 +155,7 @@ class VannaFlaskApp:
         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.
             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".
@@ -140,6 +177,7 @@ class VannaFlaskApp:
         """
         self.flask_app = Flask(__name__)
         self.vn = vn
+        self.auth = auth
         self.cache = cache
         self.allow_llm_to_see_data = allow_llm_to_see_data
         self.logo = logo
@@ -160,32 +198,50 @@ class VannaFlaskApp:
         log = logging.getLogger("werkzeug")
         log.setLevel(logging.ERROR)
 
+        @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"])
-        def get_config():
+        @self.requires_auth
+        def get_config(user: any):
+            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,
+            }
+
+            config = self.auth.override_config_for_user(user, 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,
-                    },
+                    "config": config
                 }
             )
 
         @self.flask_app.route("/api/v0/generate_questions", methods=["GET"])
-        def generate_questions():
+        @self.requires_auth
+        def generate_questions(user: any):
             # If self has an _model attribute and model=='chinook'
             if hasattr(self.vn, "_model") and self.vn._model == "chinook":
                 return jsonify(
@@ -240,7 +296,8 @@ class VannaFlaskApp:
                 )
 
         @self.flask_app.route("/api/v0/generate_sql", methods=["GET"])
-        def generate_sql():
+        @self.requires_auth
+        def generate_sql(user: any):
             question = flask.request.args.get("question")
 
             if question is None:
@@ -261,8 +318,9 @@ class VannaFlaskApp:
             )
 
         @self.flask_app.route("/api/v0/run_sql", methods=["GET"])
+        @self.requires_auth
         @self.requires_cache(["sql"])
-        def run_sql(id: str, sql: str):
+        def run_sql(user: any, id: str, sql: str):
             try:
                 if not vn.run_sql_is_set:
                     return jsonify(
@@ -274,7 +332,7 @@ class VannaFlaskApp:
 
                 df = vn.run_sql(sql=sql)
 
-                cache.set(id=id, field="df", value=df)
+                self.cache.set(id=id, field="df", value=df)
 
                 return jsonify(
                     {
@@ -288,8 +346,9 @@ class VannaFlaskApp:
                 return jsonify({"type": "sql_error", "error": str(e)})
 
         @self.flask_app.route("/api/v0/fix_sql", methods=["POST"])
+        @self.requires_auth
         @self.requires_cache(["question", "sql"])
-        def fix_sql(id: str, question:str, sql: str):
+        def fix_sql(user: any, id: str, question:str, sql: str):
             error = flask.request.json.get("error")
 
             if error is None:
@@ -311,14 +370,15 @@ class VannaFlaskApp:
 
 
         @self.flask_app.route('/api/v0/update_sql', methods=['POST'])
+        @self.requires_auth
         @self.requires_cache([])
-        def update_sql(id: str):
+        def update_sql(user: any, 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)
+            self.cache.set(id=id, field='sql', value=sql)
 
             return jsonify(
                 {
@@ -328,8 +388,9 @@ class VannaFlaskApp:
                 })
 
         @self.flask_app.route("/api/v0/download_csv", methods=["GET"])
+        @self.requires_auth
         @self.requires_cache(["df"])
-        def download_csv(id: str, df):
+        def download_csv(user: any, id: str, df):
             csv = df.to_csv()
 
             return Response(
@@ -339,8 +400,9 @@ class VannaFlaskApp:
             )
 
         @self.flask_app.route("/api/v0/generate_plotly_figure", methods=["GET"])
+        @self.requires_auth
         @self.requires_cache(["df", "question", "sql"])
-        def generate_plotly_figure(id: str, df, question, sql):
+        def generate_plotly_figure(user: any, id: str, df, question, sql):
             chart_instructions = flask.request.args.get('chart_instructions')
 
             if chart_instructions is not None:
@@ -355,7 +417,7 @@ class VannaFlaskApp:
                 fig = vn.get_plotly_figure(plotly_code=code, df=df, dark_mode=False)
                 fig_json = fig.to_json()
 
-                cache.set(id=id, field="fig_json", value=fig_json)
+                self.cache.set(id=id, field="fig_json", value=fig_json)
 
                 return jsonify(
                     {
@@ -373,7 +435,8 @@ class VannaFlaskApp:
                 return jsonify({"type": "error", "error": str(e)})
 
         @self.flask_app.route("/api/v0/get_training_data", methods=["GET"])
-        def get_training_data():
+        @self.requires_auth
+        def get_training_data(user: any):
             df = vn.get_training_data()
 
             if df is None or len(df) == 0:
@@ -393,7 +456,8 @@ class VannaFlaskApp:
             )
 
         @self.flask_app.route("/api/v0/remove_training_data", methods=["POST"])
-        def remove_training_data():
+        @self.requires_auth
+        def remove_training_data(user: any):
             # Get id from the JSON body
             id = flask.request.json.get("id")
 
@@ -408,7 +472,8 @@ class VannaFlaskApp:
                 )
 
         @self.flask_app.route("/api/v0/train", methods=["POST"])
-        def add_training_data():
+        @self.requires_auth
+        def add_training_data(user: any):
             question = flask.request.json.get("question")
             sql = flask.request.json.get("sql")
             ddl = flask.request.json.get("ddl")
@@ -425,8 +490,9 @@ class VannaFlaskApp:
                 return jsonify({"type": "error", "error": str(e)})
 
         @self.flask_app.route("/api/v0/generate_followup_questions", methods=["GET"])
+        @self.requires_auth
         @self.requires_cache(["df", "question", "sql"])
-        def generate_followup_questions(id: str, df, question, sql):
+        def generate_followup_questions(user: any, id: str, df, question, sql):
             if self.allow_llm_to_see_data:
                 followup_questions = vn.generate_followup_questions(
                     question=question, sql=sql, df=df
@@ -434,7 +500,7 @@ class VannaFlaskApp:
                 if followup_questions is not None and len(followup_questions) > 5:
                     followup_questions = followup_questions[:5]
 
-                cache.set(id=id, field="followup_questions", value=followup_questions)
+                self.cache.set(id=id, field="followup_questions", value=followup_questions)
 
                 return jsonify(
                     {
@@ -445,7 +511,7 @@ class VannaFlaskApp:
                     }
                 )
             else:
-                cache.set(id=id, field="followup_questions", value=[])
+                self.cache.set(id=id, field="followup_questions", value=[])
                 return jsonify(
                     {
                         "type": "question_list",
@@ -456,10 +522,14 @@ class VannaFlaskApp:
                 )
 
         @self.flask_app.route("/api/v0/generate_summary", methods=["GET"])
+        @self.requires_auth
         @self.requires_cache(["df", "question"])
-        def generate_summary(id: str, df, question):
+        def generate_summary(user: any, id: str, df, question):
             if self.allow_llm_to_see_data:
                 summary = vn.generate_summary(question=question, df=df)
+
+                self.cache.set(id=id, field="summary", value=summary)
+
                 return jsonify(
                     {
                         "type": "text",
@@ -477,10 +547,12 @@ class VannaFlaskApp:
                 )
 
         @self.flask_app.route("/api/v0/load_question", methods=["GET"])
+        @self.requires_auth
         @self.requires_cache(
-            ["question", "sql", "df", "fig_json"]
+            ["question", "sql", "df", "fig_json"],
+            optional_fields=["summary"]
         )
-        def load_question(id: str, question, sql, df, fig_json):
+        def load_question(user: any, id: str, question, sql, df, fig_json, summary):
             try:
                 return jsonify(
                     {
@@ -488,8 +560,9 @@ class VannaFlaskApp:
                         "id": id,
                         "question": question,
                         "sql": sql,
-                        "df": df.head(10).to_json(orient="records"),
+                        "df": df.head(10).to_json(orient="records", date_format="iso"),
                         "fig": fig_json,
+                        "summary": summary,
                     }
                 )
 
@@ -497,7 +570,8 @@ class VannaFlaskApp:
                 return jsonify({"type": "error", "error": str(e)})
 
         @self.flask_app.route("/api/v0/get_question_history", methods=["GET"])
-        def get_question_history():
+        @self.requires_auth
+        def get_question_history(user: any):
             return jsonify(
                 {
                     "type": "question_history",
@@ -525,7 +599,7 @@ class VannaFlaskApp:
         # Proxy the /vanna.svg file to the remote server
         @self.flask_app.route("/vanna.svg")
         def proxy_vanna_svg():
-            remote_url = f"https://vanna.ai/img/vanna.svg"
+            remote_url = "https://vanna.ai/img/vanna.svg"
             response = requests.get(remote_url, stream=True)
 
             # Check if the request to the remote URL was successful