Merge pull request #2 from vanna-ai/init

generate question

docs/vanna.html

@@ -57,6 +57,9 @@ <h2>API Documentation</h2>
             <li>
                     <a class="function" href="#generate_explanation">generate_explanation</a>
             </li>
+            <li>
+                    <a class="function" href="#generate_question">generate_question</a>
+            </li>
     </ul>
 
 
@@ -182,7 +185,7 @@ <h1 class="modulename">
                     <div class="attr function">
 
         <span class="def">def</span>
-        <span class="name">generate_sql</span><span class="signature pdoc-code condensed">(<span class="param"><span class="n">question</span><span class="p">:</span> <span class="nb">str</span></span><span class="return-annotation">) -> <span class="nb">str</span> <span class="o">|</span> <span class="kc">None</span>:</span></span>
+        <span class="name">generate_sql</span><span class="signature pdoc-code condensed">(<span class="param"><span class="n">question</span><span class="p">:</span> <span class="nb">str</span></span><span class="return-annotation">) -> <span class="nb">str</span>:</span></span>
 
 
     </div>
@@ -203,7 +206,7 @@ <h1 class="modulename">
                     <div class="attr function">
 
         <span class="def">def</span>
-        <span class="name">generate_plotly_code</span><span class="signature pdoc-code multiline">(<span class="param">	<span class="n">question</span><span class="p">:</span> <span class="nb">str</span> <span class="o">|</span> <span class="kc">None</span>,</span><span class="param">	<span class="n">sql</span><span class="p">:</span> <span class="nb">str</span> <span class="o">|</span> <span class="kc">None</span>,</span><span class="param">	<span class="n">df</span><span class="p">:</span> <span class="n">pandas</span><span class="o">.</span><span class="n">core</span><span class="o">.</span><span class="n">frame</span><span class="o">.</span><span class="n">DataFrame</span></span><span class="return-annotation">) -> <span class="nb">str</span> <span class="o">|</span> <span class="kc">None</span>:</span></span>
+        <span class="name">generate_plotly_code</span><span class="signature pdoc-code multiline">(<span class="param">	<span class="n">question</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">str</span><span class="p">]</span>,</span><span class="param">	<span class="n">sql</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="nb">str</span><span class="p">]</span>,</span><span class="param">	<span class="n">df</span><span class="p">:</span> <span class="n">pandas</span><span class="o">.</span><span class="n">core</span><span class="o">.</span><span class="n">frame</span><span class="o">.</span><span class="n">DataFrame</span></span><span class="return-annotation">) -> <span class="nb">str</span>:</span></span>
 
 
     </div>
@@ -226,7 +229,7 @@ <h1 class="modulename">
                     <div class="attr function">
 
         <span class="def">def</span>
-        <span class="name">get_plotly_figure</span><span class="signature pdoc-code multiline">(<span class="param">	<span class="n">plotly_code</span><span class="p">:</span> <span class="nb">str</span>,</span><span class="param">	<span class="n">df</span><span class="p">:</span> <span class="n">pandas</span><span class="o">.</span><span class="n">core</span><span class="o">.</span><span class="n">frame</span><span class="o">.</span><span class="n">DataFrame</span>,</span><span class="param">	<span class="n">dark_mode</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span></span><span class="return-annotation">) -> <span class="n">plotly</span><span class="o">.</span><span class="n">graph_objs</span><span class="o">.</span><span class="n">_figure</span><span class="o">.</span><span class="n">Figure</span> <span class="o">|</span> <span class="kc">None</span>:</span></span>
+        <span class="name">get_plotly_figure</span><span class="signature pdoc-code multiline">(<span class="param">	<span class="n">plotly_code</span><span class="p">:</span> <span class="nb">str</span>,</span><span class="param">	<span class="n">df</span><span class="p">:</span> <span class="n">pandas</span><span class="o">.</span><span class="n">core</span><span class="o">.</span><span class="n">frame</span><span class="o">.</span><span class="n">DataFrame</span>,</span><span class="param">	<span class="n">dark_mode</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span></span><span class="return-annotation">) -> <span class="n">plotly</span><span class="o">.</span><span class="n">graph_objs</span><span class="o">.</span><span class="n">_figure</span><span class="o">.</span><span class="n">Figure</span>:</span></span>
 
 
     </div>
@@ -256,19 +259,13 @@ <h1 class="modulename">
 
             <div class="docstring"><p>Get the results of an SQL query using the Vanna.AI API.</p>
 
-<h6 id="parameters">Parameters</h6>
-
-<ul>
-<li><strong>cs</strong>:  The Snowflake cursor to use.</li>
-<li><strong>default_database</strong>:  The default database to use (executed as "USE DATABASE {default_database};")</li>
-<li><strong>sql</strong>:  The SQL query to run.</li>
-</ul>
-
-<h6 id="returns">Returns</h6>
+<p>Args:
+    cs (pyodbc.Cursor): The cursor to use.
+    default_database (str): The default database to use.
+    sql (str): The SQL query to execute.</p>
 
-<blockquote>
-  <p>The results of the SQL query.</p>
-</blockquote>
+<p>Returns:
+    pd.DataFrame: The results of the SQL query.</p>
 </div>
 
 
@@ -277,7 +274,7 @@ <h6 id="returns">Returns</h6>
                     <div class="attr function">
 
         <span class="def">def</span>
-        <span class="name">generate_explanation</span><span class="signature pdoc-code condensed">(<span class="param"><span class="n">sql</span><span class="p">:</span> <span class="nb">str</span></span><span class="return-annotation">) -> <span class="nb">str</span> <span class="o">|</span> <span class="kc">None</span>:</span></span>
+        <span class="name">generate_explanation</span><span class="signature pdoc-code condensed">(<span class="param"><span class="n">sql</span><span class="p">:</span> <span class="nb">str</span></span><span class="return-annotation">) -> <span class="nb">str</span>:</span></span>
 
 
     </div>
@@ -293,17 +290,40 @@ <h6 id="returns">Returns</h6>
 
 <p>Generate an explanation of an SQL query using the Vanna.AI API.</p>
 
-<h6 id="parameters">Parameters</h6>
+<p>Args:
+    sql (str): The SQL query to generate an explanation for.</p>
+
+<p>Returns:
+    str or None: The explanation, or None if an error occurred.</p>
+</div>
+
+
+                </section>
+                <section id="generate_question">
+                    <div class="attr function">
+
+        <span class="def">def</span>
+        <span class="name">generate_question</span><span class="signature pdoc-code condensed">(<span class="param"><span class="n">sql</span><span class="p">:</span> <span class="nb">str</span></span><span class="return-annotation">) -> <span class="nb">str</span>:</span></span>
+
+
+    </div>
+    <a class="headerlink" href="#generate_question"></a>
+
+            <div class="docstring"><h2 id="example">Example</h2>
+
+<div class="pdoc-code codehilite">
+<pre><span></span><code><span class="n">vn</span><span class="o">.</span><span class="n">generate_question</span><span class="p">(</span><span class="n">sql</span><span class="o">=</span><span class="s2">&quot;SELECT * FROM students WHERE name = &#39;John Doe&#39;&quot;</span><span class="p">)</span>
+<span class="c1"># &#39;AI Response&#39;</span>
+</code></pre>
+</div>
 
-<ul>
-<li><strong>sql</strong>:  The SQL query to explain.</li>
-</ul>
+<p>Generate a question from an SQL query using the Vanna.AI API.</p>
 
-<h6 id="returns">Returns</h6>
+<p>Args:
+    sql (str): The SQL query to generate a question for.</p>
 
-<blockquote>
-  <p>The explanation of the SQL query, or None if an error occurred.</p>
-</blockquote>
+<p>Returns:
+    str or None: The question, or None if an error occurred.</p>
 </div>
 
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "vanna"
-version = "0.0.2"
+version = "0.0.3"
 authors = [
   { name="Zain Hoda", email="[email protected]" },
 ]

src/vanna/__init__.py

@@ -143,7 +143,7 @@ def remove_sql(question: str) -> bool:
 
     return status.success
 
-def generate_sql(question: str) -> str | None:
+def generate_sql(question: str) -> str:
     """
     Generate an SQL query using the Vanna.AI API.
 
@@ -165,7 +165,7 @@ def generate_sql(question: str) -> str | None:
 
     return sql_answer.sql
 
-def generate_plotly_code(question: str | None, sql: str | None, df: pd.DataFrame) -> str | None:
+def generate_plotly_code(question: Union[str, None], sql: Union[str, None], df: pd.DataFrame) -> str:
     """
     Generate Plotly code using the Vanna.AI API.
 
@@ -195,7 +195,7 @@ def generate_plotly_code(question: str | None, sql: str | None, df: pd.DataFrame
 
     return plotly_code.plotly_code
 
-def get_plotly_figure(plotly_code: str, df: pd.DataFrame, dark_mode: bool = True) -> plotly.graph_objs.Figure | None:
+def get_plotly_figure(plotly_code: str, df: pd.DataFrame, dark_mode: bool = True) -> plotly.graph_objs.Figure:
     """
     Get a Plotly figure from a dataframe and Plotly code.
 
@@ -223,15 +223,13 @@ def get_results(cs, default_database: str, sql: str) -> pd.DataFrame:
     """
     Get the results of an SQL query using the Vanna.AI API.
 
-    :param cs: The Snowflake cursor to use.
-    :type cs: snowflake.connector.cursor.SnowflakeCursor
-    :param default_database: The default database to use (executed as "USE DATABASE {default_database};")
-    :type default_database: str
-    :param sql: The SQL query to run.
-    :type sql: str
-    
-    :return: The results of the SQL query.
-    :rtype: pd.DataFrame
+    Args:
+        cs (pyodbc.Cursor): The cursor to use.
+        default_database (str): The default database to use.
+        sql (str): The SQL query to execute.
+
+    Returns:
+        pd.DataFrame: The results of the SQL query.
     """
     cs.execute(f"USE DATABASE {default_database}")
 
@@ -245,7 +243,7 @@ def get_results(cs, default_database: str, sql: str) -> pd.DataFrame:
     return df
 
 
-def generate_explanation(sql: str) -> str | None:
+def generate_explanation(sql: str) -> str:
     """
 
     ## Example
@@ -256,11 +254,11 @@ def generate_explanation(sql: str) -> str | None:
     
     Generate an explanation of an SQL query using the Vanna.AI API.
 
-    :param sql: The SQL query to explain.
-    :type sql: str
-
-    :return: The explanation of the SQL query, or None if an error occurred.
-    :rtype: str or None
+    Args:
+        sql (str): The SQL query to generate an explanation for.
+    
+    Returns:
+        str or None: The explanation, or None if an error occurred.
 
     """
     params = [SQLAnswer(
@@ -279,3 +277,38 @@ def generate_explanation(sql: str) -> str | None:
     explanation = Explanation(**d['result'])
 
     return explanation.explanation
+
+def generate_question(sql: str) -> str:
+    """
+
+    ## Example
+    ```python
+    vn.generate_question(sql="SELECT * FROM students WHERE name = 'John Doe'")
+    # 'AI Response'
+    ```
+    
+    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 = __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