@mastra/lance 1.0.0-beta.10 → 1.0.0-beta.11

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, MastraStorage, 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';
@@ -1603,24 +1603,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;
@@ -1630,6 +1612,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 });
@@ -1761,7 +1761,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(
         {
@@ -1838,7 +1838,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) {
@@ -2328,6 +2328,7 @@ var LanceVectorStore = class _LanceVectorStore extends MastraVector {
   }
   async query({
     tableName,
+    indexName,
     queryVector,
     filter,
     includeVector = false,
@@ -2335,12 +2336,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");
@@ -2351,25 +2353,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);
@@ -2399,7 +2406,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
       );
@@ -2434,13 +2441,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");
@@ -2451,18 +2459,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]);
@@ -2479,7 +2490,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(
@@ -2487,7 +2533,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
       );
@@ -2584,7 +2635,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,
@@ -2593,13 +2654,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");
       }
@@ -2612,19 +2672,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";
@@ -2633,8 +2707,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,
@@ -2643,7 +2724,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,
@@ -2657,7 +2738,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
       );
@@ -2955,7 +3036,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) {