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

vanna/flask/auth.py

@@ -0,0 +1,55 @@
+from abc import ABC, abstractmethod
+
+import flask
+
+
+class AuthInterface(ABC):
+    @abstractmethod
+    def get_user(self, flask_request) -> any:
+        pass
+
+    @abstractmethod
+    def is_logged_in(self, user: any) -> bool:
+        pass
+
+    @abstractmethod
+    def override_config_for_user(self, user: any, config: dict) -> dict:
+        pass
+
+    @abstractmethod
+    def login_form(self) -> str:
+        pass
+
+    @abstractmethod
+    def login_handler(self, flask_request) -> str:
+        pass
+
+    @abstractmethod
+    def callback_handler(self, flask_request) -> str:
+        pass
+
+    @abstractmethod
+    def logout_handler(self, flask_request) -> str:
+        pass
+
+class NoAuth(AuthInterface):
+    def get_user(self, flask_request) -> any:
+        return {}
+
+    def is_logged_in(self, user: any) -> bool:
+        return True
+
+    def override_config_for_user(self, user: any, config: dict) -> dict:
+        return config
+
+    def login_form(self) -> str:
+        return ''
+
+    def login_handler(self, flask_request) -> str:
+        return 'No login required'
+
+    def callback_handler(self, flask_request) -> str:
+        return 'No login required'
+
+    def logout_handler(self, flask_request) -> str:
+        return 'No login required'

vanna/google/__init__.py

@@ -0,0 +1 @@
+from .gemini_chat import GoogleGeminiChat

vanna/google/gemini_chat.py

@@ -0,0 +1,52 @@
+import os
+from ..base import VannaBase
+
+
+class GoogleGeminiChat(VannaBase):
+    def __init__(self, config=None):
+        VannaBase.__init__(self, config=config)
+
+        # default temperature - can be overrided using config
+        self.temperature = 0.7      
+
+        if "temperature" in config:
+            self.temperature = config["temperature"]
+
+        if "model_name" in config:
+            model_name = config["model_name"]
+        else:
+            model_name = "gemini-1.0-pro"
+
+        self.google_api_key = None
+
+        if "api_key" in config or os.getenv("GOOGLE_API_KEY"):
+            """
+            If Google api_key is provided through config
+            or set as an environment variable, assign it.
+            """
+            import google.generativeai as genai
+
+            genai.configure(api_key=config["api_key"])
+            self.chat_model = genai.GenerativeModel(model_name)
+        else:
+            # Authenticate using VertexAI
+            from vertexai.preview.generative_models import GenerativeModel
+            self.chat_model = GenerativeModel("gemini-pro")
+
+    def system_message(self, message: str) -> any:
+        return message
+
+    def user_message(self, message: str) -> any:
+        return message
+
+    def assistant_message(self, message: str) -> any:
+        return message
+
+    def submit_prompt(self, prompt, **kwargs) -> str:
+        response = self.chat_model.generate_content(
+            prompt,
+            generation_config={
+                "temperature": self.temperature,
+            },
+        )
+        return response.text

vanna/remote.py

@@ -34,11 +34,13 @@ from .types import (
   UserOTP,
   Visibility,
 )
+from .vannadb import VannaDB_VectorStore
 
 
-class VannaDefault(VannaBase):
+class VannaDefault(VannaDB_VectorStore):
     def __init__(self, model: str, api_key: str, config=None):
         VannaBase.__init__(self, config=config)
+        VannaDB_VectorStore.__init__(self, vanna_model=model, vanna_api_key=api_key, config=config)
 
         self._model = model
         self._api_key = api_key
@@ -48,50 +50,6 @@ class VannaDefault(VannaBase):
             if config is None or "endpoint" not in config
             else config["endpoint"]
         )
-        self._unauthenticated_endpoint = (
-            "https://ask.vanna.ai/unauthenticated_rpc"
-            if config is None or "unauthenticated_endpoint" not in config
-            else config["unauthenticated_endpoint"]
-        )
-
-    def _unauthenticated_rpc_call(self, method, params):
-        headers = {
-            "Content-Type": "application/json",
-        }
-        data = {
-            "method": method,
-            "params": [self._dataclass_to_dict(obj) for obj in params],
-        }
-
-        response = requests.post(
-            self._unauthenticated_endpoint, headers=headers, data=json.dumps(data)
-        )
-        return response.json()
-
-    def _rpc_call(self, method, params):
-        if method != "list_orgs":
-            headers = {
-                "Content-Type": "application/json",
-                "Vanna-Key": self._api_key,
-                "Vanna-Org": self._model,
-            }
-        else:
-            headers = {
-                "Content-Type": "application/json",
-                "Vanna-Key": self._api_key,
-                "Vanna-Org": "demo-tpc-h",
-            }
-
-        data = {
-            "method": method,
-            "params": [self._dataclass_to_dict(obj) for obj in params],
-        }
-
-        response = requests.post(self._endpoint, headers=headers, data=json.dumps(data))
-        return response.json()
-
-    def _dataclass_to_dict(self, obj):
-        return dataclasses.asdict(obj)
 
     def system_message(self, message: str) -> any:
         return {"role": "system", "content": message}
