@mastra/lance 0.2.0 → 0.2.1

package/src/storage/domains/operations/index.ts

@@ -0,0 +1,489 @@
+import type { Connection } from '@lancedb/lancedb';
+import { ErrorCategory, ErrorDomain, MastraError } from '@mastra/core/error';
+import { StoreOperations } from '@mastra/core/storage';
+import type { TABLE_NAMES, StorageColumn } from '@mastra/core/storage';
+import { Utf8, Int32, Float32, Binary, Schema, Field, Float64 } from 'apache-arrow';
+import type { DataType } from 'apache-arrow';
+import { getPrimaryKeys, getTableSchema, processResultWithTypeConversion, validateKeyTypes } from '../utils';
+
+export class StoreOperationsLance extends StoreOperations {
+  client: Connection;
+  constructor({ client }: { client: Connection }) {
+    super();
+    this.client = client;
+  }
+
+  protected getDefaultValue(type: StorageColumn['type']): string {
+    switch (type) {
+      case 'text':
+        return "''";
+      case 'timestamp':
+        return 'CURRENT_TIMESTAMP';
+      case 'integer':
+      case 'bigint':
+        return '0';
+      case 'jsonb':
+        return "'{}'";
+      case 'uuid':
+        return "''";
+      default:
+        return super.getDefaultValue(type);
+    }
+  }
+
+  async hasColumn(tableName: TABLE_NAMES, columnName: string): Promise<boolean> {
+    const table = await this.client.openTable(tableName);
+    const schema = await table.schema();
+    return schema.fields.some(field => field.name === columnName);
+  }
+
+  private translateSchema(schema: Record<string, StorageColumn>): Schema {
+    const fields = Object.entries(schema).map(([name, column]) => {
+      // Convert string type to Arrow DataType
+      let arrowType: DataType;
+      switch (column.type.toLowerCase()) {
+        case 'text':
+        case 'uuid':
+          arrowType = new Utf8();
+          break;
+        case 'int':
+        case 'integer':
+          arrowType = new Int32();
+          break;
+        case 'bigint':
+          arrowType = new Float64();
+          break;
+        case 'float':
+          arrowType = new Float32();
+          break;
+        case 'jsonb':
+        case 'json':
+          arrowType = new Utf8();
+          break;
+        case 'binary':
+          arrowType = new Binary();
+          break;
+        case 'timestamp':
+          arrowType = new Float64();
+          break;
+        default:
+          // Default to string for unknown types
+          arrowType = new Utf8();
+      }
+
+      // Create a field with the appropriate arrow type
+      return new Field(name, arrowType, column.nullable ?? true);
+    });
+
+    return new Schema(fields);
+  }
+
+  async createTable({
+    tableName,
+    schema,
+  }: {
+    tableName: TABLE_NAMES;
+    schema: Record<string, StorageColumn>;
+  }): Promise<void> {
+    try {
+      if (!this.client) {
+        throw new Error('LanceDB client not initialized. Call LanceStorage.create() first.');
+      }
+      if (!tableName) {
+        throw new Error('tableName is required for createTable.');
+      }
+      if (!schema) {
+        throw new Error('schema is required for createTable.');
+      }
+    } catch (error) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_LANCE_STORAGE_CREATE_TABLE_INVALID_ARGS',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.USER,
+          details: { tableName },
+        },
+        error,
+      );
+    }
+
+    try {
+      const arrowSchema = this.translateSchema(schema);
+      await this.client.createEmptyTable(tableName, arrowSchema);
+    } catch (error: any) {
+      if (error.message?.includes('already exists')) {
+        this.logger.debug(`Table '${tableName}' already exists, skipping create`);
+        return;
+      }
+      throw new MastraError(
+        {
+          id: 'STORAGE_LANCE_STORAGE_CREATE_TABLE_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { tableName },
+        },
+        error,
+      );
+    }
+  }
+
+  async dropTable({ tableName }: { tableName: TABLE_NAMES }): Promise<void> {
+    try {
+      if (!this.client) {
+        throw new Error('LanceDB client not initialized. Call LanceStorage.create() first.');
+      }
+      if (!tableName) {
+        throw new Error('tableName is required for dropTable.');
+      }
+    } catch (validationError: any) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_LANCE_STORAGE_DROP_TABLE_INVALID_ARGS',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.USER,
+          text: validationError.message,
+          details: { tableName },
+        },
+        validationError,
+      );
+    }
+
+    try {
+      await this.client.dropTable(tableName);
+    } catch (error: any) {
+      if (error.toString().includes('was not found') || error.message?.includes('Table not found')) {
+        this.logger.debug(`Table '${tableName}' does not exist, skipping drop`);
+        return;
+      }
+      throw new MastraError(
+        {
+          id: 'STORAGE_LANCE_STORAGE_DROP_TABLE_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { tableName },
+        },
+        error,
+      );
+    }
+  }
+
+  async alterTable({
+    tableName,
+    schema,
+    ifNotExists,
+  }: {
+    tableName: string;
+    schema: Record<string, StorageColumn>;
+    ifNotExists: string[];
+  }): Promise<void> {
+    try {
+      if (!this.client) {
+        throw new Error('LanceDB client not initialized. Call LanceStorage.create() first.');
+      }
+      if (!tableName) {
+        throw new Error('tableName is required for alterTable.');
+      }
+      if (!schema) {
+        throw new Error('schema is required for alterTable.');
+      }
+      if (!ifNotExists || ifNotExists.length === 0) {
+        this.logger.debug('No columns specified to add in alterTable, skipping.');
+        return;
+      }
+    } catch (validationError: any) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_LANCE_STORAGE_ALTER_TABLE_INVALID_ARGS',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.USER,
+          text: validationError.message,
+          details: { tableName },
+        },
+        validationError,
+      );
+    }
+
+    try {
+      const table = await this.client.openTable(tableName);
+      const currentSchema = await table.schema();
+      const existingFields = new Set(currentSchema.fields.map((f: any) => f.name));
+
+      const typeMap: Record<string, string> = {
+        text: 'string',
+        integer: 'int',
+        bigint: 'bigint',
+        timestamp: 'timestamp',
+        jsonb: 'string',
+        uuid: 'string',
+      };
+
+      // Find columns to add
+      const columnsToAdd = ifNotExists
+        .filter(col => schema[col] && !existingFields.has(col))
+        .map(col => {
+          const colDef = schema[col];
+          return {
+            name: col,
+            valueSql: colDef?.nullable
+              ? `cast(NULL as ${typeMap[colDef.type ?? 'text']})`
+              : `cast(${this.getDefaultValue(colDef?.type ?? 'text')} as ${typeMap[colDef?.type ?? 'text']})`,
+          };
+        });
+
+      if (columnsToAdd.length > 0) {
+        await table.addColumns(columnsToAdd);
+        this.logger?.info?.(`Added columns [${columnsToAdd.map(c => c.name).join(', ')}] to table ${tableName}`);
+      }
+    } catch (error: any) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_LANCE_STORAGE_ALTER_TABLE_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { tableName },
+        },
+        error,
+      );
+    }
+  }
+
+  async clearTable({ tableName }: { tableName: TABLE_NAMES }): Promise<void> {
+    try {
+      if (!this.client) {
+        throw new Error('LanceDB client not initialized. Call LanceStorage.create() first.');
+      }
+      if (!tableName) {
+        throw new Error('tableName is required for clearTable.');
+      }
+    } catch (validationError: any) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_LANCE_STORAGE_CLEAR_TABLE_INVALID_ARGS',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.USER,
+          text: validationError.message,
+          details: { tableName },
+        },
+        validationError,
+      );
+    }
+
+    try {
+      const table = await this.client.openTable(tableName);
+
+      // delete function always takes a predicate as an argument, so we use '1=1' to delete all records because it is always true.
+      await table.delete('1=1');
+    } catch (error: any) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_LANCE_STORAGE_CLEAR_TABLE_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { tableName },
+        },
+        error,
+      );
+    }
+  }
+
+  async insert({ tableName, record }: { tableName: string; record: Record<string, any> }): Promise<void> {
+    try {
+      if (!this.client) {
+        throw new Error('LanceDB client not initialized. Call LanceStorage.create() first.');
+      }
+      if (!tableName) {
+        throw new Error('tableName is required for insert.');
+      }
+      if (!record || Object.keys(record).length === 0) {
+        throw new Error('record is required and cannot be empty for insert.');
+      }
+    } catch (validationError: any) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_LANCE_STORAGE_INSERT_INVALID_ARGS',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.USER,
+          text: validationError.message,
+          details: { tableName },
+        },
+        validationError,
+      );
+    }
+
+    try {
+      const table = await this.client.openTable(tableName);
+
+      const primaryId = getPrimaryKeys(tableName as TABLE_NAMES);
+
+      const processedRecord = { ...record };
+
+      for (const key in processedRecord) {
+        if (
+          processedRecord[key] !== null &&
+          typeof processedRecord[key] === 'object' &&
+          !(processedRecord[key] instanceof Date)
+        ) {
+          this.logger.debug('Converting object to JSON string: ', processedRecord[key]);
+          processedRecord[key] = JSON.stringify(processedRecord[key]);
+        }
+      }
+      console.log(await table.schema());
+
+      await table.mergeInsert(primaryId).whenMatchedUpdateAll().whenNotMatchedInsertAll().execute([processedRecord]);
+    } catch (error: any) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_LANCE_STORAGE_INSERT_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { tableName },
+        },
+        error,
+      );
+    }
+  }
+
+  async batchInsert({ tableName, records }: { tableName: string; records: Record<string, any>[] }): Promise<void> {
+    try {
+      if (!this.client) {
+        throw new Error('LanceDB client not initialized. Call LanceStorage.create() first.');
+      }
+      if (!tableName) {
+        throw new Error('tableName is required for batchInsert.');
+      }
+      if (!records || records.length === 0) {
+        throw new Error('records array is required and cannot be empty for batchInsert.');
+      }
+    } catch (validationError: any) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_LANCE_STORAGE_BATCH_INSERT_INVALID_ARGS',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.USER,
+          text: validationError.message,
+          details: { tableName },
+        },
+        validationError,
+      );
+    }
+
+    try {
+      const table = await this.client.openTable(tableName);
+
+      const primaryId = getPrimaryKeys(tableName as TABLE_NAMES);
+
+      const processedRecords = records.map(record => {
+        const processedRecord = { ...record };
+
+        // Convert values based on schema type
+        for (const key in processedRecord) {
+          // Skip null/undefined values
+          if (processedRecord[key] == null) continue;
+
+          if (
+            processedRecord[key] !== null &&
+            typeof processedRecord[key] === 'object' &&
+            !(processedRecord[key] instanceof Date)
+          ) {
+            processedRecord[key] = JSON.stringify(processedRecord[key]);
+          }
+        }
+
+        return processedRecord;
+      });
+
+      console.log(processedRecords);
+
+      await table.mergeInsert(primaryId).whenMatchedUpdateAll().whenNotMatchedInsertAll().execute(processedRecords);
+    } catch (error: any) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_LANCE_STORAGE_BATCH_INSERT_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { tableName },
+        },
+        error,
+      );
+    }
+  }
+
+  async load({ tableName, keys }: { tableName: TABLE_NAMES; keys: Record<string, any> }): Promise<any> {
+    try {
+      if (!this.client) {
+        throw new Error('LanceDB client not initialized. Call LanceStorage.create() first.');
+      }
+      if (!tableName) {
+        throw new Error('tableName is required for load.');
+      }
+      if (!keys || Object.keys(keys).length === 0) {
+        throw new Error('keys are required and cannot be empty for load.');
+      }
+    } catch (validationError: any) {
+      throw new MastraError(
+        {
+          id: 'STORAGE_LANCE_STORAGE_LOAD_INVALID_ARGS',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.USER,
+          text: validationError.message,
+          details: { tableName },
+        },
+        validationError,
+      );
+    }
+
+    try {
+      const table = await this.client.openTable(tableName);
+      const tableSchema = await getTableSchema({ tableName, client: this.client });
+      const query = table.query();
+
+      // Build filter condition with 'and' between all conditions
+      if (Object.keys(keys).length > 0) {
+        // Validate key types against schema
+        validateKeyTypes(keys, tableSchema);
+
+        const filterConditions = Object.entries(keys)
+          .map(([key, value]) => {
+            // Check if key is in camelCase and wrap it in backticks if it is
+            const isCamelCase = /^[a-z][a-zA-Z]*$/.test(key) && /[A-Z]/.test(key);
+            const quotedKey = isCamelCase ? `\`${key}\`` : key;
+
+            // Handle different types appropriately
+            if (typeof value === 'string') {
+              return `${quotedKey} = '${value}'`;
+            } else if (value === null) {
+              return `${quotedKey} IS NULL`;
+            } else {
+              // For numbers, booleans, etc.
+              return `${quotedKey} = ${value}`;
+            }
+          })
+          .join(' AND ');
+
+        this.logger.debug('where clause generated: ' + filterConditions);
+        query.where(filterConditions);
+      }
+
+      const result = await query.limit(1).toArray();
+
+      if (result.length === 0) {
+        this.logger.debug('No record found');
+        return null;
+      }
+      // Process the result with type conversions
+      return processResultWithTypeConversion(result[0], tableSchema);
+    } catch (error: any) {
+      // If it's already a MastraError (e.g. from validateKeyTypes if we change it later), rethrow
+      if (error instanceof MastraError) throw error;
+      throw new MastraError(
+        {
+          id: 'STORAGE_LANCE_STORAGE_LOAD_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { tableName, keyCount: Object.keys(keys).length, firstKey: Object.keys(keys)[0] ?? '' },
+        },
+        error,
+      );
+    }
+  }
+}

