feat(query-generation): add better prompts for visualization

Author: akshatdubeysf

src/components/db-query/nodes/generate-description.node.ts

@@ -84,8 +84,8 @@ Return a short bulleted list where each bullet is one condition, filter, or piec
       if (token) {
         output += token;
         config.writer?.({
-          type: LLMStreamEventType.Message,
-          data: {message: token},
+          type: LLMStreamEventType.ToolStatus,
+          data: {thinkingToken: token},
         });
       }
     }

src/components/db-query/tools/get-data-as-dataset.tool.ts

@@ -53,8 +53,9 @@ export class GetDataAsDatasetTool implements IGraphTool {
     }) as AnyObject[string];
     return graph.asTool({
       name: this.key,
-      description: `Query tool for generating SQL queries for a users request. Use it to find data from the database based on the user's request.
-                Note that it does not return the query, instead only a dataset ID that is not relevant to the user. 
+      description: `Query tool for generating SQL queries for a users request. Use it only when the user needs raw tabular data from the database.
+                Do not use this tool if the user's request involves trends, growth, decline, comparisons, distributions, patterns, or any form of analytical insight — use the 'generate-visualization' tool instead.
+                Note that it does not return the query, instead only a dataset ID that is not relevant to the user.
                 It internally fires an event that renders a grid for the dataset on the UI for the user to see.`,
       schema,
     });

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

@@ -51,8 +51,9 @@ export class GenerateVisualizationTool implements IGraphTool {
         ),
       datasetId: z
         .string()
+        .optional()
         .describe(
-          `ID of the dataset that needs to be visualized. Use the dataset ID from 'get-data-as-dataset' or 'improve-dataset' tool`,
+          `ID of the dataset that needs to be visualized. Use the dataset ID from 'get-data-as-dataset' or 'improve-dataset' tool if available. If not provided, the tool will internally fetch the data.`,
         ),
       type: z
         .string()
@@ -63,12 +64,11 @@ export class GenerateVisualizationTool implements IGraphTool {
     }) as AnyObject[string];
     return graph.asTool({
       name: this.key,
-      description: `
-It takes in a prompt and an optional dataset ID that needs to be visualized.
-Must use this over 'get-data-as-dataset' tool if the user request would be better served with a visualization rather than a tabular representation of the data, even if the user does not explicitly ask for it.
-No need to call 'get-data-as-dataset' tool before this, if the dataset ID is not provided, the tool will internally fetch the data to be visualized.
+      description: `Generates a visualization for the user's request. It takes in a prompt and an optional dataset ID.
+If the user's request involves trends, growth, decline, comparisons, distributions, patterns, correlations, or any analytical insight, ALWAYS use this tool instead of 'get-data-as-dataset'.
+No need to call 'get-data-as-dataset' tool before this — if the dataset ID is not provided, this tool will internally fetch the data to be visualized.
 It does not return anything, instead it fires an event internally that renders the visualization on the UI for the user to see.
-It supports followings types of visualizations: ${visualizations.map(v => v.name).join(', ')}.`,
+It supports the following types of visualizations: ${visualizations.map(v => v.name).join(', ')}.`,
       schema,
     });
   }

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

@@ -57,7 +57,7 @@ You are an expert data visualization assistant. Your task is to create a line ch
   }) as z.AnyZodObject;
 
   constructor(
-    @inject(AiIntegrationBindings.CheapLLM)
+    @inject(AiIntegrationBindings.SmartNonThinkingLLM)
     private readonly llm: LLMProvider,
   ) {}