@@ -102,299 +60,6 @@ class VannaDefault(VannaBase):
     def assistant_message(self, message: str) -> any:
         return {"role": "assistant", "content": message}
 
-    def get_training_data(self, **kwargs) -> pd.DataFrame:
-        """
-        Get the training data for the current model
-
-        **Example:**
-        ```python
-        training_data = vn.get_training_data()
-        ```
-
-        Returns:
-            pd.DataFrame or None: The training data, or None if an error occurred.
-
-        """
-        params = []
-
-        d = self._rpc_call(method="get_training_data", params=params)
-
-        if "result" not in d:
-            return None
-
-        # Load the result into a dataclass
-        training_data = DataFrameJSON(**d["result"])
-
-        df = pd.read_json(StringIO(training_data.data))
-
-        return df
-
-    def remove_training_data(self, id: str, **kwargs) -> bool:
-        """
-        Remove training data from the model
-
-        **Example:**
-        ```python
-        vn.remove_training_data(id="1-ddl")
-        ```
-
-        Args:
-            id (str): The ID of the training data to remove.
-        """
-        params = [StringData(data=id)]
-
-        d = self._rpc_call(method="remove_training_data", params=params)
-
-        if "result" not in d:
-            raise Exception(f"Error removing training data")
-
-        status = Status(**d["result"])
-
-        if not status.success:
-            raise Exception(f"Error removing training data: {status.message}")
-
-        return status.success
-
-    def generate_questions(self) -> list[str]:
-        """
-        **Example:**
-        ```python
-        vn.generate_questions()
-        # ['What is the average salary of employees?', 'What is the total salary of employees?', ...]
-        ```
-
-        Generate questions using the Vanna.AI API.
-
-        Returns:
-            List[str] or None: The questions, or None if an error occurred.
-        """
-        d = self._rpc_call(method="generate_questions", params=[])
-
-        if "result" not in d:
-            return None
-
-        # Load the result into a dataclass
-        question_string_list = QuestionStringList(**d["result"])
-
-        return question_string_list.questions
-
-    def add_ddl(self, ddl: str, **kwargs) -> str:
-        """
-        Adds a DDL statement to the model's training data
-
-        **Example:**
-        ```python
-        vn.add_ddl(
-            ddl="CREATE TABLE employees (id INT, name VARCHAR(255), salary INT)"
-        )
-        ```
-
-        Args:
-            ddl (str): The DDL statement to store.
-
-        Returns:
-            str: The ID of the DDL statement.
-        """
-        params = [StringData(data=ddl)]
-
-        d = self._rpc_call(method="add_ddl", params=params)
-
-        if "result" not in d:
-            raise Exception("Error adding DDL", d)
-
-        status = StatusWithId(**d["result"])
-
-        return status.id
-
-    def add_documentation(self, documentation: str, **kwargs) -> str:
-        """
-        Adds documentation to the model's training data
-
-        **Example:**
-        ```python
-        vn.add_documentation(
-            documentation="Our organization's definition of sales is the discount price of an item multiplied by the quantity sold."
-        )
-        ```
-
-        Args:
-            documentation (str): The documentation string to store.
-
-        Returns:
-            str: The ID of the documentation string.
-        """
-        params = [StringData(data=documentation)]
-
-        d = self._rpc_call(method="add_documentation", params=params)
-
-        if "result" not in d:
-            raise Exception("Error adding documentation", d)
-
-        status = StatusWithId(**d["result"])
-
-        return status.id
-
-    def add_question_sql(self, question: str, sql: str, **kwargs) -> str:
-        """
-        Adds a question and its corresponding SQL query to the model's training data. The preferred way to call this is to use [`vn.train(sql=...)`][vanna.train].
-
-        **Example:**
-        ```python
-        vn.add_sql(
-            question="What is the average salary of employees?",
-            sql="SELECT AVG(salary) FROM employees"
-        )
-        ```
-
-        Args:
-            question (str): The question to store.
-            sql (str): The SQL query to store.
-            tag (Union[str, None]): A tag to associate with the question and SQL query.
-
-        Returns:
-            str: The ID of the question and SQL query.
-        """
-        if "tag" in kwargs:
-            tag = kwargs["tag"]
-        else:
-            tag = "Manually Trained"
-
-        params = [QuestionSQLPair(question=question, sql=sql, tag=tag)]
-
-        d = self._rpc_call(method="add_sql", params=params)
-
-        if "result" not in d:
-            raise Exception("Error adding question and SQL pair", d)
-
-        status = StatusWithId(**d["result"])
-
-        return status.id
-
-    def generate_embedding(self, data: str, **kwargs) -> list[float]:
-        """
-        Not necessary for remote models as embeddings are generated on the server side.
-        """
-        pass
-
-    def generate_plotly_code(
-        self, question: str = None, sql: str = None, df_metadata: str = None, **kwargs
-    ) -> str:
-        """
-        **Example:**
-        ```python
-        vn.generate_plotly_code(
-            question="What is the average salary of employees?",
-            sql="SELECT AVG(salary) FROM employees",
-            df_metadata=df.dtypes
-        )
-        # fig = px.bar(df, x="name", y="salary")
-        ```
-        Generate Plotly code using the Vanna.AI API.
-
-        Args:
-            question (str): The question to generate Plotly code for.
-            sql (str): The SQL query to generate Plotly code for.
-            df (pd.DataFrame): The dataframe to generate Plotly code for.
-            chart_instructions (str): Optional instructions for how to plot the chart.
-
-        Returns:
-            str or None: The Plotly code, or None if an error occurred.
-        """
-        if kwargs is not None and "chart_instructions" in kwargs:
-            if question is not None:
-                question = (
-                    question
-                    + " -- When plotting, follow these instructions: "
-                    + kwargs["chart_instructions"]
-                )
-            else:
-                question = (
-                    "When plotting, follow these instructions: "
-                    + kwargs["chart_instructions"]
-                )
-
-        params = [
-            DataResult(
-                question=question,
-                sql=sql,
-                table_markdown=df_metadata,
-                error=None,
-                correction_attempts=0,
-            )
-        ]
-
-        d = self._rpc_call(method="generate_plotly_code", params=params)
-
-        if "result" not in d:
-            return None
-
-        # Load the result into a dataclass
-        plotly_code = PlotlyResult(**d["result"])
-
-        return plotly_code.plotly_code
-
-    def generate_question(self, sql: str, **kwargs) -> str:
-        """
-
-        **Example:**
-        ```python
-        vn.generate_question(sql="SELECT * FROM students WHERE name = 'John Doe'")
-        # 'What is the name of the student?'
-        ```
-
-        Generate a question from an SQL query using the Vanna.AI API.
-
-        Args:
-            sql (str): The SQL query to generate a question for.
-
-        Returns:
-            str or None: The question, or None if an error occurred.
-
-        """
-        params = [
-            SQLAnswer(
-                raw_answer="",
-                prefix="",
-                postfix="",
-                sql=sql,
-            )
-        ]
-
-        d = self._rpc_call(method="generate_question", params=params)
-
-        if "result" not in d:
-            return None
-
-        # Load the result into a dataclass
-        question = Question(**d["result"])
-
-        return question.question
-
-    def get_sql_prompt(
-        self,
-        question: str,
-        question_sql_list: list,
-        ddl_list: list,
-        doc_list: list,
-        **kwargs,
-    ):
-        """
-        Not necessary for remote models as prompts are generated on the server side.
-        """
-
-    def get_followup_questions_prompt(
-        self,
-        question: str,
-        df: pd.DataFrame,
-        question_sql_list: list,
-        ddl_list: list,
-        doc_list: list,
-        **kwargs,
-    ):
-        """
-        Not necessary for remote models as prompts are generated on the server side.
-        """
-
     def submit_prompt(self, prompt, **kwargs) -> str:
         # JSON-ify the prompt
         json_prompt = json.dumps(prompt)
