@superatomai/sdk-node 0.0.17-mds → 0.0.18-mds

package/dist/index.js

@@ -5082,6 +5082,38 @@ var StreamBuffer = class {
     this.callback = void 0;
   }
 };
+var TaggedStreamBuffer = class extends StreamBuffer {
+  constructor(callId, parent, sourceName, intent) {
+    super();
+    this.started = false;
+    this.callId = callId;
+    this.parent = parent;
+    this.sourceName = sourceName;
+    this.intent = intent;
+  }
+  hasCallback() {
+    return this.parent.hasCallback();
+  }
+  getFullText() {
+    return this.parent.getFullText();
+  }
+  write(chunk) {
+    if (!this.parent.hasCallback()) return;
+    if (!this.started) {
+      this.started = true;
+      const escapedName = this.sourceName.replace(/\|/g, "-");
+      const escapedIntent = this.intent.replace(/\|/g, " ").replace(/"/g, """).substring(0, 200);
+      this.parent.write(`__SB_${this.callId}_START__${escapedName}|${escapedIntent}__SB_END__`);
+    }
+    this.parent.write(`__SB_${this.callId}_MSG__${chunk}__SB_END__`);
+  }
+  flush() {
+    this.parent.flush();
+  }
+  dispose() {
+    this.parent.flush();
+  }
+};
 function streamDelay(ms = STREAM_DELAY_MS) {
   return new Promise((resolve) => setTimeout(resolve, ms));
 }
