@deepagents/text2sql 0.12.1 → 0.13.1

package/dist/index.js

@@ -392,6 +392,7 @@ import z2 from "zod";
 import { toState } from "@deepagents/agent";
 
 // packages/context/dist/index.js
+import { mergeWith } from "lodash-es";
 import { encode } from "gpt-tokenizer";
 import { generateId } from "ai";
 import pluralize from "pluralize";
@@ -967,6 +968,8 @@ var ContextEngine = class {
   #branch = null;
   #chatData = null;
   #initialized = false;
+  /** Initial metadata to merge on first initialization */
+  #initialMetadata;
   constructor(options) {
     if (!options.chatId) {
       throw new Error("chatId is required");
@@ -978,6 +981,7 @@ var ContextEngine = class {
     this.#chatId = options.chatId;
     this.#userId = options.userId;
     this.#branchName = "main";
+    this.#initialMetadata = options.metadata;
   }
   /**
    * Initialize the chat and branch if they don't exist.
@@ -990,6 +994,15 @@ var ContextEngine = class {
       id: this.#chatId,
       userId: this.#userId
     });
+    if (this.#initialMetadata) {
+      this.#chatData = await this.#store.updateChat(this.#chatId, {
+        metadata: {
+          ...this.#chatData.metadata,
+          ...this.#initialMetadata
+        }
+      });
+      this.#initialMetadata = void 0;
+    }
     this.#branch = await this.#store.getActiveBranch(this.#chatId);
     this.#initialized = true;
   }
@@ -1463,6 +1476,36 @@ var ContextEngine = class {
     }
     this.#chatData = await this.#store.updateChat(this.#chatId, storeUpdates);
   }
+  /**
+   * Track token usage for the current chat.
+   * Accumulates usage metrics in chat.metadata.usage.
+   *
+   * @param usage - Token usage from AI SDK (LanguageModelUsage)
+   *
+   * @example
+   * ```ts
+   * // In onFinish callback
+   * const usage = await result.totalUsage;
+   * await context.trackUsage(usage);
+   * ```
+   */
+  async trackUsage(usage) {
+    await this.#ensureInitialized();
+    const freshChatData = await this.#store.getChat(this.#chatId);
+    const currentUsage = freshChatData?.metadata?.usage ?? {};
+    const updatedUsage = mergeWith(
+      {},
+      currentUsage,
+      usage,
+      (a, b) => typeof a === "number" || typeof b === "number" ? (a ?? 0) + (b ?? 0) : void 0
+    );
+    this.#chatData = await this.#store.updateChat(this.#chatId, {
+      metadata: {
+        ...freshChatData?.metadata,
+        usage: updatedUsage
+      }
+    });
+  }
   /**
    * Consolidate context fragments (no-op for now).
    *
@@ -1742,6 +1785,7 @@ var errorRecoveryGuardrail = {
         "My response format was invalid. Let me try again with a properly formatted response."
       );
     }
+    console.dir({ part }, { depth: null });
     return logAndFail(
       "Unknown error",
       `An error occurred: ${errorText}. Let me try a different approach.`
@@ -1929,12 +1973,20 @@ var SqliteContextStore = class extends ContextStore {
   }
   async listChats(options) {
     const params = [];
-    let whereClause = "";
+    const whereClauses = [];
     let limitClause = "";
     if (options?.userId) {
-      whereClause = "WHERE c.userId = ?";
+      whereClauses.push("c.userId = ?");
       params.push(options.userId);
     }
+    if (options?.metadata) {
+      whereClauses.push(`json_extract(c.metadata, '$.' || ?) = ?`);
+      params.push(options.metadata.key);
+      params.push(
+        typeof options.metadata.value === "boolean" ? options.metadata.value ? 1 : 0 : options.metadata.value
+      );
+    }
+    const whereClause = whereClauses.length > 0 ? `WHERE ${whereClauses.join(" AND ")}` : "";
     if (options?.limit !== void 0) {
       limitClause = " LIMIT ?";
       params.push(options.limit);
@@ -1948,6 +2000,7 @@ var SqliteContextStore = class extends ContextStore {
           c.id,
           c.userId,
           c.title,
+          c.metadata,
           c.createdAt,
           c.updatedAt,
           COUNT(DISTINCT m.id) as messageCount,
@@ -1963,6 +2016,7 @@ var SqliteContextStore = class extends ContextStore {
       id: row.id,
       userId: row.userId,
       title: row.title ?? void 0,
+      metadata: row.metadata ? JSON.parse(row.metadata) : void 0,
       messageCount: row.messageCount,
       branchCount: row.branchCount,
       createdAt: row.createdAt,
@@ -1992,11 +2046,17 @@ var SqliteContextStore = class extends ContextStore {
   // Message Operations (Graph Nodes)
   // ==========================================================================
   async addMessage(message2) {
-    const existingParent = message2.parentId === message2.id ? this.#db.prepare("SELECT parentId FROM messages WHERE id = ?").get(message2.id) : void 0;
-    const parentId = message2.parentId === message2.id ? existingParent?.parentId ?? null : message2.parentId;
     this.#db.prepare(
       `INSERT INTO messages (id, chatId, parentId, name, type, data, createdAt)
-         VALUES (?, ?, ?, ?, ?, ?, ?)
+         VALUES (
+           ?1,
+           ?2,
+           CASE WHEN ?3 = ?1 THEN (SELECT parentId FROM messages WHERE id = ?1) ELSE ?3 END,
+           ?4,
+           ?5,
+           ?6,
+           ?7
+         )
          ON CONFLICT(id) DO UPDATE SET
            name = excluded.name,
            type = excluded.type,
@@ -2004,7 +2064,7 @@ var SqliteContextStore = class extends ContextStore {
     ).run(
       message2.id,
       message2.chatId,
-      parentId,
+      message2.parentId,
       message2.name,
       message2.type ?? null,
       JSON.stringify(message2.data),
@@ -2565,10 +2625,10 @@ var repairToolCall = async ({
   if (NoSuchToolError.isInstance(error)) {
     return null;
   }
-  const tool3 = tools3[toolCall.toolName];
+  const tool4 = tools3[toolCall.toolName];
   const { output } = await generateText({
     model: groq("openai/gpt-oss-20b"),
-    output: Output.object({ schema: tool3.inputSchema }),
+    output: Output.object({ schema: tool4.inputSchema }),
     prompt: [
       `The model tried to call the tool "${toolCall.toolName}" with the following inputs:`,
       JSON.stringify(toolCall.input),
@@ -2699,6 +2759,7 @@ var fragments = [
 var developer_agent_default = { tools, fragments };
 
 // packages/text2sql/src/lib/agents/result-tools.ts
+import "ai";
 import { createBashTool as createBashTool2 } from "bash-tool";
 import chalk3 from "chalk";
 import {
@@ -2762,12 +2823,20 @@ function createSqlCommand(adapter) {
         try {
           const rows = await adapter.execute(query);
           const rowsArray = Array.isArray(rows) ? rows : [];
-          const filePath = `/results/${v7()}.json`;
-          await ctx.fs.writeFile(filePath, JSON.stringify(rowsArray, null, 2));
+          const content = JSON.stringify(rowsArray, null, 2);
+          const filename = `${v7()}.json`;
+          const isolatedPath = `/results/${filename}`;
+          const sharedPath = `/artifacts/${filename}`;
+          await Promise.all([
+            ctx.fs.writeFile(isolatedPath, content),
+            // Current turn's isolated copy
+            ctx.fs.writeFile(sharedPath, content)
+            // Shared copy for cross-turn access
+          ]);
           const columns = rowsArray.length > 0 ? Object.keys(rowsArray[0]) : [];
           return {
             stdout: [
-              filePath,
+              `results stored in ${sharedPath}`,
               `columns: ${columns.join(", ") || "(none)"}`,
               `rows: ${rowsArray.length}`
             ].join("\n") + "\n",
@@ -3085,12 +3154,12 @@ var suggestionsAgent = agent2({
 });
 
 // packages/text2sql/src/lib/agents/text2sql.agent.ts
-import { tool as tool2 } from "ai";
+import { tool as tool3 } from "ai";
 import z5 from "zod";
 import { toState as toState2 } from "@deepagents/agent";
 import { scratchpad_tool } from "@deepagents/toolbox";
 var tools2 = {
-  validate_query: tool2({
+  validate_query: tool3({
     description: `Validate SQL query syntax before execution. Use this to check if your SQL is valid before running db_query. This helps catch errors early and allows you to correct the query if needed.`,
     inputSchema: z5.object({
       sql: z5.string().describe("The SQL query to validate.")
@@ -3104,7 +3173,7 @@ var tools2 = {
       return "Query is valid.";
     }
   }),
-  db_query: tool2({
+  db_query: tool3({
     description: `Internal tool to fetch data from the store's database. Write a SQL query to retrieve the information needed to answer the user's question. The results will be returned as data that you can then present to the user in natural language.`,
     inputSchema: z5.object({
       reasoning: z5.string().describe(
@@ -3441,6 +3510,10 @@ function reasoningFramework() {
       principle({
         title: "Inhibit your response",
         description: "Only take an action after all the above reasoning is completed. Once you've taken an action, you cannot take it back."
+      }),
+      principle({
+        title: "Continuous self-monitoring",
+        description: "Constantly evaluate your own reasoning process for any gaps, biases, or errors. Apply the above principles iteratively as needed."
       })
     )
   ];
@@ -3654,7 +3727,31 @@ Action: Ask user: "Top by what metric\u2014total revenue, number of orders, or m
         "Maintain consistency with previous query structure when possible."
       ],
       notes: "If reference is ambiguous, ask which previous result or entity the user means."
-    })
+    }),
+    fragment(
+      "Bash tool usage",
+      workflow({
+        task: "Query execution",
+        steps: [
+          'Execute SQL through bash tool: sql run "SELECT ..."',
+          "Read the output: file path, column names, and row count.",
+          "Use column names to construct jq filters: cat <path> | jq '.[] | {col1, col2}'",
+          "For large results, slice first: cat <path> | jq '.[:10]'"
+        ]
+      }),
+      hint(
+        `You cannot access sql through a tool, it'll fail so the proper way to access it is through the bash tool using "sql run" and "sql validate" commands.`
+      ),
+      hint(
+        "The sql command outputs: file path, column names (comma-separated), and row count. Use column names to construct precise jq queries."
+      ),
+      hint(
+        'This is virtual bash environment and "sql" commands proxy to the database hence you cannot access sql files directly.'
+      ),
+      hint(
+        "If a query fails, the sql command returns an error message in stderr."
+      )
+    )
   ];
   if (date === "strict") {
     baseTeachings.push(
@@ -3687,14 +3784,12 @@ var Text2Sql = class {
   #config;
   constructor(config) {
     this.#config = {
+      teachingsOptions: config.teachingsOptions,
       adapter: config.adapter,
-      store: config.store,
-      instructions: [
-        ...guidelines(config.teachingsOptions),
-        ...config.instructions ?? []
-      ],
+      context: config.context,
       tools: config.tools ?? {},
       model: config.model,
+      transform: config.transform,
       introspection: new JsonCache(
         "introspection-" + config.version
       )
@@ -3706,7 +3801,7 @@ var Text2Sql = class {
       input,
       adapter: this.#config.adapter,
       schemaFragments,
-      instructions: this.#config.instructions,
+      instructions: [],
       model: this.#config.model
     });
     return result.sql;
@@ -3766,40 +3861,11 @@ var Text2Sql = class {
       })
     ];
   }
-  async chat(messages, params) {
-    const schemaFragments = await this.index();
-    const context = new ContextEngine({
-      store: this.#config.store,
-      chatId: params.chatId,
-      userId: params.userId
-    }).set(
-      ...schemaFragments,
-      ...this.#buildRenderingInstructions(),
-      fragment(
-        "Bash tool usage",
-        workflow({
-          task: "Query execution",
-          steps: [
-            'Execute SQL through bash tool: sql run "SELECT ..."',
-            "Read the output: file path, column names, and row count.",
-            "Use column names to construct jq filters: cat <path> | jq '.[] | {col1, col2}'",
-            "For large results, slice first: cat <path> | jq '.[:10]'"
-          ]
-        }),
-        hint(
-          `You cannot access sql through a tool, it'll fail so the proper way to access it is through the bash tool using "sql run" and "sql validate" commands.`
-        ),
-        hint(
-          "The sql command outputs: file path, column names (comma-separated), and row count. Use column names to construct precise jq queries."
-        ),
-        hint(
-          'This is virtual bash environment and "sql" commands proxy to the database hence you cannot access sql files directly.'
-        ),
-        hint(
-          "If a query fails, the sql command returns an error message in stderr."
-        )
-      ),
-      ...this.#config.instructions
+  async chat(messages) {
+    const context = this.#config.context(
+      ...guidelines(this.#config.teachingsOptions),
+      ...await this.index(),
+      ...this.#buildRenderingInstructions()
     );
     const userMsg = messages.at(-1);
     if (userMsg) {
@@ -3810,7 +3876,7 @@ var Text2Sql = class {
     const skillMounts = context.getSkillMounts();
     const { bash } = await createResultTools({
       adapter: this.#config.adapter,
-      chatId: params.chatId,
+      chatId: context.chatId,
       messageId,
       skillMounts
     });
@@ -3825,7 +3891,10 @@ var Text2Sql = class {
       guardrails: [errorRecoveryGuardrail],
       maxGuardrailRetries: 3
     });
-    const result = await chatAgent.stream({});
+    const result = await chatAgent.stream(
+      {},
+      { transform: this.#config.transform }
+    );
     return result.toUIMessageStream({
       onError: (error) => this.#formatError(error),
       sendStart: true,
@@ -3837,23 +3906,16 @@ var Text2Sql = class {
       onFinish: async ({ responseMessage }) => {
         context.set(assistant(responseMessage));
         await context.save();
+        await context.trackUsage(await result.totalUsage);
       }
     });
   }
-  /**
-   * Developer chat interface - power-user mode for SQL generation.
-   * Uses db_query tool for direct SQL execution (LLM writes SQL).
-   */
-  async developer(messages, params) {
-    const schemaFragments = await this.index();
-    const context = new ContextEngine({
-      store: this.#config.store,
-      chatId: params.chatId,
-      userId: params.userId
-    }).set(
+  async developer(messages) {
+    const context = this.#config.context(
+      ...guidelines(this.#config.teachingsOptions),
       ...developer_agent_default.fragments,
-      ...this.#config.instructions,
-      ...schemaFragments
+      ...await this.index(),
+      ...this.#buildRenderingInstructions()
     );
     const userMsg = messages.at(-1);
     if (userMsg) {
@@ -3879,6 +3941,7 @@ var Text2Sql = class {
       onFinish: async ({ responseMessage }) => {
         context.set(assistant(responseMessage));
         await context.save();
+        await context.trackUsage(await result.totalUsage);
       }
     });
   }