@@ -410,46 +75,3 @@ class VannaDefault(VannaBase):
         results = StringData(**d["result"])
 
         return results.data
-
-    def get_similar_question_sql(self, question: str, **kwargs) -> list:
-        """
-        Not necessary for remote models as similar questions are generated on the server side.
-        """
-
-    def get_related_ddl(self, question: str, **kwargs) -> list:
-        """
-        Not necessary for remote models as related DDL statements are generated on the server side.
-        """
-
-    def get_related_documentation(self, question: str, **kwargs) -> list:
-        """
-        Not necessary for remote models as related documentation is generated on the server side.
-        """
-
-    def generate_sql(self, question: str, **kwargs) -> str:
-        """
-        **Example:**
-        ```python
-        vn.generate_sql_from_question(question="What is the average salary of employees?")
-        # SELECT AVG(salary) FROM employees
-        ```
-
-        Generate an SQL query using the Vanna.AI API.
-
-        Args:
-            question (str): The question to generate an SQL query for.
-
-        Returns:
-            str or None: The SQL query, or None if an error occurred.
-        """
-        params = [Question(question=question)]
-
-        d = self._rpc_call(method="generate_sql_from_question", params=params)
-
-        if "result" not in d:
-            return None
-
-        # Load the result into a dataclass
-        sql_answer = SQLAnswer(**d["result"])
-
-        return sql_answer.sql