@@ -7010,16 +7042,6 @@ var SourceAgent = class {
     const startTime = Date.now();
     const { intent, aggregation = "raw" } = input;
     logger.info(`[SourceAgent:${this.tool.name}] Starting | intent: "${intent}" | aggregation: ${aggregation}`);
-    if (this.streamBuffer.hasCallback()) {
-      const sourceType = this.extractSourceType();
-      const sourceLabel = sourceType !== "unknown" ? ` (${sourceType})` : "";
-      this.streamBuffer.write(`
-
-\u{1F517} **Connecting to ${this.tool.name}**${sourceLabel}
-
-`);
-      await streamDelay();
-    }
     try {
       const prompts = await this.buildPrompt(intent, aggregation);
       logger.logLLMPrompt(`sourceAgent:${this.tool.name}`, "system", extractPromptText(prompts.system));
@@ -7029,38 +7051,46 @@ var SourceAgent = class {
       let resultData = [];
       let queryExecuted;
       let totalRowsMatched = 0;
+      let querySucceeded = false;
+      let successfulQueries = 0;
       const schemaSearchFn = this.tool.schemaSearchFn;
       const toolHandler = async (toolName, toolInput) => {
         if (toolName.endsWith("_search_schema") && schemaSearchFn) {
           const keywords = toolInput.keywords || [];
           logger.info(`[SourceAgent:${this.tool.name}] Schema search: ${keywords.join(", ")}`);
           if (this.streamBuffer.hasCallback()) {
-            this.streamBuffer.write(`\u{1F50D} **Searching schema for:** ${keywords.join(", ")}
+            this.streamBuffer.write(`\u{1F50D} **Looking up tables:** ${keywords.join(", ")}
 
 `);
             await streamDelay();
           }
           const result = schemaSearchFn(keywords);
-          if (this.streamBuffer.hasCallback()) {
-            const matchCount = (result.match(/^• /gm) || []).length;
-            this.streamBuffer.write(`\u{1F4CB} Found ${matchCount} matching table(s)
-
-`);
-          }
           return `Schema search results:
 
 ${result}`;
         }
+        if (successfulQueries >= 2) {
+          return `\u26D4 Maximum successful queries reached (2). Use the data already retrieved.`;
+        }
         this.attempts++;
         if (this.attempts > this.config.maxRetries) {
           throw new Error(`Max retry attempts (${this.config.maxRetries}) reached for ${this.tool.name}`);
         }
         if (this.attempts > 1 && this.streamBuffer.hasCallback()) {
-          this.streamBuffer.write(`
+          if (querySucceeded) {
+            const reason = toolInput.reason || "";
+            this.streamBuffer.write(`
+
+\u{1F4DD} **Follow-up query${reason ? `:** ${reason}` : "**"}
+
+`);
+          } else {
+            this.streamBuffer.write(`
 
-\u{1F504} **Query failed, retrying with corrected query** (attempt ${this.attempts}/${this.config.maxRetries})
+\u{1F504} **Retrying query** (attempt ${this.attempts}/${this.config.maxRetries})
 
 `);
+          }
           await streamDelay();
         }
         const cappedInput = { ...toolInput };
@@ -7076,13 +7106,6 @@ ${result}`;
 ${queryDisplay}
 \`\`\`
 
-`);
-          } else {
-            this.streamBuffer.write(`\u{1F4DD} **Query parameters:**
-\`\`\`json
-${JSON.stringify(cappedInput, null, 2)}
-\`\`\`
-
 `);
           }
           await streamDelay();
@@ -7096,17 +7119,36 @@ ${JSON.stringify(cappedInput, null, 2)}
           if (result && result.error) {
             const errorMsg = typeof result.error === "string" ? result.error : JSON.stringify(result.error);
             logger.warn(`[SourceAgent:${this.tool.name}] Tool returned error (attempt ${this.attempts}/${this.config.maxRetries}): ${errorMsg}`);
+            if (this.streamBuffer.hasCallback()) {
+              this.streamBuffer.write(`\u274C **Query failed:** ${errorMsg}
+
+`);
+            }
             return `\u274C ERROR: ${errorMsg}
 
 Analyze the error and try again with a corrected query.`;
           }
           resultData = result.data || [];
           totalRowsMatched = result.metadata?.totalCount || result.count || resultData.length;
+          querySucceeded = true;
+          successfulQueries++;
           if (this.streamBuffer.hasCallback()) {
-            const totalInfo = totalRowsMatched > resultData.length ? ` (${totalRowsMatched} total matched)` : "";
-            this.streamBuffer.write(`\u2705 Retrieved ${resultData.length} rows${totalInfo}
+            const totalInfo = totalRowsMatched > resultData.length ? ` of ${totalRowsMatched} total` : "";
+            this.streamBuffer.write(`\u2705 **${resultData.length} rows${totalInfo} from ${this.tool.name}**
+
+`);
+            if (resultData.length > 0) {
+              try {
+                const previewData = resultData.slice(0, this.config.maxRowsPerSource);
+                this.streamBuffer.write(`<DataTable>${JSON.stringify(previewData)}</DataTable>
 
 `);
+              } catch {
+                this.streamBuffer.write(`_Data preview not available_
+
+`);
+              }
+            }
           }
           const formattedResult = formatToolResultForLLM(result, {
             toolName: this.tool.name,
@@ -7128,19 +7170,28 @@ Analyze the error and try again with a corrected query.`;
             sourceType: this.extractSourceType()
           };
           const formatted = typeof formattedResult === "string" ? formattedResult : JSON.stringify(formattedResult);
-          return `\u2705 Query executed successfully. ${resultData.length} rows returned (${totalRowsMatched} total matched). Data is ready \u2014 do NOT call the tool again.
+          const followUpNote = successfulQueries < 2 ? "You may make ONE follow-up query if this data is incomplete for the intent. Otherwise, STOP." : "Maximum queries reached. Use the data you have.";
+          return `\u2705 Query executed successfully. ${resultData.length} rows returned (${totalRowsMatched} total matched). ${followUpNote}
 
 ${formatted}`;
         } catch (execError) {
           const errorMsg = execError instanceof Error ? execError.message : typeof execError === "object" && execError !== null ? execError.message || execError.error || JSON.stringify(execError) : String(execError);
           logger.warn(`[SourceAgent:${this.tool.name}] Tool execution failed (attempt ${this.attempts}/${this.config.maxRetries}): ${errorMsg}`);
+          if (this.streamBuffer.hasCallback()) {
+            this.streamBuffer.write(`\u274C **Query failed:** ${errorMsg}
+
+`);
+          }
           return `\u274C ERROR: ${errorMsg}
 
 Analyze the error and try again with a corrected query.`;
         }
       };
       const hasSchemaSearch = !!schemaSearchFn;
-      const maxIterations = this.config.maxRetries + 2 + (hasSchemaSearch ? 2 : 0);
+      const queryIterations = this.config.maxRetries + 1;
+      const schemaIterations = hasSchemaSearch ? 3 : 0;
+      const responseIterations = 2;
+      const maxIterations = queryIterations + schemaIterations + responseIterations;
       await LLM.streamWithTools(
         { sys: prompts.system, user: prompts.user },
         tools,
@@ -7193,7 +7244,7 @@ Analyze the error and try again with a corrected query.`;
       if (this.streamBuffer.hasCallback()) {
         this.streamBuffer.write(`
 
-\u274C **${this.tool.name} failed:** ${errorMsg}
+\u274C **Could not retrieve data from ${this.tool.name}:** ${errorMsg}
 
 `);
       }
@@ -7224,26 +7275,43 @@ Analyze the error and try again with a corrected query.`;
     const sourceName = this.tool.name;
     const sourceType = this.extractSourceType();
     const fullSchema = this.tool.fullSchema || this.tool.description || "No schema available";
+    const databaseRules = await promptLoader.loadDatabaseRulesForType(sourceType);
+    const rowLimitSyntax = {
+      "mssql": `Use SELECT TOP ${this.config.maxRowsPerSource} in every SELECT statement`,
+      "postgres": `Add LIMIT ${this.config.maxRowsPerSource} at the end of every query`,
+      "mysql": `Add LIMIT ${this.config.maxRowsPerSource} at the end of every query`,
+      "excel": `Add LIMIT ${this.config.maxRowsPerSource} at the end of every query`,
+      "csv": `Add LIMIT ${this.config.maxRowsPerSource} at the end of every query`
+    };
+    const rowLimitInstruction = rowLimitSyntax[sourceType] || `Limit results to ${this.config.maxRowsPerSource} rows`;
     const hasSchemaSearch = !!this.tool.schemaSearchFn;
     const schemaTier = this.tool.schemaTier || "";
     let schemaSearchInstructions = "";
     if (hasSchemaSearch && schemaTier === "very_large") {
-      schemaSearchInstructions = `## Schema Search
+      schemaSearchInstructions = `## Schema Search \u2014 REQUIRED
 This source has a very large schema. The schema above shows only table names \u2014 no column details.
-Before writing any query, you MUST use the search_schema tool to find exact table and column names.
-1. Search with keywords related to the requested data
-2. Review the returned column details and sample values carefully
-3. Then write your SQL query using the exact names from the search results
-You may search multiple times with different keywords if the first search doesn't find what you need.`;
+**You MUST use search_schema BEFORE writing any query.** Do NOT guess column names.
+1. FIRST: Call search_schema with keywords related to the requested data
+2. Review the returned column details, types, and sample values carefully
+3. THEN write your SQL query using the EXACT column names from the search results
+You may search multiple times with different keywords if the first search doesn't find what you need.
+**Never write a query without searching first \u2014 guessing column names wastes time on failed queries.**`;
     } else if (hasSchemaSearch) {
-      schemaSearchInstructions = `## Schema Search
-This source has a large schema. The top tables are shown in detail above, but not all tables are listed with columns.
-If you need data from a table that is only listed in the catalog (no columns shown), use the search_schema tool to look up its column details before writing a query.`;
+      schemaSearchInstructions = `## Schema Search \u2014 REQUIRED
+This source has a large schema. The top tables are shown in detail above, but not all tables have full column details.
+**You MUST use search_schema BEFORE writing any query** to verify the exact column names for the tables you plan to query.
+1. FIRST: Call search_schema with the table names or keywords you intend to use
+2. Review the returned column names, types, and sample values
+3. THEN write your SQL query using the EXACT column names from the search results
+Even if a table appears in the detailed schema above, search_schema returns sample values that help you write correct WHERE filters.
+**Never skip the search step \u2014 it prevents failed queries and retries.**`;
     }
     const prompts = await promptLoader.loadPrompts("agent-source-query", {
       SOURCE_NAME: sourceName,
       SOURCE_TYPE: sourceType,
       FULL_SCHEMA: fullSchema,
+      DATABASE_RULES: databaseRules,
+      ROW_LIMIT_SYNTAX: rowLimitInstruction,
       SCHEMA_SEARCH_INSTRUCTIONS: schemaSearchInstructions,
       MAX_ROWS: String(this.config.maxRowsPerSource),
       AGGREGATION_MODE: aggregation,
@@ -7298,6 +7366,10 @@ If you need data from a table that is only listed in the catalog (no columns sho
         required.push(key);
       }
     });
+    properties["reason"] = {
+      type: "string",
+      description: "string (optional) - Brief reason for this query. Required for follow-up queries after a successful first query. Explain what additional data you need and why."
+    };
     return {
       name: this.tool.id,
       description: this.tool.description || `Query ${this.tool.name}`,
@@ -7386,6 +7458,7 @@ var MainAgent = class {
     const tools = [...sourceToolDefs, ...directToolDefs];
     const sourceResults = [];
     const executedTools = [];
+    let sourceCallCounter = 0;
     const toolHandler = async (toolName, toolInput) => {
       const externalTool = this.externalTools.find((t) => t.id === toolName);
       if (!externalTool) {
@@ -7400,7 +7473,15 @@ var MainAgent = class {
         aggregation: toolInput.aggregation || "raw"
       };
       logger.info(`[MainAgent] Dispatching SourceAgent for "${externalTool.name}" | intent: "${sourceInput.intent}"`);
-      const sourceAgent = new SourceAgent(externalTool, this.config, this.streamBuffer);
+      sourceCallCounter++;
+      const callId = `src_${sourceCallCounter}`;
+      const taggedBuffer = new TaggedStreamBuffer(
+        callId,
+        this.streamBuffer,
+        externalTool.name,
+        sourceInput.intent
+      );
+      const sourceAgent = new SourceAgent(externalTool, this.config, taggedBuffer);
       const result = await sourceAgent.execute(sourceInput);
       sourceResults.push(result);
       if (result.success) {
@@ -7630,10 +7711,28 @@ Execution time: ${metadata.executionTimeMs}ms
       output += "No data returned.";
       return output;
     }
-    output += `### Results (${data.length} rows)
+    const maxRowsForLLM = Math.min(data.length, 10);
+    const truncatedData = data.slice(0, maxRowsForLLM).map((row) => {
+      const truncatedRow = {};
+      for (const [key, value] of Object.entries(row)) {
+        if (typeof value === "string" && value.length > 200) {
+          truncatedRow[key] = value.substring(0, 200) + "...";
+        } else {
+          truncatedRow[key] = value;
+        }
+      }
+      return truncatedRow;
+    });
+    const truncationNote = data.length > maxRowsForLLM ? `
+(showing ${maxRowsForLLM} of ${data.length} rows)` : "";
+    output += `### Results (${data.length} rows${truncationNote})
 `;
     output += "```json\n";
-    output += JSON.stringify(data, null, 2);
+    try {
+      output += JSON.stringify(truncatedData, null, 2);
+    } catch {
+      output += "[Data could not be serialized]";
+    }
     output += "\n```\n";
     return output;
   }
@@ -7653,7 +7752,9 @@ var DEFAULT_AGENT_CONFIG = {
   sourceAgentModel: "",
   // will use the provider's default model
   maxRetries: 3,
-  maxIterations: 10
+  // 3 retries = 4 total query attempts (1 initial + 3 retries for SQL errors)
+  maxIterations: 15
+  // must accommodate: schema searches (2-3) + query attempts (4) + LLM responses + final
 };
 
 // src/userResponse/utils/component-props-processor.ts