4297 added a prototype for AI assisted annotation feature for the for…

…ecast chart (#4374)
DARPA-ASKEM · Aug 9, 2024 · a544c56 · a544c56
1 parent b99c9ba
commit a544c56
Show file tree

Hide file tree

Showing 2 changed files with 334 additions and 42 deletions.
diff --git a/packages/client/hmi-client/src/services/charts.ts b/packages/client/hmi-client/src/services/charts.ts
@@ -1,5 +1,8 @@
+import API from '@/api/api';
+import { b64DecodeUnicode } from '@/utils/binary';
 import { percentile } from '@/utils/math';
 import { isEmpty } from 'lodash';
+import { v4 as uuidv4 } from 'uuid';
 
 const VEGALITE_SCHEMA = 'https://vega.github.io/schema/vega-lite/v5.json';
 
@@ -30,6 +33,13 @@ export interface ForecastChartLayer {
 	groupField?: string;
 }
 
+export interface ForecastChartAnnotation {
+	id: string;
+	layerSpec: any;
+	isLLMGenerated: boolean;
+	metadata: any;
+}
+
 export interface HistogramChartOptions extends BaseChartOptions {
 	maxBins?: number;
 	variables: { field: string; label?: string; width: number; color: string }[];
@@ -42,7 +52,7 @@ export interface ErrorChartOptions extends Omit<BaseChartOptions, 'height' | 'yA
 	variables: { field: string; label?: string }[];
 }
 
-export const createErrorChart = (dataset: Record<string, any>[], options: ErrorChartOptions) => {
+export function createErrorChart(dataset: Record<string, any>[], options: ErrorChartOptions) {
 	const axisColor = '#EEE';
 	const labelColor = '#667085';
 	const labelFontWeight = 'normal';
@@ -204,9 +214,9 @@ export const createErrorChart = (dataset: Record<string, any>[], options: ErrorC
 			]
 		}
 	};
-};
+}
 
-export const createHistogramChart = (dataset: Record<string, any>[], options: HistogramChartOptions) => {
+export function createHistogramChart(dataset: Record<string, any>[], options: HistogramChartOptions) {
 	const maxBins = options.maxBins ?? 10;
 	const axisColor = '#EEE';
 	const labelColor = '#667085';
@@ -290,7 +300,7 @@ export const createHistogramChart = (dataset: Record<string, any>[], options: Hi
 		}
 	};
 	return spec;
-};
+}
 
 /**
  * Generate Vegalite specs for simulation/forecast charts. The chart can contain:
@@ -311,12 +321,12 @@ export const createHistogramChart = (dataset: Record<string, any>[], options: Hi
  *
  * Then we use the new 'var' and 'value' columns to render timeseries
  * */
