feat(tools): rename query builder tools to dataset-based tools

BREAKING CHANGE: Tool keys renamed.
generate-query is now get-data-as-dataset.
improve-query is now improve-dataset.
Descriptions updated to dataset-centric language.

Author: akshatdubeysf

src/__tests__/db-query/unit/db-query.graph.unit.ts

@@ -34,8 +34,9 @@ describe(`DbQueryGraph Unit`, function () {
     stubMap[DbQueryNodes.GenerateChecklist].callsFake(async () => ({
       validationChecklist: '1. Test check',
     }));
-    stubMap[DbQueryNodes.GenerateDescription].callsFake(async (state: Record<string, unknown>) =>
-      state.description ? {} : {description: 'Test description'},
+    stubMap[DbQueryNodes.GenerateDescription].callsFake(
+      async (state: Record<string, unknown>) =>
+        state.description ? {} : {description: 'Test description'},
     );
     // VerifyChecklist runs in parallel with SqlGeneration — must return partial state
     stubMap[DbQueryNodes.VerifyChecklist].callsFake(async () => ({}));
@@ -164,7 +165,9 @@ describe(`DbQueryGraph Unit`, function () {
     expect(stubMap[DbQueryNodes.IsImprovement].calledOnce).to.be.true();
     expect(stubMap[DbQueryNodes.CheckCache].calledOnce).to.be.true();
     expect(stubMap[DbQueryNodes.GetTables].calledOnce).to.be.true();
-    expect(stubMap[DbQueryNodes.SqlGeneration].getCalls().length).to.be.eql(MAX_ATTEMPTS);
+    expect(stubMap[DbQueryNodes.SqlGeneration].getCalls().length).to.be.eql(
+      MAX_ATTEMPTS,
+    );
     expect(
       stubMap[DbQueryNodes.SyntacticValidator].getCalls().length,
     ).to.be.eql(MAX_ATTEMPTS);

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

@@ -1,8 +1,5 @@
 import {expect, sinon} from '@loopback/testlab';
-import {
-  EvaluationResult,
-  SemanticValidatorNode,
-} from '../../../../components';
+import {EvaluationResult, SemanticValidatorNode} from '../../../../components';
 import {LLMProvider} from '../../../../types';
 
 describe('SemanticValidatorNode Unit', function () {

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

@@ -55,8 +55,7 @@ describe('SqlGenerationNode Unit', function () {
 
   it('should generate SQL query based on the provided prompt', async () => {
     llmStub.resolves({
-      content:
-        '<think>thinking about it</think>SELECT * FROM employees;',
+      content: '<think>thinking about it</think>SELECT * FROM employees;',
     });
 
     const state = {
@@ -145,8 +144,7 @@ It should have no other character or symbol or character that is not part of SQL
 
   it('should generate SQL query based on the provided prompt with a single feedback from some validation stage', async () => {
     llmStub.resolves({
-      content:
-        '<think>thinking about it</think>SELECT * FROM employees;',
+      content: '<think>thinking about it</think>SELECT * FROM employees;',
     });
 
     const state = {
@@ -249,8 +247,7 @@ It should have no other character or symbol or character that is not part of SQL
 
   it('should generate SQL query based on the provided prompt with a multiple feedbacks from from previous loops', async () => {
     llmStub.resolves({
-      content:
-        '<think>thinking about it</think>SELECT * FROM employees;',
+      content: '<think>thinking about it</think>SELECT * FROM employees;',
     });
 
     const state = {
@@ -361,8 +358,7 @@ It should have no other character or symbol or character that is not part of SQL
 
   it('should generate SQL query with sample queries when no feedbacks but has sample SQL', async () => {
     llmStub.resolves({
-      content:
-        '<think>thinking about it</think>SELECT * FROM employees;',
+      content: '<think>thinking about it</think>SELECT * FROM employees;',
     });
 
     const state = {
@@ -419,8 +415,7 @@ It should have no other character or symbol or character that is not part of SQL
 
   it('should generate SQL query with baseline sample queries when no feedbacks and not from cache', async () => {
     llmStub.resolves({
-      content:
-        '<think>thinking about it</think>SELECT * FROM employees;',
+      content: '<think>thinking about it</think>SELECT * FROM employees;',
     });
 
     const state = {
@@ -514,8 +509,7 @@ It should have no other character or symbol or character that is not part of SQL
     it('should use cheap LLM when OPTIMIZE_CACHED_QUERIES is true and sampleSql exists', async () => {
       process.env.OPTIMIZE_CACHED_QUERIES = 'true';
       cheapLLMStub.resolves({
-        content:
-          'SELECT * FROM employees WHERE id = 1;',
+        content: 'SELECT * FROM employees WHERE id = 1;',
       });
 
       const state = {
@@ -575,8 +569,7 @@ It should have no other character or symbol or character that is not part of SQL
     it('should use smart LLM when OPTIMIZE_CACHED_QUERIES is false and sampleSql exists', async () => {
       process.env.OPTIMIZE_CACHED_QUERIES = 'false';
       smartLLMStub.resolves({
-        content:
-          'SELECT * FROM employees WHERE id = 1;',
+        content: 'SELECT * FROM employees WHERE id = 1;',
       });
 
       const state = {
@@ -636,8 +629,7 @@ It should have no other character or symbol or character that is not part of SQL
     it('should use cheap LLM for single table schemas regardless of cache', async () => {
       process.env.OPTIMIZE_CACHED_QUERIES = 'false';
       cheapLLMStub.resolves({
-        content:
-          'SELECT * FROM employees;',
+        content: 'SELECT * FROM employees;',
       });
 
       const state = {
@@ -751,8 +743,7 @@ It should have no other character or symbol or character that is not part of SQL
     it('should default to true for OPTIMIZE_CACHED_QUERIES when env var is not set', async () => {
       delete process.env.OPTIMIZE_CACHED_QUERIES;
       cheapLLMStub.resolves({
-        content:
-          'SELECT * FROM employees WHERE id = 1;',
+        content: 'SELECT * FROM employees WHERE id = 1;',
       });
 
       const state = {
@@ -812,8 +803,7 @@ It should have no other character or symbol or character that is not part of SQL
     it('should use smart LLM when sampleSql is null despite optimization being enabled', async () => {
       process.env.OPTIMIZE_CACHED_QUERIES = 'true';
       smartLLMStub.resolves({
-        content:
-          'SELECT * FROM employees, departments;',
+        content: 'SELECT * FROM employees, departments;',
       });
 
       const state = {

src/components/db-query/db-query.component.ts

@@ -37,8 +37,8 @@ import {SchemaStore} from './services/schema.store';
 import {TableSearchService} from './services/search/table-search.service';
 import {
   AskAboutDatasetTool,
-  GenerateQueryTool,
-  ImproveQueryTool,
+  GetDataAsDatasetTool,
+  ImproveDatasetTool,
 } from './tools';
 import {PgWithRlsConnector} from './connectors/pg';
 
@@ -72,8 +72,8 @@ export class DbQueryComponent implements Component {
       DbQueryGraph,
       // tools
       AskAboutDatasetTool,
-      GenerateQueryTool,
-      ImproveQueryTool,
+      GetDataAsDatasetTool,
+      ImproveDatasetTool,
       // nodes
       IsImprovementNode,
       GetTablesNode,

src/components/db-query/db-query.graph.ts

@@ -59,56 +59,53 @@ export class DbQueryGraph extends BaseGraph<DbQueryState> {
       .addNode(DbQueryNodes.PostCacheAndTables, async () => ({}))
       .addNode(DbQueryNodes.PreValidation, async () => ({}))
       // PostValidation: merges syntactic + semantic results into status/feedbacks
-      .addNode(
-        DbQueryNodes.PostValidation,
-        async (state: DbQueryState) => {
-          // Syntactic failures take priority
-          if (
-            state.syntacticStatus &&
-            state.syntacticStatus !== EvaluationResult.Pass
-          ) {
-            return {
-              status: state.syntacticStatus,
-              feedbacks: [
-                ...(state.feedbacks ?? []),
-                ...(state.syntacticFeedback ? [state.syntacticFeedback] : []),
-              ],
-              syntacticStatus: undefined,
-              syntacticFeedback: undefined,
-              semanticStatus: undefined,
-              semanticFeedback: undefined,
-            };
-          }
-          // Semantic failure
-          if (
-            state.semanticStatus &&
-            state.semanticStatus !== EvaluationResult.Pass
-          ) {
-            return {
-              status: state.semanticStatus,
-              feedbacks: [
-                ...(state.feedbacks ?? []),
-                ...(state.semanticFeedback ? [state.semanticFeedback] : []),
-              ],
-              syntacticStatus: undefined,
-              syntacticFeedback: undefined,
-              semanticStatus: undefined,
-              semanticFeedback: undefined,
-            };
-          }
-          // Both passed — clear internal validator feedbacks
+      .addNode(DbQueryNodes.PostValidation, async (state: DbQueryState) => {
+        // Syntactic failures take priority
+        if (
+          state.syntacticStatus &&
+          state.syntacticStatus !== EvaluationResult.Pass
+        ) {
           return {
-            status: EvaluationResult.Pass,
-            feedbacks: (state.feedbacks ?? []).filter(
-              f => !f.startsWith('Query Validation Failed'),
-            ),
+            status: state.syntacticStatus,
+            feedbacks: [
+              ...(state.feedbacks ?? []),
+              ...(state.syntacticFeedback ? [state.syntacticFeedback] : []),
+            ],
             syntacticStatus: undefined,
             syntacticFeedback: undefined,
             semanticStatus: undefined,
             semanticFeedback: undefined,
           };
-        },
-      )
+        }
+        // Semantic failure
+        if (
+          state.semanticStatus &&
+          state.semanticStatus !== EvaluationResult.Pass
+        ) {
+          return {
+            status: state.semanticStatus,
+            feedbacks: [
+              ...(state.feedbacks ?? []),
+              ...(state.semanticFeedback ? [state.semanticFeedback] : []),
+            ],
+            syntacticStatus: undefined,
+            syntacticFeedback: undefined,
+            semanticStatus: undefined,
+            semanticFeedback: undefined,
+          };
+        }
+        // Both passed — clear internal validator feedbacks
+        return {
+          status: EvaluationResult.Pass,
+          feedbacks: (state.feedbacks ?? []).filter(
+            f => !f.startsWith('Query Validation Failed'),
+          ),
+          syntacticStatus: undefined,
+          syntacticFeedback: undefined,
+          semanticStatus: undefined,
+          semanticFeedback: undefined,
+        };
+      })
       // === EDGES ===
       // Parallel fan-out: cache check and table selection
       .addEdge(START, DbQueryNodes.IsImprovement)

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

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

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

@@ -101,9 +101,7 @@ Keep these feedbacks in mind while validating the new query.
         semanticStatus: EvaluationResult.Pass,
       } as DbQueryState;
     } else {
-      const reason = invalidMatch
-        ? invalidMatch[1].trim()
-        : response.trim();
+      const reason = invalidMatch ? invalidMatch[1].trim() : response.trim();
       config.writer?.({
         type: LLMStreamEventType.Log,
         data: `Query Validation Failed by LLM: ${reason}`,

src/components/db-query/nodes/sql-generation.node.ts

@@ -91,8 +91,7 @@ In the last attempt, you generated this SQL query -
       (process.env.OPTIMIZE_CACHED_QUERIES ?? 'true') === 'true';
 
     const isSingleTable =
-      state.schema.tables &&
-      Object.keys(state.schema.tables).length === 1;
+      state.schema.tables && Object.keys(state.schema.tables).length === 1;
 
     // Use cheap LLM for validation fix retries — the query is close, just needs small corrections
     const isValidationFixRetry =

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

@@ -37,7 +37,7 @@ export function getTableNodeTests(cases: GetTableNodeTestCase[]) {
         } as unknown as RunnableConfig,
       );
       testCase.expectedTables.forEach(table => {
-        expect(result.schema.tables).to.have.property(table);
+        expect(result.schema?.tables).to.have.property(table);
       });
     },
   }));

src/components/db-query/tools/ask-about-dataset.tool.ts

@@ -82,7 +82,7 @@ export class AskAboutDatasetTool implements IGraphTool {
       {
         name: this.key,
         description:
-          'Tool for answering questions about a query, note that it does not run the query and can only answer questions about the query itself, not the data it returns on execution. Call this only if you have a valid dataset ID available.',
+          'Tool for answering questions about an existing dataset, note that it can only answer questions about the dataset definition, not the data it contains. Call this only if you have a valid dataset ID available.',
         schema,
       },
     );