package/src/storage/domains/scores/index.ts

@@ -0,0 +1,221 @@
+import type { Connection } from '@lancedb/lancedb';
+import { ErrorCategory, ErrorDomain, MastraError } from '@mastra/core/error';
+import type { ScoreRowData } from '@mastra/core/scores';
+import { ScoresStorage, TABLE_SCORERS } from '@mastra/core/storage';
+import type { PaginationInfo, StoragePagination } from '@mastra/core/storage';
+import { getTableSchema, processResultWithTypeConversion } from '../utils';
+
+export class StoreScoresLance extends ScoresStorage {
+  private client: Connection;
+  constructor({ client }: { client: Connection }) {
+    super();
+    this.client = client;
+  }
+
+  async saveScore(score: ScoreRowData): Promise<{ score: ScoreRowData }> {
+    try {
+      const table = await this.client.openTable(TABLE_SCORERS);
+      // Fetch schema fields for mastra_scorers
+      const schema = await getTableSchema({ tableName: TABLE_SCORERS, client: this.client });
+      const allowedFields = new Set(schema.fields.map((f: any) => f.name));
+      // Filter out fields not in schema
+      const filteredScore: Record<string, any> = {};
+      (Object.keys(score) as (keyof ScoreRowData)[]).forEach(key => {
+        if (allowedFields.has(key)) {
+          filteredScore[key] = score[key];
+        }
+      });
+      // Convert any object fields to JSON strings for storage
+      for (const key in filteredScore) {
+        if (
+          filteredScore[key] !== null &&
+          typeof filteredScore[key] === 'object' &&
+          !(filteredScore[key] instanceof Date)
+        ) {
+          filteredScore[key] = JSON.stringify(filteredScore[key]);
+        }
+      }
+
+      console.log('Saving score to LanceStorage:', filteredScore);
+
+      await table.add([filteredScore], { mode: 'append' });
+      return { score };
+    } catch (error: any) {
+      throw new MastraError(
+        {
+          id: 'LANCE_STORAGE_SAVE_SCORE_FAILED',
+          text: 'Failed to save score in LanceStorage',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { error: error?.message },
+        },
+        error,
+      );
+    }
+  }
+
+  async getScoreById({ id }: { id: string }): Promise<ScoreRowData | null> {
+    try {
+      const table = await this.client.openTable(TABLE_SCORERS);
+
+      const query = table.query().where(`id = '${id}'`).limit(1);
+
+      const records = await query.toArray();
+
+      if (records.length === 0) return null;
+      const schema = await getTableSchema({ tableName: TABLE_SCORERS, client: this.client });
+      return processResultWithTypeConversion(records[0], schema) as ScoreRowData;
+    } catch (error: any) {
+      throw new MastraError(
+        {
+          id: 'LANCE_STORAGE_GET_SCORE_BY_ID_FAILED',
+          text: 'Failed to get score by id in LanceStorage',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { error: error?.message },
+        },
+        error,
+      );
+    }
+  }
+
+  async getScoresByScorerId({
+    scorerId,
+    pagination,
+  }: {
+    scorerId: string;
+    pagination: StoragePagination;
+  }): Promise<{ pagination: PaginationInfo; scores: ScoreRowData[] }> {
+    try {
+      const table = await this.client.openTable(TABLE_SCORERS);
+      // Use zero-based pagination (default page = 0)
+      const { page = 0, perPage = 10 } = pagination || {};
+      const offset = page * perPage;
+      // Query for scores with the given scorerId
+      // Must use backticks for field names to handle camelCase
+      const query = table.query().where(`\`scorerId\` = '${scorerId}'`).limit(perPage);
+      if (offset > 0) query.offset(offset);
+      const records = await query.toArray();
+      const schema = await getTableSchema({ tableName: TABLE_SCORERS, client: this.client });
+      const scores = processResultWithTypeConversion(records, schema) as ScoreRowData[];
+
+      const allRecords = await table.query().where(`\`scorerId\` = '${scorerId}'`).toArray();
+      const total = allRecords.length;
+
+      return {
+        pagination: {
+          page,
+          perPage,
+          total,
+          hasMore: offset + scores.length < total,
+        },
+        scores,
+      };
+    } catch (error: any) {
+      throw new MastraError(
+        {
+          id: 'LANCE_STORAGE_GET_SCORES_BY_SCORER_ID_FAILED',
+          text: 'Failed to get scores by scorerId in LanceStorage',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { error: error?.message },
+        },
+        error,
+      );
+    }
+  }
+
+  async getScoresByRunId({
+    runId,
+    pagination,
+  }: {
+    runId: string;
+    pagination: StoragePagination;
+  }): Promise<{ pagination: PaginationInfo; scores: ScoreRowData[] }> {
+    try {
+      const table = await this.client.openTable(TABLE_SCORERS);
+      const { page = 0, perPage = 10 } = pagination || {};
+      const offset = page * perPage;
+      // Query for scores with the given runId
+      const query = table.query().where(`\`runId\` = '${runId}'`).limit(perPage);
+      if (offset > 0) query.offset(offset);
+      const records = await query.toArray();
+      const schema = await getTableSchema({ tableName: TABLE_SCORERS, client: this.client });
+      const scores = processResultWithTypeConversion(records, schema) as ScoreRowData[];
+      // Get total count for pagination
+      const allRecords = await table.query().where(`\`runId\` = '${runId}'`).toArray();
+      const total = allRecords.length;
+      return {
+        pagination: {
+          page,
+          perPage,
+          total,
+          hasMore: offset + scores.length < total,
+        },
+        scores,
+      };
+    } catch (error: any) {
+      throw new MastraError(
+        {
+          id: 'LANCE_STORAGE_GET_SCORES_BY_RUN_ID_FAILED',
+          text: 'Failed to get scores by runId in LanceStorage',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { error: error?.message },
+        },
+        error,
+      );
+    }
+  }
+
+  async getScoresByEntityId({
+    entityId,
+    entityType,
+    pagination,
+  }: {
+    pagination: StoragePagination;
+    entityId: string;
+    entityType: string;
+  }): Promise<{ pagination: PaginationInfo; scores: ScoreRowData[] }> {
+    try {
+      const table = await this.client.openTable(TABLE_SCORERS);
+      const { page = 0, perPage = 10 } = pagination || {};
+      const offset = page * perPage;
+      // Query for scores with the given entityId and entityType
+      const query = table
+        .query()
+        .where(`\`entityId\` = '${entityId}' AND \`entityType\` = '${entityType}'`)
+        .limit(perPage);
+      if (offset > 0) query.offset(offset);
+      const records = await query.toArray();
+      const schema = await getTableSchema({ tableName: TABLE_SCORERS, client: this.client });
+      const scores = processResultWithTypeConversion(records, schema) as ScoreRowData[];
+      // Get total count for pagination
+      const allRecords = await table
+        .query()
+        .where(`\`entityId\` = '${entityId}' AND \`entityType\` = '${entityType}'`)
+        .toArray();
+      const total = allRecords.length;
+      return {
+        pagination: {
+          page,
+          perPage,
+          total,
+          hasMore: offset + scores.length < total,
+        },
+        scores,
+      };
+    } catch (error: any) {
+      throw new MastraError(
+        {
+          id: 'LANCE_STORAGE_GET_SCORES_BY_ENTITY_ID_FAILED',
+          text: 'Failed to get scores by entityId and entityType in LanceStorage',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { error: error?.message },
+        },
+        error,
+      );
+    }
+  }
+}