feat(query-generation): handle missing token count, handle duplicate context

Author: akshatdubeysf

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

@@ -129,10 +129,8 @@ export class DbQueryGraph extends BaseGraph<DbQueryState> {
       // GetColumns → GenerateChecklist (fast pass) → parallel fan-out
       .addEdge(DbQueryNodes.GetColumns, DbQueryNodes.GenerateChecklist)
       .addEdge(DbQueryNodes.GenerateChecklist, DbQueryNodes.SqlGeneration)
-      .addEdge(DbQueryNodes.GenerateChecklist, DbQueryNodes.GenerateDescription)
       .addEdge(DbQueryNodes.GenerateChecklist, DbQueryNodes.VerifyChecklist)
-      // All three fan-in to PreValidation
-      .addEdge(DbQueryNodes.GenerateDescription, DbQueryNodes.PreValidation)
+      // Both fan-in to PreValidation
       .addEdge(DbQueryNodes.VerifyChecklist, DbQueryNodes.PreValidation)
       // SqlGeneration routes to validation or failure
       .addConditionalEdges(
@@ -146,12 +144,14 @@ export class DbQueryGraph extends BaseGraph<DbQueryState> {
           Failed: DbQueryNodes.Failed,
         },
       )
-      // Parallel fan-out: both validators run concurrently
+      // Parallel fan-out: validators and description generation run concurrently
       .addEdge(DbQueryNodes.PreValidation, DbQueryNodes.SyntacticValidator)
       .addEdge(DbQueryNodes.PreValidation, DbQueryNodes.SemanticValidator)
+      .addEdge(DbQueryNodes.PreValidation, DbQueryNodes.GenerateDescription)
       // Fan-in at PostValidation
       .addEdge(DbQueryNodes.SyntacticValidator, DbQueryNodes.PostValidation)
       .addEdge(DbQueryNodes.SemanticValidator, DbQueryNodes.PostValidation)
+      .addEdge(DbQueryNodes.GenerateDescription, DbQueryNodes.PostValidation)
       .addConditionalEdges(
         DbQueryNodes.PostValidation,
         (state: DbQueryState) => {
@@ -164,9 +164,7 @@ export class DbQueryGraph extends BaseGraph<DbQueryState> {
           return 'Failed';
         },
         {
-          // SaveDataset fans-in from both PostValidation and GenerateDescription
           Accepted: DbQueryNodes.SaveDataset,
-          // FixSQL goes through GenerateChecklist (no-op) to re-trigger both SqlGeneration and GenerateDescription
           FixSQL: DbQueryNodes.GenerateChecklist,
           ReselectTables: DbQueryNodes.GetTables,
           Failed: DbQueryNodes.Failed,

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

@@ -27,14 +27,19 @@ export class GenerateDescriptionNode implements IGraphNode<DbQueryState> {
 
   prompt = PromptTemplate.fromTemplate(`
 <instructions>
-You are an AI assistant that summarizes what data a query would fetch to answer the user's question.
-Write a concise, bulleted summary in plain english. No SQL, no technical jargon, no table/column names.
+You are an AI assistant that describes what a SQL query does in plain english.
+Analyze the actual query below and write a concise, bulleted summary of the data it retrieves and any filters/conditions it applies.
+Write in plain english. No SQL, no technical jargon, no table/column names.
 </instructions>
 
 <user-question>
 {prompt}
 </user-question>
 
+<sql-query>
+{sql}
+</sql-query>
+
 <database-schema>
 {schema}
 </database-schema>
@@ -56,7 +61,7 @@ Return a short bulleted list where each bullet is one condition, filter, or piec
     const generateDesc =
       this.config.nodes?.sqlGenerationNode?.generateDescription !== false;
 
-    if (!generateDesc || state.description) {
+    if (!generateDesc || !state.sql) {
       return {} as DbQueryState;
     }
 
@@ -68,6 +73,7 @@ Return a short bulleted list where each bullet is one condition, filter, or piec
     const chain = RunnableSequence.from([this.prompt, this.llm]);
     const stream = await chain.stream({
       prompt: state.prompt,
+      sql: state.sql,
       schema: this.schemaHelper.asString(state.schema),
       checks: [
         '<must-follow-rules>',

src/components/db-query/services/db-schema-helper.service.ts

@@ -28,19 +28,19 @@ export class DbSchemaHelperService {
     private readonly config: DbQueryConfig,
   ) {}
   getTablesContext(schema: DatabaseSchema) {
-    const tableContexts: string[] = [];
+    const contextSet = new Set<string>();
     Object.keys(schema.tables).forEach(table => {
       if (schema.tables[table].context) {
         for (const item of schema.tables[table].context) {
           if (typeof item === 'string' && item.trim().length > 0) {
-            tableContexts.push(item.trim());
+            contextSet.add(item.trim());
           } else if (typeof item === 'object') {
             const tableSet = new Set(
               Object.keys(schema.tables).map(t => t.split('.').pop() ?? t),
             );
             Object.keys(item).forEach(withTable => {
               if (tableSet.has(`${withTable}`)) {
-                tableContexts.push(item[withTable].trim());
+                contextSet.add(item[withTable].trim());
               }
             });
           } else {
@@ -49,7 +49,7 @@ export class DbSchemaHelperService {
         }
       }
     });
-    return tableContexts;
+    return [...contextSet];
   }
 
   asString(schema: DatabaseSchema): string {

src/services/token-counter.service.ts

@@ -37,10 +37,10 @@ export class TokenCounter {
       outputTokens: 0,
     };
     if (usageMetadata) {
-      this.inputs += usageMetadata.input_tokens;
-      this.outputs += usageMetadata.output_tokens;
-      prev.inputTokens += usageMetadata.input_tokens;
-      prev.outputTokens += usageMetadata.output_tokens;
+      this.inputs += usageMetadata.input_tokens ?? 0;
+      this.outputs += usageMetadata.output_tokens ?? 0;
+      prev.inputTokens += usageMetadata.input_tokens ?? 0;
+      prev.outputTokens += usageMetadata.output_tokens ?? 0;
       this.countMap.set(llmName, prev);
     }
     return {