@mastra/lance 1.0.0-beta.7 → 1.0.0-beta.9

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/dist/index.js CHANGED
@@ -1,7 +1,8 @@
1
1
  import { connect, Index } from '@lancedb/lancedb';
2
2
  import { MastraError, ErrorCategory, ErrorDomain } from '@mastra/core/error';
3
- import { MastraStorage, createStorageErrorId, createVectorErrorId, StoreOperations, MemoryStorage, TABLE_THREADS, TABLE_MESSAGES, normalizePerPage, calculatePagination, TABLE_RESOURCES, ScoresStorage, TABLE_SCORERS, WorkflowsStorage, TABLE_WORKFLOW_SNAPSHOT, ensureDate } from '@mastra/core/storage';
3
+ 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';
4
4
  import { MessageList } from '@mastra/core/agent';
5
+ import { MastraBase } from '@mastra/core/base';
5
6
  import { Utf8, Float64, Binary, Float32, Int32, Field, Schema } from 'apache-arrow';
6
7
  import { saveScorePayloadSchema } from '@mastra/core/evals';
7
8
  import { MastraVector } from '@mastra/core/vector';
@@ -122,622 +123,503 @@ async function getTableSchema({
122
123
  }
123
124
  }
124
125
 
125
- // src/storage/domains/memory/index.ts
126
- var StoreMemoryLance = class extends MemoryStorage {
126
+ // src/storage/db/index.ts
127
+ function resolveLanceConfig(config) {
128
+ return config.client;
129
+ }
130
+ var LanceDB = class extends MastraBase {
127
131
  client;
128
- operations;
129
- constructor({ client, operations }) {
130
- super();
132
+ constructor({ client }) {
133
+ super({ name: "lance-db" });
131
134
  this.client = client;
132
- this.operations = operations;
133
135
  }
134
- // Utility to escape single quotes in SQL strings
135
- escapeSql(str) {
136
- return str.replace(/'/g, "''");
136
+ getDefaultValue(type) {
137
+ switch (type) {
138
+ case "text":
139
+ return "''";
140
+ case "timestamp":
141
+ return "CURRENT_TIMESTAMP";
142
+ case "integer":
143
+ case "bigint":
144
+ return "0";
145
+ case "jsonb":
146
+ return "'{}'";
147
+ case "uuid":
148
+ return "''";
149
+ default:
150
+ return getDefaultValue(type);
151
+ }
137
152
  }
138
- async getThreadById({ threadId }) {
153
+ async hasColumn(tableName, columnName) {
154
+ const table = await this.client.openTable(tableName);
155
+ const schema = await table.schema();
156
+ return schema.fields.some((field) => field.name === columnName);
157
+ }
158
+ translateSchema(schema) {
159
+ const fields = Object.entries(schema).map(([name, column]) => {
160
+ let arrowType;
161
+ switch (column.type.toLowerCase()) {
162
+ case "text":
163
+ case "uuid":
164
+ arrowType = new Utf8();
165
+ break;
166
+ case "int":
167
+ case "integer":
168
+ arrowType = new Int32();
169
+ break;
170
+ case "bigint":
171
+ arrowType = new Float64();
172
+ break;
173
+ case "float":
174
+ arrowType = new Float32();
175
+ break;
176
+ case "jsonb":
177
+ case "json":
178
+ arrowType = new Utf8();
179
+ break;
180
+ case "binary":
181
+ arrowType = new Binary();
182
+ break;
183
+ case "timestamp":
184
+ arrowType = new Float64();
185
+ break;
186
+ default:
187
+ arrowType = new Utf8();
188
+ }
189
+ return new Field(name, arrowType, column.nullable ?? true);
190
+ });
191
+ return new Schema(fields);
192
+ }
193
+ async createTable({
194
+ tableName,
195
+ schema
196
+ }) {
139
197
  try {
140
- const thread = await this.operations.load({ tableName: TABLE_THREADS, keys: { id: threadId } });
141
- if (!thread) {
142
- return null;
198
+ if (!this.client) {
199
+ throw new Error("LanceDB client not initialized. Call LanceStorage.create() first.");
200
+ }
201
+ if (!tableName) {
202
+ throw new Error("tableName is required for createTable.");
203
+ }
204
+ if (!schema) {
205
+ throw new Error("schema is required for createTable.");
143
206
  }
144
- return {
145
- ...thread,
146
- createdAt: new Date(thread.createdAt),
147
- updatedAt: new Date(thread.updatedAt)
148
- };
149
207
  } catch (error) {
150
208
  throw new MastraError(
151
209
  {
152
- id: createStorageErrorId("LANCE", "GET_THREAD_BY_ID", "FAILED"),
210
+ id: createStorageErrorId("LANCE", "CREATE_TABLE", "INVALID_ARGS"),
153
211
  domain: ErrorDomain.STORAGE,
154
- category: ErrorCategory.THIRD_PARTY
212
+ category: ErrorCategory.USER,
213
+ details: { tableName }
155
214
  },
156
215
  error
157
216
  );
158
217
  }
159
- }
160
- /**
161
- * Saves a thread to the database. This function doesn't overwrite existing threads.
162
- * @param thread - The thread to save
163
- * @returns The saved thread
164
- */
165
- async saveThread({ thread }) {
166
218
  try {
167
- const record = { ...thread, metadata: JSON.stringify(thread.metadata) };
168
- const table = await this.client.openTable(TABLE_THREADS);
169
- await table.add([record], { mode: "append" });
170
- return thread;
219
+ const arrowSchema = this.translateSchema(schema);
220
+ await this.client.createEmptyTable(tableName, arrowSchema);
171
221
  } catch (error) {
222
+ if (error.message?.includes("already exists")) {
223
+ this.logger.debug(`Table '${tableName}' already exists, skipping create`);
224
+ return;
225
+ }
172
226
  throw new MastraError(
173
227
  {
174
- id: createStorageErrorId("LANCE", "SAVE_THREAD", "FAILED"),
228
+ id: createStorageErrorId("LANCE", "CREATE_TABLE", "FAILED"),
175
229
  domain: ErrorDomain.STORAGE,
176
- category: ErrorCategory.THIRD_PARTY
230
+ category: ErrorCategory.THIRD_PARTY,
231
+ details: { tableName }
177
232
  },
178
233
  error
179
234
  );
180
235
  }
181
236
  }
182
- async updateThread({
183
- id,
184
- title,
185
- metadata
186
- }) {
187
- const maxRetries = 5;
188
- for (let attempt = 0; attempt < maxRetries; attempt++) {
189
- try {
190
- const current = await this.getThreadById({ threadId: id });
191
- if (!current) {
192
- throw new Error(`Thread with id ${id} not found`);
193
- }
194
- const mergedMetadata = { ...current.metadata, ...metadata };
195
- const record = {
196
- id,
197
- title,
198
- metadata: JSON.stringify(mergedMetadata),
199
- updatedAt: (/* @__PURE__ */ new Date()).getTime()
200
- };
201
- const table = await this.client.openTable(TABLE_THREADS);
202
- await table.mergeInsert("id").whenMatchedUpdateAll().whenNotMatchedInsertAll().execute([record]);
203
- const updatedThread = await this.getThreadById({ threadId: id });
204
- if (!updatedThread) {
205
- throw new Error(`Failed to retrieve updated thread ${id}`);
206
- }
207
- return updatedThread;
208
- } catch (error) {
209
- if (error.message?.includes("Commit conflict") && attempt < maxRetries - 1) {
210
- const delay = Math.pow(2, attempt) * 10;
211
- await new Promise((resolve) => setTimeout(resolve, delay));
212
- continue;
213
- }
214
- throw new MastraError(
215
- {
216
- id: createStorageErrorId("LANCE", "UPDATE_THREAD", "FAILED"),
217
- domain: ErrorDomain.STORAGE,
218
- category: ErrorCategory.THIRD_PARTY
219
- },
220
- error
221
- );
222
- }
223
- }
224
- throw new MastraError(
225
- {
226
- id: createStorageErrorId("LANCE", "UPDATE_THREAD", "FAILED"),
227
- domain: ErrorDomain.STORAGE,
228
- category: ErrorCategory.THIRD_PARTY
229
- },
230
- new Error("All retries exhausted")
231
- );
232
- }
233
- async deleteThread({ threadId }) {
237
+ async dropTable({ tableName }) {
234
238
  try {
235
- const table = await this.client.openTable(TABLE_THREADS);
236
- await table.delete(`id = '${threadId}'`);
237
- const messagesTable = await this.client.openTable(TABLE_MESSAGES);
238
- await messagesTable.delete(`thread_id = '${threadId}'`);
239
- } catch (error) {
239
+ if (!this.client) {
240
+ throw new Error("LanceDB client not initialized. Call LanceStorage.create() first.");
241
+ }
242
+ if (!tableName) {
243
+ throw new Error("tableName is required for dropTable.");
244
+ }
245
+ } catch (validationError) {
240
246
  throw new MastraError(
241
247
  {
242
- id: createStorageErrorId("LANCE", "DELETE_THREAD", "FAILED"),
248
+ id: createStorageErrorId("LANCE", "DROP_TABLE", "INVALID_ARGS"),
243
249
  domain: ErrorDomain.STORAGE,
244
- category: ErrorCategory.THIRD_PARTY
250
+ category: ErrorCategory.USER,
251
+ text: validationError.message,
252
+ details: { tableName }
245
253
  },
246
- error
254
+ validationError
247
255
  );
248
256
  }
249
- }
250
- normalizeMessage(message) {
251
- const { thread_id, ...rest } = message;
252
- return {
253
- ...rest,
254
- threadId: thread_id,
255
- content: typeof message.content === "string" ? (() => {
256
- try {
257
- return JSON.parse(message.content);
258
- } catch {
259
- return message.content;
260
- }
261
- })() : message.content
262
- };
263
- }
264
- async listMessagesById({ messageIds }) {
265
- if (messageIds.length === 0) return { messages: [] };
266
257
  try {
267
- const table = await this.client.openTable(TABLE_MESSAGES);
268
- const quotedIds = messageIds.map((id) => `'${id}'`).join(", ");
269
- const allRecords = await table.query().where(`id IN (${quotedIds})`).toArray();
270
- const messages = processResultWithTypeConversion(
271
- allRecords,
272
- await getTableSchema({ tableName: TABLE_MESSAGES, client: this.client })
273
- );
274
- const list = new MessageList().add(
275
- messages.map(this.normalizeMessage),
276
- "memory"
277
- );
278
- return { messages: list.get.all.db() };
258
+ await this.client.dropTable(tableName);
279
259
  } catch (error) {
260
+ if (error.toString().includes("was not found") || error.message?.includes("Table not found")) {
261
+ this.logger.debug(`Table '${tableName}' does not exist, skipping drop`);
262
+ return;
263
+ }
280
264
  throw new MastraError(
281
265
  {
282
- id: createStorageErrorId("LANCE", "LIST_MESSAGES_BY_ID", "FAILED"),
266
+ id: createStorageErrorId("LANCE", "DROP_TABLE", "FAILED"),
283
267
  domain: ErrorDomain.STORAGE,
284
268
  category: ErrorCategory.THIRD_PARTY,
285
- details: {
286
- messageIds: JSON.stringify(messageIds)
287
- }
269
+ details: { tableName }
288
270
  },
289
271
  error
290
272
  );
291
273
  }
292
274
  }
293
- async listMessages(args) {
294
- const { threadId, resourceId, include, filter, perPage: perPageInput, page = 0, orderBy } = args;
295
- const threadIds = Array.isArray(threadId) ? threadId : [threadId];
296
- if (threadIds.length === 0 || threadIds.some((id) => !id.trim())) {
275
+ async alterTable({
276
+ tableName,
277
+ schema,
278
+ ifNotExists
279
+ }) {
280
+ try {
281
+ if (!this.client) {
282
+ throw new Error("LanceDB client not initialized. Call LanceStorage.create() first.");
283
+ }
284
+ if (!tableName) {
285
+ throw new Error("tableName is required for alterTable.");
286
+ }
287
+ if (!schema) {
288
+ throw new Error("schema is required for alterTable.");
289
+ }
290
+ if (!ifNotExists || ifNotExists.length === 0) {
291
+ this.logger.debug("No columns specified to add in alterTable, skipping.");
292
+ return;
293
+ }
294
+ } catch (validationError) {
297
295
  throw new MastraError(
298
296
  {
299
- id: createStorageErrorId("LANCE", "LIST_MESSAGES", "INVALID_THREAD_ID"),
297
+ id: createStorageErrorId("LANCE", "ALTER_TABLE", "INVALID_ARGS"),
300
298
  domain: ErrorDomain.STORAGE,
301
- category: ErrorCategory.THIRD_PARTY,
302
- details: { threadId: Array.isArray(threadId) ? threadId.join(",") : threadId }
299
+ category: ErrorCategory.USER,
300
+ text: validationError.message,
301
+ details: { tableName }
303
302
  },
304
- new Error("threadId must be a non-empty string or array of non-empty strings")
303
+ validationError
305
304
  );
306
305
  }
307
- const perPage = normalizePerPage(perPageInput, 40);
308
- const { offset, perPage: perPageForResponse } = calculatePagination(page, perPageInput, perPage);
309
306
  try {
310
- if (page < 0) {
311
- throw new MastraError(
312
- {
313
- id: createStorageErrorId("LANCE", "LIST_MESSAGES", "INVALID_PAGE"),
314
- domain: ErrorDomain.STORAGE,
315
- category: ErrorCategory.USER,
316
- details: { page }
317
- },
318
- new Error("page must be >= 0")
319
- );
307
+ const table = await this.client.openTable(tableName);
308
+ const currentSchema = await table.schema();
309
+ const existingFields = new Set(currentSchema.fields.map((f) => f.name));
310
+ const typeMap = {
311
+ text: "string",
312
+ integer: "int",
313
+ bigint: "bigint",
314
+ timestamp: "timestamp",
315
+ jsonb: "string",
316
+ uuid: "string"
317
+ };
318
+ const columnsToAdd = ifNotExists.filter((col) => schema[col] && !existingFields.has(col)).map((col) => {
319
+ const colDef = schema[col];
320
+ return {
321
+ name: col,
322
+ valueSql: colDef?.nullable ? `cast(NULL as ${typeMap[colDef.type ?? "text"]})` : `cast(${this.getDefaultValue(colDef?.type ?? "text")} as ${typeMap[colDef?.type ?? "text"]})`
323
+ };
324
+ });
325
+ if (columnsToAdd.length > 0) {
326
+ await table.addColumns(columnsToAdd);
327
+ this.logger?.info?.(`Added columns [${columnsToAdd.map((c) => c.name).join(", ")}] to table ${tableName}`);
320
328
  }
321
- const { field, direction } = this.parseOrderBy(orderBy, "ASC");
322
- const table = await this.client.openTable(TABLE_MESSAGES);
323
- const threadCondition = threadIds.length === 1 ? `thread_id = '${this.escapeSql(threadIds[0])}'` : `thread_id IN (${threadIds.map((t) => `'${this.escapeSql(t)}'`).join(", ")})`;
324
- const conditions = [threadCondition];
325
- if (resourceId) {
326
- conditions.push(`\`resourceId\` = '${this.escapeSql(resourceId)}'`);
329
+ } catch (error) {
330
+ throw new MastraError(
331
+ {
332
+ id: createStorageErrorId("LANCE", "ALTER_TABLE", "FAILED"),
333
+ domain: ErrorDomain.STORAGE,
334
+ category: ErrorCategory.THIRD_PARTY,
335
+ details: { tableName }
336
+ },
337
+ error
338
+ );
339
+ }
340
+ }
341
+ async clearTable({ tableName }) {
342
+ try {
343
+ if (!this.client) {
344
+ throw new Error("LanceDB client not initialized. Call LanceStorage.create() first.");
327
345
  }
328
- if (filter?.dateRange?.start) {
329
- const startTime = filter.dateRange.start instanceof Date ? filter.dateRange.start.getTime() : new Date(filter.dateRange.start).getTime();
330
- conditions.push(`\`createdAt\` >= ${startTime}`);
346
+ if (!tableName) {
347
+ throw new Error("tableName is required for clearTable.");
331
348
  }
332
- if (filter?.dateRange?.end) {
333
- const endTime = filter.dateRange.end instanceof Date ? filter.dateRange.end.getTime() : new Date(filter.dateRange.end).getTime();
334
- conditions.push(`\`createdAt\` <= ${endTime}`);
349
+ } catch (validationError) {
350
+ throw new MastraError(
351
+ {
352
+ id: createStorageErrorId("LANCE", "CLEAR_TABLE", "INVALID_ARGS"),
353
+ domain: ErrorDomain.STORAGE,
354
+ category: ErrorCategory.USER,
355
+ text: validationError.message,
356
+ details: { tableName }
357
+ },
358
+ validationError
359
+ );
360
+ }
361
+ try {
362
+ const table = await this.client.openTable(tableName);
363
+ await table.delete("1=1");
364
+ } catch (error) {
365
+ throw new MastraError(
366
+ {
367
+ id: createStorageErrorId("LANCE", "CLEAR_TABLE", "FAILED"),
368
+ domain: ErrorDomain.STORAGE,
369
+ category: ErrorCategory.THIRD_PARTY,
370
+ details: { tableName }
371
+ },
372
+ error
373
+ );
374
+ }
375
+ }
376
+ async insert({ tableName, record }) {
377
+ try {
378
+ if (!this.client) {
379
+ throw new Error("LanceDB client not initialized. Call LanceStorage.create() first.");
335
380
  }
336
- const whereClause = conditions.join(" AND ");
337
- const total = await table.countRows(whereClause);
338
- const query = table.query().where(whereClause);
339
- let allRecords = await query.toArray();
340
- allRecords.sort((a, b) => {
341
- const aValue = field === "createdAt" ? a.createdAt : a[field];
342
- const bValue = field === "createdAt" ? b.createdAt : b[field];
343
- if (aValue == null && bValue == null) return 0;
344
- if (aValue == null) return direction === "ASC" ? -1 : 1;
345
- if (bValue == null) return direction === "ASC" ? 1 : -1;
346
- if (typeof aValue === "string" && typeof bValue === "string") {
347
- return direction === "ASC" ? aValue.localeCompare(bValue) : bValue.localeCompare(aValue);
348
- }
349
- return direction === "ASC" ? aValue - bValue : bValue - aValue;
350
- });
351
- const paginatedRecords = allRecords.slice(offset, offset + perPage);
352
- const messages = paginatedRecords.map((row) => this.normalizeMessage(row));
353
- if (total === 0 && messages.length === 0 && (!include || include.length === 0)) {
354
- return {
355
- messages: [],
356
- total: 0,
357
- page,
358
- perPage: perPageForResponse,
359
- hasMore: false
360
- };
381
+ if (!tableName) {
382
+ throw new Error("tableName is required for insert.");
361
383
  }
362
- const messageIds = new Set(messages.map((m) => m.id));
363
- if (include && include.length > 0) {
364
- const threadIds2 = [...new Set(include.map((item) => item.threadId || threadId))];
365
- const allThreadMessages = [];
366
- for (const tid of threadIds2) {
367
- const threadQuery = table.query().where(`thread_id = '${tid}'`);
368
- let threadRecords = await threadQuery.toArray();
369
- allThreadMessages.push(...threadRecords);
370
- }
371
- allThreadMessages.sort((a, b) => a.createdAt - b.createdAt);
372
- const contextMessages = this.processMessagesWithContext(allThreadMessages, include);
373
- const includedMessages = contextMessages.map((row) => this.normalizeMessage(row));
374
- for (const includeMsg of includedMessages) {
375
- if (!messageIds.has(includeMsg.id)) {
376
- messages.push(includeMsg);
377
- messageIds.add(includeMsg.id);
378
- }
379
- }
384
+ if (!record || Object.keys(record).length === 0) {
385
+ throw new Error("record is required and cannot be empty for insert.");
380
386
  }
381
- const list = new MessageList().add(messages, "memory");
382
- let finalMessages = list.get.all.db();
383
- finalMessages = finalMessages.sort((a, b) => {
384
- const aValue = field === "createdAt" ? new Date(a.createdAt).getTime() : a[field];
385
- const bValue = field === "createdAt" ? new Date(b.createdAt).getTime() : b[field];
386
- if (aValue == null && bValue == null) return 0;
387
- if (aValue == null) return direction === "ASC" ? -1 : 1;
388
- if (bValue == null) return direction === "ASC" ? 1 : -1;
389
- if (typeof aValue === "string" && typeof bValue === "string") {
390
- return direction === "ASC" ? aValue.localeCompare(bValue) : bValue.localeCompare(aValue);
387
+ } catch (validationError) {
388
+ throw new MastraError(
389
+ {
390
+ id: createStorageErrorId("LANCE", "INSERT", "INVALID_ARGS"),
391
+ domain: ErrorDomain.STORAGE,
392
+ category: ErrorCategory.USER,
393
+ text: validationError.message,
394
+ details: { tableName }
395
+ },
396
+ validationError
397
+ );
398
+ }
399
+ try {
400
+ const table = await this.client.openTable(tableName);
401
+ const primaryId = getPrimaryKeys(tableName);
402
+ const processedRecord = { ...record };
403
+ for (const key in processedRecord) {
404
+ if (processedRecord[key] !== null && typeof processedRecord[key] === "object" && !(processedRecord[key] instanceof Date)) {
405
+ this.logger.debug("Converting object to JSON string: ", processedRecord[key]);
406
+ processedRecord[key] = JSON.stringify(processedRecord[key]);
391
407
  }
392
- return direction === "ASC" ? aValue - bValue : bValue - aValue;
393
- });
394
- const returnedThreadMessageIds = new Set(finalMessages.filter((m) => m.threadId === threadId).map((m) => m.id));
395
- const allThreadMessagesReturned = returnedThreadMessageIds.size >= total;
396
- const fetchedAll = perPageInput === false || allThreadMessagesReturned;
397
- const hasMore = !fetchedAll && offset + perPage < total;
398
- return {
399
- messages: finalMessages,
400
- total,
401
- page,
402
- perPage: perPageForResponse,
403
- hasMore
404
- };
408
+ }
409
+ await table.mergeInsert(primaryId).whenMatchedUpdateAll().whenNotMatchedInsertAll().execute([processedRecord]);
405
410
  } catch (error) {
406
- const mastraError = new MastraError(
411
+ throw new MastraError(
407
412
  {
408
- id: createStorageErrorId("LANCE", "LIST_MESSAGES", "FAILED"),
413
+ id: createStorageErrorId("LANCE", "INSERT", "FAILED"),
409
414
  domain: ErrorDomain.STORAGE,
410
415
  category: ErrorCategory.THIRD_PARTY,
411
- details: {
412
- threadId: Array.isArray(threadId) ? threadId.join(",") : threadId,
413
- resourceId: resourceId ?? ""
414
- }
416
+ details: { tableName }
415
417
  },
416
418
  error
417
419
  );
418
- this.logger?.error?.(mastraError.toString());
419
- this.logger?.trackException?.(mastraError);
420
- return {
421
- messages: [],
422
- total: 0,
423
- page,
424
- perPage: perPageForResponse,
425
- hasMore: false
426
- };
427
420
  }
428
421
  }
429
- async saveMessages(args) {
422
+ async batchInsert({ tableName, records }) {
430
423
  try {
431
- const { messages } = args;
432
- if (messages.length === 0) {
433
- return { messages: [] };
424
+ if (!this.client) {
425
+ throw new Error("LanceDB client not initialized. Call LanceStorage.create() first.");
434
426
  }
435
- const threadId = messages[0]?.threadId;
436
- if (!threadId) {
437
- throw new Error("Thread ID is required");
427
+ if (!tableName) {
428
+ throw new Error("tableName is required for batchInsert.");
438
429
  }
439
- for (const message of messages) {
440
- if (!message.id) {
441
- throw new Error("Message ID is required");
442
- }
443
- if (!message.threadId) {
444
- throw new Error("Thread ID is required for all messages");
445
- }
446
- if (message.resourceId === null || message.resourceId === void 0) {
447
- throw new Error("Resource ID cannot be null or undefined");
448
- }
449
- if (!message.content) {
450
- throw new Error("Message content is required");
451
- }
430
+ if (!records || records.length === 0) {
431
+ throw new Error("records array is required and cannot be empty for batchInsert.");
452
432
  }
453
- const transformedMessages = messages.map((message) => {
454
- const { threadId: threadId2, type, ...rest } = message;
455
- return {
456
- ...rest,
457
- thread_id: threadId2,
458
- type: type ?? "v2",
459
- content: JSON.stringify(message.content)
460
- };
433
+ } catch (validationError) {
434
+ throw new MastraError(
435
+ {
436
+ id: createStorageErrorId("LANCE", "BATCH_INSERT", "INVALID_ARGS"),
437
+ domain: ErrorDomain.STORAGE,
438
+ category: ErrorCategory.USER,
439
+ text: validationError.message,
440
+ details: { tableName }
441
+ },
442
+ validationError
443
+ );
444
+ }
445
+ try {
446
+ const table = await this.client.openTable(tableName);
447
+ const primaryId = getPrimaryKeys(tableName);
448
+ const processedRecords = records.map((record) => {
449
+ const processedRecord = { ...record };
450
+ for (const key in processedRecord) {
451
+ if (processedRecord[key] == null) continue;
452
+ if (processedRecord[key] !== null && typeof processedRecord[key] === "object" && !(processedRecord[key] instanceof Date)) {
453
+ processedRecord[key] = JSON.stringify(processedRecord[key]);
454
+ }
455
+ }
456
+ return processedRecord;
461
457
  });
462
- const table = await this.client.openTable(TABLE_MESSAGES);
463
- await table.mergeInsert("id").whenMatchedUpdateAll().whenNotMatchedInsertAll().execute(transformedMessages);
464
- const threadsTable = await this.client.openTable(TABLE_THREADS);
465
- const currentTime = (/* @__PURE__ */ new Date()).getTime();
466
- const updateRecord = { id: threadId, updatedAt: currentTime };
467
- await threadsTable.mergeInsert("id").whenMatchedUpdateAll().whenNotMatchedInsertAll().execute([updateRecord]);
468
- const list = new MessageList().add(messages, "memory");
469
- return { messages: list.get.all.db() };
458
+ await table.mergeInsert(primaryId).whenMatchedUpdateAll().whenNotMatchedInsertAll().execute(processedRecords);
470
459
  } catch (error) {
471
460
  throw new MastraError(
472
461
  {
473
- id: createStorageErrorId("LANCE", "SAVE_MESSAGES", "FAILED"),
462
+ id: createStorageErrorId("LANCE", "BATCH_INSERT", "FAILED"),
474
463
  domain: ErrorDomain.STORAGE,
475
- category: ErrorCategory.THIRD_PARTY
464
+ category: ErrorCategory.THIRD_PARTY,
465
+ details: { tableName }
476
466
  },
477
467
  error
478
468
  );
479
469
  }
480
470
  }
481
- async listThreadsByResourceId(args) {
471
+ async load({ tableName, keys }) {
482
472
  try {
483
- const { resourceId, page = 0, perPage: perPageInput, orderBy } = args;
484
- const perPage = normalizePerPage(perPageInput, 100);
485
- if (page < 0) {
486
- throw new MastraError(
487
- {
488
- id: createStorageErrorId("LANCE", "LIST_THREADS_BY_RESOURCE_ID", "INVALID_PAGE"),
489
- domain: ErrorDomain.STORAGE,
490
- category: ErrorCategory.USER,
491
- details: { page }
492
- },
493
- new Error("page must be >= 0")
494
- );
473
+ if (!this.client) {
474
+ throw new Error("LanceDB client not initialized. Call LanceStorage.create() first.");
495
475
  }
496
- const { offset, perPage: perPageForResponse } = calculatePagination(page, perPageInput, perPage);
497
- const { field, direction } = this.parseOrderBy(orderBy);
498
- const table = await this.client.openTable(TABLE_THREADS);
499
- const total = await table.countRows(`\`resourceId\` = '${this.escapeSql(resourceId)}'`);
500
- const query = table.query().where(`\`resourceId\` = '${this.escapeSql(resourceId)}'`);
501
- const records = await query.toArray();
502
- records.sort((a, b) => {
503
- const aValue = ["createdAt", "updatedAt"].includes(field) ? new Date(a[field]).getTime() : a[field];
504
- const bValue = ["createdAt", "updatedAt"].includes(field) ? new Date(b[field]).getTime() : b[field];
505
- if (aValue == null && bValue == null) return 0;
506
- if (aValue == null) return direction === "ASC" ? -1 : 1;
507
- if (bValue == null) return direction === "ASC" ? 1 : -1;
508
- if (typeof aValue === "string" && typeof bValue === "string") {
509
- return direction === "ASC" ? aValue.localeCompare(bValue) : bValue.localeCompare(aValue);
510
- }
511
- return direction === "ASC" ? aValue - bValue : bValue - aValue;
512
- });
513
- const paginatedRecords = records.slice(offset, offset + perPage);
514
- const schema = await getTableSchema({ tableName: TABLE_THREADS, client: this.client });
515
- const threads = paginatedRecords.map(
516
- (record) => processResultWithTypeConversion(record, schema)
476
+ if (!tableName) {
477
+ throw new Error("tableName is required for load.");
478
+ }
479
+ if (!keys || Object.keys(keys).length === 0) {
480
+ throw new Error("keys are required and cannot be empty for load.");
481
+ }
482
+ } catch (validationError) {
483
+ throw new MastraError(
484
+ {
485
+ id: createStorageErrorId("LANCE", "LOAD", "INVALID_ARGS"),
486
+ domain: ErrorDomain.STORAGE,
487
+ category: ErrorCategory.USER,
488
+ text: validationError.message,
489
+ details: { tableName }
490
+ },
491
+ validationError
517
492
  );
518
- return {
519
- threads,
520
- total,
521
- page,
522
- perPage: perPageForResponse,
523
- hasMore: offset + perPage < total
524
- };
493
+ }
494
+ try {
495
+ const table = await this.client.openTable(tableName);
496
+ const tableSchema = await getTableSchema({ tableName, client: this.client });
497
+ const query = table.query();
498
+ if (Object.keys(keys).length > 0) {
499
+ validateKeyTypes(keys, tableSchema);
500
+ const filterConditions = Object.entries(keys).map(([key, value]) => {
501
+ const isCamelCase = /^[a-z][a-zA-Z]*$/.test(key) && /[A-Z]/.test(key);
502
+ const quotedKey = isCamelCase ? `\`${key}\`` : key;
503
+ if (typeof value === "string") {
504
+ return `${quotedKey} = '${value}'`;
505
+ } else if (value === null) {
506
+ return `${quotedKey} IS NULL`;
507
+ } else {
508
+ return `${quotedKey} = ${value}`;
509
+ }
510
+ }).join(" AND ");
511
+ this.logger.debug("where clause generated: " + filterConditions);
512
+ query.where(filterConditions);
513
+ }
514
+ const result = await query.limit(1).toArray();
515
+ if (result.length === 0) {
516
+ this.logger.debug("No record found");
517
+ return null;
518
+ }
519
+ return processResultWithTypeConversion(result[0], tableSchema);
525
520
  } catch (error) {
521
+ if (error instanceof MastraError) throw error;
526
522
  throw new MastraError(
527
523
  {
528
- id: createStorageErrorId("LANCE", "LIST_THREADS_BY_RESOURCE_ID", "FAILED"),
524
+ id: createStorageErrorId("LANCE", "LOAD", "FAILED"),
529
525
  domain: ErrorDomain.STORAGE,
530
- category: ErrorCategory.THIRD_PARTY
526
+ category: ErrorCategory.THIRD_PARTY,
527
+ details: { tableName, keyCount: Object.keys(keys).length, firstKey: Object.keys(keys)[0] ?? "" }
531
528
  },
532
529
  error
533
530
  );
534
531
  }
535
532
  }
536
- /**
537
- * Processes messages to include context messages based on withPreviousMessages and withNextMessages
538
- * @param records - The sorted array of records to process
539
- * @param include - The array of include specifications with context parameters
540
- * @returns The processed array with context messages included
541
- */
542
- processMessagesWithContext(records, include) {
543
- const messagesWithContext = include.filter((item) => item.withPreviousMessages || item.withNextMessages);
544
- if (messagesWithContext.length === 0) {
545
- return records;
546
- }
547
- const messageIndexMap = /* @__PURE__ */ new Map();
548
- records.forEach((message, index) => {
549
- messageIndexMap.set(message.id, index);
550
- });
551
- const additionalIndices = /* @__PURE__ */ new Set();
552
- for (const item of messagesWithContext) {
553
- const messageIndex = messageIndexMap.get(item.id);
554
- if (messageIndex !== void 0) {
555
- if (item.withPreviousMessages) {
556
- const startIdx = Math.max(0, messageIndex - item.withPreviousMessages);
557
- for (let i = startIdx; i < messageIndex; i++) {
558
- additionalIndices.add(i);
559
- }
560
- }
561
- if (item.withNextMessages) {
562
- const endIdx = Math.min(records.length - 1, messageIndex + item.withNextMessages);
563
- for (let i = messageIndex + 1; i <= endIdx; i++) {
564
- additionalIndices.add(i);
565
- }
566
- }
567
- }
568
- }
569
- if (additionalIndices.size === 0) {
570
- return records;
571
- }
572
- const originalMatchIds = new Set(include.map((item) => item.id));
573
- const allIndices = /* @__PURE__ */ new Set();
574
- records.forEach((record, index) => {
575
- if (originalMatchIds.has(record.id)) {
576
- allIndices.add(index);
577
- }
578
- });
579
- additionalIndices.forEach((index) => {
580
- allIndices.add(index);
533
+ };
534
+
535
+ // src/storage/domains/memory/index.ts
536
+ var StoreMemoryLance = class extends MemoryStorage {
537
+ client;
538
+ #db;
539
+ constructor(config) {
540
+ super();
541
+ const client = resolveLanceConfig(config);
542
+ this.client = client;
543
+ this.#db = new LanceDB({ client });
544
+ }
545
+ async init() {
546
+ await this.#db.createTable({ tableName: TABLE_THREADS, schema: TABLE_SCHEMAS[TABLE_THREADS] });
547
+ await this.#db.createTable({ tableName: TABLE_MESSAGES, schema: TABLE_SCHEMAS[TABLE_MESSAGES] });
548
+ await this.#db.createTable({ tableName: TABLE_RESOURCES, schema: TABLE_SCHEMAS[TABLE_RESOURCES] });
549
+ await this.#db.alterTable({
550
+ tableName: TABLE_MESSAGES,
551
+ schema: TABLE_SCHEMAS[TABLE_MESSAGES],
552
+ ifNotExists: ["resourceId"]
581
553
  });
582
- return Array.from(allIndices).sort((a, b) => a - b).map((index) => records[index]);
583
554
  }
584
- /**
585
- * Parse message data from LanceDB record format to MastraDBMessage format
586
- */
587
- parseMessageData(data) {
588
- const { thread_id, ...rest } = data;
589
- return {
590
- ...rest,
591
- threadId: thread_id,
592
- content: typeof data.content === "string" ? (() => {
593
- try {
594
- return JSON.parse(data.content);
595
- } catch {
596
- return data.content;
597
- }
598
- })() : data.content,
599
- createdAt: new Date(data.createdAt),
600
- updatedAt: new Date(data.updatedAt)
601
- };
555
+ async dangerouslyClearAll() {
556
+ await this.#db.clearTable({ tableName: TABLE_THREADS });
557
+ await this.#db.clearTable({ tableName: TABLE_MESSAGES });
558
+ await this.#db.clearTable({ tableName: TABLE_RESOURCES });
602
559
  }
603
- async updateMessages(args) {
604
- const { messages } = args;
605
- this.logger.debug("Updating messages", { count: messages.length });
606
- if (!messages.length) {
607
- return [];
560
+ async deleteMessages(messageIds) {
561
+ if (!messageIds || messageIds.length === 0) {
562
+ return;
608
563
  }
609
- const updatedMessages = [];
610
- const affectedThreadIds = /* @__PURE__ */ new Set();
564
+ this.logger.debug("Deleting messages", { count: messageIds.length });
611
565
  try {
612
- for (const updateData of messages) {
613
- const { id, ...updates } = updateData;
614
- const existingMessage = await this.operations.load({ tableName: TABLE_MESSAGES, keys: { id } });
615
- if (!existingMessage) {
616
- this.logger.warn("Message not found for update", { id });
617
- continue;
618
- }
619
- const existingMsg = this.parseMessageData(existingMessage);
620
- const originalThreadId = existingMsg.threadId;
621
- affectedThreadIds.add(originalThreadId);
622
- const updatePayload = {};
623
- if ("role" in updates && updates.role !== void 0) updatePayload.role = updates.role;
624
- if ("type" in updates && updates.type !== void 0) updatePayload.type = updates.type;
625
- if ("resourceId" in updates && updates.resourceId !== void 0) updatePayload.resourceId = updates.resourceId;
626
- if ("threadId" in updates && updates.threadId !== void 0 && updates.threadId !== null) {
627
- updatePayload.thread_id = updates.threadId;
628
- affectedThreadIds.add(updates.threadId);
629
- }
630
- if (updates.content) {
631
- const existingContent = existingMsg.content;
632
- let newContent = { ...existingContent };
633
- if (updates.content.metadata !== void 0) {
634
- newContent.metadata = {
635
- ...existingContent.metadata || {},
636
- ...updates.content.metadata || {}
637
- };
638
- }
639
- if (updates.content.content !== void 0) {
640
- newContent.content = updates.content.content;
641
- }
642
- if ("parts" in updates.content && updates.content.parts !== void 0) {
643
- newContent.parts = updates.content.parts;
644
- }
645
- updatePayload.content = JSON.stringify(newContent);
646
- }
647
- await this.operations.insert({ tableName: TABLE_MESSAGES, record: { id, ...updatePayload } });
648
- const updatedMessage = await this.operations.load({ tableName: TABLE_MESSAGES, keys: { id } });
649
- if (updatedMessage) {
650
- updatedMessages.push(this.parseMessageData(updatedMessage));
566
+ const threadIds = /* @__PURE__ */ new Set();
567
+ for (const messageId of messageIds) {
568
+ const message = await this.#db.load({ tableName: TABLE_MESSAGES, keys: { id: messageId } });
569
+ if (message?.thread_id) {
570
+ threadIds.add(message.thread_id);
651
571
  }
652
572
  }
653
- for (const threadId of affectedThreadIds) {
654
- await this.operations.insert({
655
- tableName: TABLE_THREADS,
656
- record: { id: threadId, updatedAt: Date.now() }
657
- });
573
+ const messagesTable = await this.client.openTable(TABLE_MESSAGES);
574
+ const idConditions = messageIds.map((id) => `id = '${this.escapeSql(id)}'`).join(" OR ");
575
+ await messagesTable.delete(idConditions);
576
+ const now = (/* @__PURE__ */ new Date()).getTime();
577
+ const threadsTable = await this.client.openTable(TABLE_THREADS);
578
+ for (const threadId of threadIds) {
579
+ const thread = await this.getThreadById({ threadId });
580
+ if (thread) {
581
+ const record = {
582
+ id: threadId,
583
+ resourceId: thread.resourceId,
584
+ title: thread.title,
585
+ metadata: JSON.stringify(thread.metadata),
586
+ createdAt: new Date(thread.createdAt).getTime(),
587
+ updatedAt: now
588
+ };
589
+ await threadsTable.mergeInsert("id").whenMatchedUpdateAll().whenNotMatchedInsertAll().execute([record]);
590
+ }
658
591
  }
659
- return updatedMessages;
660
592
  } catch (error) {
661
593
  throw new MastraError(
662
594
  {
663
- id: createStorageErrorId("LANCE", "UPDATE_MESSAGES", "FAILED"),
595
+ id: createStorageErrorId("LANCE", "DELETE_MESSAGES", "FAILED"),
664
596
  domain: ErrorDomain.STORAGE,
665
597
  category: ErrorCategory.THIRD_PARTY,
666
- details: { count: messages.length }
598
+ details: { count: messageIds.length }
667
599
  },
668
600
  error
669
601
  );
670
602
  }
671
603
  }
672
- async getResourceById({ resourceId }) {
604
+ // Utility to escape single quotes in SQL strings
605
+ escapeSql(str) {
606
+ return str.replace(/'/g, "''");
607
+ }
608
+ async getThreadById({ threadId }) {
673
609
  try {
674
- const resource = await this.operations.load({ tableName: TABLE_RESOURCES, keys: { id: resourceId } });
675
- if (!resource) {
610
+ const thread = await this.#db.load({ tableName: TABLE_THREADS, keys: { id: threadId } });
611
+ if (!thread) {
676
612
  return null;
677
613
  }
678
- let createdAt;
679
- let updatedAt;
680
- try {
681
- if (resource.createdAt instanceof Date) {
682
- createdAt = resource.createdAt;
683
- } else if (typeof resource.createdAt === "string") {
684
- createdAt = new Date(resource.createdAt);
685
- } else if (typeof resource.createdAt === "number") {
686
- createdAt = new Date(resource.createdAt);
687
- } else {
688
- createdAt = /* @__PURE__ */ new Date();
689
- }
690
- if (isNaN(createdAt.getTime())) {
691
- createdAt = /* @__PURE__ */ new Date();
692
- }
693
- } catch {
694
- createdAt = /* @__PURE__ */ new Date();
695
- }
696
- try {
697
- if (resource.updatedAt instanceof Date) {
698
- updatedAt = resource.updatedAt;
699
- } else if (typeof resource.updatedAt === "string") {
700
- updatedAt = new Date(resource.updatedAt);
701
- } else if (typeof resource.updatedAt === "number") {
702
- updatedAt = new Date(resource.updatedAt);
703
- } else {
704
- updatedAt = /* @__PURE__ */ new Date();
705
- }
706
- if (isNaN(updatedAt.getTime())) {
707
- updatedAt = /* @__PURE__ */ new Date();
708
- }
709
- } catch {
710
- updatedAt = /* @__PURE__ */ new Date();
711
- }
712
- let workingMemory = resource.workingMemory;
713
- if (workingMemory === null || workingMemory === void 0) {
714
- workingMemory = void 0;
715
- } else if (workingMemory === "") {
716
- workingMemory = "";
717
- } else if (typeof workingMemory === "object") {
718
- workingMemory = JSON.stringify(workingMemory);
719
- }
720
- let metadata = resource.metadata;
721
- if (metadata === "" || metadata === null || metadata === void 0) {
722
- metadata = void 0;
723
- } else if (typeof metadata === "string") {
724
- try {
725
- metadata = JSON.parse(metadata);
726
- } catch {
727
- metadata = metadata;
728
- }
729
- }
730
614
  return {
731
- ...resource,
732
- createdAt,
733
- updatedAt,
734
- workingMemory,
735
- metadata
615
+ ...thread,
616
+ createdAt: new Date(thread.createdAt),
617
+ updatedAt: new Date(thread.updatedAt)
736
618
  };
737
619
  } catch (error) {
738
620
  throw new MastraError(
739
621
  {
740
- id: createStorageErrorId("LANCE", "GET_RESOURCE_BY_ID", "FAILED"),
622
+ id: createStorageErrorId("LANCE", "GET_THREAD_BY_ID", "FAILED"),
741
623
  domain: ErrorDomain.STORAGE,
742
624
  category: ErrorCategory.THIRD_PARTY
743
625
  },
@@ -745,23 +627,21 @@ var StoreMemoryLance = class extends MemoryStorage {
745
627
  );
746
628
  }
747
629
  }
748
- async saveResource({ resource }) {
630
+ /**
631
+ * Saves a thread to the database. This function doesn't overwrite existing threads.
632
+ * @param thread - The thread to save
633
+ * @returns The saved thread
634
+ */
635
+ async saveThread({ thread }) {
749
636
  try {
750
- const record = {
751
- ...resource,
752
- metadata: resource.metadata ? JSON.stringify(resource.metadata) : "",
753
- createdAt: resource.createdAt.getTime(),
754
- // Store as timestamp (milliseconds)
755
- updatedAt: resource.updatedAt.getTime()
756
- // Store as timestamp (milliseconds)
757
- };
758
- const table = await this.client.openTable(TABLE_RESOURCES);
637
+ const record = { ...thread, metadata: JSON.stringify(thread.metadata) };
638
+ const table = await this.client.openTable(TABLE_THREADS);
759
639
  await table.add([record], { mode: "append" });
760
- return resource;
640
+ return thread;
761
641
  } catch (error) {
762
642
  throw new MastraError(
763
643
  {
764
- id: createStorageErrorId("LANCE", "SAVE_RESOURCE", "FAILED"),
644
+ id: createStorageErrorId("LANCE", "SAVE_THREAD", "FAILED"),
765
645
  domain: ErrorDomain.STORAGE,
766
646
  category: ErrorCategory.THIRD_PARTY
767
647
  },
@@ -769,44 +649,32 @@ var StoreMemoryLance = class extends MemoryStorage {
769
649
  );
770
650
  }
771
651
  }
772
- async updateResource({
773
- resourceId,
774
- workingMemory,
652
+ async updateThread({
653
+ id,
654
+ title,
775
655
  metadata
776
656
  }) {
777
- const maxRetries = 3;
778
- for (let attempt = 0; attempt < maxRetries; attempt++) {
779
- try {
780
- const existingResource = await this.getResourceById({ resourceId });
781
- if (!existingResource) {
782
- const newResource = {
783
- id: resourceId,
784
- workingMemory,
785
- metadata: metadata || {},
786
- createdAt: /* @__PURE__ */ new Date(),
787
- updatedAt: /* @__PURE__ */ new Date()
788
- };
789
- return this.saveResource({ resource: newResource });
657
+ const maxRetries = 5;
658
+ for (let attempt = 0; attempt < maxRetries; attempt++) {
659
+ try {
660
+ const current = await this.getThreadById({ threadId: id });
661
+ if (!current) {
662
+ throw new Error(`Thread with id ${id} not found`);
790
663
  }
791
- const updatedResource = {
792
- ...existingResource,
793
- workingMemory: workingMemory !== void 0 ? workingMemory : existingResource.workingMemory,
794
- metadata: {
795
- ...existingResource.metadata,
796
- ...metadata
797
- },
798
- updatedAt: /* @__PURE__ */ new Date()
799
- };
664
+ const mergedMetadata = { ...current.metadata, ...metadata };
800
665
  const record = {
801
- id: resourceId,
802
- workingMemory: updatedResource.workingMemory || "",
803
- metadata: updatedResource.metadata ? JSON.stringify(updatedResource.metadata) : "",
804
- updatedAt: updatedResource.updatedAt.getTime()
805
- // Store as timestamp (milliseconds)
666
+ id,
667
+ title,
668
+ metadata: JSON.stringify(mergedMetadata),
669
+ updatedAt: (/* @__PURE__ */ new Date()).getTime()
806
670
  };
807
- const table = await this.client.openTable(TABLE_RESOURCES);
671
+ const table = await this.client.openTable(TABLE_THREADS);
808
672
  await table.mergeInsert("id").whenMatchedUpdateAll().whenNotMatchedInsertAll().execute([record]);
809
- return updatedResource;
673
+ const updatedThread = await this.getThreadById({ threadId: id });
674
+ if (!updatedThread) {
675
+ throw new Error(`Failed to retrieve updated thread ${id}`);
676
+ }
677
+ return updatedThread;
810
678
  } catch (error) {
811
679
  if (error.message?.includes("Commit conflict") && attempt < maxRetries - 1) {
812
680
  const delay = Math.pow(2, attempt) * 10;
@@ -815,7 +683,7 @@ var StoreMemoryLance = class extends MemoryStorage {
815
683
  }
816
684
  throw new MastraError(
817
685
  {
818
- id: createStorageErrorId("LANCE", "UPDATE_RESOURCE", "FAILED"),
686
+ id: createStorageErrorId("LANCE", "UPDATE_THREAD", "FAILED"),
819
687
  domain: ErrorDomain.STORAGE,
820
688
  category: ErrorCategory.THIRD_PARTY
821
689
  },
@@ -823,418 +691,630 @@ var StoreMemoryLance = class extends MemoryStorage {
823
691
  );
824
692
  }
825
693
  }
826
- throw new Error("Unexpected end of retry loop");
827
- }
828
- };
829
- var StoreOperationsLance = class extends StoreOperations {
830
- client;
831
- constructor({ client }) {
832
- super();
833
- this.client = client;
834
- }
835
- getDefaultValue(type) {
836
- switch (type) {
837
- case "text":
838
- return "''";
839
- case "timestamp":
840
- return "CURRENT_TIMESTAMP";
841
- case "integer":
842
- case "bigint":
843
- return "0";
844
- case "jsonb":
845
- return "'{}'";
846
- case "uuid":
847
- return "''";
848
- default:
849
- return super.getDefaultValue(type);
850
- }
851
- }
852
- async hasColumn(tableName, columnName) {
853
- const table = await this.client.openTable(tableName);
854
- const schema = await table.schema();
855
- return schema.fields.some((field) => field.name === columnName);
856
- }
857
- translateSchema(schema) {
858
- const fields = Object.entries(schema).map(([name, column]) => {
859
- let arrowType;
860
- switch (column.type.toLowerCase()) {
861
- case "text":
862
- case "uuid":
863
- arrowType = new Utf8();
864
- break;
865
- case "int":
866
- case "integer":
867
- arrowType = new Int32();
868
- break;
869
- case "bigint":
870
- arrowType = new Float64();
871
- break;
872
- case "float":
873
- arrowType = new Float32();
874
- break;
875
- case "jsonb":
876
- case "json":
877
- arrowType = new Utf8();
878
- break;
879
- case "binary":
880
- arrowType = new Binary();
881
- break;
882
- case "timestamp":
883
- arrowType = new Float64();
884
- break;
885
- default:
886
- arrowType = new Utf8();
887
- }
888
- return new Field(name, arrowType, column.nullable ?? true);
889
- });
890
- return new Schema(fields);
694
+ throw new MastraError(
695
+ {
696
+ id: createStorageErrorId("LANCE", "UPDATE_THREAD", "FAILED"),
697
+ domain: ErrorDomain.STORAGE,
698
+ category: ErrorCategory.THIRD_PARTY
699
+ },
700
+ new Error("All retries exhausted")
701
+ );
891
702
  }
892
- async createTable({
893
- tableName,
894
- schema
895
- }) {
703
+ async deleteThread({ threadId }) {
896
704
  try {
897
- if (!this.client) {
898
- throw new Error("LanceDB client not initialized. Call LanceStorage.create() first.");
899
- }
900
- if (!tableName) {
901
- throw new Error("tableName is required for createTable.");
902
- }
903
- if (!schema) {
904
- throw new Error("schema is required for createTable.");
905
- }
705
+ const table = await this.client.openTable(TABLE_THREADS);
706
+ await table.delete(`id = '${threadId}'`);
707
+ const messagesTable = await this.client.openTable(TABLE_MESSAGES);
708
+ await messagesTable.delete(`thread_id = '${threadId}'`);
906
709
  } catch (error) {
907
710
  throw new MastraError(
908
711
  {
909
- id: createStorageErrorId("LANCE", "CREATE_TABLE", "INVALID_ARGS"),
712
+ id: createStorageErrorId("LANCE", "DELETE_THREAD", "FAILED"),
910
713
  domain: ErrorDomain.STORAGE,
911
- category: ErrorCategory.USER,
912
- details: { tableName }
714
+ category: ErrorCategory.THIRD_PARTY
913
715
  },
914
716
  error
915
717
  );
916
718
  }
719
+ }
720
+ normalizeMessage(message) {
721
+ const { thread_id, ...rest } = message;
722
+ return {
723
+ ...rest,
724
+ threadId: thread_id,
725
+ content: typeof message.content === "string" ? (() => {
726
+ try {
727
+ return JSON.parse(message.content);
728
+ } catch {
729
+ return message.content;
730
+ }
731
+ })() : message.content
732
+ };
733
+ }
734
+ async listMessagesById({ messageIds }) {
735
+ if (messageIds.length === 0) return { messages: [] };
917
736
  try {
918
- const arrowSchema = this.translateSchema(schema);
919
- await this.client.createEmptyTable(tableName, arrowSchema);
737
+ const table = await this.client.openTable(TABLE_MESSAGES);
738
+ const quotedIds = messageIds.map((id) => `'${id}'`).join(", ");
739
+ const allRecords = await table.query().where(`id IN (${quotedIds})`).toArray();
740
+ const messages = processResultWithTypeConversion(
741
+ allRecords,
742
+ await getTableSchema({ tableName: TABLE_MESSAGES, client: this.client })
743
+ );
744
+ const list = new MessageList().add(
745
+ messages.map(this.normalizeMessage),
746
+ "memory"
747
+ );
748
+ return { messages: list.get.all.db() };
920
749
  } catch (error) {
921
- if (error.message?.includes("already exists")) {
922
- this.logger.debug(`Table '${tableName}' already exists, skipping create`);
923
- return;
924
- }
925
750
  throw new MastraError(
926
751
  {
927
- id: createStorageErrorId("LANCE", "CREATE_TABLE", "FAILED"),
752
+ id: createStorageErrorId("LANCE", "LIST_MESSAGES_BY_ID", "FAILED"),
928
753
  domain: ErrorDomain.STORAGE,
929
754
  category: ErrorCategory.THIRD_PARTY,
930
- details: { tableName }
755
+ details: {
756
+ messageIds: JSON.stringify(messageIds)
757
+ }
931
758
  },
932
759
  error
933
760
  );
934
761
  }
935
762
  }
936
- async dropTable({ tableName }) {
937
- try {
938
- if (!this.client) {
939
- throw new Error("LanceDB client not initialized. Call LanceStorage.create() first.");
940
- }
941
- if (!tableName) {
942
- throw new Error("tableName is required for dropTable.");
943
- }
944
- } catch (validationError) {
763
+ async listMessages(args) {
764
+ const { threadId, resourceId, include, filter, perPage: perPageInput, page = 0, orderBy } = args;
765
+ const threadIds = Array.isArray(threadId) ? threadId : [threadId];
766
+ if (threadIds.length === 0 || threadIds.some((id) => !id.trim())) {
945
767
  throw new MastraError(
946
768
  {
947
- id: createStorageErrorId("LANCE", "DROP_TABLE", "INVALID_ARGS"),
769
+ id: createStorageErrorId("LANCE", "LIST_MESSAGES", "INVALID_THREAD_ID"),
948
770
  domain: ErrorDomain.STORAGE,
949
- category: ErrorCategory.USER,
950
- text: validationError.message,
951
- details: { tableName }
771
+ category: ErrorCategory.THIRD_PARTY,
772
+ details: { threadId: Array.isArray(threadId) ? threadId.join(",") : threadId }
952
773
  },
953
- validationError
774
+ new Error("threadId must be a non-empty string or array of non-empty strings")
954
775
  );
955
776
  }
777
+ const perPage = normalizePerPage(perPageInput, 40);
778
+ const { offset, perPage: perPageForResponse } = calculatePagination(page, perPageInput, perPage);
956
779
  try {
957
- await this.client.dropTable(tableName);
958
- } catch (error) {
959
- if (error.toString().includes("was not found") || error.message?.includes("Table not found")) {
960
- this.logger.debug(`Table '${tableName}' does not exist, skipping drop`);
961
- return;
780
+ if (page < 0) {
781
+ throw new MastraError(
782
+ {
783
+ id: createStorageErrorId("LANCE", "LIST_MESSAGES", "INVALID_PAGE"),
784
+ domain: ErrorDomain.STORAGE,
785
+ category: ErrorCategory.USER,
786
+ details: { page }
787
+ },
788
+ new Error("page must be >= 0")
789
+ );
790
+ }
791
+ const { field, direction } = this.parseOrderBy(orderBy, "ASC");
792
+ const table = await this.client.openTable(TABLE_MESSAGES);
793
+ const threadCondition = threadIds.length === 1 ? `thread_id = '${this.escapeSql(threadIds[0])}'` : `thread_id IN (${threadIds.map((t) => `'${this.escapeSql(t)}'`).join(", ")})`;
794
+ const conditions = [threadCondition];
795
+ if (resourceId) {
796
+ conditions.push(`\`resourceId\` = '${this.escapeSql(resourceId)}'`);
797
+ }
798
+ if (filter?.dateRange?.start) {
799
+ const startTime = filter.dateRange.start instanceof Date ? filter.dateRange.start.getTime() : new Date(filter.dateRange.start).getTime();
800
+ conditions.push(`\`createdAt\` >= ${startTime}`);
801
+ }
802
+ if (filter?.dateRange?.end) {
803
+ const endTime = filter.dateRange.end instanceof Date ? filter.dateRange.end.getTime() : new Date(filter.dateRange.end).getTime();
804
+ conditions.push(`\`createdAt\` <= ${endTime}`);
805
+ }
806
+ const whereClause = conditions.join(" AND ");
807
+ const total = await table.countRows(whereClause);
808
+ const query = table.query().where(whereClause);
809
+ let allRecords = await query.toArray();
810
+ allRecords.sort((a, b) => {
811
+ const aValue = field === "createdAt" ? a.createdAt : a[field];
812
+ const bValue = field === "createdAt" ? b.createdAt : b[field];
813
+ if (aValue == null && bValue == null) return 0;
814
+ if (aValue == null) return direction === "ASC" ? -1 : 1;
815
+ if (bValue == null) return direction === "ASC" ? 1 : -1;
816
+ if (typeof aValue === "string" && typeof bValue === "string") {
817
+ return direction === "ASC" ? aValue.localeCompare(bValue) : bValue.localeCompare(aValue);
818
+ }
819
+ return direction === "ASC" ? aValue - bValue : bValue - aValue;
820
+ });
821
+ const paginatedRecords = allRecords.slice(offset, offset + perPage);
822
+ const messages = paginatedRecords.map((row) => this.normalizeMessage(row));
823
+ if (total === 0 && messages.length === 0 && (!include || include.length === 0)) {
824
+ return {
825
+ messages: [],
826
+ total: 0,
827
+ page,
828
+ perPage: perPageForResponse,
829
+ hasMore: false
830
+ };
831
+ }
832
+ const messageIds = new Set(messages.map((m) => m.id));
833
+ if (include && include.length > 0) {
834
+ const threadIds2 = [...new Set(include.map((item) => item.threadId || threadId))];
835
+ const allThreadMessages = [];
836
+ for (const tid of threadIds2) {
837
+ const threadQuery = table.query().where(`thread_id = '${tid}'`);
838
+ let threadRecords = await threadQuery.toArray();
839
+ allThreadMessages.push(...threadRecords);
840
+ }
841
+ allThreadMessages.sort((a, b) => a.createdAt - b.createdAt);
842
+ const contextMessages = this.processMessagesWithContext(allThreadMessages, include);
843
+ const includedMessages = contextMessages.map((row) => this.normalizeMessage(row));
844
+ for (const includeMsg of includedMessages) {
845
+ if (!messageIds.has(includeMsg.id)) {
846
+ messages.push(includeMsg);
847
+ messageIds.add(includeMsg.id);
848
+ }
849
+ }
962
850
  }
963
- throw new MastraError(
851
+ const list = new MessageList().add(messages, "memory");
852
+ let finalMessages = list.get.all.db();
853
+ finalMessages = finalMessages.sort((a, b) => {
854
+ const aValue = field === "createdAt" ? new Date(a.createdAt).getTime() : a[field];
855
+ const bValue = field === "createdAt" ? new Date(b.createdAt).getTime() : b[field];
856
+ if (aValue == null && bValue == null) return 0;
857
+ if (aValue == null) return direction === "ASC" ? -1 : 1;
858
+ if (bValue == null) return direction === "ASC" ? 1 : -1;
859
+ if (typeof aValue === "string" && typeof bValue === "string") {
860
+ return direction === "ASC" ? aValue.localeCompare(bValue) : bValue.localeCompare(aValue);
861
+ }
862
+ return direction === "ASC" ? aValue - bValue : bValue - aValue;
863
+ });
864
+ const returnedThreadMessageIds = new Set(finalMessages.filter((m) => m.threadId === threadId).map((m) => m.id));
865
+ const allThreadMessagesReturned = returnedThreadMessageIds.size >= total;
866
+ const fetchedAll = perPageInput === false || allThreadMessagesReturned;
867
+ const hasMore = !fetchedAll && offset + perPage < total;
868
+ return {
869
+ messages: finalMessages,
870
+ total,
871
+ page,
872
+ perPage: perPageForResponse,
873
+ hasMore
874
+ };
875
+ } catch (error) {
876
+ const mastraError = new MastraError(
964
877
  {
965
- id: createStorageErrorId("LANCE", "DROP_TABLE", "FAILED"),
878
+ id: createStorageErrorId("LANCE", "LIST_MESSAGES", "FAILED"),
966
879
  domain: ErrorDomain.STORAGE,
967
880
  category: ErrorCategory.THIRD_PARTY,
968
- details: { tableName }
881
+ details: {
882
+ threadId: Array.isArray(threadId) ? threadId.join(",") : threadId,
883
+ resourceId: resourceId ?? ""
884
+ }
969
885
  },
970
886
  error
971
887
  );
888
+ this.logger?.error?.(mastraError.toString());
889
+ this.logger?.trackException?.(mastraError);
890
+ return {
891
+ messages: [],
892
+ total: 0,
893
+ page,
894
+ perPage: perPageForResponse,
895
+ hasMore: false
896
+ };
972
897
  }
973
898
  }
974
- async alterTable({
975
- tableName,
976
- schema,
977
- ifNotExists
978
- }) {
899
+ async saveMessages(args) {
979
900
  try {
980
- if (!this.client) {
981
- throw new Error("LanceDB client not initialized. Call LanceStorage.create() first.");
982
- }
983
- if (!tableName) {
984
- throw new Error("tableName is required for alterTable.");
901
+ const { messages } = args;
902
+ if (messages.length === 0) {
903
+ return { messages: [] };
985
904
  }
986
- if (!schema) {
987
- throw new Error("schema is required for alterTable.");
905
+ const threadId = messages[0]?.threadId;
906
+ if (!threadId) {
907
+ throw new Error("Thread ID is required");
988
908
  }
989
- if (!ifNotExists || ifNotExists.length === 0) {
990
- this.logger.debug("No columns specified to add in alterTable, skipping.");
991
- return;
909
+ for (const message of messages) {
910
+ if (!message.id) {
911
+ throw new Error("Message ID is required");
912
+ }
913
+ if (!message.threadId) {
914
+ throw new Error("Thread ID is required for all messages");
915
+ }
916
+ if (message.resourceId === null || message.resourceId === void 0) {
917
+ throw new Error("Resource ID cannot be null or undefined");
918
+ }
919
+ if (!message.content) {
920
+ throw new Error("Message content is required");
921
+ }
992
922
  }
993
- } catch (validationError) {
994
- throw new MastraError(
995
- {
996
- id: createStorageErrorId("LANCE", "ALTER_TABLE", "INVALID_ARGS"),
997
- domain: ErrorDomain.STORAGE,
998
- category: ErrorCategory.USER,
999
- text: validationError.message,
1000
- details: { tableName }
1001
- },
1002
- validationError
1003
- );
1004
- }
1005
- try {
1006
- const table = await this.client.openTable(tableName);
1007
- const currentSchema = await table.schema();
1008
- const existingFields = new Set(currentSchema.fields.map((f) => f.name));
1009
- const typeMap = {
1010
- text: "string",
1011
- integer: "int",
1012
- bigint: "bigint",
1013
- timestamp: "timestamp",
1014
- jsonb: "string",
1015
- uuid: "string"
1016
- };
1017
- const columnsToAdd = ifNotExists.filter((col) => schema[col] && !existingFields.has(col)).map((col) => {
1018
- const colDef = schema[col];
923
+ const transformedMessages = messages.map((message) => {
924
+ const { threadId: threadId2, type, ...rest } = message;
1019
925
  return {
1020
- name: col,
1021
- valueSql: colDef?.nullable ? `cast(NULL as ${typeMap[colDef.type ?? "text"]})` : `cast(${this.getDefaultValue(colDef?.type ?? "text")} as ${typeMap[colDef?.type ?? "text"]})`
926
+ ...rest,
927
+ thread_id: threadId2,
928
+ type: type ?? "v2",
929
+ content: JSON.stringify(message.content)
1022
930
  };
1023
931
  });
1024
- if (columnsToAdd.length > 0) {
1025
- await table.addColumns(columnsToAdd);
1026
- this.logger?.info?.(`Added columns [${columnsToAdd.map((c) => c.name).join(", ")}] to table ${tableName}`);
1027
- }
932
+ const table = await this.client.openTable(TABLE_MESSAGES);
933
+ await table.mergeInsert("id").whenMatchedUpdateAll().whenNotMatchedInsertAll().execute(transformedMessages);
934
+ const threadsTable = await this.client.openTable(TABLE_THREADS);
935
+ const currentTime = (/* @__PURE__ */ new Date()).getTime();
936
+ const updateRecord = { id: threadId, updatedAt: currentTime };
937
+ await threadsTable.mergeInsert("id").whenMatchedUpdateAll().whenNotMatchedInsertAll().execute([updateRecord]);
938
+ const list = new MessageList().add(messages, "memory");
939
+ return { messages: list.get.all.db() };
1028
940
  } catch (error) {
1029
941
  throw new MastraError(
1030
942
  {
1031
- id: createStorageErrorId("LANCE", "ALTER_TABLE", "FAILED"),
943
+ id: createStorageErrorId("LANCE", "SAVE_MESSAGES", "FAILED"),
1032
944
  domain: ErrorDomain.STORAGE,
1033
- category: ErrorCategory.THIRD_PARTY,
1034
- details: { tableName }
945
+ category: ErrorCategory.THIRD_PARTY
1035
946
  },
1036
947
  error
1037
948
  );
1038
949
  }
1039
950
  }
1040
- async clearTable({ tableName }) {
951
+ async listThreadsByResourceId(args) {
1041
952
  try {
1042
- if (!this.client) {
1043
- throw new Error("LanceDB client not initialized. Call LanceStorage.create() first.");
1044
- }
1045
- if (!tableName) {
1046
- throw new Error("tableName is required for clearTable.");
953
+ const { resourceId, page = 0, perPage: perPageInput, orderBy } = args;
954
+ const perPage = normalizePerPage(perPageInput, 100);
955
+ if (page < 0) {
956
+ throw new MastraError(
957
+ {
958
+ id: createStorageErrorId("LANCE", "LIST_THREADS_BY_RESOURCE_ID", "INVALID_PAGE"),
959
+ domain: ErrorDomain.STORAGE,
960
+ category: ErrorCategory.USER,
961
+ details: { page }
962
+ },
963
+ new Error("page must be >= 0")
964
+ );
1047
965
  }
1048
- } catch (validationError) {
1049
- throw new MastraError(
1050
- {
1051
- id: createStorageErrorId("LANCE", "CLEAR_TABLE", "INVALID_ARGS"),
1052
- domain: ErrorDomain.STORAGE,
1053
- category: ErrorCategory.USER,
1054
- text: validationError.message,
1055
- details: { tableName }
1056
- },
1057
- validationError
966
+ const { offset, perPage: perPageForResponse } = calculatePagination(page, perPageInput, perPage);
967
+ const { field, direction } = this.parseOrderBy(orderBy);
968
+ const table = await this.client.openTable(TABLE_THREADS);
969
+ const total = await table.countRows(`\`resourceId\` = '${this.escapeSql(resourceId)}'`);
970
+ const query = table.query().where(`\`resourceId\` = '${this.escapeSql(resourceId)}'`);
971
+ const records = await query.toArray();
972
+ records.sort((a, b) => {
973
+ const aValue = ["createdAt", "updatedAt"].includes(field) ? new Date(a[field]).getTime() : a[field];
974
+ const bValue = ["createdAt", "updatedAt"].includes(field) ? new Date(b[field]).getTime() : b[field];
975
+ if (aValue == null && bValue == null) return 0;
976
+ if (aValue == null) return direction === "ASC" ? -1 : 1;
977
+ if (bValue == null) return direction === "ASC" ? 1 : -1;
978
+ if (typeof aValue === "string" && typeof bValue === "string") {
979
+ return direction === "ASC" ? aValue.localeCompare(bValue) : bValue.localeCompare(aValue);
980
+ }
981
+ return direction === "ASC" ? aValue - bValue : bValue - aValue;
982
+ });
983
+ const paginatedRecords = records.slice(offset, offset + perPage);
984
+ const schema = await getTableSchema({ tableName: TABLE_THREADS, client: this.client });
985
+ const threads = paginatedRecords.map(
986
+ (record) => processResultWithTypeConversion(record, schema)
1058
987
  );
1059
- }
1060
- try {
1061
- const table = await this.client.openTable(tableName);
1062
- await table.delete("1=1");
988
+ return {
989
+ threads,
990
+ total,
991
+ page,
992
+ perPage: perPageForResponse,
993
+ hasMore: offset + perPage < total
994
+ };
1063
995
  } catch (error) {
1064
996
  throw new MastraError(
1065
997
  {
1066
- id: createStorageErrorId("LANCE", "CLEAR_TABLE", "FAILED"),
998
+ id: createStorageErrorId("LANCE", "LIST_THREADS_BY_RESOURCE_ID", "FAILED"),
1067
999
  domain: ErrorDomain.STORAGE,
1068
- category: ErrorCategory.THIRD_PARTY,
1069
- details: { tableName }
1000
+ category: ErrorCategory.THIRD_PARTY
1070
1001
  },
1071
1002
  error
1072
1003
  );
1073
1004
  }
1074
1005
  }
1075
- async insert({ tableName, record }) {
1076
- try {
1077
- if (!this.client) {
1078
- throw new Error("LanceDB client not initialized. Call LanceStorage.create() first.");
1079
- }
1080
- if (!tableName) {
1081
- throw new Error("tableName is required for insert.");
1006
+ /**
1007
+ * Processes messages to include context messages based on withPreviousMessages and withNextMessages
1008
+ * @param records - The sorted array of records to process
1009
+ * @param include - The array of include specifications with context parameters
1010
+ * @returns The processed array with context messages included
1011
+ */
1012
+ processMessagesWithContext(records, include) {
1013
+ const messagesWithContext = include.filter((item) => item.withPreviousMessages || item.withNextMessages);
1014
+ if (messagesWithContext.length === 0) {
1015
+ return records;
1016
+ }
1017
+ const messageIndexMap = /* @__PURE__ */ new Map();
1018
+ records.forEach((message, index) => {
1019
+ messageIndexMap.set(message.id, index);
1020
+ });
1021
+ const additionalIndices = /* @__PURE__ */ new Set();
1022
+ for (const item of messagesWithContext) {
1023
+ const messageIndex = messageIndexMap.get(item.id);
1024
+ if (messageIndex !== void 0) {
1025
+ if (item.withPreviousMessages) {
1026
+ const startIdx = Math.max(0, messageIndex - item.withPreviousMessages);
1027
+ for (let i = startIdx; i < messageIndex; i++) {
1028
+ additionalIndices.add(i);
1029
+ }
1030
+ }
1031
+ if (item.withNextMessages) {
1032
+ const endIdx = Math.min(records.length - 1, messageIndex + item.withNextMessages);
1033
+ for (let i = messageIndex + 1; i <= endIdx; i++) {
1034
+ additionalIndices.add(i);
1035
+ }
1036
+ }
1082
1037
  }
1083
- if (!record || Object.keys(record).length === 0) {
1084
- throw new Error("record is required and cannot be empty for insert.");
1038
+ }
1039
+ if (additionalIndices.size === 0) {
1040
+ return records;
1041
+ }
1042
+ const originalMatchIds = new Set(include.map((item) => item.id));
1043
+ const allIndices = /* @__PURE__ */ new Set();
1044
+ records.forEach((record, index) => {
1045
+ if (originalMatchIds.has(record.id)) {
1046
+ allIndices.add(index);
1085
1047
  }
1086
- } catch (validationError) {
1087
- throw new MastraError(
1088
- {
1089
- id: createStorageErrorId("LANCE", "INSERT", "INVALID_ARGS"),
1090
- domain: ErrorDomain.STORAGE,
1091
- category: ErrorCategory.USER,
1092
- text: validationError.message,
1093
- details: { tableName }
1094
- },
1095
- validationError
1096
- );
1048
+ });
1049
+ additionalIndices.forEach((index) => {
1050
+ allIndices.add(index);
1051
+ });
1052
+ return Array.from(allIndices).sort((a, b) => a - b).map((index) => records[index]);
1053
+ }
1054
+ /**
1055
+ * Parse message data from LanceDB record format to MastraDBMessage format
1056
+ */
1057
+ parseMessageData(data) {
1058
+ const { thread_id, ...rest } = data;
1059
+ return {
1060
+ ...rest,
1061
+ threadId: thread_id,
1062
+ content: typeof data.content === "string" ? (() => {
1063
+ try {
1064
+ return JSON.parse(data.content);
1065
+ } catch {
1066
+ return data.content;
1067
+ }
1068
+ })() : data.content,
1069
+ createdAt: new Date(data.createdAt),
1070
+ updatedAt: new Date(data.updatedAt)
1071
+ };
1072
+ }
1073
+ async updateMessages(args) {
1074
+ const { messages } = args;
1075
+ this.logger.debug("Updating messages", { count: messages.length });
1076
+ if (!messages.length) {
1077
+ return [];
1097
1078
  }
1079
+ const updatedMessages = [];
1080
+ const affectedThreadIds = /* @__PURE__ */ new Set();
1098
1081
  try {
1099
- const table = await this.client.openTable(tableName);
1100
- const primaryId = getPrimaryKeys(tableName);
1101
- const processedRecord = { ...record };
1102
- for (const key in processedRecord) {
1103
- if (processedRecord[key] !== null && typeof processedRecord[key] === "object" && !(processedRecord[key] instanceof Date)) {
1104
- this.logger.debug("Converting object to JSON string: ", processedRecord[key]);
1105
- processedRecord[key] = JSON.stringify(processedRecord[key]);
1082
+ for (const updateData of messages) {
1083
+ const { id, ...updates } = updateData;
1084
+ const existingMessage = await this.#db.load({ tableName: TABLE_MESSAGES, keys: { id } });
1085
+ if (!existingMessage) {
1086
+ this.logger.warn("Message not found for update", { id });
1087
+ continue;
1088
+ }
1089
+ const existingMsg = this.parseMessageData(existingMessage);
1090
+ const originalThreadId = existingMsg.threadId;
1091
+ affectedThreadIds.add(originalThreadId);
1092
+ const updatePayload = {};
1093
+ if ("role" in updates && updates.role !== void 0) updatePayload.role = updates.role;
1094
+ if ("type" in updates && updates.type !== void 0) updatePayload.type = updates.type;
1095
+ if ("resourceId" in updates && updates.resourceId !== void 0) updatePayload.resourceId = updates.resourceId;
1096
+ if ("threadId" in updates && updates.threadId !== void 0 && updates.threadId !== null) {
1097
+ updatePayload.thread_id = updates.threadId;
1098
+ affectedThreadIds.add(updates.threadId);
1099
+ }
1100
+ if (updates.content) {
1101
+ const existingContent = existingMsg.content;
1102
+ let newContent = { ...existingContent };
1103
+ if (updates.content.metadata !== void 0) {
1104
+ newContent.metadata = {
1105
+ ...existingContent.metadata || {},
1106
+ ...updates.content.metadata || {}
1107
+ };
1108
+ }
1109
+ if (updates.content.content !== void 0) {
1110
+ newContent.content = updates.content.content;
1111
+ }
1112
+ if ("parts" in updates.content && updates.content.parts !== void 0) {
1113
+ newContent.parts = updates.content.parts;
1114
+ }
1115
+ updatePayload.content = JSON.stringify(newContent);
1116
+ }
1117
+ await this.#db.insert({ tableName: TABLE_MESSAGES, record: { id, ...updatePayload } });
1118
+ const updatedMessage = await this.#db.load({ tableName: TABLE_MESSAGES, keys: { id } });
1119
+ if (updatedMessage) {
1120
+ updatedMessages.push(this.parseMessageData(updatedMessage));
1106
1121
  }
1107
1122
  }
1108
- await table.mergeInsert(primaryId).whenMatchedUpdateAll().whenNotMatchedInsertAll().execute([processedRecord]);
1123
+ for (const threadId of affectedThreadIds) {
1124
+ await this.#db.insert({
1125
+ tableName: TABLE_THREADS,
1126
+ record: { id: threadId, updatedAt: Date.now() }
1127
+ });
1128
+ }
1129
+ return updatedMessages;
1109
1130
  } catch (error) {
1110
1131
  throw new MastraError(
1111
1132
  {
1112
- id: createStorageErrorId("LANCE", "INSERT", "FAILED"),
1133
+ id: createStorageErrorId("LANCE", "UPDATE_MESSAGES", "FAILED"),
1113
1134
  domain: ErrorDomain.STORAGE,
1114
1135
  category: ErrorCategory.THIRD_PARTY,
1115
- details: { tableName }
1136
+ details: { count: messages.length }
1116
1137
  },
1117
1138
  error
1118
1139
  );
1119
1140
  }
1120
1141
  }
1121
- async batchInsert({ tableName, records }) {
1142
+ async getResourceById({ resourceId }) {
1122
1143
  try {
1123
- if (!this.client) {
1124
- throw new Error("LanceDB client not initialized. Call LanceStorage.create() first.");
1144
+ const resource = await this.#db.load({ tableName: TABLE_RESOURCES, keys: { id: resourceId } });
1145
+ if (!resource) {
1146
+ return null;
1125
1147
  }
1126
- if (!tableName) {
1127
- throw new Error("tableName is required for batchInsert.");
1148
+ let createdAt;
1149
+ let updatedAt;
1150
+ try {
1151
+ if (resource.createdAt instanceof Date) {
1152
+ createdAt = resource.createdAt;
1153
+ } else if (typeof resource.createdAt === "string") {
1154
+ createdAt = new Date(resource.createdAt);
1155
+ } else if (typeof resource.createdAt === "number") {
1156
+ createdAt = new Date(resource.createdAt);
1157
+ } else {
1158
+ createdAt = /* @__PURE__ */ new Date();
1159
+ }
1160
+ if (isNaN(createdAt.getTime())) {
1161
+ createdAt = /* @__PURE__ */ new Date();
1162
+ }
1163
+ } catch {
1164
+ createdAt = /* @__PURE__ */ new Date();
1128
1165
  }
1129
- if (!records || records.length === 0) {
1130
- throw new Error("records array is required and cannot be empty for batchInsert.");
1166
+ try {
1167
+ if (resource.updatedAt instanceof Date) {
1168
+ updatedAt = resource.updatedAt;
1169
+ } else if (typeof resource.updatedAt === "string") {
1170
+ updatedAt = new Date(resource.updatedAt);
1171
+ } else if (typeof resource.updatedAt === "number") {
1172
+ updatedAt = new Date(resource.updatedAt);
1173
+ } else {
1174
+ updatedAt = /* @__PURE__ */ new Date();
1175
+ }
1176
+ if (isNaN(updatedAt.getTime())) {
1177
+ updatedAt = /* @__PURE__ */ new Date();
1178
+ }
1179
+ } catch {
1180
+ updatedAt = /* @__PURE__ */ new Date();
1131
1181
  }
1132
- } catch (validationError) {
1133
- throw new MastraError(
1134
- {
1135
- id: createStorageErrorId("LANCE", "BATCH_INSERT", "INVALID_ARGS"),
1136
- domain: ErrorDomain.STORAGE,
1137
- category: ErrorCategory.USER,
1138
- text: validationError.message,
1139
- details: { tableName }
1140
- },
1141
- validationError
1142
- );
1143
- }
1144
- try {
1145
- const table = await this.client.openTable(tableName);
1146
- const primaryId = getPrimaryKeys(tableName);
1147
- const processedRecords = records.map((record) => {
1148
- const processedRecord = { ...record };
1149
- for (const key in processedRecord) {
1150
- if (processedRecord[key] == null) continue;
1151
- if (processedRecord[key] !== null && typeof processedRecord[key] === "object" && !(processedRecord[key] instanceof Date)) {
1152
- processedRecord[key] = JSON.stringify(processedRecord[key]);
1153
- }
1182
+ let workingMemory = resource.workingMemory;
1183
+ if (workingMemory === null || workingMemory === void 0) {
1184
+ workingMemory = void 0;
1185
+ } else if (workingMemory === "") {
1186
+ workingMemory = "";
1187
+ } else if (typeof workingMemory === "object") {
1188
+ workingMemory = JSON.stringify(workingMemory);
1189
+ }
1190
+ let metadata = resource.metadata;
1191
+ if (metadata === "" || metadata === null || metadata === void 0) {
1192
+ metadata = void 0;
1193
+ } else if (typeof metadata === "string") {
1194
+ try {
1195
+ metadata = JSON.parse(metadata);
1196
+ } catch {
1197
+ metadata = metadata;
1154
1198
  }
1155
- return processedRecord;
1156
- });
1157
- await table.mergeInsert(primaryId).whenMatchedUpdateAll().whenNotMatchedInsertAll().execute(processedRecords);
1199
+ }
1200
+ return {
1201
+ ...resource,
1202
+ createdAt,
1203
+ updatedAt,
1204
+ workingMemory,
1205
+ metadata
1206
+ };
1158
1207
  } catch (error) {
1159
1208
  throw new MastraError(
1160
1209
  {
1161
- id: createStorageErrorId("LANCE", "BATCH_INSERT", "FAILED"),
1210
+ id: createStorageErrorId("LANCE", "GET_RESOURCE_BY_ID", "FAILED"),
1162
1211
  domain: ErrorDomain.STORAGE,
1163
- category: ErrorCategory.THIRD_PARTY,
1164
- details: { tableName }
1212
+ category: ErrorCategory.THIRD_PARTY
1165
1213
  },
1166
1214
  error
1167
1215
  );
1168
1216
  }
1169
1217
  }
1170
- async load({ tableName, keys }) {
1171
- try {
1172
- if (!this.client) {
1173
- throw new Error("LanceDB client not initialized. Call LanceStorage.create() first.");
1174
- }
1175
- if (!tableName) {
1176
- throw new Error("tableName is required for load.");
1177
- }
1178
- if (!keys || Object.keys(keys).length === 0) {
1179
- throw new Error("keys are required and cannot be empty for load.");
1180
- }
1181
- } catch (validationError) {
1182
- throw new MastraError(
1183
- {
1184
- id: createStorageErrorId("LANCE", "LOAD", "INVALID_ARGS"),
1185
- domain: ErrorDomain.STORAGE,
1186
- category: ErrorCategory.USER,
1187
- text: validationError.message,
1188
- details: { tableName }
1189
- },
1190
- validationError
1191
- );
1192
- }
1218
+ async saveResource({ resource }) {
1193
1219
  try {
1194
- const table = await this.client.openTable(tableName);
1195
- const tableSchema = await getTableSchema({ tableName, client: this.client });
1196
- const query = table.query();
1197
- if (Object.keys(keys).length > 0) {
1198
- validateKeyTypes(keys, tableSchema);
1199
- const filterConditions = Object.entries(keys).map(([key, value]) => {
1200
- const isCamelCase = /^[a-z][a-zA-Z]*$/.test(key) && /[A-Z]/.test(key);
1201
- const quotedKey = isCamelCase ? `\`${key}\`` : key;
1202
- if (typeof value === "string") {
1203
- return `${quotedKey} = '${value}'`;
1204
- } else if (value === null) {
1205
- return `${quotedKey} IS NULL`;
1206
- } else {
1207
- return `${quotedKey} = ${value}`;
1208
- }
1209
- }).join(" AND ");
1210
- this.logger.debug("where clause generated: " + filterConditions);
1211
- query.where(filterConditions);
1212
- }
1213
- const result = await query.limit(1).toArray();
1214
- if (result.length === 0) {
1215
- this.logger.debug("No record found");
1216
- return null;
1217
- }
1218
- return processResultWithTypeConversion(result[0], tableSchema);
1220
+ const record = {
1221
+ ...resource,
1222
+ metadata: resource.metadata ? JSON.stringify(resource.metadata) : "",
1223
+ createdAt: resource.createdAt.getTime(),
1224
+ // Store as timestamp (milliseconds)
1225
+ updatedAt: resource.updatedAt.getTime()
1226
+ // Store as timestamp (milliseconds)
1227
+ };
1228
+ const table = await this.client.openTable(TABLE_RESOURCES);
1229
+ await table.add([record], { mode: "append" });
1230
+ return resource;
1219
1231
  } catch (error) {
1220
- if (error instanceof MastraError) throw error;
1221
1232
  throw new MastraError(
1222
1233
  {
1223
- id: createStorageErrorId("LANCE", "LOAD", "FAILED"),
1234
+ id: createStorageErrorId("LANCE", "SAVE_RESOURCE", "FAILED"),
1224
1235
  domain: ErrorDomain.STORAGE,
1225
- category: ErrorCategory.THIRD_PARTY,
1226
- details: { tableName, keyCount: Object.keys(keys).length, firstKey: Object.keys(keys)[0] ?? "" }
1236
+ category: ErrorCategory.THIRD_PARTY
1227
1237
  },
1228
1238
  error
1229
1239
  );
1230
1240
  }
1231
1241
  }
1242
+ async updateResource({
1243
+ resourceId,
1244
+ workingMemory,
1245
+ metadata
1246
+ }) {
1247
+ const maxRetries = 3;
1248
+ for (let attempt = 0; attempt < maxRetries; attempt++) {
1249
+ try {
1250
+ const existingResource = await this.getResourceById({ resourceId });
1251
+ if (!existingResource) {
1252
+ const newResource = {
1253
+ id: resourceId,
1254
+ workingMemory,
1255
+ metadata: metadata || {},
1256
+ createdAt: /* @__PURE__ */ new Date(),
1257
+ updatedAt: /* @__PURE__ */ new Date()
1258
+ };
1259
+ return this.saveResource({ resource: newResource });
1260
+ }
1261
+ const updatedResource = {
1262
+ ...existingResource,
1263
+ workingMemory: workingMemory !== void 0 ? workingMemory : existingResource.workingMemory,
1264
+ metadata: {
1265
+ ...existingResource.metadata,
1266
+ ...metadata
1267
+ },
1268
+ updatedAt: /* @__PURE__ */ new Date()
1269
+ };
1270
+ const record = {
1271
+ id: resourceId,
1272
+ workingMemory: updatedResource.workingMemory || "",
1273
+ metadata: updatedResource.metadata ? JSON.stringify(updatedResource.metadata) : "",
1274
+ updatedAt: updatedResource.updatedAt.getTime()
1275
+ // Store as timestamp (milliseconds)
1276
+ };
1277
+ const table = await this.client.openTable(TABLE_RESOURCES);
1278
+ await table.mergeInsert("id").whenMatchedUpdateAll().whenNotMatchedInsertAll().execute([record]);
1279
+ return updatedResource;
1280
+ } catch (error) {
1281
+ if (error.message?.includes("Commit conflict") && attempt < maxRetries - 1) {
1282
+ const delay = Math.pow(2, attempt) * 10;
1283
+ await new Promise((resolve) => setTimeout(resolve, delay));
1284
+ continue;
1285
+ }
1286
+ throw new MastraError(
1287
+ {
1288
+ id: createStorageErrorId("LANCE", "UPDATE_RESOURCE", "FAILED"),
1289
+ domain: ErrorDomain.STORAGE,
1290
+ category: ErrorCategory.THIRD_PARTY
1291
+ },
1292
+ error
1293
+ );
1294
+ }
1295
+ }
1296
+ throw new Error("Unexpected end of retry loop");
1297
+ }
1232
1298
  };
1233
1299
  var StoreScoresLance = class extends ScoresStorage {
1234
1300
  client;
1235
- constructor({ client }) {
1301
+ #db;
1302
+ constructor(config) {
1236
1303
  super();
1304
+ const client = resolveLanceConfig(config);
1237
1305
  this.client = client;
1306
+ this.#db = new LanceDB({ client });
1307
+ }
1308
+ async init() {
1309
+ await this.#db.createTable({ tableName: TABLE_SCORERS, schema: SCORERS_SCHEMA });
1310
+ await this.#db.alterTable({
1311
+ tableName: TABLE_SCORERS,
1312
+ schema: SCORERS_SCHEMA,
1313
+ ifNotExists: ["spanId", "requestContext"]
1314
+ });
1315
+ }
1316
+ async dangerouslyClearAll() {
1317
+ await this.#db.clearTable({ tableName: TABLE_SCORERS });
1238
1318
  }
1239
1319
  async saveScore(score) {
1240
1320
  let validatedScore;
@@ -1248,7 +1328,7 @@ var StoreScoresLance = class extends ScoresStorage {
1248
1328
  domain: ErrorDomain.STORAGE,
1249
1329
  category: ErrorCategory.USER,
1250
1330
  details: {
1251
- scorer: score.scorer?.id ?? "unknown",
1331
+ scorer: typeof score.scorer?.id === "string" ? score.scorer.id : String(score.scorer?.id ?? "unknown"),
1252
1332
  entityId: score.entityId ?? "unknown",
1253
1333
  entityType: score.entityType ?? "unknown",
1254
1334
  traceId: score.traceId ?? "",
@@ -1518,6 +1598,9 @@ var StoreScoresLance = class extends ScoresStorage {
1518
1598
  }
1519
1599
  }
1520
1600
  };
1601
+ function escapeSql(str) {
1602
+ return str.replace(/'/g, "''");
1603
+ }
1521
1604
  function parseWorkflowRun(row) {
1522
1605
  let parsedSnapshot = row.snapshot;
1523
1606
  if (typeof parsedSnapshot === "string") {
@@ -1538,42 +1621,88 @@ function parseWorkflowRun(row) {
1538
1621
  }
1539
1622
  var StoreWorkflowsLance = class extends WorkflowsStorage {
1540
1623
  client;
1541
- constructor({ client }) {
1624
+ #db;
1625
+ constructor(config) {
1542
1626
  super();
1627
+ const client = resolveLanceConfig(config);
1543
1628
  this.client = client;
1629
+ this.#db = new LanceDB({ client });
1630
+ }
1631
+ async init() {
1632
+ const schema = TABLE_SCHEMAS[TABLE_WORKFLOW_SNAPSHOT];
1633
+ await this.#db.createTable({ tableName: TABLE_WORKFLOW_SNAPSHOT, schema });
1634
+ await this.#db.alterTable({
1635
+ tableName: TABLE_WORKFLOW_SNAPSHOT,
1636
+ schema,
1637
+ ifNotExists: ["resourceId"]
1638
+ });
1544
1639
  }
1545
- updateWorkflowResults({
1546
- // workflowName,
1547
- // runId,
1548
- // stepId,
1549
- // result,
1550
- // requestContext,
1640
+ async dangerouslyClearAll() {
1641
+ await this.#db.clearTable({ tableName: TABLE_WORKFLOW_SNAPSHOT });
1642
+ }
1643
+ async updateWorkflowResults({
1644
+ workflowName,
1645
+ runId,
1646
+ stepId,
1647
+ result,
1648
+ requestContext
1551
1649
  }) {
1552
- throw new Error("Method not implemented.");
1650
+ let snapshot = await this.loadWorkflowSnapshot({ workflowName, runId });
1651
+ if (!snapshot) {
1652
+ snapshot = {
1653
+ context: {},
1654
+ activePaths: [],
1655
+ timestamp: Date.now(),
1656
+ suspendedPaths: {},
1657
+ activeStepsPath: {},
1658
+ resumeLabels: {},
1659
+ serializedStepGraph: [],
1660
+ status: "pending",
1661
+ value: {},
1662
+ waitingPaths: {},
1663
+ runId,
1664
+ requestContext: {}
1665
+ };
1666
+ }
1667
+ snapshot.context[stepId] = result;
1668
+ snapshot.requestContext = { ...snapshot.requestContext, ...requestContext };
1669
+ await this.persistWorkflowSnapshot({ workflowName, runId, snapshot });
1670
+ return snapshot.context;
1553
1671
  }
1554
- updateWorkflowState({
1555
- // workflowName,
1556
- // runId,
1557
- // opts,
1672
+ async updateWorkflowState({
1673
+ workflowName,
1674
+ runId,
1675
+ opts
1558
1676
  }) {
1559
- throw new Error("Method not implemented.");
1677
+ const snapshot = await this.loadWorkflowSnapshot({ workflowName, runId });
1678
+ if (!snapshot) {
1679
+ return void 0;
1680
+ }
1681
+ if (!snapshot.context) {
1682
+ throw new Error(`Snapshot not found for runId ${runId}`);
1683
+ }
1684
+ const updatedSnapshot = { ...snapshot, ...opts };
1685
+ await this.persistWorkflowSnapshot({ workflowName, runId, snapshot: updatedSnapshot });
1686
+ return updatedSnapshot;
1560
1687
  }
1561
1688
  async persistWorkflowSnapshot({
1562
1689
  workflowName,
1563
1690
  runId,
1564
1691
  resourceId,
1565
- snapshot
1692
+ snapshot,
1693
+ createdAt,
1694
+ updatedAt
1566
1695
  }) {
1567
1696
  try {
1568
1697
  const table = await this.client.openTable(TABLE_WORKFLOW_SNAPSHOT);
1569
- const query = table.query().where(`workflow_name = '${workflowName}' AND run_id = '${runId}'`);
1698
+ const query = table.query().where(`workflow_name = '${escapeSql(workflowName)}' AND run_id = '${escapeSql(runId)}'`);
1570
1699
  const records = await query.toArray();
1571
- let createdAt;
1572
- const now = Date.now();
1700
+ let createdAtValue;
1701
+ const now = createdAt?.getTime() ?? Date.now();
1573
1702
  if (records.length > 0) {
1574
- createdAt = records[0].createdAt ?? now;
1703
+ createdAtValue = records[0].createdAt ?? now;
1575
1704
  } else {
1576
- createdAt = now;
1705
+ createdAtValue = now;
1577
1706
  }
1578
1707
  const { status, value, ...rest } = snapshot;
1579
1708
  const record = {
@@ -1582,8 +1711,8 @@ var StoreWorkflowsLance = class extends WorkflowsStorage {
1582
1711
  resourceId,
1583
1712
  snapshot: JSON.stringify({ status, value, ...rest }),
1584
1713
  // this is to ensure status is always just before value, for when querying the db by status
1585
- createdAt,
1586
- updatedAt: now
1714
+ createdAt: createdAtValue,
1715
+ updatedAt: updatedAt ?? now
1587
1716
  };
1588
1717
  await table.mergeInsert(["workflow_name", "run_id"]).whenMatchedUpdateAll().whenNotMatchedInsertAll().execute([record]);
1589
1718
  } catch (error) {
@@ -1604,7 +1733,7 @@ var StoreWorkflowsLance = class extends WorkflowsStorage {
1604
1733
  }) {
1605
1734
  try {
1606
1735
  const table = await this.client.openTable(TABLE_WORKFLOW_SNAPSHOT);
1607
- const query = table.query().where(`workflow_name = '${workflowName}' AND run_id = '${runId}'`);
1736
+ const query = table.query().where(`workflow_name = '${escapeSql(workflowName)}' AND run_id = '${escapeSql(runId)}'`);
1608
1737
  const records = await query.toArray();
1609
1738
  return records.length > 0 ? JSON.parse(records[0].snapshot) : null;
1610
1739
  } catch (error) {
@@ -1622,9 +1751,9 @@ var StoreWorkflowsLance = class extends WorkflowsStorage {
1622
1751
  async getWorkflowRunById(args) {
1623
1752
  try {
1624
1753
  const table = await this.client.openTable(TABLE_WORKFLOW_SNAPSHOT);
1625
- let whereClause = `run_id = '${args.runId}'`;
1754
+ let whereClause = `run_id = '${escapeSql(args.runId)}'`;
1626
1755
  if (args.workflowName) {
1627
- whereClause += ` AND workflow_name = '${args.workflowName}'`;
1756
+ whereClause += ` AND workflow_name = '${escapeSql(args.workflowName)}'`;
1628
1757
  }
1629
1758
  const query = table.query().where(whereClause);
1630
1759
  const records = await query.toArray();
@@ -1646,7 +1775,7 @@ var StoreWorkflowsLance = class extends WorkflowsStorage {
1646
1775
  async deleteWorkflowRunById({ runId, workflowName }) {
1647
1776
  try {
1648
1777
  const table = await this.client.openTable(TABLE_WORKFLOW_SNAPSHOT);
1649
- const whereClause = `run_id = '${runId.replace(/'/g, "''")}' AND workflow_name = '${workflowName.replace(/'/g, "''")}'`;
1778
+ const whereClause = `run_id = '${escapeSql(runId)}' AND workflow_name = '${escapeSql(workflowName)}'`;
1650
1779
  await table.delete(whereClause);
1651
1780
  } catch (error) {
1652
1781
  throw new MastraError(
@@ -1666,14 +1795,14 @@ var StoreWorkflowsLance = class extends WorkflowsStorage {
1666
1795
  let query = table.query();
1667
1796
  const conditions = [];
1668
1797
  if (args?.workflowName) {
1669
- conditions.push(`workflow_name = '${args.workflowName.replace(/'/g, "''")}'`);
1798
+ conditions.push(`workflow_name = '${escapeSql(args.workflowName)}'`);
1670
1799
  }
1671
1800
  if (args?.status) {
1672
1801
  const escapedStatus = args.status.replace(/\\/g, "\\\\").replace(/'/g, "''").replace(/%/g, "\\%").replace(/_/g, "\\_");
1673
1802
  conditions.push(`\`snapshot\` LIKE '%"status":"${escapedStatus}","value"%'`);
1674
1803
  }
1675
1804
  if (args?.resourceId) {
1676
- conditions.push(`\`resourceId\` = '${args.resourceId}'`);
1805
+ conditions.push(`\`resourceId\` = '${escapeSql(args.resourceId)}'`);
1677
1806
  }
1678
1807
  if (args?.fromDate instanceof Date) {
1679
1808
  conditions.push(`\`createdAt\` >= ${args.fromDate.getTime()}`);
@@ -1762,12 +1891,10 @@ var LanceStorage = class _LanceStorage extends MastraStorage {
1762
1891
  const instance = new _LanceStorage(id, name, storageOptions?.disableInit);
1763
1892
  try {
1764
1893
  instance.lanceClient = await connect(uri, connectionOptions);
1765
- const operations = new StoreOperationsLance({ client: instance.lanceClient });
1766
1894
  instance.stores = {
1767
- operations: new StoreOperationsLance({ client: instance.lanceClient }),
1768
1895
  workflows: new StoreWorkflowsLance({ client: instance.lanceClient }),
1769
1896
  scores: new StoreScoresLance({ client: instance.lanceClient }),
1770
- memory: new StoreMemoryLance({ client: instance.lanceClient, operations })
1897
+ memory: new StoreMemoryLance({ client: instance.lanceClient })
1771
1898
  };
1772
1899
  return instance;
1773
1900
  } catch (e) {
@@ -1784,221 +1911,43 @@ var LanceStorage = class _LanceStorage extends MastraStorage {
1784
1911
  }
1785
1912
  }
1786
1913
  /**
1787
- * @internal
1788
- * Private constructor to enforce using the create factory method
1789
- */
1790
- constructor(id, name, disableInit) {
1791
- super({ id, name, disableInit });
1792
- const operations = new StoreOperationsLance({ client: this.lanceClient });
1793
- this.stores = {
1794
- operations: new StoreOperationsLance({ client: this.lanceClient }),
1795
- workflows: new StoreWorkflowsLance({ client: this.lanceClient }),
1796
- scores: new StoreScoresLance({ client: this.lanceClient }),
1797
- memory: new StoreMemoryLance({ client: this.lanceClient, operations })
1798
- };
1799
- }
1800
- async createTable({
1801
- tableName,
1802
- schema
1803
- }) {
1804
- return this.stores.operations.createTable({ tableName, schema });
1805
- }
1806
- async dropTable({ tableName }) {
1807
- return this.stores.operations.dropTable({ tableName });
1808
- }
1809
- async alterTable({
1810
- tableName,
1811
- schema,
1812
- ifNotExists
1813
- }) {
1814
- return this.stores.operations.alterTable({ tableName, schema, ifNotExists });
1815
- }
1816
- async clearTable({ tableName }) {
1817
- return this.stores.operations.clearTable({ tableName });
1818
- }
1819
- async insert({ tableName, record }) {
1820
- return this.stores.operations.insert({ tableName, record });
1821
- }
1822
- async batchInsert({ tableName, records }) {
1823
- return this.stores.operations.batchInsert({ tableName, records });
1824
- }
1825
- async load({ tableName, keys }) {
1826
- return this.stores.operations.load({ tableName, keys });
1827
- }
1828
- async getThreadById({ threadId }) {
1829
- return this.stores.memory.getThreadById({ threadId });
1830
- }
1831
- /**
1832
- * Saves a thread to the database. This function doesn't overwrite existing threads.
1833
- * @param thread - The thread to save
1834
- * @returns The saved thread
1914
+ * Creates a new instance of LanceStorage from a pre-configured LanceDB connection.
1915
+ * Use this when you need to configure the connection before initialization.
1916
+ *
1917
+ * @param id The unique identifier for this storage instance
1918
+ * @param name The name for this storage instance
1919
+ * @param client Pre-configured LanceDB connection
1920
+ * @param options Storage options including disableInit
1921
+ *
1922
+ * @example
1923
+ * ```typescript
1924
+ * import { connect } from '@lancedb/lancedb';
1925
+ *
1926
+ * const client = await connect('/path/to/db', {
1927
+ * // Custom connection options
1928
+ * });
1929
+ *
1930
+ * const store = LanceStorage.fromClient('my-id', 'MyStorage', client);
1931
+ * ```
1835
1932
  */
1836
- async saveThread({ thread }) {
1837
- return this.stores.memory.saveThread({ thread });
1838
- }
1839
- async updateThread({
1840
- id,
1841
- title,
1842
- metadata
1843
- }) {
1844
- return this.stores.memory.updateThread({ id, title, metadata });
1845
- }
1846
- async deleteThread({ threadId }) {
1847
- return this.stores.memory.deleteThread({ threadId });
1848
- }
1849
- get supports() {
1850
- return {
1851
- selectByIncludeResourceScope: true,
1852
- resourceWorkingMemory: true,
1853
- hasColumn: true,
1854
- createTable: true,
1855
- deleteMessages: false,
1856
- listScoresBySpan: true
1933
+ static fromClient(id, name, client, options) {
1934
+ const instance = new _LanceStorage(id, name, options?.disableInit);
1935
+ instance.lanceClient = client;
1936
+ instance.stores = {
1937
+ workflows: new StoreWorkflowsLance({ client }),
1938
+ scores: new StoreScoresLance({ client }),
1939
+ memory: new StoreMemoryLance({ client })
1857
1940
  };
1858
- }
1859
- async getResourceById({ resourceId }) {
1860
- return this.stores.memory.getResourceById({ resourceId });
1861
- }
1862
- async saveResource({ resource }) {
1863
- return this.stores.memory.saveResource({ resource });
1864
- }
1865
- async updateResource({
1866
- resourceId,
1867
- workingMemory,
1868
- metadata
1869
- }) {
1870
- return this.stores.memory.updateResource({ resourceId, workingMemory, metadata });
1941
+ return instance;
1871
1942
  }
1872
1943
  /**
1873
- * Processes messages to include context messages based on withPreviousMessages and withNextMessages
1874
- * @param records - The sorted array of records to process
1875
- * @param include - The array of include specifications with context parameters
1876
- * @returns The processed array with context messages included
1944
+ * @internal
1945
+ * Private constructor to enforce using the create factory method.
1946
+ * Note: stores is initialized in create() after the lanceClient is connected.
1877
1947
  */
1878
- processMessagesWithContext(records, include) {
1879
- const messagesWithContext = include.filter((item) => item.withPreviousMessages || item.withNextMessages);
1880
- if (messagesWithContext.length === 0) {
1881
- return records;
1882
- }
1883
- const messageIndexMap = /* @__PURE__ */ new Map();
1884
- records.forEach((message, index) => {
1885
- messageIndexMap.set(message.id, index);
1886
- });
1887
- const additionalIndices = /* @__PURE__ */ new Set();
1888
- for (const item of messagesWithContext) {
1889
- const messageIndex = messageIndexMap.get(item.id);
1890
- if (messageIndex !== void 0) {
1891
- if (item.withPreviousMessages) {
1892
- const startIdx = Math.max(0, messageIndex - item.withPreviousMessages);
1893
- for (let i = startIdx; i < messageIndex; i++) {
1894
- additionalIndices.add(i);
1895
- }
1896
- }
1897
- if (item.withNextMessages) {
1898
- const endIdx = Math.min(records.length - 1, messageIndex + item.withNextMessages);
1899
- for (let i = messageIndex + 1; i <= endIdx; i++) {
1900
- additionalIndices.add(i);
1901
- }
1902
- }
1903
- }
1904
- }
1905
- if (additionalIndices.size === 0) {
1906
- return records;
1907
- }
1908
- const originalMatchIds = new Set(include.map((item) => item.id));
1909
- const allIndices = /* @__PURE__ */ new Set();
1910
- records.forEach((record, index) => {
1911
- if (originalMatchIds.has(record.id)) {
1912
- allIndices.add(index);
1913
- }
1914
- });
1915
- additionalIndices.forEach((index) => {
1916
- allIndices.add(index);
1917
- });
1918
- return Array.from(allIndices).sort((a, b) => a - b).map((index) => records[index]);
1919
- }
1920
- async listMessagesById({ messageIds }) {
1921
- return this.stores.memory.listMessagesById({ messageIds });
1922
- }
1923
- async saveMessages(args) {
1924
- return this.stores.memory.saveMessages(args);
1925
- }
1926
- async updateMessages(_args) {
1927
- return this.stores.memory.updateMessages(_args);
1928
- }
1929
- async listWorkflowRuns(args) {
1930
- return this.stores.workflows.listWorkflowRuns(args);
1931
- }
1932
- async getWorkflowRunById(args) {
1933
- return this.stores.workflows.getWorkflowRunById(args);
1934
- }
1935
- async deleteWorkflowRunById({ runId, workflowName }) {
1936
- return this.stores.workflows.deleteWorkflowRunById({ runId, workflowName });
1937
- }
1938
- async updateWorkflowResults({
1939
- workflowName,
1940
- runId,
1941
- stepId,
1942
- result,
1943
- requestContext
1944
- }) {
1945
- return this.stores.workflows.updateWorkflowResults({ workflowName, runId, stepId, result, requestContext });
1946
- }
1947
- async updateWorkflowState({
1948
- workflowName,
1949
- runId,
1950
- opts
1951
- }) {
1952
- return this.stores.workflows.updateWorkflowState({ workflowName, runId, opts });
1953
- }
1954
- async persistWorkflowSnapshot({
1955
- workflowName,
1956
- runId,
1957
- resourceId,
1958
- snapshot
1959
- }) {
1960
- return this.stores.workflows.persistWorkflowSnapshot({ workflowName, runId, resourceId, snapshot });
1961
- }
1962
- async loadWorkflowSnapshot({
1963
- workflowName,
1964
- runId
1965
- }) {
1966
- return this.stores.workflows.loadWorkflowSnapshot({ workflowName, runId });
1967
- }
1968
- async getScoreById({ id: _id }) {
1969
- return this.stores.scores.getScoreById({ id: _id });
1970
- }
1971
- async listScoresByScorerId({
1972
- scorerId,
1973
- source,
1974
- entityId,
1975
- entityType,
1976
- pagination
1977
- }) {
1978
- return this.stores.scores.listScoresByScorerId({ scorerId, source, pagination, entityId, entityType });
1979
- }
1980
- async saveScore(score) {
1981
- return this.stores.scores.saveScore(score);
1982
- }
1983
- async listScoresByRunId({
1984
- runId,
1985
- pagination
1986
- }) {
1987
- return this.stores.scores.listScoresByRunId({ runId, pagination });
1988
- }
1989
- async listScoresByEntityId({
1990
- entityId,
1991
- entityType,
1992
- pagination
1993
- }) {
1994
- return this.stores.scores.listScoresByEntityId({ entityId, entityType, pagination });
1995
- }
1996
- async listScoresBySpan({
1997
- traceId,
1998
- spanId,
1999
- pagination
2000
- }) {
2001
- return this.stores.scores.listScoresBySpan({ traceId, spanId, pagination });
1948
+ constructor(id, name, disableInit) {
1949
+ super({ id, name, disableInit });
1950
+ this.stores = {};
2002
1951
  }
2003
1952
  };
2004
1953
  var LanceFilterTranslator = class extends BaseFilterTranslator {
@@ -3221,6 +3170,6 @@ var LanceVectorStore = class _LanceVectorStore extends MastraVector {
3221
3170
  }
3222
3171
  };
3223
3172
 
3224
- export { LanceStorage, LanceVectorStore };
3173
+ export { LanceStorage, LanceVectorStore, StoreMemoryLance, StoreScoresLance, StoreWorkflowsLance };
3225
3174
  //# sourceMappingURL=index.js.map
3226
3175
  //# sourceMappingURL=index.js.map