feat(visualization): better visualization graph

Author: akshatdubeysf

src/__tests__/db-query/unit/connectors/pg.connector.unit.ts

@@ -25,7 +25,7 @@ describe(`PgConnector Unit`, () => {
     expect(ddl).to.be.equal(`-- Model representing an employee in the system.
 CREATE TABLE public.employees (
  -- Unique identifier for the employee record
-  id UUID NOT NULL,
+  id UUID,
  -- Name of the employee
   name TEXT NOT NULL,
  -- Unique code for the employee, used for identification

src/__tests__/db-query/unit/nodes/semantic-validator.node.unit.ts

@@ -72,6 +72,7 @@ describe('SemanticValidatorNode Unit', function () {
       fromCache: false,
       resultArray: undefined,
       description: undefined,
+      directCall: false,
     };
     llmStub.resolves({
       content: 'valid',
@@ -110,6 +111,7 @@ describe('SemanticValidatorNode Unit', function () {
       sampleSql: '',
       fromCache: false,
       resultArray: undefined,
+      directCall: false,
       description: undefined,
     };
     llmStub.resolves({
@@ -193,6 +195,7 @@ valid
       sampleSql: '',
       fromCache: false,
       description: undefined,
+      directCall: false,
       resultArray: undefined,
     };
     llmStub.resolves({

src/__tests__/db-query/unit/nodes/sql-generation.node.unit.ts

@@ -87,6 +87,7 @@ describe('SqlGenerationNode Unit', function () {
       datasetId: undefined,
       fromCache: false,
       resultArray: undefined,
+      directCall: false,
       description: undefined,
     };
 
@@ -180,6 +181,7 @@ Do not use any DB concepts like enum numbers, joins, CTEs, subqueries etc. in th
       datasetId: undefined,
       fromCache: false,
       resultArray: undefined,
+      directCall: false,
       description: undefined,
     };
 
@@ -291,6 +293,7 @@ Do not use any DB concepts like enum numbers, joins, CTEs, subqueries etc. in th
       datasetId: undefined,
       fromCache: false,
       resultArray: undefined,
+      directCall: false,
       description: undefined,
     };
 
@@ -402,6 +405,7 @@ Do not use any DB concepts like enum numbers, joins, CTEs, subqueries etc. in th
       datasetId: undefined,
       fromCache: true,
       resultArray: undefined,
+      directCall: false,
       description: undefined,
     };
 
@@ -454,6 +458,7 @@ Do not use any DB concepts like enum numbers, joins, CTEs, subqueries etc. in th
       datasetId: undefined,
       fromCache: false,
       resultArray: undefined,
+      directCall: false,
       description: undefined,
     };
 

src/components/db-query/connectors/pg/pg.connector.ts

@@ -85,7 +85,7 @@ export class PgConnector implements IDbConnector {
       const columnDefinitions = Object.keys(columns).map(columnName => {
         const column = columns[columnName];
         const dataType = this.mapColumnToDbType(columnName, column);
-        const notNull = column.required || column.id ? ' NOT NULL' : '';
+        const notNull = column.required ? ' NOT NULL' : '';
         let descriptionString = '';
         if (column.description) {
           descriptionString = ` -- ${column.description.replace(/'/g, "''")}\n`;

src/components/db-query/nodes/check-cache.node.ts

@@ -148,15 +148,18 @@ If no queries are relevant, return '${CacheResults.NotRelevant}' and nothing els
         return state;
       }
       const datasetId = relevantDocs[indexNum].metadata.datasetId;
-      config.writer?.({
-        type: LLMStreamEventType.ToolStatus,
-        data: {
-          status: ToolStatus.Completed,
+      if (!state.directCall) {
+        config.writer?.({
+          type: LLMStreamEventType.ToolStatus,
           data: {
-            datasetId,
+            status: ToolStatus.Completed,
+            data: {
+              datasetId,
+            },
           },
-        },
-      });
+        });
+      }
+
       return {
         ...state,
         fromCache: true,

src/components/db-query/nodes/save-dataset-node.ts

@@ -90,15 +90,17 @@ export class SaveDataSetNode implements IGraphNode<DbQueryState> {
       votes: 0,
     });
 
-    config.writer?.({
-      type: LLMStreamEventType.ToolStatus,
-      data: {
-        status: ToolStatus.Completed,
+    if (!state.directCall) {
+      config.writer?.({
+        type: LLMStreamEventType.ToolStatus,
         data: {
-          datasetId: dataset.id,
+          status: ToolStatus.Completed,
+          data: {
+            datasetId: dataset.id,
+          },
         },
-      },
-    });
+      });
+    }
 
     let result: undefined | AnyObject[] = undefined;
     if (this.config.readAccessForAI && dataset.id) {

src/components/db-query/state.ts

@@ -17,6 +17,7 @@ export const DbQueryGraphStateAnnotation = Annotation.Root({
   done: Annotation<boolean | undefined>,
   resultArray: Annotation<AnyObject[string][] | undefined>,
   description: Annotation<string | undefined>,
+  directCall: Annotation<boolean | undefined>,
 });
 
 export type DbQueryState = typeof DbQueryGraphStateAnnotation.State;

src/components/db-query/testing/get-table.node.builder.ts

@@ -24,6 +24,7 @@ export function getTableNodeTests(cases: GetTableNodeTestCase[]) {
           fromCache: undefined,
           done: false,
           resultArray: undefined,
+          directCall: false,
           description: undefined,
         },
         {

src/components/visualization/nodes.enum.ts

@@ -2,4 +2,5 @@ export enum VisualizationGraphNodes {
   GetDatasetData = 'get_dataset_data',
   SelectVisualisation = 'select_visualization',
   RenderVisualization = 'render_visualization',
+  CallQueryGeneration = 'call_query_generation',
 }

src/components/visualization/nodes/call-query-generation.node.ts

@@ -0,0 +1,56 @@
+import {service} from '@loopback/core';
+import {graphNode} from '../../../decorators';
+import {IGraphNode, LLMStreamEventType, RunnableConfig} from '../../../graphs';
+import {DbQueryGraph, POST_DATASET_TAG} from '../../db-query';
+import {VisualizationGraphNodes} from '../nodes.enum';
+import {VisualizationGraphState} from '../state';
+
+@graphNode(VisualizationGraphNodes.CallQueryGeneration, {
+  [POST_DATASET_TAG]: true,
+})
+export class CallQueryGenerationNode
+  implements IGraphNode<VisualizationGraphState>
+{
+  constructor(
+    @service(DbQueryGraph)
+    private readonly queryPipeline: DbQueryGraph,
+  ) {}
+  async execute(
+    state: VisualizationGraphState,
+    config: RunnableConfig,
+  ): Promise<VisualizationGraphState> {
+    if (state.datasetId) {
+      return state;
+    }
+
+    const queryGraph = await this.queryPipeline.build();
+
+    const result = await queryGraph.invoke(
+      {
+        datasetId: state.datasetId,
+        directCall: true,
+        prompt: `Generate a query to fetch data for visualization based on the following user prompt: ${state.prompt}.${state.visualizer?.context ? ` Ensure that the query structure satisfies the following context: ${state.visualizer.context}` : ''}`,
+      },
+      config,
+    );
+
+    if (!result.datasetId) {
+      config.writer?.({
+        type: LLMStreamEventType.Error,
+        data: {
+          status: `Failed to create dataset for visualization: ${result.replyToUser ?? 'Unknown error'}`,
+        },
+      });
+      return {
+        ...state,
+        error:
+          result.replyToUser ?? 'Failed to create dataset for visualization',
+      };
+    }
+
+    return {
+      ...state,
+      datasetId: result.datasetId,
+    };
+  }
+}