From dc2161251d76ec82c862136dbbe1997edeb7eefa Mon Sep 17 00:00:00 2001
From: dloman118 <99347459+dloman118@users.noreply.github.com>
Date: Thu, 4 Jan 2024 20:50:31 -0500
Subject: [PATCH] make data_table_output False

---
 .../library/nba/game_recap.yaml                |  1 +
 .../openassistants/contrib/sqlalchemy_query.py | 18 ++++++++++--------
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/examples/fast-api-server/library/nba/game_recap.yaml b/examples/fast-api-server/library/nba/game_recap.yaml
index d29f7ea..614fff6 100644
--- a/examples/fast-api-server/library/nba/game_recap.yaml
+++ b/examples/fast-api-server/library/nba/game_recap.yaml
@@ -15,6 +15,7 @@ parameters:
       - team
 type: DuckDBQueryFunction
 dataset: 'dummy-data/'
+data_table_output: False
 sqls:
   - |
     SELECT date,home_team,away_team,home_team_abbr,away_team_abbr,home_team_winner,away_team_winner,home_team_score,away_team_score,venue,location,home_team_overall_record,away_team_overall_record,
diff --git a/packages/openassistants/openassistants/contrib/sqlalchemy_query.py b/packages/openassistants/openassistants/contrib/sqlalchemy_query.py
index f7221d3..f30adf4 100644
--- a/packages/openassistants/openassistants/contrib/sqlalchemy_query.py
+++ b/packages/openassistants/openassistants/contrib/sqlalchemy_query.py
@@ -1,6 +1,6 @@
 import abc
 import asyncio
-from typing import Annotated, Any, List, Literal, Sequence
+from typing import Annotated, Any, List, Literal, Sequence, Optional
 
 import jsonschema
 import pandas as pd
@@ -77,6 +77,7 @@ class QueryFunction(BaseFunction, abc.ABC):
     sqls: List[str]
     visualizations: List[str]
     summarization: str
+    data_table_output: Optional[bool] = True
     suggested_follow_ups: Annotated[List[SuggestedPrompt], Field(default_factory=list)]
 
     @abc.abstractmethod
@@ -159,14 +160,15 @@ async def execute(
         results: List[FunctionOutput] = []
 
         dataframes = await self._execute_sqls(deps)
-        results.extend(
-            [
-                DataFrameOutput(dataframe=SerializedDataFrame.from_pd(df))
-                for df in dataframes
-            ]
-        )
+        if self.data_table_output:
+            results.extend(
+                [
+                    DataFrameOutput(dataframe=SerializedDataFrame.from_pd(df))
+                    for df in dataframes
+                ]
+            )
 
-        yield results
+            yield results
 
         visualizations = await self._execute_visualizations(dataframes, deps)