vanna/vannadb/vannadb_vector.py

@@ -7,14 +7,17 @@ import requests
 
 from ..base import VannaBase
 from ..types import (
-    DataFrameJSON,
-    Question,
-    QuestionSQLPair,
-    Status,
-    StatusWithId,
-    StringData,
-    TrainingData,
+  DataFrameJSON,
+  NewOrganization,
+  OrganizationList,
+  Question,
+  QuestionSQLPair,
+  Status,
+  StatusWithId,
+  StringData,
+  TrainingData,
 )
+from ..utils import sanitize_model_name
 
 
 class VannaDB_VectorStore(VannaBase):
@@ -29,27 +32,8 @@ class VannaDB_VectorStore(VannaBase):
             if config is None or "endpoint" not in config
             else config["endpoint"]
         )
-        self._unauthenticated_endpoint = (
-            "https://ask.vanna.ai/unauthenticated_rpc"
-            if config is None or "unauthenticated_endpoint" not in config
-            else config["unauthenticated_endpoint"]
-        )
         self.related_training_data = {}
 
-    def _unauthenticated_rpc_call(self, method, params):
-        headers = {
-            "Content-Type": "application/json",
-        }
-        data = {
-            "method": method,
-            "params": [self._dataclass_to_dict(obj) for obj in params],
-        }
-
-        response = requests.post(
-            self._unauthenticated_endpoint, headers=headers, data=json.dumps(data)
-        )
-        return response.json()
-
     def _rpc_call(self, method, params):
         if method != "list_orgs":
             headers = {
@@ -75,6 +59,53 @@ class VannaDB_VectorStore(VannaBase):
     def _dataclass_to_dict(self, obj):
         return dataclasses.asdict(obj)
 
+    def create_model(self, model: str, **kwargs) -> bool:
+        """
+        **Example:**
+        ```python
+        success = vn.create_model("my_model")
+        ```
+        Create a new model.
+
+        Args:
+            model (str): The name of the model to create.
+
+        Returns:
+            bool: True if the model was created, False otherwise.
+        """
+        model = sanitize_model_name(model)
+        params = [NewOrganization(org_name=model, db_type="")]
+
+        d = self._rpc_call(method="create_org", params=params)
+
+        if "result" not in d:
+            return False
+
+        status = Status(**d["result"])
+
+        return status.success
+
+    def get_models(self) -> list:
+        """
+        **Example:**
+        ```python
+        models = vn.get_models()
+        ```
+
+        List the models that belong to the user.
+
+        Returns:
+            List[str]: A list of model names.
+        """
+        d = self._rpc_call(method="list_my_models", params=[])
+
+        if "result" not in d:
+            return []
+
+        orgs = OrganizationList(**d["result"])
+
+        return orgs.organizations
+
     def generate_embedding(self, data: str, **kwargs) -> list[float]:
         # This is done server-side
         pass
@@ -141,7 +172,7 @@ class VannaDB_VectorStore(VannaBase):
         d = self._rpc_call(method="remove_training_data", params=params)
 
         if "result" not in d:
-            raise Exception(f"Error removing training data")
+            raise Exception("Error removing training data")
 
         status = Status(**d["result"])
 

{vanna-0.3.3.dist-info → vanna-0.4.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vanna
-Version: 0.3.3
+Version: 0.4.0
 Summary: Generate SQL queries from natural language
 Author-email: Zain Hoda <zain@vanna.ai>
 Requires-Python: >=3.9
@@ -28,11 +28,15 @@ Requires-Dist: chromadb ; extra == "all"
 Requires-Dist: anthropic ; extra == "all"
 Requires-Dist: zhipuai ; extra == "all"
 Requires-Dist: marqo ; extra == "all"
+Requires-Dist: google-generativeai ; extra == "all"
+Requires-Dist: google-cloud-aiplatform ; extra == "all"
 Requires-Dist: anthropic ; extra == "anthropic"
 Requires-Dist: google-cloud-bigquery ; extra == "bigquery"
 Requires-Dist: chromadb ; extra == "chromadb"
 Requires-Dist: duckdb ; extra == "duckdb"
 Requires-Dist: google-generativeai ; extra == "gemini"
+Requires-Dist: google-generativeai ; extra == "google"
+Requires-Dist: google-cloud-aiplatform ; extra == "google"
 Requires-Dist: marqo ; extra == "marqo"
 Requires-Dist: mistralai ; extra == "mistralai"
 Requires-Dist: PyMySQL ; extra == "mysql"
@@ -50,6 +54,7 @@ Provides-Extra: bigquery
 Provides-Extra: chromadb
 Provides-Extra: duckdb
 Provides-Extra: gemini
+Provides-Extra: google
 Provides-Extra: marqo
 Provides-Extra: mistralai
 Provides-Extra: mysql