vanna 0.0.10__tar.gz → 0.0.11__tar.gz

{vanna-0.0.10/src/vanna.egg-info → vanna-0.0.11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vanna
-Version: 0.0.10
+Version: 0.0.11
 Summary: Generate SQL queries from natural language
 Author-email: Zain Hoda <zain@vanna.ai>
 Project-URL: Homepage, https://github.com/vanna-ai/vanna-py

{vanna-0.0.10 → vanna-0.0.11}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "vanna"
-version = "0.0.10"
+version = "0.0.11"
 authors = [
   { name="Zain Hoda", email="zain@vanna.ai" },
 ]

{vanna-0.0.10 → vanna-0.0.11}/src/vanna/__init__.py

@@ -204,6 +204,7 @@ def add_user_to_dataset(dataset: str, email: str, is_admin: bool) -> bool:
     Args:
         dataset (str): The name of the dataset to add the user to.
         email (str): The email address of the user to add.
+        is_admin (bool): Whether or not the user should be an admin.
 
     Returns:
         bool: True if the user was added successfully, False otherwise.
@@ -258,14 +259,14 @@ def _set_org(org: str) -> None:
         d = __unauthenticated_rpc_call(method="check_org_exists", params=[Organization(name=org, user=None, connection=None)])
 
         if 'result' not in d:
-            raise Exception("Failed to check if organization exists")
+            raise Exception("Failed to check if dataset exists")
 
         status = Status(**d['result'])
 
         if status.success:
             raise Exception(f"An organization with the name {org} already exists")
 
-        create = input(f"Would you like to create organization '{org}'? (y/n): ")
+        create = input(f"Would you like to create dataset '{org}'? (y/n): ")
 
         if create.lower() == 'y':
             db_type = input("What type of database would you like to use? (Snowflake, BigQuery, Postgres, etc.): ")
@@ -274,7 +275,7 @@ def _set_org(org: str) -> None:
                 __org = org
             else:
                 __org = None
-                raise Exception("Failed to create organization")
+                raise Exception("Failed to create dataset")
     else:
         __org = org
 
@@ -509,6 +510,41 @@ def generate_sql(question: str) -> str:
 
     return sql_answer.sql
 
+def generate_followup_questions(question: str, df: pd.DataFrame) -> List[str]:
+    """
+    ## Example
+    ```python
+    vn.generate_followup_questions(question="What is the average salary of employees?", df=df)
+    # ['What is the average salary of employees in the Sales department?', 'What is the average salary of employees in the Engineering department?', ...]
+    ```
+
+    Generate follow-up questions using the Vanna.AI API.
+
+    Args:
+        question (str): The question to generate follow-up questions for.
+        df (pd.DataFrame): The DataFrame to generate follow-up questions for.
+
+    Returns:
+        List[str] or None: The follow-up questions, or None if an error occurred.
+    """
+    params = [DataResult(
+        question=question,
+        sql=None,
+        table_markdown=df.head().to_markdown(),
+        error=None,
+        correction_attempts=0,
+    )]
+
+    d = __rpc_call(method="generate_followup_questions", params=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 generate_questions() -> List[str]:
     """
     ## Example

{vanna-0.0.10 → vanna-0.0.11/src/vanna.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vanna
-Version: 0.0.10
+Version: 0.0.11
 Summary: Generate SQL queries from natural language
 Author-email: Zain Hoda <zain@vanna.ai>
 Project-URL: Homepage, https://github.com/vanna-ai/vanna-py