@mastra/lance 1.0.0-beta.9 → 1.0.0

package/dist/index.js

@@ -1,6 +1,6 @@
 import { connect, Index } from '@lancedb/lancedb';
 import { MastraError, ErrorCategory, ErrorDomain } from '@mastra/core/error';
-import { MemoryStorage, TABLE_SCHEMAS, TABLE_THREADS, TABLE_MESSAGES, TABLE_RESOURCES, createStorageErrorId, normalizePerPage, calculatePagination, ScoresStorage, SCORERS_SCHEMA, TABLE_SCORERS, WorkflowsStorage, TABLE_WORKFLOW_SNAPSHOT, MastraStorage, createVectorErrorId, getDefaultValue, ensureDate } from '@mastra/core/storage';
+import { MemoryStorage, TABLE_SCHEMAS, TABLE_THREADS, TABLE_MESSAGES, TABLE_RESOURCES, createStorageErrorId, normalizePerPage, calculatePagination, ScoresStorage, SCORERS_SCHEMA, TABLE_SCORERS, WorkflowsStorage, ensureDate, TABLE_WORKFLOW_SNAPSHOT, MastraCompositeStore, createVectorErrorId, getDefaultValue } from '@mastra/core/storage';
 import { MessageList } from '@mastra/core/agent';
 import { MastraBase } from '@mastra/core/base';
 import { Utf8, Float64, Binary, Float32, Int32, Field, Schema } from 'apache-arrow';
@@ -797,11 +797,13 @@ var StoreMemoryLance = class extends MemoryStorage {
       }
       if (filter?.dateRange?.start) {
         const startTime = filter.dateRange.start instanceof Date ? filter.dateRange.start.getTime() : new Date(filter.dateRange.start).getTime();
-        conditions.push(`\`createdAt\` >= ${startTime}`);
+        const startOp = filter.dateRange.startExclusive ? ">" : ">=";
+        conditions.push(`\`createdAt\` ${startOp} ${startTime}`);
       }
       if (filter?.dateRange?.end) {
         const endTime = filter.dateRange.end instanceof Date ? filter.dateRange.end.getTime() : new Date(filter.dateRange.end).getTime();
-        conditions.push(`\`createdAt\` <= ${endTime}`);
+        const endOp = filter.dateRange.endExclusive ? "<" : "<=";
+        conditions.push(`\`createdAt\` ${endOp} ${endTime}`);
       }
       const whereClause = conditions.join(" AND ");
       const total = await table.countRows(whereClause);
@@ -948,27 +950,63 @@ var StoreMemoryLance = class extends MemoryStorage {
       );
     }
   }
