add spatial join with openassistant 0.0.4

Author: lixun910

examples/demo-app/package.json

@@ -25,7 +25,7 @@
     "@loaders.gl/csv": "^4.3.2",
     "@loaders.gl/json": "^4.3.2",
     "@loaders.gl/parquet": "^4.3.2",
-    "@openassistant/ui": "^0.0.1",
+    "@openassistant/ui": "^0.0.4",
     "@types/classnames": "^2.3.1",
     "@types/keymirror": "^0.1.1",
     "classnames": "^2.2.1",

src/ai-assistant/package.json

@@ -36,10 +36,10 @@
     "@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.3",
-    "@openassistant/echarts": "^0.0.3",
-    "@openassistant/geoda": "^0.0.3",
-    "@openassistant/ui": "^0.0.3",
+    "@openassistant/core": "^0.0.4",
+    "@openassistant/echarts": "^0.0.4",
+    "@openassistant/geoda": "^0.0.4",
+    "@openassistant/ui": "^0.0.4",
     "color-interpolate": "^1.0.5"
   },
   "devDependencies": {

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

@@ -4,7 +4,11 @@
 import React, {useEffect} from 'react';
 import styled, {withTheme} from 'styled-components';
 import {MessageModel, useAssistant} from '@openassistant/core';
-import {dataClassifyFunctionDefinition} from '@openassistant/geoda';
+import {
+  dataClassifyFunctionDefinition,
+  spatialCountFunctionDefinition,
+  SpatialJoinGeometries
+} from '@openassistant/geoda';
 import {histogramFunctionDefinition, scatterplotFunctionDefinition} from '@openassistant/echarts';
 import {AiAssistant} from '@openassistant/ui';
 import '@openassistant/echarts/dist/index.css';
@@ -34,9 +38,11 @@ import {updateLayerColorFunctionDefinition} from '../tools/layer-style-function'
 import {SelectedKeplerGlActions} from './ai-assistant-manager';
 import {
   getDatasetContext,
+  getGeometriesFromDataset,
   getScatterplotValuesFromDataset,
   getValuesFromDataset,
-  highlightRows
+  highlightRows,
+  saveAsDataset
 } from '../tools/utils';
 
 export type AiAssistantComponentProps = {
@@ -72,6 +78,21 @@ function AiAssistantComponentFactory() {
     mapStyle,
     visState
   }: AiAssistantComponentProps) => {
+    // get values from dataset, used by LLM functions
+    const getValuesCallback = (datasetName: string, variableName: string): number[] =>
+      getValuesFromDataset(visState.datasets, datasetName, variableName);
+
+    // highlight rows, used by LLM functions
+    const highlightRowsCallback = (datasetName: string, selectedRowIndices: number[]) =>
+      highlightRows(
+        visState.datasets,
+        visState.layers,
+        datasetName,
+        selectedRowIndices,
+        keplerGlActions.layerSetIsValid
+      );
+
+    // define LLM functions
     const functions = [
       basemapFunctionDefinition({mapStyleChange: keplerGlActions.mapStyleChange, mapStyle}),
       loadUrlFunctionDefinition({
@@ -95,49 +116,36 @@ function AiAssistantComponentFactory() {
         setFilterPlot: keplerGlActions.setFilterPlot
       }),
       histogramFunctionDefinition({
-        getValues: (datasetName: string, variableName: string): number[] =>
-          getValuesFromDataset(visState.datasets, datasetName, variableName),
-        onSelected: (datasetName: string, selectedRowIndices: number[]) =>
-          highlightRows(
-            visState.datasets,
-            visState.layers,
-            datasetName,
-            selectedRowIndices,
-            keplerGlActions.layerSetIsValid
-          )
+        getValues: getValuesCallback,
+        onSelected: highlightRowsCallback
       }),
       scatterplotFunctionDefinition({
-        getValues: (
-          datasetName: string,
-          xVariableName: string,
-          yVariableName: string
-        ): Promise<{x: number[]; y: number[]}> =>
-          Promise.resolve(
-            getScatterplotValuesFromDataset(
-              visState.datasets,
-              datasetName,
-              xVariableName,
-              yVariableName
-            )
-          ),
-        onSelected: (datasetName: string, selectedRowIndices: number[]) =>
-          highlightRows(
-            visState.datasets,
-            visState.layers,
-            datasetName,
-            selectedRowIndices,
-            keplerGlActions.layerSetIsValid
-          )
+        getValues: async (datasetName: string, xVar: string, yVar: string) =>
+          getScatterplotValuesFromDataset(visState.datasets, datasetName, xVar, yVar),
+        onSelected: highlightRowsCallback
       }),
       dataClassifyFunctionDefinition({
-        getValues: (datasetName: string, variableName: string): number[] =>
-          getValuesFromDataset(visState.datasets, datasetName, variableName)
+        getValues: getValuesCallback
+      }),
+      spatialCountFunctionDefinition({
+        getValues: getValuesCallback,
+        getGeometries: (datasetName: string): SpatialJoinGeometries =>
+          getGeometriesFromDataset(
+            visState.datasets,
+            visState.layers,
+            visState.layerData,
+            datasetName
+          ),
+        saveAsDataset: (datasetName: string, data: Record<string, number[]>) =>
+          saveAsDataset(visState.datasets, datasetName, data, keplerGlActions.addDataToMap)
       })
     ];
 
+    // enable voice and screen capture
     const enableVoiceAndScreenCapture =
       aiAssistant.config.provider === 'openai' || aiAssistant.config.provider === 'google';
 
+    // define assistant props
     const assistantProps = {
       name: ASSISTANT_NAME,
       description: ASSISTANT_DESCRIPTION,
@@ -151,12 +159,14 @@ function AiAssistantComponentFactory() {
 
     const {initializeAssistant, addAdditionalContext} = useAssistant(assistantProps);
 
+    // initialize assistant with context
     const initializeAssistantWithContext = async () => {
       await initializeAssistant();
       const context = getDatasetContext(visState.datasets, visState.layers);
       addAdditionalContext({context});
     };
 
+    // initialize assistant with context
     useEffect(() => {
       initializeAssistantWithContext();
       // re-initialize assistant when datasets, filters or layers change

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

@@ -41,6 +41,7 @@ export type SelectedKeplerGlActions = {
   setFilter: ActionHandler<typeof setFilter>;
   setFilterPlot: ActionHandler<typeof setFilterPlot>;
   layerSetIsValid: ActionHandler<typeof layerSetIsValid>;
+  addTableColumn: ActionHandler<typeof addTableColumn>;
 };
 
 export type AiAssistantManagerState = {

src/ai-assistant/src/constants.ts

@@ -2,7 +2,7 @@
 // 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.';
+  '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.';
 
 export const WELCOME_MESSAGE = `Hi, I am Kepler.gl AI Assistant!\nHere are some tasks I can help you with:\n\n${TASK_LIST}`;
 
@@ -14,6 +14,8 @@ When responding to user queries:
    - Identify the appropriate function to call
    - Determine all required parameters
    - If parameters are missing, ask the user to provide them
+   - 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
 
 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

@@ -5,7 +5,17 @@ import interpolate from 'color-interpolate';
 
 import {Layer} from '@kepler.gl/layers';
 import {Datasets, KeplerTable} from '@kepler.gl/table';
+import {SpatialJoinGeometries} from '@openassistant/geoda';
+import {ALL_FIELD_TYPES} from '@kepler.gl/constants';
+import {AddDataToMapPayload, ProtoDataset, ProtoDatasetField} from '@kepler.gl/types';
 
+/**
+ * Check if the dataset exists
+ * @param datasets The kepler.gl datasets
+ * @param datasetName The name of the dataset
+ * @param functionName The name of the function
+ * @returns The result of the check
+ */
 export function checkDatasetNotExists(
   datasets: Datasets,
   datasetName: string,
@@ -26,6 +36,13 @@ export function checkDatasetNotExists(
   return null;
 }
 
+/**
+ * Check if the field exists
+ * @param dataset The kepler.gl dataset
+ * @param fieldName The name of the field
+ * @param functionName The name of the function
+ * @returns The result of the check
+ */
 export function checkFieldNotExists(dataset: KeplerTable, fieldName: string, functionName: string) {
   const field = dataset.fields.find(f => f.name === fieldName);
   if (!field) {
@@ -43,6 +60,12 @@ export function checkFieldNotExists(dataset: KeplerTable, fieldName: string, fun
   return null;
 }
 
+/**
+ * Interpolate the colors from the original colors with the given number of colors
+ * @param originalColors The original colors
+ * @param numberOfColors The number of colors
+ * @returns The interpolated colors
+ */
 export function interpolateColor(originalColors: string[], numberOfColors: number) {
   if (originalColors.length === numberOfColors) {
     return originalColors;
@@ -98,6 +121,14 @@ export function getScatterplotValuesFromDataset(
   return {x: xValues, y: yValues};
 }
 
+/**
+ * Highlight the rows in a dataset
+ * @param datasets The kepler.gl datasets
+ * @param layers The kepler.gl layers
+ * @param datasetName The name of the dataset
+ * @param selectedRowIndices The indices of the rows to highlight
+ * @param layerSetIsValid The function to set the layer validity
+ */
 export function highlightRows(
   datasets: Datasets,
   layers: Layer[],
@@ -117,11 +148,17 @@ export function highlightRows(
     selectLayers.forEach(layer => {
       layer.formatLayerData(datasets);
       // trigger a re-render using layerSetIsValid() to update the top layer
-      layerSetIsValid(selectLayers[0], true);
+      layerSetIsValid(layer, true);
     });
   }
 }
 
+/**
+ * Get the dataset context, which is used to provide the dataset information to the AI assistant
+ * @param datasets The kepler.gl datasets
+ * @param layers The kepler.gl layers
+ * @returns The dataset context
+ */
 export function getDatasetContext(datasets: Datasets, layers: Layer[]) {
   const context = 'Please remember the following dataset context:';
   const dataMeta = Object.values(datasets).map(dataset => ({
@@ -130,7 +167,128 @@ export function getDatasetContext(datasets: Datasets, layers: Layer[]) {
     fields: dataset.fields.map(field => ({[field.name]: field.type})),
     layers: layers
       .filter(layer => layer.config.dataId === dataset.id)
-      .map(layer => ({id: layer.id, label: layer.config.label, type: layer.type}))
+      .map(layer => ({
+        id: layer.id,
+        label: layer.config.label,
+        type: layer.type,
+        geometryMode: layer.config.columnMode,
+        // get the valid geometry columns as string
+        geometryColumns: Object.fromEntries(
+          Object.entries(layer.config.columns)
+            .filter(([_, value]) => value !== null)
+            .map(([key, value]) => [
+              key,
+              typeof value === 'object' && value !== null
+                ? Object.fromEntries(Object.entries(value).filter(([_, v]) => v !== null))
+                : value
+            ])
+        )
+      }))
   }));
   return `${context}\n${JSON.stringify(dataMeta)}`;
 }
+
+/**
+ * Get the geometries from a dataset
+ * @param datasets The kepler.gl datasets
+ * @param layers The kepler.gl layers
+ * @param layerData The layer data
+ * @param datasetName The name of the dataset
+ * @returns The geometries
+ */
+export function getGeometriesFromDataset(
+  datasets: Datasets,
+  layers: Layer[],
+  layerData: any[],
+  datasetName: string
+): SpatialJoinGeometries {
+  const datasetId = Object.keys(datasets).find(dataId => datasets[dataId].label === datasetName);
+  if (!datasetId) return [];
+  const dataset = datasets[datasetId];
+
+  // get the index of the layer
+  const layerIndex = layers.findIndex(layer => layer.config.dataId === dataset.id);
+  if (layerIndex === -1) return [];
+
+  const geometries = layerData[layerIndex];
+
+  return geometries?.data;
+}
+
+/**
+ * Save the data as a new dataset by joining it with the left dataset
+ * @param datasets The kepler.gl datasets
+ * @param datasetName The name of the left dataset
+ * @param data The data to save
+ * @param addDataToMap The function to add the data to the map
+ */
+export function saveAsDataset(
+  datasets: Datasets,
+  datasetName: string,
+  data: Record<string, number[]>,
+  addDataToMap: (data: AddDataToMapPayload) => void
+) {
+  // find datasetId from datasets
+  const datasetId = Object.keys(datasets).find(dataId => datasets[dataId].label === datasetName);
+  if (!datasetId) return;
+
+  const leftDataset = datasets[datasetId];
+  const numRows = leftDataset.length;
+
+  const fields: ProtoDatasetField[] = [
+    // New fields from data
+    ...Object.keys(data).map((fieldName, index) => ({
+      name: fieldName,
+      id: `${fieldName}_${index}`,
+      displayName: fieldName,
+      type: determineFieldType(data[fieldName][0])
+    })),
+    // Existing fields from leftDataset
+    ...leftDataset.fields.map((field, index) => ({
+      name: field.name,
+      id: field.id || `${field.name}_${index}`,
+      displayName: field.displayName,
+      type: field.type
+    }))
+  ];
+
+  // Pre-calculate data values array
+  const dataValues = Object.values(data);
+
+  const rows = Array(numRows)
+    .fill(null)
+    .map((_, rowIdx) => [
+      // New data values
+      ...dataValues.map(col => col[rowIdx]),
+      // Existing dataset values
+      ...leftDataset.fields.map(field => leftDataset.getValue(field.name, rowIdx))
+    ]);
+
+  // create new dataset
+  const newDatasetName = `${datasetName}_joined`;
+  const newDataset: ProtoDataset = {
+    info: {
+      id: newDatasetName,
+      label: newDatasetName
+    },
+    data: {
+      fields,
+      rows
+    }
+  };
+
+  addDataToMap({datasets: [newDataset], options: {autoCreateLayers: true, centerMap: true}});
+}
+
+/**
+ * Helper function to determine field type
+ * @param value The value to determine the field type
+ * @returns The field type
+ */
+function determineFieldType(value: unknown): keyof typeof ALL_FIELD_TYPES {
+  return typeof value === 'number'
+    ? Number.isInteger(value)
+      ? ALL_FIELD_TYPES.integer
+      : ALL_FIELD_TYPES.real
+    : ALL_FIELD_TYPES.string;
+}