feat(query-generation): reduce cognitive complexity

Author: akshatdubeysf

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

@@ -1,5 +1,6 @@
 import {expect, sinon} from '@loopback/testlab';
 import {EvaluationResult, SemanticValidatorNode} from '../../../../components';
+import {DbSchemaHelperService} from '../../../../components/db-query/services';
 import {LLMProvider} from '../../../../types';
 
 describe('SemanticValidatorNode Unit', function () {
@@ -9,8 +10,11 @@ describe('SemanticValidatorNode Unit', function () {
   beforeEach(() => {
     llmStub = sinon.stub();
     const llm = llmStub as unknown as LLMProvider;
+    const schemaHelper = {
+      asString: sinon.stub().returns(''),
+    } as unknown as DbSchemaHelperService;
 
-    node = new SemanticValidatorNode(llm, llm, {models: []});
+    node = new SemanticValidatorNode(llm, llm, {models: []}, schemaHelper);
   });
 
   afterEach(() => {
@@ -84,11 +88,12 @@ describe('SemanticValidatorNode Unit', function () {
     sinon.assert.calledOnce(llmStub);
 
     const prompt = llmStub.firstCall.args[0];
-    // Verify the prompt contains the checklist and SQL but not schema or user prompt
+    // Verify the prompt contains the user question, checklist, SQL, and schema
     expect(prompt.value).to.containEql(state.sql);
+    expect(prompt.value).to.containEql(state.prompt);
     expect(prompt.value).to.containEql('1. Query selects from users table');
-    expect(prompt.value).to.not.containEql('<database-schema>');
-    expect(prompt.value).to.not.containEql('<user-question>');
+    expect(prompt.value).to.containEql('<database-schema>');
+    expect(prompt.value).to.containEql('<user-question>');
   });
 
   it('should include feedbacks in the prompt', async () => {

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

@@ -7,6 +7,7 @@ import {IGraphNode, LLMStreamEventType} from '../../../graphs';
 import {AiIntegrationBindings} from '../../../keys';
 import {LLMProvider} from '../../../types';
 import {stripThinkingTokens} from '../../../utils';
+import {AIMessage} from '@langchain/core/messages';
 import {DbQueryAIExtensionBindings} from '../keys';
 import {DbQueryNodes} from '../nodes.enum';
 import {DbSchemaHelperService} from '../services';
@@ -68,15 +69,14 @@ If no rules are relevant, return: none
     state: DbQueryState,
     config: LangGraphRunnableConfig,
   ): Promise<DbQueryState> {
-    // Skip if checklist was already generated (e.g. retry paths)
+    const empty = {} as DbQueryState;
     if (state.validationChecklist) {
-      return {} as DbQueryState;
+      return empty;
     }
 
-    // Skip for small schemas (1-2 tables) — context is already small enough
     const tableCount = Object.keys(state.schema?.tables ?? {}).length;
     if (tableCount <= 2) {
-      return {} as DbQueryState;
+      return empty;
     }
 
     const allChecks = [
@@ -85,14 +85,32 @@ If no rules are relevant, return: none
     ];
 
     if (allChecks.length === 0) {
-      return {} as DbQueryState;
+      return empty;
     }
 
     config.writer?.({
       type: LLMStreamEventType.Log,
       data: 'Filtering validation checklist for semantic validation.',
     });
 
+    const mergedIndexes = await this.runParallelChecklist(state, allChecks);
+
+    if (mergedIndexes.size === 0) {
+      return empty;
+    }
+
+    const validationChecklist = Array.from(mergedIndexes)
+      .sort((a, b) => a - b)
+      .map(i => allChecks[i - 1])
+      .join('\n');
+
+    return {validationChecklist} as DbQueryState;
+  }
+
+  private async runParallelChecklist(
+    state: DbQueryState,
+    allChecks: string[],
+  ): Promise<Set<number>> {
     const indexedChecks = allChecks
       .map((check, i) => `${i + 1}. ${check}`)
       .join('\n');
@@ -108,33 +126,25 @@ If no rules are relevant, return: none
       indexedChecks,
     };
 
-    // Run N parallel calls and union the results
     const results = await Promise.all(
       Array.from({length: parallelism}, () => chain.invoke(invokeArgs)),
     );
 
     const mergedIndexes = new Set<number>();
     for (const output of results) {
-      const response = stripThinkingTokens(output).trim();
-      if (!response) continue;
-      const indexStr = response;
-      if (indexStr === 'none') continue;
-      indexStr
-        .split(',')
-        .map(s => parseInt(s.trim(), 10))
-        .filter(n => !isNaN(n) && n >= 1 && n <= allChecks.length)
-        .forEach(n => mergedIndexes.add(n));
-    }
-
-    if (mergedIndexes.size === 0) {
-      return {} as DbQueryState;
+      this.parseIndexes(output, allChecks.length).forEach(n =>
+        mergedIndexes.add(n),
+      );
     }
+    return mergedIndexes;
+  }
 
-    const validationChecklist = Array.from(mergedIndexes)
-      .sort((a, b) => a - b)
-      .map(i => allChecks[i - 1])
-      .join('\n');
-
-    return {validationChecklist} as DbQueryState;
+  private parseIndexes(output: AIMessage, maxIndex: number): number[] {
+    const response = stripThinkingTokens(output).trim();
+    if (!response || response === 'none') return [];
+    return response
+      .split(',')
+      .map(s => Number.parseInt(s.trim(), 10))
+      .filter(n => !Number.isNaN(n) && n >= 1 && n <= maxIndex);
   }
 }

src/components/db-query/nodes/get-tables.node.ts

@@ -209,12 +209,13 @@ Use these if they are relevant to the table selection, otherwise ignore them, th
   }
 
   private _filterByPermissions(tables: string[]): string[] {
-    if (!this.permissionHelper) {
+    const permHelper = this.permissionHelper;
+    if (!permHelper) {
       return tables;
     }
     return tables.filter(t => {
       const name = t.toLowerCase().slice(t.indexOf('.') + 1);
-      return this.permissionHelper!.findMissingPermissions([name]).length === 0;
+      return permHelper.findMissingPermissions([name]).length === 0;
     });
   }
 

src/components/db-query/nodes/semantic-validator.node.ts

@@ -1,14 +1,15 @@
 import {PromptTemplate} from '@langchain/core/prompts';
 import {RunnableSequence} from '@langchain/core/runnables';
 import {LangGraphRunnableConfig} from '@langchain/langgraph';
-import {inject} from '@loopback/context';
+import {inject, service} from '@loopback/core';
 import {graphNode} from '../../../decorators';
 import {IGraphNode, LLMStreamEventType} from '../../../graphs';
 import {AiIntegrationBindings} from '../../../keys';
 import {LLMProvider} from '../../../types';
 import {stripThinkingTokens} from '../../../utils';
 import {DbQueryAIExtensionBindings} from '../keys';
 import {DbQueryNodes} from '../nodes.enum';
+import {DbSchemaHelperService} from '../services';
 import {DbQueryState} from '../state';
 import {DbQueryConfig, EvaluationResult} from '../types';
 
@@ -21,6 +22,8 @@ export class SemanticValidatorNode implements IGraphNode<DbQueryState> {
     private readonly cheapllm: LLMProvider,
     @inject(DbQueryAIExtensionBindings.Config)
     private readonly config: DbQueryConfig,
+    @service(DbSchemaHelperService)
+    private readonly schemaHelper: DbSchemaHelperService,
   ) {}
 
   prompt = PromptTemplate.fromTemplate(`
@@ -31,10 +34,18 @@ Go through each checklist item and verify it against the SQL query.
 DO NOT make up issues that do not exist in the query.
 </instructions>
 
+<user-question>
+{userPrompt}
+</user-question>
+
 <sql-query>
 {query}
 </sql-query>
 
+<database-schema>
+{schema}
+</database-schema>
+
 <validation-checklist>
 {checklist}
 </validation-checklist>
@@ -86,13 +97,15 @@ Keep these feedbacks in mind while validating the new query.
     const llm = useSmartLLM ? this.smartllm : this.cheapllm;
     const chain = RunnableSequence.from([this.prompt, llm]);
     const output = await chain.invoke({
+      userPrompt: state.prompt,
       query: state.sql,
+      schema: this.schemaHelper.asString(state.schema),
       checklist: state.validationChecklist ?? 'No checklist provided.',
       feedbacks: await this.getFeedbacks(state),
     });
     const response = stripThinkingTokens(output);
 
-    const invalidMatch = response.match(/<invalid>(.*?)<\/invalid>/s);
+    const invalidMatch = /<invalid>(.*?)<\/invalid>/s.exec(response);
     const isValid =
       response.includes('<valid/>') || response.includes('<valid />');
 

src/components/db-query/nodes/verify-checklist.node.ts

@@ -7,6 +7,7 @@ import {IGraphNode, LLMStreamEventType} from '../../../graphs';
 import {AiIntegrationBindings} from '../../../keys';
 import {LLMProvider} from '../../../types';
 import {stripThinkingTokens} from '../../../utils';
+import {AIMessage} from '@langchain/core/messages';
 import {DbQueryAIExtensionBindings} from '../keys';
 import {DbQueryNodes} from '../nodes.enum';
 import {DbSchemaHelperService} from '../services';
@@ -93,15 +94,15 @@ If no rules are relevant:
     state: DbQueryState,
     config: LangGraphRunnableConfig,
   ): Promise<DbQueryState> {
-    // Skip on retry — checklist was already verified on the first pass
+    const empty = {} as DbQueryState;
+
     if (state.feedbacks?.length) {
-      return {} as DbQueryState;
+      return empty;
     }
 
-    // Skip for small schemas (1-2 tables) — context is already small enough
     const tableCount = Object.keys(state.schema?.tables ?? {}).length;
     if (tableCount <= 2) {
-      return {} as DbQueryState;
+      return empty;
     }
 
     const allChecks = [
@@ -110,14 +111,34 @@ If no rules are relevant:
     ];
 
     if (allChecks.length === 0) {
-      return {} as DbQueryState;
+      return empty;
     }
 
     config.writer?.({
       type: LLMStreamEventType.Log,
       data: 'Verifying validation checklist with chain-of-thought.',
     });
 
+    const output = await this.invokeVerification(state, allChecks);
+    const verifiedIndexes = this.parseVerifiedIndexes(output, allChecks.length);
+
+    if (verifiedIndexes.length === 0) {
+      return empty;
+    }
+
+    const validationChecklist = this.mergeWithExisting(
+      state.validationChecklist,
+      verifiedIndexes,
+      allChecks,
+    );
+
+    return {validationChecklist} as DbQueryState;
+  }
+
+  private async invokeVerification(
+    state: DbQueryState,
+    allChecks: string[],
+  ): Promise<AIMessage> {
     const indexedChecks = allChecks
       .map((check, i) => `${i + 1}. ${check}`)
       .join('\n');
@@ -130,42 +151,40 @@ If no rules are relevant:
           ? this.evaluationOutputInstructions
           : this.simpleOutputInstructions),
     );
+
     const chain = RunnableSequence.from([promptTemplate, this.llm]);
-    const output = await chain.invoke({
+    return chain.invoke({
       prompt: state.prompt,
       tables: Object.keys(state.schema?.tables ?? {}).join(', '),
       schema: this.schemaHelper.asString(state.schema),
       indexedChecks,
     });
+  }
 
+  private parseVerifiedIndexes(output: AIMessage, maxIndex: number): number[] {
     const response = stripThinkingTokens(output).trim();
-    const resultMatch = response.match(/<result>(.*?)<\/result>/s);
+    const resultMatch = /<result>(.*?)<\/result>/s.exec(response);
     const indexStr = resultMatch ? resultMatch[1].trim() : response;
 
-    if (indexStr === 'none' || !indexStr) {
-      return {} as DbQueryState;
-    }
+    if (!indexStr || indexStr === 'none') return [];
 
-    const verifiedIndexes = indexStr
+    return indexStr
       .split(',')
-      .map(s => parseInt(s.trim(), 10))
-      .filter(n => !isNaN(n) && n >= 1 && n <= allChecks.length);
-
-    if (verifiedIndexes.length === 0) {
-      return {} as DbQueryState;
-    }
+      .map(s => Number.parseInt(s.trim(), 10))
+      .filter(n => !Number.isNaN(n) && n >= 1 && n <= maxIndex);
+  }
 
-    // Merge with existing checklist — union of both passes
+  private mergeWithExisting(
+    existing: string | undefined,
+    verifiedIndexes: number[],
+    allChecks: string[],
+  ): string {
     const existingChecks = new Set(
-      (state.validationChecklist ?? '').split('\n').filter(c => c.length > 0),
+      (existing ?? '').split('\n').filter(c => c.length > 0),
     );
-    const verifiedChecks = verifiedIndexes.map(i => allChecks[i - 1]);
-    for (const check of verifiedChecks) {
+    for (const check of verifiedIndexes.map(i => allChecks[i - 1])) {
       existingChecks.add(check);
     }
-
-    const validationChecklist = Array.from(existingChecks).join('\n');
-
-    return {validationChecklist} as DbQueryState;
+    return Array.from(existingChecks).join('\n');
   }
 }

src/components/db-query/services/knowledge-graph/db-knowledge-graph.service.ts

@@ -18,9 +18,10 @@ import {Concept, GraphEdge, GraphNode, KnowledgeGraph} from './types';
 const debug = require('debug')('ai-integration:knowledge-graph');
 
 @injectable({scope: BindingScope.SINGLETON})
-export class DbKnowledgeGraphService
-  implements KnowledgeGraph<string, DatabaseSchema>
-{
+export class DbKnowledgeGraphService implements KnowledgeGraph<
+  string,
+  DatabaseSchema
+> {
   edges: Map<string, GraphEdge[]>;
   nodes: Map<string, GraphNode>;
   private vectorWeight: number;

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

@@ -48,14 +48,14 @@ export class GetDataAsDatasetTool implements IGraphTool {
       prompt: z
         .string()
         .describe(
-          `The user's request describing what data they need from the database.`,
+          `Prompt from the user that will be used for generating an SQL query and create a dataset from it.`,
         ),
     }) as AnyObject[string];
     return graph.asTool({
       name: this.key,
-      description: `Tool for fetching data from the database and returning it as a dataset based on the user's request.
-                Use this whenever the user wants to retrieve, look up, or explore data from the database.
-                It returns a dataset ID and renders a data grid on the UI for the user to see the results.`,
+      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. 
+                It internally fires an event that renders a grid for the dataset on the UI for the user to see.`,
       schema,
     });
   }

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

@@ -8,9 +8,7 @@ import {VisualizationGraphState} from '../state';
 @graphNode(VisualizationGraphNodes.CallQueryGeneration, {
   [POST_DATASET_TAG]: true,
 })
-export class CallQueryGenerationNode
-  implements IGraphNode<VisualizationGraphState>
-{
+export class CallQueryGenerationNode implements IGraphNode<VisualizationGraphState> {
   constructor(
     @service(DbQueryGraph)
     private readonly queryPipeline: DbQueryGraph,