-  async listThreadsByResourceId(args) {
+  async listThreads(args) {
+    const { page = 0, perPage: perPageInput, orderBy, filter } = args;
+    try {
+      this.validatePaginationInput(page, perPageInput ?? 100);
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: createStorageErrorId("LANCE", "LIST_THREADS", "INVALID_PAGE"),
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.USER,
+          details: { page, ...perPageInput !== void 0 && { perPage: perPageInput } }
+        },
+        error instanceof Error ? error : new Error("Invalid pagination parameters")
+      );
+    }
+    const perPage = normalizePerPage(perPageInput, 100);
+    try {
+      this.validateMetadataKeys(filter?.metadata);
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: createStorageErrorId("LANCE", "LIST_THREADS", "INVALID_METADATA_KEY"),
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.USER,
+          details: { metadataKeys: filter?.metadata ? Object.keys(filter.metadata).join(", ") : "" }
+        },
+        error instanceof Error ? error : new Error("Invalid metadata key")
+      );
+    }
     try {
-      const { resourceId, page = 0, perPage: perPageInput, orderBy } = args;
-      const perPage = normalizePerPage(perPageInput, 100);
-      if (page < 0) {
-        throw new MastraError(
-          {
-            id: createStorageErrorId("LANCE", "LIST_THREADS_BY_RESOURCE_ID", "INVALID_PAGE"),
-            domain: ErrorDomain.STORAGE,
-            category: ErrorCategory.USER,
-            details: { page }
-          },
-          new Error("page must be >= 0")
-        );
-      }
       const { offset, perPage: perPageForResponse } = calculatePagination(page, perPageInput, perPage);
       const { field, direction } = this.parseOrderBy(orderBy);
       const table = await this.client.openTable(TABLE_THREADS);
-      const total = await table.countRows(`\`resourceId\` = '${this.escapeSql(resourceId)}'`);
-      const query = table.query().where(`\`resourceId\` = '${this.escapeSql(resourceId)}'`);
-      const records = await query.toArray();
+      const whereClauses = [];
+      if (filter?.resourceId) {
+        whereClauses.push(`\`resourceId\` = '${this.escapeSql(filter.resourceId)}'`);
+      }
+      const whereClause = whereClauses.length > 0 ? whereClauses.join(" AND ") : "";
+      const query = whereClause ? table.query().where(whereClause) : table.query();
+      let records = await query.toArray();
+      if (filter?.metadata && Object.keys(filter.metadata).length > 0) {
+        records = records.filter((record) => {
+          if (!record.metadata) return false;
+          let recordMeta;
+          if (typeof record.metadata === "string") {
+            try {
+              recordMeta = JSON.parse(record.metadata);
+            } catch {
+              return false;
+            }
+          } else {
+            recordMeta = record.metadata;
+          }
+          return Object.entries(filter.metadata).every(([key, value]) => recordMeta[key] === value);
+        });
+      }
+      const total = records.length;
       records.sort((a, b) => {
         const aValue = ["createdAt", "updatedAt"].includes(field) ? new Date(a[field]).getTime() : a[field];
         const bValue = ["createdAt", "updatedAt"].includes(field) ? new Date(b[field]).getTime() : b[field];
@@ -995,7 +1033,7 @@ var StoreMemoryLance = class extends MemoryStorage {
     } catch (error) {
       throw new MastraError(
         {
-          id: createStorageErrorId("LANCE", "LIST_THREADS_BY_RESOURCE_ID", "FAILED"),
+          id: createStorageErrorId("LANCE", "LIST_THREADS", "FAILED"),
           domain: ErrorDomain.STORAGE,
           category: ErrorCategory.THIRD_PARTY
         },
@@ -1601,24 +1639,6 @@ var StoreScoresLance = class extends ScoresStorage {
 function escapeSql(str) {
   return str.replace(/'/g, "''");
 }
-function parseWorkflowRun(row) {
-  let parsedSnapshot = row.snapshot;
-  if (typeof parsedSnapshot === "string") {
-    try {
-      parsedSnapshot = JSON.parse(row.snapshot);
-    } catch (e) {
-      console.warn(`Failed to parse snapshot for workflow ${row.workflow_name}: ${e}`);
-    }
-  }
-  return {
-    workflowName: row.workflow_name,
-    runId: row.run_id,
-    snapshot: parsedSnapshot,
-    createdAt: ensureDate(row.createdAt),
-    updatedAt: ensureDate(row.updatedAt),
-    resourceId: row.resourceId
-  };
-}
 var StoreWorkflowsLance = class extends WorkflowsStorage {
   client;
   #db;
@@ -1628,6 +1648,24 @@ var StoreWorkflowsLance = class extends WorkflowsStorage {
     this.client = client;
     this.#db = new LanceDB({ client });
   }
+  parseWorkflowRun(row) {
+    let parsedSnapshot = row.snapshot;
+    if (typeof parsedSnapshot === "string") {
+      try {
+        parsedSnapshot = JSON.parse(row.snapshot);
+      } catch (e) {
+        this.logger.warn(`Failed to parse snapshot for workflow ${row.workflow_name}: ${e}`);
+      }
+    }
+    return {
+      workflowName: row.workflow_name,
+      runId: row.run_id,
+      snapshot: parsedSnapshot,
+      createdAt: ensureDate(row.createdAt),
+      updatedAt: ensureDate(row.updatedAt),
+      resourceId: row.resourceId
+    };
+  }
   async init() {
     const schema = TABLE_SCHEMAS[TABLE_WORKFLOW_SNAPSHOT];
     await this.#db.createTable({ tableName: TABLE_WORKFLOW_SNAPSHOT, schema });
@@ -1759,7 +1797,7 @@ var StoreWorkflowsLance = class extends WorkflowsStorage {
       const records = await query.toArray();
       if (records.length === 0) return null;
       const record = records[0];
-      return parseWorkflowRun(record);
+      return this.parseWorkflowRun(record);
     } catch (error) {
       throw new MastraError(
         {
@@ -1836,7 +1874,7 @@ var StoreWorkflowsLance = class extends WorkflowsStorage {
       }
       const records = await query.toArray();
       return {
-        runs: records.map((record) => parseWorkflowRun(record)),
+        runs: records.map((record) => this.parseWorkflowRun(record)),
         total: total || records.length
       };
     } catch (error) {
@@ -1854,7 +1892,7 @@ var StoreWorkflowsLance = class extends WorkflowsStorage {
 };
 
 // src/storage/index.ts
-var LanceStorage = class _LanceStorage extends MastraStorage {
+var LanceStorage = class _LanceStorage extends MastraCompositeStore {
   stores;
   lanceClient;
   /**
@@ -2326,6 +2364,7 @@ var LanceVectorStore = class _LanceVectorStore extends MastraVector {
   }
   async query({
     tableName,
+    indexName,
     queryVector,
     filter,
     includeVector = false,
@@ -2333,12 +2372,13 @@ var LanceVectorStore = class _LanceVectorStore extends MastraVector {
     columns = [],
     includeAllColumns = false
   }) {
+    const resolvedTableName = tableName ?? indexName;
     try {
       if (!this.lanceClient) {
         throw new Error("LanceDB client not initialized. Use LanceVectorStore.create() to create an instance");
       }
-      if (!tableName) {
-        throw new Error("tableName is required");
+      if (!resolvedTableName) {
+        throw new Error("tableName or indexName is required");
       }
       if (!queryVector) {
         throw new Error("queryVector is required");
@@ -2349,25 +2389,30 @@ var LanceVectorStore = class _LanceVectorStore extends MastraVector {
           id: createVectorErrorId("LANCE", "QUERY", "INVALID_ARGS"),
           domain: ErrorDomain.STORAGE,
           category: ErrorCategory.USER,
-          text: "LanceDB client not initialized. Use LanceVectorStore.create() to create an instance",
-          details: { tableName }
+          text: error instanceof Error ? error.message : "Invalid query arguments",
+          details: { tableName: resolvedTableName }
         },
         error
       );
     }
     try {
-      const table = await this.lanceClient.openTable(tableName);
-      const selectColumns = [...columns];
-      if (!selectColumns.includes("id")) {
-        selectColumns.push("id");
+      const tables = await this.lanceClient.tableNames();
+      if (!tables.includes(resolvedTableName)) {
+        this.logger.debug(`Table ${resolvedTableName} does not exist. Returning empty results.`);
+        return [];
       }
+      const table = await this.lanceClient.openTable(resolvedTableName);
       let query = table.search(queryVector);
       if (filter && Object.keys(filter).length > 0) {
         const whereClause = this.filterTranslator(filter);
         this.logger.debug(`Where clause generated: ${whereClause}`);
         query = query.where(whereClause);
       }
-      if (!includeAllColumns && selectColumns.length > 0) {
+      if (!includeAllColumns && columns.length > 0) {
+        const selectColumns = [...columns];
+        if (!selectColumns.includes("id")) {
+          selectColumns.push("id");
+        }
         query = query.select(selectColumns);
       }
       query = query.limit(topK);
@@ -2397,7 +2442,7 @@ var LanceVectorStore = class _LanceVectorStore extends MastraVector {
           id: createVectorErrorId("LANCE", "QUERY", "FAILED"),
           domain: ErrorDomain.STORAGE,
           category: ErrorCategory.THIRD_PARTY,
-          details: { tableName, includeVector, columnsCount: columns?.length, includeAllColumns }
+          details: { tableName: resolvedTableName, includeVector, columnsCount: columns?.length, includeAllColumns }
         },
         error
       );
@@ -2432,13 +2477,14 @@ var LanceVectorStore = class _LanceVectorStore extends MastraVector {
     const translator = new LanceFilterTranslator();
     return translator.translate(prefixedFilter);
   }
-  async upsert({ tableName, vectors, metadata = [], ids = [] }) {
+  async upsert({ tableName, indexName, vectors, metadata = [], ids = [] }) {
+    const resolvedTableName = tableName ?? indexName;
     try {
       if (!this.lanceClient) {
         throw new Error("LanceDB client not initialized. Use LanceVectorStore.create() to create an instance");
       }
-      if (!tableName) {
-        throw new Error("tableName is required");
+      if (!resolvedTableName) {
+        throw new Error("tableName or indexName is required");
       }
       if (!vectors || !Array.isArray(vectors) || vectors.length === 0) {
         throw new Error("vectors array is required and must not be empty");
@@ -2449,18 +2495,21 @@ var LanceVectorStore = class _LanceVectorStore extends MastraVector {
           id: createVectorErrorId("LANCE", "UPSERT", "INVALID_ARGS"),
           domain: ErrorDomain.STORAGE,
           category: ErrorCategory.USER,
-          text: "LanceDB client not initialized. Use LanceVectorStore.create() to create an instance",
-          details: { tableName }
+          text: error instanceof Error ? error.message : "Invalid upsert arguments",
+          details: { tableName: resolvedTableName }
         },
         error
       );
     }
     try {
       const tables = await this.lanceClient.tableNames();
-      if (!tables.includes(tableName)) {
-        throw new Error(`Table ${tableName} does not exist`);
+      const tableExists = tables.includes(resolvedTableName);
+      let table = null;
+      if (!tableExists) {
+        this.logger.debug(`Table ${resolvedTableName} does not exist. Creating it with the first upsert data.`);
+      } else {
+        table = await this.lanceClient.openTable(resolvedTableName);
       }
-      const table = await this.lanceClient.openTable(tableName);
       const vectorIds = ids.length === vectors.length ? ids : vectors.map((_, i) => ids[i] || crypto.randomUUID());
       const data = vectors.map((vector, i) => {
         const id = String(vectorIds[i]);
@@ -2477,7 +2526,42 @@ var LanceVectorStore = class _LanceVectorStore extends MastraVector {
         }
         return rowData;
       });
-      await table.add(data, { mode: "overwrite" });
+      if (table !== null) {
+        const rowCount = await table.countRows();
+        const schema = await table.schema();
+        const existingColumns = new Set(schema.fields.map((f) => f.name));
+        const dataColumns = new Set(Object.keys(data[0] || {}));
+        const extraColumns = [...dataColumns].filter((col) => !existingColumns.has(col));
+        const missingSchemaColumns = [...existingColumns].filter((col) => !dataColumns.has(col));
+        const hasSchemaMismatch = extraColumns.length > 0 || missingSchemaColumns.length > 0;
+        if (rowCount === 0 && extraColumns.length > 0) {
+          this.logger.warn(
+            `Table ${resolvedTableName} is empty and data has extra columns ${extraColumns.join(", ")}. Recreating with new schema.`
+          );
+          await this.lanceClient.dropTable(resolvedTableName);
+          await this.lanceClient.createTable(resolvedTableName, data);
+        } else if (hasSchemaMismatch) {
+          if (extraColumns.length > 0) {
+            this.logger.warn(
+              `Table ${resolvedTableName} has ${rowCount} rows. Columns ${extraColumns.join(", ")} will be dropped from upsert.`
+            );
+          }
+          const schemaFieldNames = schema.fields.map((f) => f.name);
+          const normalizedData = data.map((row) => {
+            const normalized = {};
+            for (const col of schemaFieldNames) {
+              normalized[col] = col in row ? row[col] : null;
+            }
+            return normalized;
+          });
+          await table.mergeInsert("id").whenMatchedUpdateAll().whenNotMatchedInsertAll().execute(normalizedData);
+        } else {
+          await table.mergeInsert("id").whenMatchedUpdateAll().whenNotMatchedInsertAll().execute(data);
+        }
+      } else {
+        this.logger.debug(`Creating table ${resolvedTableName} with initial data`);
+        await this.lanceClient.createTable(resolvedTableName, data);
+      }
       return vectorIds;
     } catch (error) {
       throw new MastraError(
@@ -2485,7 +2569,12 @@ var LanceVectorStore = class _LanceVectorStore extends MastraVector {
           id: createVectorErrorId("LANCE", "UPSERT", "FAILED"),
           domain: ErrorDomain.STORAGE,
           category: ErrorCategory.THIRD_PARTY,
-          details: { tableName, vectorCount: vectors.length, metadataCount: metadata.length, idsCount: ids.length }
+          details: {
+            tableName: resolvedTableName,
+            vectorCount: vectors.length,
+            metadataCount: metadata.length,
+            idsCount: ids.length
+          }
         },
         error
       );
@@ -2582,7 +2671,17 @@ var LanceVectorStore = class _LanceVectorStore extends MastraVector {
     }
   }
   /**
-   * indexName is actually a column name in a table in lanceDB
+   * Creates a vector index on a table.
+   *
+   * The behavior of `indexName` depends on whether `tableName` is provided:
+   * - With `tableName`: `indexName` is the column to index (advanced use case)
+   * - Without `tableName`: `indexName` becomes the table name, and 'vector' is used as the column (Memory compatibility)
+   *
+   * @param tableName - Optional table name. If not provided, defaults to indexName.
+   * @param indexName - The index/column name, or table name if tableName is not provided.
+   * @param dimension - Vector dimension size.
+   * @param metric - Distance metric: 'cosine', 'euclidean', or 'dotproduct'.
+   * @param indexConfig - Optional index configuration.
    */
   async createIndex({
     tableName,
@@ -2591,13 +2690,12 @@ var LanceVectorStore = class _LanceVectorStore extends MastraVector {
     metric = "cosine",
     indexConfig = {}
   }) {
+    const resolvedTableName = tableName ?? indexName;
+    const columnToIndex = tableName ? indexName : "vector";
     try {
       if (!this.lanceClient) {
         throw new Error("LanceDB client not initialized. Use LanceVectorStore.create() to create an instance");
       }
-      if (!tableName) {
-        throw new Error("tableName is required");
-      }
       if (!indexName) {
         throw new Error("indexName is required");
       }
@@ -2610,19 +2708,33 @@ var LanceVectorStore = class _LanceVectorStore extends MastraVector {
           id: createVectorErrorId("LANCE", "CREATE_INDEX", "INVALID_ARGS"),
           domain: ErrorDomain.STORAGE,
           category: ErrorCategory.USER,
-          details: { tableName: tableName || "", indexName, dimension, metric }
+          details: { tableName: resolvedTableName, indexName, dimension, metric }
         },
         err
       );
     }
     try {
       const tables = await this.lanceClient.tableNames();
-      if (!tables.includes(tableName)) {
-        throw new Error(
-          `Table ${tableName} does not exist. Please create the table first by calling createTable() method.`
+      let table;
+      if (!tables.includes(resolvedTableName)) {
+        this.logger.debug(
+          `Table ${resolvedTableName} does not exist. Creating empty table with dimension ${dimension}.`
         );
+        const initVector = new Array(dimension).fill(0);
+        table = await this.lanceClient.createTable(resolvedTableName, [{ id: "__init__", vector: initVector }]);
+        try {
+          await table.delete("id = '__init__'");
+        } catch (deleteError) {
+          this.logger.warn(
+            `Failed to delete initialization row from ${resolvedTableName}. Subsequent queries may include '__init__' row.`,
+            deleteError
+          );
+        }
+        this.logger.debug(`Table ${resolvedTableName} created. Index creation deferred until data is available.`);
+        return;
+      } else {
+        table = await this.lanceClient.openTable(resolvedTableName);
       }
-      const table = await this.lanceClient.openTable(tableName);
       let metricType;
       if (metric === "euclidean") {
         metricType = "l2";
@@ -2631,8 +2743,15 @@ var LanceVectorStore = class _LanceVectorStore extends MastraVector {
       } else if (metric === "cosine") {
         metricType = "cosine";
       }
+      const rowCount = await table.countRows();
+      if (rowCount < 256) {
+        this.logger.warn(
+          `Table ${resolvedTableName} has ${rowCount} rows, which is below the 256 row minimum for index creation. Skipping index creation.`
+        );
+        return;
+      }
       if (indexConfig.type === "ivfflat") {
-        await table.createIndex(indexName, {
+        await table.createIndex(columnToIndex, {
           config: Index.ivfPq({
             numPartitions: indexConfig.numPartitions || 128,
             numSubVectors: indexConfig.numSubVectors || 16,
@@ -2641,7 +2760,7 @@ var LanceVectorStore = class _LanceVectorStore extends MastraVector {
         });
       } else {
         this.logger.debug("Creating HNSW PQ index with config:", indexConfig);
-        await table.createIndex(indexName, {
+        await table.createIndex(columnToIndex, {
           config: Index.hnswPq({
             m: indexConfig?.hnsw?.m || 16,
             efConstruction: indexConfig?.hnsw?.efConstruction || 100,
@@ -2655,7 +2774,7 @@ var LanceVectorStore = class _LanceVectorStore extends MastraVector {
           id: createVectorErrorId("LANCE", "CREATE_INDEX", "FAILED"),
           domain: ErrorDomain.STORAGE,
           category: ErrorCategory.THIRD_PARTY,
-          details: { tableName: tableName || "", indexName, dimension }
+          details: { tableName: resolvedTableName, indexName, dimension }
         },
         error
       );
@@ -2953,7 +3072,7 @@ var LanceVectorStore = class _LanceVectorStore extends MastraVector {
               }
               return rowData;
             });
-            await table.add(updatedRecords, { mode: "overwrite" });
+            await table.mergeInsert("id").whenMatchedUpdateAll().whenNotMatchedInsertAll().execute(updatedRecords);
             return;
           }
         } catch (err) {