feat(component): minor performance improvements

Author: akshatdubeysf

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

@@ -36,7 +36,9 @@ describe('SemanticValidatorNode Unit', function () {
       .stub(schemaHelper, 'getTablesContext')
       .returns(['employee salary must be converted to USD']);
 
-    node = new SemanticValidatorNode(llm, schemaHelper, ['test context']);
+    node = new SemanticValidatorNode(llm, llm, {models: []}, schemaHelper, [
+      'test context',
+    ]);
   });
 
   afterEach(() => {

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

@@ -36,6 +36,7 @@ describe('SqlGenerationNode Unit', function () {
       .returns(['Table employees contains employee information']);
 
     node = new SqlGenerationNode(
+      llm,
       llm,
       {
         db: {

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

@@ -12,13 +12,17 @@ import {DbQueryAIExtensionBindings} from '../keys';
 import {DbQueryNodes} from '../nodes.enum';
 import {DbSchemaHelperService} from '../services';
 import {DbQueryState} from '../state';
-import {EvaluationResult} from '../types';
+import {DbQueryConfig, EvaluationResult} from '../types';
 
 @graphNode(DbQueryNodes.SemanticValidator)
 export class SemanticValidatorNode implements IGraphNode<DbQueryState> {
   constructor(
     @inject(AiIntegrationBindings.SmartLLM)
-    private readonly llm: LLMProvider,
+    private readonly smartllm: LLMProvider,
+    @inject(AiIntegrationBindings.CheapLLM)
+    private readonly cheapllm: LLMProvider,
+    @inject(DbQueryAIExtensionBindings.Config)
+    private readonly config: DbQueryConfig,
     @service(DbSchemaHelperService)
     private readonly schemaHelper: DbSchemaHelperService,
     @inject(DbQueryAIExtensionBindings.GlobalContext, {optional: true})
@@ -87,7 +91,10 @@ Keep these feedbacks in mind while validating the new query.
       type: LLMStreamEventType.Log,
       data: `Validating the query semantically.`,
     });
-    const chain = RunnableSequence.from([this.prompt, this.llm]);
+    const useSmartLLM =
+      this.config.nodes?.semanticValidatorNode?.useSmartLLM ?? false;
+    const llm = useSmartLLM ? this.smartllm : this.cheapllm;
+    const chain = RunnableSequence.from([this.prompt, llm]);
     const output = await chain.invoke({
       query: state.sql,
       prompt: state.prompt,

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

@@ -86,6 +86,8 @@ In the last attempt, you generated this SQL query -
   constructor(
     @inject(AiIntegrationBindings.SmartLLM)
     private readonly sqlLLM: LLMProvider,
+    @inject(AiIntegrationBindings.CheapLLM)
+    private readonly cheapllm: LLMProvider,
     @inject(DbQueryAIExtensionBindings.Config)
     private readonly config: DbQueryConfig,
     @service(DbSchemaHelperService)
@@ -97,10 +99,17 @@ In the last attempt, you generated this SQL query -
     state: DbQueryState,
     config: LangGraphRunnableConfig,
   ): Promise<DbQueryState> {
-    const chain = RunnableSequence.from([
-      this.sqlGenerationPrompt,
-      this.sqlLLM,
-    ]);
+    let llm = this.sqlLLM;
+
+    if (
+      this.config.nodes?.sqlGenerationNode?.generateDescription !== false &&
+      state.schema.tables &&
+      Object.keys(state.schema.tables).length === 1
+    ) {
+      llm = this.cheapllm;
+    }
+
+    const chain = RunnableSequence.from([this.sqlGenerationPrompt, llm]);
 
     config.writer?.({
       type: LLMStreamEventType.Log,

src/components/db-query/types.ts

@@ -115,10 +115,14 @@ export type DbQueryConfig = {
   nodes?: {
     sqlGenerationNode?: {
       generateDescription?: boolean;
+      useSmartLLMForSingleTableQueries?: boolean;
     };
     getTablesNode?: {
       useSmartLLM?: boolean;
     };
+    semanticValidatorNode?: {
+      useSmartLLM?: boolean;
+    };
   };
   columnSelection?: boolean;
 };