@deepagents/text2sql 0.12.1 → 0.13.1

package/dist/lib/adapters/groundings/index.js

@@ -226,6 +226,7 @@ import z from "zod";
 import "@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";
@@ -801,6 +802,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");
@@ -812,6 +815,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.
@@ -824,6 +828,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;
   }
@@ -1297,6 +1310,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).
    *
@@ -1592,12 +1635,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);
@@ -1611,6 +1662,7 @@ var SqliteContextStore = class extends ContextStore {
           c.id,
           c.userId,
           c.title,
+          c.metadata,
           c.createdAt,
           c.updatedAt,
           COUNT(DISTINCT m.id) as messageCount,
@@ -1626,6 +1678,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,
@@ -1655,11 +1708,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,
@@ -1667,7 +1726,7 @@ var SqliteContextStore = class extends ContextStore {
     ).run(
       message2.id,
       message2.chatId,
-      parentId,
+      message2.parentId,
       message2.name,
       message2.type ?? null,
       JSON.stringify(message2.data),