-export const createForecastChart = (
+export function createForecastChart(
 	samplingLayer: ForecastChartLayer | null,
 	statisticsLayer: ForecastChartLayer | null,
 	groundTruthLayer: ForecastChartLayer | null,
 	options: ForecastChartOptions
-) => {
+) {
 	const axisColor = '#EEE';
 	const labelColor = '#667085';
 	const labelFontWeight = 'normal';
@@ -506,7 +516,222 @@ export const createForecastChart = (
 	}
 
 	return spec;
-};
+}
+
+export function createForecastChartAnnotation(axis: 'x' | 'y', datum: number, label: string) {
+	const layerSpec = {
+		description: `At ${axis} ${datum}, add a label '${label}'.`,
+		encoding: {
+			[axis]: { datum }
+		},
+		layer: [
+			{
+				mark: {
+					type: 'rule',
+					strokeDash: [4, 4]
+				}
+			},
+			{
+				mark: {
+					type: 'text',
+					align: 'left',
+					dx: 5,
+					dy: -5
+				},
+				encoding: {
+					text: { value: label }
+				}
+			}
+		]
+	};
+	const annotation: ForecastChartAnnotation = {
+		id: uuidv4(),
+		layerSpec,
+		isLLMGenerated: false,
+		metadata: { axis, datum, label }
+	};
+	return annotation;
+}
+
+export async function generateForecastChartAnnotation(
+	request: string,
+	timeField: string,
+	variables: string[],
+	axisTitle: { x: string; y: string }
+) {
+	const prompt = `
+	  You are an agent who is an expert in Vega-Lite chart specs. Provide a Vega-Lite layer JSON object for the annotation that can be added to an existing chart spec to satisfy the provided user request.
+
+    - The Vega-Lite schema version you must use is https://vega.github.io/schema/vega-lite/v5.json.
+    - Assume that you don’t know the exact data points from the data.
+    - You must give me the single layer object that renders all the necessary drawing objects, including multiple layers within the top layer object if needed.
+    - When adding a label, also draw a constant line perpendicular to the axis to which you are adding the label.
+
+    Assuming you are adding the annotations to the following chart spec,
+    ---- Example Chart Spec Start -----
+    {
+      "data": {"url": "data/samples.csv"},
+      "transform": [
+        {
+          "fold": ["price", "cost", "tax", "profit"],
+          "as": ["variableField", "valueField"]
+        }
+      ],
+      "layer": [
+        {
+          "mark": "line",
+          "encoding": {
+            "x": {"field": "date", "type": "quantitative", "axis": {"title": "Day"}},
+            "y": {"field": "valueField", "type": "quantitative", "axis": {"title": "Dollars"}}
+          }
+        }
+      ]
+    }
+    ---- Example Chart Spec End -----
+    Here are some example requests and the answers:
+
+    Request:
+    At day 200, add a label 'important'
+    Answer:
+    {
+      "description": "At day 200, add a label 'important'",
+      "layer": [
+        {
+          "mark": {
+            "type": "rule",
+            "strokeDash": [4, 4]
+          },
+					"encoding": {
+						"x": { "datum": 200, "axis": { "title": ""} }
+					}
+        },
+        {
+          "mark": {
+            "type": "text",
+            "align": "left",
+            "dx": 5,
+            "dy": -5
+          },
+          "encoding": {
+						"x": { "datum": 200, "axis": { "title": ""} }
+            "text": {"value": "important"}
+          }
+        }
+      ]
+    }
+
+    Request:
+    Add a label 'expensive' at price 20
+    Answer:
+    {
+      "description": "Add a label 'expensive' at price 20",
+      "layer": [
+        {
+          "mark": {
+            "type": "rule",
+            "strokeDash": [4, 4]
+          },
+          "encoding": {
+            "y": { "datum": 20, "axis": { "title": ""} }
+          }
+        },
+        {
+          "mark": {
+            "type": "text",
+            "align": "left",
+            "dx": 5,
+            "dy": -5
+          },
+          "encoding": {
+            "y": { "datum": 20, "axis": { "title": ""} },
+            "text": {"value": "expensive"}
+          }
+        }
+      ]
+    }
+
+    Request:
+    Add a vertical line for the day where the price exceeds 100.
+    Answer:
+    {
+      "description": "Add a vertical line for the day where the price exceeds 100.",
+      "transform": [
+        {"filter": "datum.valueField > 100"},
+        {"aggregate": [{"op": "min", "field": "date", "as": "min_date"}]}
+      ],
+      "layer": [
+        {
+          "mark": {
+            "type": "rule",
+            "strokeDash": [4, 4]
+          },
+          "encoding": {
+            "x": {"field": "min_date", "type": "quantitative", "axis": { "title": ""}}
+          }
+        },
+        {
+          "mark": {
+            "type": "text",
+            "align": "left",
+            "dx": 5,
+            "dy": -10
+          },
+          "encoding": {
+            "x": {"field": "min_date", "type": "quantitative", "axis": { "title": ""}},
+            "text": {"value": "Price > 100"}
+          }
+        }
+      ]
+    }
+
+    Here is the information of the existing target chart spec where you need to add the annotations:
+    - The existing chart follows a similar pattern as the above Example Chart Spec like:
+        {
+          ...
+          "transform": [
+            {
+              "fold": ${JSON.stringify(variables)},
+              "as": ["variableField', "valueField"]
+            }
+          ],
+          "layer": [
+            {
+              ...
+              "encoding": {
+                "x": {"field": "${timeField}", "type": "quantitative", "axis": {"title": "${axisTitle.x}"}},
+                "y": {"field": "valueField", "type": "quantitative", "axis": {"title": "${axisTitle.y}"}}
+              }
+            }
+            ...
+          ]
+        }
+    - Assume all unknown variables except the time field are for the y-axis and are renamed to the valueField.
+
+     Give me the layer object to be added to the existing chart spec based on the following user request.
+		 Please return only a JSON object as a response. Make sure to return plain JSON object that can be parsed as JSON. Do not include code block.
+
+		Request:
+    ${request}
+    Answer
+    {
+	`;
+	// FIXME: Use dedicated endpoint for annotation generation that's configured with JSON response_format instead of using the summary endpoint which is for text output
+	const { data } = await API.post(`/gollm/generate-summary?mode=SYNC`, prompt, {
+		headers: {
+			'Content-Type': 'application/json'
+		}
+	});
+	const str = b64DecodeUnicode(data.output);
+	const result = JSON.parse(str);
+	const layerSpec = JSON.parse(result.response);
+	const annotation: ForecastChartAnnotation = {
+		id: uuidv4(),
+		layerSpec,
+		isLLMGenerated: true,
+		metadata: { llmRequest: request }
+	};
+	return annotation;
+}
 
 /// /////////////////////////////////////////////////////////////////////////////
 // Optimize charts
@@ -681,7 +906,7 @@ export function createInterventionChartMarkers(data: { name: string; value: numb
 	return [markerSpec, labelSpec];
 }
 
-export const createInterventionChart = (interventionsData: { name: string; value: number; time: number }[]) => {
+export function createInterventionChart(interventionsData: { name: string; value: number; time: number }[]) {
 	const spec: any = {
 		$schema: VEGALITE_SCHEMA,
 		width: 400,
@@ -706,4 +931,4 @@ export const createInterventionChart = (interventionsData: { name: string; value
 		});
 	}
 	return spec;
-};
+}