vanna 0.3.2__py3-none-any.whl → 0.3.4__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 (
@@ -220,19 +220,14 @@ class ZhipuAI_Chat(VannaBase):
         if len(prompt) == 0:
             raise Exception("Prompt is empty")
 
-        client = ZhipuAI(api_key=self.api_key)  # 填写您自己的APIKey
+        client = ZhipuAI(api_key=self.api_key)
         response = client.chat.completions.create(
-            model="glm-4",  # 填写需要调用的模型名称
+            model="glm-4",
             max_tokens=max_tokens,
             temperature=temperature,
             top_p=top_p,
             stop=stop,
             messages=prompt,
         )
-        # print(prompt)
-
-        # print(response)
-
-        # print(f"Cost {response.usage.total_tokens} token")
 
         return response.choices[0].message.content

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:
@@ -147,19 +159,21 @@ class VannaBase(ABC):
             return False
 
     def generate_followup_questions(
-        self, question: str, sql: str, df: pd.DataFrame, **kwargs
+        self, question: str, sql: str, df: pd.DataFrame, n_questions: int = 5, **kwargs
     ) -> list:
         """
         **Example:**
         ```python
-        vn.generate_followup_questions("What are the top 10 customers by sales?", df)
+        vn.generate_followup_questions("What are the top 10 customers by sales?", sql, df)
         ```
 
         Generate a list of followup questions that you can ask Vanna.AI.
 
         Args:
             question (str): The question that was asked.
+            sql (str): The LLM-generated SQL query.
             df (pd.DataFrame): The results of the SQL query.
+            n_questions (int): Number of follow-up questions to generate.
 
         Returns:
             list: A list of followup questions that you can ask Vanna.AI.
@@ -170,7 +184,7 @@ class VannaBase(ABC):
                 f"You are a helpful data assistant. The user asked the question: '{question}'\n\nThe SQL query for this question was: {sql}\n\nThe following is a pandas DataFrame with the results of the query: \n{df.to_markdown()}\n\n"
             ),
             self.user_message(
-                "Generate a list of followup questions that the user might ask about this data. Respond with a list of questions, one per line. Do not answer with any explanations -- just the questions. Remember that there should be an unambiguous SQL query that can be generated from the question. Prefer questions that are answerable outside of the context of this conversation. Prefer questions that are slight modifications of the SQL query that was generated that allow digging deeper into the data. Each question will be turned into a button that the user can click to generate a new SQL query so don't use 'example' type questions. Each question must have a one-to-one correspondence with an instantiated SQL query."
+                f"Generate a list of {n_questions} followup questions that the user might ask about this data. Respond with a list of questions, one per line. Do not answer with any explanations -- just the questions. Remember that there should be an unambiguous SQL query that can be generated from the question. Prefer questions that are answerable outside of the context of this conversation. Prefer questions that are slight modifications of the SQL query that was generated that allow digging deeper into the data. Each question will be turned into a button that the user can click to generate a new SQL query so don't use 'example' type questions. Each question must have a one-to-one correspondence with an instantiated SQL query."
             ),
         ]
 
@@ -361,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 (
@@ -380,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 (
@@ -396,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 (
@@ -888,6 +902,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):
         """
@@ -1236,7 +1338,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....")
@@ -1304,12 +1406,14 @@ class VannaBase(ABC):
         table_column = df.columns[
             df.columns.str.lower().str.contains("table_name")
         ].to_list()[0]
-        column_column = df.columns[
-            df.columns.str.lower().str.contains("column_name")
-        ].to_list()[0]
-        data_type_column = df.columns[
-            df.columns.str.lower().str.contains("data_type")
-        ].to_list()[0]
+        columns = [database_column,
+                    schema_column,
+                    table_column]
+        candidates = ["column_name",
+                      "data_type",
+                      "comment"]
+        matches = df.columns.str.lower().str.contains("|".join(candidates), regex=True)
+        columns += df.columns[matches].to_list()
 
         plan = TrainingPlan([])
 
@@ -1330,15 +1434,7 @@ class VannaBase(ABC):
                         f'{database_column} == "{database}" and {schema_column} == "{schema}" and {table_column} == "{table}"'
                     )
                     doc = f"The following columns are in the {table} table in the {database} database:\n\n"
