add SQL query using duckdb in Assistant

src/ai-assistant/package.json

@@ -31,12 +31,14 @@
     "umd"
   ],
   "dependencies": {
+    "@duckdb/duckdb-wasm": "^1.29.0",
     "@kepler.gl/components": "3.1.0-alpha.1",
     "@kepler.gl/constants": "3.1.0-alpha.1",
     "@kepler.gl/layers": "3.1.0-alpha.1",
     "@kepler.gl/types": "3.1.0-alpha.1",
     "@kepler.gl/utils": "3.1.0-alpha.1",
     "@openassistant/core": "^0.0.5",
+    "@openassistant/duckdb": "^0.0.5",
     "@openassistant/echarts": "^0.0.5",
     "@openassistant/geoda": "^0.0.5",
     "@openassistant/ui": "^0.0.5",

src/ai-assistant/src/components/ai-assistant-component.tsx

@@ -11,6 +11,7 @@ import {
 } from '@openassistant/geoda';
 import {histogramFunctionDefinition, scatterplotFunctionDefinition} from '@openassistant/echarts';
 import {AiAssistant} from '@openassistant/ui';
+import {queryDuckDBFunctionDefinition} from '@openassistant/duckdb';
 import '@openassistant/echarts/dist/index.css';
 import '@openassistant/ui/dist/index.css';
 
@@ -32,7 +33,6 @@ import {
   PROMPT_IDEAS,
   WELCOME_MESSAGE
 } from '../constants';
-import {filterFunctionDefinition} from '../tools/filter-function';
 import {addLayerFunctionDefinition} from '../tools/layer-creation-function';
 import {updateLayerColorFunctionDefinition} from '../tools/layer-style-function';
 import {SelectedKeplerGlActions} from './ai-assistant-manager';
@@ -42,6 +42,7 @@ import {
   getScatterplotValuesFromDataset,
   getValuesFromDataset,
   highlightRows,
+  highlightRowsByColumnValues,
   saveAsDataset
 } from '../tools/utils';
 
@@ -93,6 +94,20 @@ function AiAssistantComponentFactory() {
         keplerGlActions.layerSetIsValid
       );
 
+    const highlightRowsByColumnValuesOnSelected = (
+      datasetName: string,
+      columnName: string,
+      selectedValues: unknown[]
+    ) =>
+      highlightRowsByColumnValues(
+        visState.datasets,
+        visState.layers,
+        datasetName,
+        columnName,
+        selectedValues,
+        keplerGlActions.layerSetIsValid
+      );
+
     // define LLM functions
     const functions = [
       basemapFunctionDefinition({mapStyleChange: keplerGlActions.mapStyleChange, mapStyle}),
@@ -109,12 +124,11 @@ function AiAssistantComponentFactory() {
         layerVisualChannelConfigChange: keplerGlActions.layerVisualChannelConfigChange,
         layers: visState.layers
       }),
-      filterFunctionDefinition({
-        datasets: visState.datasets,
-        filters: visState.filters,
-        createOrUpdateFilter: keplerGlActions.createOrUpdateFilter,
-        setFilter: keplerGlActions.setFilter,
-        setFilterPlot: keplerGlActions.setFilterPlot
+      queryDuckDBFunctionDefinition({
+        // duckDB: kepler.gl duckdb instance
+        getValues: (datasetName: string, variableName: string): number[] =>
+          getValuesFromDataset(visState.datasets, datasetName, variableName),
+        onSelected: highlightRowsByColumnValuesOnSelected
       }),
       histogramFunctionDefinition({
         getValues: getValuesCallback,

src/ai-assistant/src/constants.ts

@@ -2,11 +2,11 @@
 // Copyright contributors to the kepler.gl project
 
 export const TASK_LIST =
-  '1. Show dataset/layer/variable info.\n2. Change the basemap style.\n3. Load data from url.\n4. Create a map layer using variable.\n5. Filter the data of a variable.\n6. Create a histogram.\n7. Classify the data of a variable.\n8. Spatial join two datasets.\n9. Query the data using SQL.';
+  '1. Show dataset/layer/variable info.\n2. Change the basemap style.\n3. Load data from url.\n4. Create a map layer using variable.\n5. Create a histogram.\n6. Classify the data of a variable.\n7. Spatial join two datasets.\n8. Query the data using SQL.';
 
 export const WELCOME_MESSAGE = `Hi, I am Kepler.gl AI Assistant!\nHere are some tasks I can help you with:\n\n${TASK_LIST}`;
 
-export const INSTRUCTIONS = `You are a Kepler.gl AI Assistant that can answer questions and help with tasks of mapping and spatial data analysis.
+export const INSTRUCTIONS = `You are a Kepler.gl AI Assistant that can answer questions and help with tasks of mapping, spatial data analysis, and SQL query.
 
 When responding to user queries:
 1. Analyze if the task requires one or multiple function calls
@@ -17,6 +17,9 @@ When responding to user queries:
    - Please ask the user to confirm the parameters
    - If the user doesn't agree, try to provide variable functions to the user
    - Execute functions in a sequential order
+3. For SQL query, please help to generate select query clause using the content of the dataset:
+   - please only use the columns that are in the dataset context
+   - please try to use the aggregate functions if possible
 
 You can execute multiple functions to complete complex tasks, but execute them one at a time in a logical sequence. Always validate the success of each function call before proceeding to the next one.
 

src/ai-assistant/src/tools/utils.ts

@@ -292,3 +292,29 @@ function determineFieldType(value: unknown): keyof typeof ALL_FIELD_TYPES {
       : ALL_FIELD_TYPES.real
     : ALL_FIELD_TYPES.string;
 }
+
+export function highlightRowsByColumnValues(
+  datasets: Datasets,
+  layers: Layer[],
+  datasetName: string,
+  columnName: string,
+  selectedValues: unknown[],
+  layerSetIsValid: (layer: Layer, isValid: boolean) => void
+) {
+  const datasetId = Object.keys(datasets).find(dataId => datasets[dataId].label === datasetName);
+  if (!datasetId) return;
+  const dataset = datasets[datasetId];
+  if (dataset) {
+    // get the values of the column
+    const values = Array.from({length: dataset.length}, (_, i) => dataset.getValue(columnName, i));
+    // create a dict using the values
+    const valueDict = values.reduce((acc, value, index) => {
+      acc[value] = index;
+      return acc;
+    }, {});
+    // @ts-expect-error need to fix the type error of value here
+    const selectedIndices = selectedValues.map(value => valueDict[value]);
+    // highlight the rows
+    highlightRows(datasets, layers, datasetName, selectedIndices, layerSetIsValid);
+  }
+}