add plotting capabilities

app.py

@@ -1,13 +1,16 @@
 from gradio_huggingfacehub_search import HuggingfaceHubSearch
 from huggingface_hub import HfApi
+import matplotlib.pyplot as plt
+from typing import Tuple, Optional
 import pandas as pd
 import gradio as gr
 import duckdb
 import requests
 import llama_cpp
 import instructor
+import enum
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 BASE_DATASETS_SERVER_URL = "https://datasets-server.huggingface.co"
 view_name = "dataset_view"
@@ -21,6 +24,7 @@ llama = llama_cpp.Llama(
     chat_format="chatml",
     n_ctx=2048,
     verbose=False,
+    temperature=0.1,
 )
 
 create = instructor.patch(
@@ -29,8 +33,23 @@ create = instructor.patch(
 )
 
 
+class OutputTypes(str, enum.Enum):
+    TABLE = "table"
+    BARCHART = "barchart"
+    LINECHART = "linechart"
+
+
 class SQLResponse(BaseModel):
     sql: str
+    visualization_type: Optional[OutputTypes] = Field(
+        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"
+    )
+    label_key: Optional[str] = Field(
+        None, description="The column name that contains the labels for chart responses"
+    )
 
 
 def get_dataset_ddl(dataset_id: str) -> str:
@@ -63,7 +82,7 @@ CREATE TABLE {} (
     return sql_ddl
 
 
-def generate_sql(dataset_id: str, query: str) -> str:
+def generate_query(dataset_id: str, query: str) -> str:
     ddl = get_dataset_ddl(dataset_id)
 
     system_prompt = f"""
@@ -76,6 +95,8 @@ def generate_sql(dataset_id: str, query: str) -> str:
     Please assist the user by writing a SQL query that answers the user's question.
     """
 
+    print("Calling LLM with system prompt: ", system_prompt)
+
     resp: SQLResponse = create(
         model="Hermes-2-Pro-Llama-3-8B",
         messages=[
@@ -88,15 +109,28 @@ def generate_sql(dataset_id: str, query: str) -> str:
         response_model=SQLResponse,
     )
 
-    return resp.sql
+    print("Received Response: ", resp)
 
+    return resp
 
-def query_dataset(dataset_id: str, query: str) -> tuple[pd.DataFrame, str]:
-    sql_query = generate_sql(dataset_id, query)
-    df = conn.execute(sql_query).fetchdf()
 
-    markdown_output = f"""```sql\n{sql_query}```"""
-    return df, markdown_output
+def query_dataset(dataset_id: str, query: str) -> Tuple[pd.DataFrame, str, plt.Figure]:
+    response: SQLResponse = generate_query(dataset_id, query)
+    df = conn.execute(response.sql).fetchdf()
+
+    plot = None
+
+    if response.visualization_type == OutputTypes.LINECHART:
+        plot = df.plot(
+            kind="line", x=response.data_key, y=response.label_key
+        ).get_figure()
+    elif response.visualization_type == OutputTypes.BARCHART:
+        plot = df.plot(
+            kind="bar", x=response.data_key, y=response.label_key
+        ).get_figure()
+
+    markdown_output = f"""```sql\n{response.sql}\n```"""
+    return df, markdown_output, plot
 
 
 with gr.Blocks() as demo:
@@ -105,19 +139,20 @@ with gr.Blocks() as demo:
         label="Hub Dataset ID",
         placeholder="Find your favorite dataset...",
         search_type="dataset",
-        value="jamescalam/world-cities-geo",
+        value="teknium/OpenHermes-2.5",
     )
     user_query = gr.Textbox("", label="Ask anything...")
 
     btn = gr.Button("Ask 🪄")
 
-    df = gr.DataFrame()
     sql_query = gr.Markdown(label="Output SQL Query")
+    df = gr.DataFrame()
+    plot = gr.Plot()
 
     btn.click(
         query_dataset,
         inputs=[dataset_id, user_query],
-        outputs=[df, sql_query],
+        outputs=[df, sql_query, plot],
     )