fix pickle issue by using dict instead of pydantic model

app.py

@@ -86,7 +86,7 @@ CREATE TABLE {} (
 
 
 @spaces.GPU
-def generate_query(dataset_id: str, query: str) -> str:
+def generate_query(dataset_id: str, query: str) -> dict:
     ddl = get_dataset_ddl(dataset_id)
 
     system_prompt = f"""
@@ -118,37 +118,38 @@ def generate_query(dataset_id: str, query: str) -> str:
 
     print("Received Response: ", resp)
 
-    return resp
+    return resp.model_dump()
 
 
 def query_dataset(dataset_id: str, query: str) -> Tuple[pd.DataFrame, str, plt.Figure]:
-    response: SQLResponse = generate_query(dataset_id, query)
+    response = generate_query(dataset_id, query)
 
     print("Querying Parquet...")
-    df = conn.execute(response.sql).fetchdf()
+    df = conn.execute(response.get("sql")).fetchdf()
 
     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
+    label_key = response.get("label_key")
+    data_key = response.get("data_key")
+    viz_type = response.get("visualization_type")
+    sql = response.get("sql")
 
-    if response.visualization_type == OutputTypes.LINECHART:
-        plot = df.plot(
-            kind="line", x=response.label_key, y=response.data_key
-        ).get_figure()
+    # handle incorrect data and label keys
+    if label_key and label_key not in df.columns:
+        label_key = None
+    if data_key and data_key not in df.columns:
+        data_key = None
+
+    if viz_type == OutputTypes.LINECHART:
+        plot = df.plot(kind="line", x=label_key, y=data_key).get_figure()
         plt.xticks(rotation=45, ha="right")
         plt.tight_layout()
-    elif response.visualization_type == OutputTypes.BARCHART:
-        plot = df.plot(
-            kind="bar", x=response.label_key, y=response.data_key
-        ).get_figure()
+    elif viz_type == OutputTypes.BARCHART:
+        plot = df.plot(kind="bar", x=label_key, y=data_key).get_figure()
         plt.xticks(rotation=45, ha="right")
         plt.tight_layout()
 
-    markdown_output = f"""```sql\n{response.sql}\n```"""
+    markdown_output = f"""```sql\n{sql}\n```"""
     return df, markdown_output, plot