feat(component): add tool for getting visualization config structure

Author: akshatdubeysf

src/components/db-query/tools/generate-query.tool.ts

@@ -48,7 +48,7 @@ export class GenerateQueryTool implements IGraphTool {
       prompt: z
         .string()
         .describe(
-          `Prompt from the user that will be used for generating the query. This should be exactly as the user asked without any assumptions.`,
+          `Prompt from the user that will be used for generating the query.`,
         ),
     }) as AnyObject[string];
     return graph.asTool({

src/components/visualization/nodes/select-visualization.node.ts

@@ -49,6 +49,7 @@ You need to suggest a visualisation from a list of visualisation that would best
 <instructions>
 The output should be a single string that has the name from the visualizations list and nothing else.
 If the required visualization would not be possible due to the structure of the data, return "none" followed by the changes required in the data to be able to render the visualization.
+Do not try to force fit the data to any visualization if it does not make sense. Prefer to returning none with appropriate reason instead.
 </instructions>
 <output-example-1>
 type-of-visualization

src/components/visualization/tools/generate-visualization.tool.ts

@@ -24,7 +24,11 @@ export class GenerateVisualizationTool implements IGraphTool {
     if (result.error) {
       return `Visualization could not be generated. Reason: ${result.error}`;
     }
-    return `Visualization rendered for the user.`;
+    return `Visualization rendered for the user with the following config: ${JSON.stringify(
+      result.visualizerConfig,
+      undefined,
+      2,
+    )}`;
   }
 
   getMetadata(result: Record<string, string>): AnyObject {

src/components/visualization/tools/get-visualization-context.tool.ts

@@ -0,0 +1,63 @@
+import {Context, inject} from '@loopback/core';
+import {graphTool} from '../../../decorators';
+import {IGraphTool} from '../../../graphs';
+import {VISUALIZATION_KEY} from '../keys';
+import {IVisualizer} from '../types';
+import {StructuredToolInterface, tool} from '@langchain/core/tools';
+import {RunnableToolLike} from '@langchain/core/runnables';
+import z from 'zod';
+import {AnyObject} from '@loopback/repository';
+
+@graphTool()
+export class GetVisualizationContextTool implements IGraphTool {
+  key = 'get-visualization-context';
+  needsReview = false;
+
+  constructor(
+    @inject.context()
+    private readonly context: Context,
+  ) {}
+
+  async build(): Promise<StructuredToolInterface | RunnableToolLike> {
+    const visualizations = await this._getVisualizations();
+    return tool(
+      ({type}: {type: string}) => {
+        const viz = visualizations.find(v => v.name === type);
+        if (!viz) {
+          throw new Error(`Visualization with type ${type} not found`);
+        }
+        return (
+          viz.context ??
+          'No additional context available for this visualization.'
+        );
+      },
+      {
+        name: this.key,
+        description:
+          'Tool to get the context information for a specific visualization type. ' +
+          'Always call this before using generate/improve query tool if a visualization is needed or expected in the future. ' +
+          'Use the response of this tool to inform the query tool about the structure expected for a visualization',
+        schema: z.object({
+          type: z
+            .string()
+            .optional()
+            .describe(
+              `Type of visualization to be generated. It can be one of the following: ${visualizations.map(v => v.name).join(', ')}. If not provided, the system will decide the best visualization based on the data and prompt.`,
+            ),
+        }) as AnyObject[string],
+      },
+    );
+  }
+
+  private async _getVisualizations() {
+    const bindings = this.context.findByTag({
+      [VISUALIZATION_KEY]: true,
+    });
+    if (bindings.length === 0) {
+      throw new Error(`Node with key ${VISUALIZATION_KEY} not found`);
+    }
+    return Promise.all(
+      bindings.map(binding => this.context.get<IVisualizer>(binding.key)),
+    );
+  }
+}

src/components/visualization/tools/index.ts

@@ -1 +1,2 @@
 export * from './generate-visualization.tool';
+export * from './get-visualization-context.tool';

src/components/visualization/types.ts

@@ -4,5 +4,6 @@ import {VisualizationGraphState} from './state';
 export interface IVisualizer {
   name: string;
   description: string;
+  context?: string;
   getConfig(state: VisualizationGraphState): Promise<AnyObject> | AnyObject;
 }

src/components/visualization/visualizer.component.ts

@@ -16,6 +16,7 @@ import {
 } from './nodes';
 import {GenerateVisualizationTool} from './tools/generate-visualization.tool';
 import {PieVisualizer, BarVisualizer, LineVisualizer} from './visualizers';
+import {GetVisualizationContextTool} from './tools';
 
 export class VisualizerComponent implements Component {
   services: ServiceOrProviderClass[] | undefined;
@@ -35,6 +36,7 @@ export class VisualizerComponent implements Component {
       VisualizationGraph,
       // tools
       GenerateVisualizationTool,
+      GetVisualizationContextTool,
       // nodes
       GetDatasetDataNode,
       SelectVisualizationNode,

src/components/visualization/visualizers/bar.visualizer.ts

@@ -33,6 +33,9 @@ You are an expert data visualization assistant. Your task is to create a bar cha
 </user-prompt>
 </inputs>`);
 
+  context?: string | undefined =
+    `A bar chart requires data with at least two columns: one for the categories (x-axis) and one for the values (y-axis). Ensure that the category column contains discrete values representing different groups or categories, while the value column contains numerical data that can be compared across these categories. Bar charts can be oriented either vertically or horizontally depending on the data representation needs.`;
+
   schema = z.object({
     categoryColumn: z
       .string()

src/components/visualization/visualizers/line.visualizer.ts

@@ -12,13 +12,17 @@ import {visualizer} from '../decorators/visualizer.decorator';
 @visualizer()
 export class LineVisualizer implements IVisualizer {
   name = 'line';
+
+  context?: string | undefined =
+    `A line chart requires data with at least two columns: one for the x-axis (typically time or sequential data) and one for the y-axis (values). Optionally, a third list of columns can be used to differentiate multiple series (lines) in the chart, this is usually the field to group the data on. Ensure that the x-axis data is continuous and ordered to accurately represent trends over time.`;
+
   description = `Renders the data in a line chart format. Best for showing trends and changes over time or continuous data.`;
   renderPrompt = PromptTemplate.fromTemplate(`
 <instructions>
 You are an expert data visualization assistant. Your task is to create a line chart config based on the provided SQL query, it's description and user prompt. Follow these steps:
 1. Analyze the SQL query results to understand the data structure.
 2. Identify the x-axis column (typically time or sequential data) and y-axis column (values) for the line chart.
-3. Determine if there are multiple series to be plotted (multiple lines).
+3. Determine if there are multiple series to be plotted (multiple lines) with combination of multiple columns, or single series based on single column.
 4. Create a configuration object for the line chart using the identified columns.
 5. Return the line chart configuration object.
 </instructions>
@@ -38,15 +42,17 @@ You are an expert data visualization assistant. Your task is to create a line ch
     xAxisColumn: z
       .string()
       .describe(
-        'Column to be used for x-axis in the line chart (typically time or sequential data)',
+        'Single column name to be used for x-axis in the line chart (typically time or sequential data)',
       ),
     yAxisColumn: z
       .string()
-      .describe('Column to be used for y-axis values in the line chart'),
-    seriesColumn: z
+      .describe(
+        'Single column name to be used for y-axis values in the line chart',
+      ),
+    seriesColumns: z
       .string()
       .describe(
-        'Optional column to group data into multiple lines/series, leave it as empty string if not needed',
+        'Optional column to group data into multiple lines/series, leave it as empty string if not needed. It can cover multiple columns separated by comma if the query needs to show multiple lines based on multiple columns. The UI supports multiple series in line chart by forming a combined key.',
       ),
   }) as z.AnyZodObject;
 
@@ -73,8 +79,15 @@ You are an expert data visualization assistant. Your task is to create a line ch
       description: state.queryDescription!,
       userPrompt: state.prompt!,
     });
-    if (settings.seriesColumn === '' || settings.seriesColumn === undefined) {
-      settings.seriesColumn = null;
+    if (
+      settings.seriesColumns === '' ||
+      settings.seriesColumns === undefined ||
+      settings.seriesColumns === null
+    ) {
+      settings.seriesColumns = null;
+    } else {
+      settings.seriesColumns =
+        settings.seriesColumns?.split(',').map((s: string) => s.trim()) ?? [];
     }
     return settings;
   }

src/components/visualization/visualizers/pie.visualizer.ts

@@ -33,6 +33,9 @@ You are an expert data visualization assistant. Your task is to create a pie cha
 </user-prompt>
 </inputs>`);
 
+  context?: string | undefined =
+    `A pie chart requires data with at least two columns: one for the labels (categories) and one for the values (numerical data). Ensure that the values are non-negative and represent parts of a whole, as pie charts are used to visualize proportions and percentages among different categories.`;
+
   schema = z.object({
     labelColumn: z
       .string()