diff --git a/core/src/sqlite_workers/sqlite_database.rs b/core/src/sqlite_workers/sqlite_database.rs index 7ae5397e9965..1046fea19789 100644 --- a/core/src/sqlite_workers/sqlite_database.rs +++ b/core/src/sqlite_workers/sqlite_database.rs @@ -54,7 +54,7 @@ impl SqliteDatabase { let time_query_start = utils::now(); // Execute the query and collect results - let mut stmt = conn.prepare(&query).unwrap(); + let mut stmt = conn.prepare(&query)?; let column_names = stmt .column_names() .into_iter() diff --git a/front/components/assistant/conversation/TablesQueryAction.tsx b/front/components/assistant/conversation/TablesQueryAction.tsx index 7722e2cd84e5..a1127350aac1 100644 --- a/front/components/assistant/conversation/TablesQueryAction.tsx +++ b/front/components/assistant/conversation/TablesQueryAction.tsx @@ -1,4 +1,5 @@ import { + Button, ChevronDownIcon, ChevronRightIcon, Chip, @@ -6,11 +7,15 @@ import { Spinner, Tooltip, } from "@dust-tt/sparkle"; +import { CloudArrowDownIcon } from "@dust-tt/sparkle"; import type { TablesQueryActionType } from "@dust-tt/types"; +import { stringify } from "csv-stringify"; import dynamic from "next/dynamic"; -import { useState } from "react"; +import { useContext, useState } from "react"; import { amber, emerald, slate } from "tailwindcss/colors"; +import { SendNotificationsContext } from "@app/components/sparkle/Notification"; + const SyntaxHighlighter = dynamic( () => import("react-syntax-highlighter").then((mod) => mod.Light), { ssr: false } @@ -21,6 +26,7 @@ export default function TablesQueryAction({ }: { tablesQueryAction: TablesQueryActionType; }) { + const sendNotification = useContext(SendNotificationsContext); const [isOutputExpanded, setIsOutputExpanded] = useState(false); // Extracting question from the params @@ -42,6 +48,39 @@ export default function TablesQueryAction({ return t.length > maxLength ? t.substring(0, maxLength) + "..." : t; }; + const handleDownload = () => { + const results = + output && + "results" in output && + Array.isArray(output?.results) && + output?.results; + + if (!results) { + return; + } + + const queryTitle = + (output && "query_title" in output && output?.query_title) ?? + "query_results"; + + stringify(results, { header: true }, (err, output) => { + if (err) { + sendNotification({ + title: "Error Downloading CSV", + type: "error", + description: `An error occurred while downloading the CSV: ${err}`, + }); + return; + } + const blob = new Blob([output], { type: "text/csv" }); + const url = URL.createObjectURL(blob); + const a = document.createElement("a"); + a.href = url; + a.download = `${queryTitle}.csv`; + a.click(); + }); + }; + return ( <> {question && ( @@ -52,6 +91,20 @@ export default function TablesQueryAction({ + {isOutputStepCompleted && + "results" in output && + Array.isArray(output?.results) && + output?.results?.length > 0 && ( +
+
+ )} )} diff --git a/front/package-lock.json b/front/package-lock.json index c1329ce1207e..6969fb375d12 100644 --- a/front/package-lock.json +++ b/front/package-lock.json @@ -32,6 +32,7 @@ "ajv": "^8.12.0", "blake3": "^2.1.7", "csv-parse": "^5.5.2", + "csv-stringify": "^6.4.5", "dd-trace": "^3.16.0", "emoji-mart": "^5.5.2", "eventsource-parser": "^1.0.0", @@ -5727,6 +5728,11 @@ "resolved": "https://registry.npmjs.org/csv-parse/-/csv-parse-5.5.2.tgz", "integrity": "sha512-YRVtvdtUNXZCMyK5zd5Wty1W6dNTpGKdqQd4EQ8tl/c6KW1aMBB1Kg1ppky5FONKmEqGJ/8WjLlTNLPne4ioVA==" }, + "node_modules/csv-stringify": { + "version": "6.4.5", + "resolved": "https://registry.npmjs.org/csv-stringify/-/csv-stringify-6.4.5.tgz", + "integrity": "sha512-SPu1Vnh8U5EnzpNOi1NDBL5jU5Rx7DVHr15DNg9LXDTAbQlAVAmEbVt16wZvEW9Fu9Qt4Ji8kmeCJ2B1+4rFTQ==" + }, "node_modules/damerau-levenshtein": { "version": "1.0.8", "resolved": "https://registry.npmjs.org/damerau-levenshtein/-/damerau-levenshtein-1.0.8.tgz", diff --git a/front/package.json b/front/package.json index b15dfcb5d76b..0046da49401f 100644 --- a/front/package.json +++ b/front/package.json @@ -40,6 +40,7 @@ "ajv": "^8.12.0", "blake3": "^2.1.7", "csv-parse": "^5.5.2", + "csv-stringify": "^6.4.5", "dd-trace": "^3.16.0", "emoji-mart": "^5.5.2", "eventsource-parser": "^1.0.0", diff --git a/types/src/front/lib/actions/registry.ts b/types/src/front/lib/actions/registry.ts index 6f66aa8cc96d..ee9c965db387 100644 --- a/types/src/front/lib/actions/registry.ts +++ b/types/src/front/lib/actions/registry.ts @@ -138,7 +138,7 @@ export const DustProdActionRegistry = createActionRegistry({ workspaceId: PRODUCTION_DUST_APPS_WORKSPACE_ID, appId: "b4f205e453", appHash: - "5f814f091a270cc56e2120a21dbf5fba7b64360868cc2dea1da53af3a2b9cc3b", + "3a4dc976412180b7aa596de606bc1b3d7c1642c91775e56db14d6ec1c2aa27bd", }, config: { MODEL: {