add examples and handle incorrect label / data key output better

README.md

@@ -22,3 +22,5 @@ Powered by [Hermes-2-Pro-Llama-3-8B](https://huggingface.co/NousResearch/Hermes-
 2. Get SQL DDL from Parquet column and data types
 3. Prompt the LLM for SQL
 4. Run the SQL against the Dataset Parquet
+
+Inspired by [Datasets-Text2SQL](https://huggingface.co/spaces/asoria/datasets-text2sql)

app.py

@@ -46,10 +46,12 @@ class SQLResponse(BaseModel):
         None, description="The type of visualization to display"
     )
     data_key: Optional[str] = Field(
-        None, description="The column name that contains the data for chart responses"
+        None,
+        description="The column name from the sql query that contains the data for chart responses",
     )
     label_key: Optional[str] = Field(
-        None, description="The column name that contains the labels for chart responses"
+        None,
+        description="The column name from the sql query that contains the labels for chart responses",
     )
 
 
@@ -95,6 +97,9 @@ def generate_query(dataset_id: str, query: str) -> str:
     ```
     
     Please assist the user by writing a SQL query that answers the user's question.
+    
+    Use Label Key as the column name for the x-axis and Data Key as the column name for the y-axis for chart responses. The 
+    label key and data key must be present in the SQL output.
     """
 
     print("Calling LLM with system prompt: ", system_prompt)
@@ -124,6 +129,12 @@ def query_dataset(dataset_id: str, query: str) -> Tuple[pd.DataFrame, str, plt.F
 
     plot = None
 
+    # handle incorrect data and label keys better
+    if response.label_key and response.label_key not in df.columns:
+        response.label_key = None
+    if response.data_key and response.data_key not in df.columns:
+        response.data_key = None
+
     if response.visualization_type == OutputTypes.LINECHART:
         plot = df.plot(
             kind="line", x=response.label_key, y=response.data_key
@@ -150,6 +161,13 @@ with gr.Blocks() as demo:
         value="gretelai/synthetic_text_to_sql",
     )
     user_query = gr.Textbox("", label="Ask anything...")
+    examples = [
+        ["Show me a preview of the data"],
+        ["Show me something interesting"],
+        ["What is the largest length of sql query context?"],
+        ["show me counts by sql_query_type in a bar chart"],
+    ]
+    gr.Examples(examples=examples, inputs=[user_query], outputs=[])
 
     btn = gr.Button("Ask 🪄")