-                    doc += df_columns_filtered_to_table[
-                        [
-                            database_column,
-                            schema_column,
-                            table_column,
-                            column_column,
-                            data_type_column,
-                        ]
-                    ].to_markdown()
+                    doc += df_columns_filtered_to_table[columns].to_markdown()
 
                     plan._plan.append(
                         TrainingPlanItem(

vanna/chromadb/chromadb_vector.py

@@ -1,5 +1,4 @@
 import json
-import uuid
 from typing import List
 
 import chromadb
@@ -8,6 +7,7 @@ from chromadb.config import Settings
 from chromadb.utils import embedding_functions
 
 from ..base import VannaBase
+from ..utils import deterministic_uuid
 
 default_ef = embedding_functions.DefaultEmbeddingFunction()
 
@@ -15,17 +15,14 @@ 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 = config.get("n_results", 10)
 
         if curr_client == "persistent":
             self.chroma_client = chromadb.PersistentClient(
@@ -42,13 +39,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]:
@@ -65,7 +68,7 @@ class ChromaDB_VectorStore(VannaBase):
             },
             ensure_ascii=False,
         )
-        id = str(uuid.uuid4()) + "-sql"
+        id = deterministic_uuid(question_sql_json) + "-sql"
         self.sql_collection.add(
             documents=question_sql_json,
             embeddings=self.generate_embedding(question_sql_json),
@@ -75,7 +78,7 @@ class ChromaDB_VectorStore(VannaBase):
         return id
 
     def add_ddl(self, ddl: str, **kwargs) -> str:
-        id = str(uuid.uuid4()) + "-ddl"
+        id = deterministic_uuid(ddl) + "-ddl"
         self.ddl_collection.add(
             documents=ddl,
             embeddings=self.generate_embedding(ddl),
@@ -84,7 +87,7 @@ class ChromaDB_VectorStore(VannaBase):
         return id
 
     def add_documentation(self, documentation: str, **kwargs) -> str:
-        id = str(uuid.uuid4()) + "-doc"
+        id = deterministic_uuid(documentation) + "-doc"
         self.documentation_collection.add(
             documents=documentation,
             embeddings=self.generate_embedding(documentation),

vanna/flask/__init__.py

@@ -525,7 +525,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

vanna/openai/openai_chat.py

@@ -41,7 +41,7 @@ class OpenAI_Chat(VannaBase):
         if config is None and client is None:
             self.client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
             return
-      
+
         if "api_key" in config:
             self.client = OpenAI(api_key=config["api_key"])
 
@@ -67,7 +67,31 @@ class OpenAI_Chat(VannaBase):
         for message in prompt:
             num_tokens += len(message["content"]) / 4
 
-        if self.config is not None and "engine" in self.config:
+        if kwargs.get("model", None) is not None:
+            model = kwargs.get("model", None)
+            print(
+                f"Using model {model} for {num_tokens} tokens (approx)"
+            )
+            response = self.client.chat.completions.create(
+                model=model,
+                messages=prompt,
+                max_tokens=self.max_tokens,
+                stop=None,
+                temperature=self.temperature,
+            )
+        elif kwargs.get("engine", None) is not None:
+            engine = kwargs.get("engine", None)
+            print(
+                f"Using model {engine} for {num_tokens} tokens (approx)"
+            )
+            response = self.client.chat.completions.create(
+                engine=engine,
+                messages=prompt,
+                max_tokens=self.max_tokens,
+                stop=None,
+                temperature=self.temperature,
+            )
+        elif self.config is not None and "engine" in self.config:
             print(
                 f"Using engine {self.config['engine']} for {num_tokens} tokens (approx)"
             )

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/utils.py

@@ -1,5 +1,8 @@
+import hashlib
 import os
 import re
+import uuid
+from typing import Union
 
 from .exceptions import ImproperlyConfigured, ValidationError
 
@@ -48,3 +51,27 @@ def sanitize_model_name(model_name):
         return model_name
     except Exception as e:
         raise ValidationError(e)
+
+
+def deterministic_uuid(content: Union[str, bytes]) -> str:
+    """Creates deterministic UUID on hash value of string or byte content.
+
+    Args:
+        content: String or byte representation of data.
+
+    Returns:
+        UUID of the content.
+    """
+    if isinstance(content, str):
+        content_bytes = content.encode("utf-8")
+    elif isinstance(content, bytes):
+        content_bytes = content
+    else:
+        raise ValueError(f"Content type {type(content)} not supported !")
+
+    hash_object = hashlib.sha256(content_bytes)
+    hash_hex = hash_object.hexdigest()
+    namespace = uuid.UUID("00000000-0000-0000-0000-000000000000")
+    content_uuid = str(uuid.uuid5(namespace, hash_hex))
+
+    return content_uuid

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.2.dist-info → vanna-0.3.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vanna
-Version: 0.3.2
+Version: 0.3.4
 Summary: Generate SQL queries from natural language
 Author-email: Zain Hoda <zain@vanna.ai>
 Requires-Python: >=3.9
@@ -26,6 +26,8 @@ Requires-Dist: openai ; extra == "all"
 Requires-Dist: mistralai ; extra == "all"
 Requires-Dist: chromadb ; extra == "all"
 Requires-Dist: anthropic ; extra == "all"
+Requires-Dist: zhipuai ; extra == "all"
+Requires-Dist: marqo ; extra == "all"
 Requires-Dist: anthropic ; extra == "anthropic"
 Requires-Dist: google-cloud-bigquery ; extra == "bigquery"
 Requires-Dist: chromadb ; extra == "chromadb"
@@ -39,6 +41,7 @@ Requires-Dist: psycopg2-binary ; extra == "postgres"
 Requires-Dist: db-dtypes ; extra == "postgres"
 Requires-Dist: snowflake-connector-python ; extra == "snowflake"
 Requires-Dist: tox ; extra == "test"
+Requires-Dist: zhipuai ; extra == "zhipuai"
 Project-URL: Bug Tracker, https://github.com/vanna-ai/vanna/issues
 Project-URL: Homepage, https://github.com/vanna-ai/vanna
 Provides-Extra: all
@@ -54,6 +57,7 @@ Provides-Extra: openai
 Provides-Extra: postgres
 Provides-Extra: snowflake
 Provides-Extra: test
+Provides-Extra: zhipuai
 
 
 

{vanna-0.3.2.dist-info → vanna-0.3.4.dist-info}/RECORD

@@ -1,18 +1,18 @@
 vanna/__init__.py,sha256=4zz2kSkVZenjwJQg-ETWsIVYdz3gio275i9DMa_aHxM,9248
 vanna/local.py,sha256=U5s8ybCRQhBUizi8I69o3jqOpTeu_6KGYY6DMwZxjG4,313
-vanna/remote.py,sha256=qqaTA4l-ikVH_1aQInae4DyTAfCCmFokQme6Jq_i1us,12803
-vanna/utils.py,sha256=Q0H4eugPYg9SVpEoTWgvmuoJZZxOVRhNzrP97E5lyak,1472
-vanna/ZhipuAI/ZhipuAI_Chat.py,sha256=C2zkK9R_HZnhOTN_lAvjiFnHSH2s_lFNfu-FEjU2ycs,8906
+vanna/remote.py,sha256=CcScbeEkqYlzBZJMeMLDlLgmaFcG0QSUQUZybFI3Y28,1856
+vanna/utils.py,sha256=cs0B_0MwhmPI2nWjVHifDYCmCR0kkddylQ2vloaPDSw,2247
+vanna/ZhipuAI/ZhipuAI_Chat.py,sha256=WtZKUBIwlNH0BGbb4lZbVR7pTWIrn7b4RLIk-7u0SuQ,8725
 vanna/ZhipuAI/ZhipuAI_embeddings.py,sha256=lUqzJg9fOx7rVFhjdkFjXcDeVGV4aAB5Ss0oERsa8pE,2849
 vanna/ZhipuAI/__init__.py,sha256=NlsijtcZp5Tj9jkOe9fNcOQND_QsGgu7otODsCLBPr0,116
 vanna/anthropic/__init__.py,sha256=85s_2mAyyPxc0T_0JEvYeAkEKWJwkwqoyUwSC5dw9Gk,43
 vanna/anthropic/anthropic_chat.py,sha256=Wk0o-NMW1uvR2fhSWxrR_2FqNh-dLprNG4uuVqpqAkY,2615
 vanna/base/__init__.py,sha256=Sl-HM1RRYzAZoSqmL1CZQmF3ZF-byYTCFQP3JZ2A5MU,28
-vanna/base/base.py,sha256=25tHTMsCnoek6X9C0hqpPyNM3H0mihcL4I8Mlq-74c0,54890
+vanna/base/base.py,sha256=89XPWy97YVx6090mNmu1zvn4k8X1pusCuAIypHHexNc,58100
 vanna/chromadb/__init__.py,sha256=-iL0nW_g4uM8nWKMuWnNePfN4nb9uk8P3WzGvezOqRg,50
-vanna/chromadb/chromadb_vector.py,sha256=fa7uj_knzSfzsVLvpSunwwu1ZJNC3GbiNZ4Yy09v4l4,8372
+vanna/chromadb/chromadb_vector.py,sha256=1n4U4XpXThCFqyJf0zAYVA7mQu9rUkjOFtYn9e04JAo,8461
 vanna/exceptions/__init__.py,sha256=N76unE7sjbGGBz6LmCrPQAugFWr9cUFv8ErJxBrCTts,717
-vanna/flask/__init__.py,sha256=UgM0Ce5pGDdadWV6ZEAXj7RXDE1E420DW1wtR-juBMw,21212
+vanna/flask/__init__.py,sha256=tpwpA8596Uyn60FAy7I5oJ81c7kgCB2JG9X044P0_SA,21211
 vanna/flask/assets.py,sha256=pOOtPV8aWtFsTuxJneFHcfrXhXh6cOSvS-Y8JO2HYrY,180924
 vanna/marqo/__init__.py,sha256=GaAWtJ0B-H5rTY607iLCCrLD7T0zMYM5qWIomEB9gLk,37
 vanna/marqo/marqo.py,sha256=W7WTtzWp4RJjZVy6OaXHqncUBIPdI4Q7qH7BRCxZ1_A,5242
@@ -21,11 +21,11 @@ vanna/mistral/mistral.py,sha256=DAEqAT9SzC91rfMM_S3SuzBZ34MrKHw9qAj6EP2MGVk,1508
 vanna/ollama/__init__.py,sha256=4xyu8aHPdnEHg5a-QAMwr5o0ns5wevsp_zkI-ndMO2k,27
 vanna/ollama/ollama.py,sha256=jfW9VQHAcmzDeo4jF3HJjOMYwAWmptknKqEJaQ0MTno,2418
 vanna/openai/__init__.py,sha256=tGkeQ7wTIPsando7QhoSHehtoQVdYLwFbKNlSmCmNeQ,86
-vanna/openai/openai_chat.py,sha256=Y3-Fhz9c6D-5vrMR1zGibavAPNPAu-hMTwqGfKhAg3Q,3852
+vanna/openai/openai_chat.py,sha256=lm-hUsQxu6Q1t06A2csC037zI4VkMk0wFbQ-_Lj74Wg,4764
 vanna/openai/openai_embeddings.py,sha256=g4pNh9LVcYP9wOoO8ecaccDFWmCUYMInebfHucAa2Gc,1260
 vanna/types/__init__.py,sha256=Qhn_YscKtJh7mFPCyCDLa2K8a4ORLMGVnPpTbv9uB2U,4957
 vanna/vannadb/__init__.py,sha256=C6UkYocmO6dmzfPKZaWojN0mI5YlZZ9VIbdcquBE58A,48
-vanna/vannadb/vannadb_vector.py,sha256=f4kddaJgTpZync7wnQi09QdODUuMtiHsK7WfKBUAmSo,5644
-vanna-0.3.2.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
-vanna-0.3.2.dist-info/METADATA,sha256=Suj6A5Wqxfyrs7d91R6g7y2lCusWtDUSctv8pUaueqg,9961
-vanna-0.3.2.dist-info/RECORD,,
+vanna/vannadb/vannadb_vector.py,sha256=9YwTO3Lh5owWQE7KPMBqLp2EkiGV0RC1sEYhslzJzgI,6168
+vanna-0.3.4.dist-info/WHEEL,sha256=EZbGkh7Ie4PoZfRQ8I0ZuP9VklN_TvcZ6DSE5Uar4z4,81
+vanna-0.3.4.dist-info/METADATA,sha256=FEg4vs5ZiSAvd5YkF5oEfFqod9n3UoNfi51Q_2WKotA,10107
+vanna-0.3.4.dist-info/RECORD,,