@mastra/lance 0.2.0 → 0.2.1

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

@@ -0,0 +1,212 @@
+import type { Connection } from '@lancedb/lancedb';
+import { MastraError, ErrorDomain, ErrorCategory } from '@mastra/core/error';
+import type { TraceType } from '@mastra/core/memory';
+import { TABLE_TRACES, TracesStorage } from '@mastra/core/storage';
+import type { PaginationInfo, StorageGetTracesPaginatedArg } from '@mastra/core/storage';
+import type { Trace } from '@mastra/core/telemetry';
+import type { StoreOperationsLance } from '../operations';
+
+export class StoreTracesLance extends TracesStorage {
+  private client: Connection;
+  private operations: StoreOperationsLance;
+  constructor({ client, operations }: { client: Connection; operations: StoreOperationsLance }) {
+    super();
+    this.client = client;
+    this.operations = operations;
+  }
+
+  async saveTrace({ trace }: { trace: TraceType }): Promise<TraceType> {
+    try {
+      const table = await this.client.openTable(TABLE_TRACES);
+      const record = {
+        ...trace,
+        attributes: JSON.stringify(trace.attributes),
+        status: JSON.stringify(trace.status),
+        events: JSON.stringify(trace.events),
+        links: JSON.stringify(trace.links),
+        other: JSON.stringify(trace.other),
+      };
+      await table.add([record], { mode: 'append' });
+      return trace;
+    } catch (error: any) {
+      throw new MastraError(
+        {
+          id: 'LANCE_STORE_SAVE_TRACE_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+        },
+        error,
+      );
+    }
+  }
+
+  async getTraceById({ traceId }: { traceId: string }): Promise<TraceType> {
+    try {
+      const table = await this.client.openTable(TABLE_TRACES);
+      const query = table.query().where(`id = '${traceId}'`);
+      const records = await query.toArray();
+      return records[0] as TraceType;
+    } catch (error: any) {
+      throw new MastraError(
+        {
+          id: 'LANCE_STORE_GET_TRACE_BY_ID_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+        },
+        error,
+      );
+    }
+  }
+
+  async getTraces({
+    name,
+    scope,
+    page = 1,
+    perPage = 10,
+    attributes,
+  }: {
+    name?: string;
+    scope?: string;
+    page: number;
+    perPage: number;
+    attributes?: Record<string, string>;
+  }): Promise<Trace[]> {
+    try {
+      const table = await this.client.openTable(TABLE_TRACES);
+      const query = table.query();
+      if (name) {
+        query.where(`name = '${name}'`);
+      }
+      if (scope) {
+        query.where(`scope = '${scope}'`);
+      }
+      if (attributes) {
+        query.where(`attributes = '${JSON.stringify(attributes)}'`);
+      }
+      // Calculate offset based on page and perPage
+      const offset = (page - 1) * perPage;
+      query.limit(perPage);
+      if (offset > 0) {
+        query.offset(offset);
+      }
+      const records = await query.toArray();
+      return records.map(record => {
+        const processed = {
+          ...record,
+          attributes: record.attributes ? JSON.parse(record.attributes) : {},
+          status: record.status ? JSON.parse(record.status) : {},
+          events: record.events ? JSON.parse(record.events) : [],
+          links: record.links ? JSON.parse(record.links) : [],
+          other: record.other ? JSON.parse(record.other) : {},
+          startTime: new Date(record.startTime),
+          endTime: new Date(record.endTime),
+          createdAt: new Date(record.createdAt),
+        };
+        if (processed.parentSpanId === null || processed.parentSpanId === undefined) {
+          processed.parentSpanId = '';
+        } else {
+          processed.parentSpanId = String(processed.parentSpanId);
+        }
+        return processed as Trace;
+      });
+    } catch (error: any) {
+      throw new MastraError(
+        {
+          id: 'LANCE_STORE_GET_TRACES_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { name: name ?? '', scope: scope ?? '' },
+        },
+        error,
+      );
+    }
+  }
+
+  async getTracesPaginated(args: StorageGetTracesPaginatedArg): Promise<PaginationInfo & { traces: Trace[] }> {
+    try {
+      const table = await this.client.openTable(TABLE_TRACES);
+      const query = table.query();
+      const conditions: string[] = [];
+      if (args.name) {
+        conditions.push(`name = '${args.name}'`);
+      }
+      if (args.scope) {
+        conditions.push(`scope = '${args.scope}'`);
+      }
+      if (args.attributes) {
+        const attributesStr = JSON.stringify(args.attributes);
+        conditions.push(`attributes LIKE '%${attributesStr.replace(/"/g, '\\"')}%'`);
+      }
+      if (args.dateRange?.start) {
+        conditions.push(`\`createdAt\` >= ${args.dateRange.start.getTime()}`);
+      }
+      if (args.dateRange?.end) {
+        conditions.push(`\`createdAt\` <= ${args.dateRange.end.getTime()}`);
+      }
+      if (conditions.length > 0) {
+        const whereClause = conditions.join(' AND ');
+        query.where(whereClause);
+      }
+      let total = 0;
+      if (conditions.length > 0) {
+        const countQuery = table.query().where(conditions.join(' AND '));
+        const allRecords = await countQuery.toArray();
+        total = allRecords.length;
+      } else {
+        total = await table.countRows();
+      }
+      const page = args.page || 0;
+      const perPage = args.perPage || 10;
+      const offset = page * perPage;
+      query.limit(perPage);
+      if (offset > 0) {
+        query.offset(offset);
+      }
+      const records = await query.toArray();
+      const traces = records.map(record => {
+        const processed = {
+          ...record,
+          attributes: record.attributes ? JSON.parse(record.attributes) : {},
+          status: record.status ? JSON.parse(record.status) : {},
+          events: record.events ? JSON.parse(record.events) : [],
+          links: record.links ? JSON.parse(record.links) : [],
+          other: record.other ? JSON.parse(record.other) : {},
+          startTime: new Date(record.startTime),
+          endTime: new Date(record.endTime),
+          createdAt: new Date(record.createdAt),
+        };
+        if (processed.parentSpanId === null || processed.parentSpanId === undefined) {
+          processed.parentSpanId = '';
+        } else {
+          processed.parentSpanId = String(processed.parentSpanId);
+        }
+        return processed as Trace;
+      });
+      return {
+        traces,
+        total,
+        page,
+        perPage,
+        hasMore: total > (page + 1) * perPage,
+      };
+    } catch (error: any) {
+      throw new MastraError(
+        {
+          id: 'LANCE_STORE_GET_TRACES_PAGINATED_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { name: args.name ?? '', scope: args.scope ?? '' },
+        },
+        error,
+      );
+    }
+  }
+
+  async batchTraceInsert({ records }: { records: Record<string, any>[] }): Promise<void> {
+    this.logger.debug('Batch inserting traces', { count: records.length });
+    await this.operations.batchInsert({
+      tableName: TABLE_TRACES,
+      records,
+    });
+  }
+}

package/src/storage/domains/utils.ts

@@ -0,0 +1,158 @@
+import type { Connection, FieldLike, SchemaLike } from '@lancedb/lancedb';
+import { ErrorCategory, ErrorDomain, MastraError } from '@mastra/core/error';
+import { TABLE_EVALS, TABLE_WORKFLOW_SNAPSHOT } from '@mastra/core/storage';
+import type { TABLE_NAMES } from '@mastra/core/storage';
+
+export function getPrimaryKeys(tableName: TABLE_NAMES): string[] {
+  let primaryId: string[] = ['id'];
+  if (tableName === TABLE_WORKFLOW_SNAPSHOT) {
+    primaryId = ['workflow_name', 'run_id'];
+  } else if (tableName === TABLE_EVALS) {
+    primaryId = ['agent_name', 'metric_name', 'run_id'];
+  }
+
+  return primaryId;
+}
+
+export function validateKeyTypes(keys: Record<string, any>, tableSchema: SchemaLike): void {
+  // Create a map of field names to their expected types
+  const fieldTypes = new Map(
+    tableSchema.fields.map((field: any) => [field.name, field.type?.toString().toLowerCase()]),
+  );
+
+  for (const [key, value] of Object.entries(keys)) {
+    const fieldType = fieldTypes.get(key);
+
+    if (!fieldType) {
+      throw new Error(`Field '${key}' does not exist in table schema`);
+    }
+
+    // Type validation
+    if (value !== null) {
+      if ((fieldType.includes('int') || fieldType.includes('bigint')) && typeof value !== 'number') {
+        throw new Error(`Expected numeric value for field '${key}', got ${typeof value}`);
+      }
+
+      if (fieldType.includes('utf8') && typeof value !== 'string') {
+        throw new Error(`Expected string value for field '${key}', got ${typeof value}`);
+      }
+
+      if (fieldType.includes('timestamp') && !(value instanceof Date) && typeof value !== 'string') {
+        throw new Error(`Expected Date or string value for field '${key}', got ${typeof value}`);
+      }
+    }
+  }
+}
+
+export function processResultWithTypeConversion(
+  rawResult: Record<string, any> | Record<string, any>[],
+  tableSchema: SchemaLike,
+): Record<string, any> | Record<string, any>[] {
+  // Build a map of field names to their schema types
+  const fieldTypeMap = new Map();
+  tableSchema.fields.forEach((field: any) => {
+    const fieldName = field.name;
+    const fieldTypeStr = field.type.toString().toLowerCase();
+    fieldTypeMap.set(fieldName, fieldTypeStr);
+  });
+
+  // Handle array case
+  if (Array.isArray(rawResult)) {
+    return rawResult.map(item => processResultWithTypeConversion(item, tableSchema));
+  }
+
+  // Handle single record case
+  const processedResult = { ...rawResult };
+
+  // Convert each field according to its schema type
+  for (const key in processedResult) {
+    const fieldTypeStr = fieldTypeMap.get(key);
+    if (!fieldTypeStr) continue;
+
+    // Skip conversion for ID fields - preserve their original format
+    // if (key === 'id') {
+    //   continue;
+    // }
+
+    // Only try to convert string values
+    if (typeof processedResult[key] === 'string') {
+      // Numeric types
+      if (fieldTypeStr.includes('int32') || fieldTypeStr.includes('float32')) {
+        if (!isNaN(Number(processedResult[key]))) {
+          processedResult[key] = Number(processedResult[key]);
+        }
+      } else if (fieldTypeStr.includes('int64')) {
+        processedResult[key] = Number(processedResult[key]);
+      } else if (fieldTypeStr.includes('utf8') && key !== 'id') {
+        try {
+          const parsed = JSON.parse(processedResult[key]);
+          if (typeof parsed === 'object') {
+            processedResult[key] = JSON.parse(processedResult[key]);
+          }
+        } catch {}
+      }
+    } else if (typeof processedResult[key] === 'bigint') {
+      // Convert BigInt values to regular numbers for application layer
+      processedResult[key] = Number(processedResult[key]);
+    } else if (fieldTypeStr.includes('float64') && ['createdAt', 'updatedAt'].includes(key)) {
+      processedResult[key] = new Date(processedResult[key]);
+    }
+
+    console.log(key, 'processedResult', processedResult);
+  }
+
+  return processedResult;
+}
+
+export async function getTableSchema({
+  tableName,
+  client,
+}: {
+  tableName: TABLE_NAMES;
+  client: Connection;
+}): Promise<SchemaLike> {
+  try {
+    if (!client) {
+      throw new Error('LanceDB client not initialized. Call LanceStorage.create() first.');
+    }
+    if (!tableName) {
+      throw new Error('tableName is required for getTableSchema.');
+    }
+  } catch (validationError: any) {
+    throw new MastraError(
+      {
+        id: 'STORAGE_LANCE_STORAGE_GET_TABLE_SCHEMA_INVALID_ARGS',
+        domain: ErrorDomain.STORAGE,
+        category: ErrorCategory.USER,
+        text: validationError.message,
+        details: { tableName },
+      },
+      validationError,
+    );
+  }
+
+  try {
+    const table = await client.openTable(tableName);
+    const rawSchema = await table.schema();
+    const fields = rawSchema.fields as FieldLike[];
+
+    // Convert schema to SchemaLike format
+    return {
+      fields,
+      metadata: new Map<string, string>(),
+      get names() {
+        return fields.map((field: FieldLike) => field.name);
+      },
+    };
+  } catch (error: any) {
+    throw new MastraError(
+      {
+        id: 'STORAGE_LANCE_STORAGE_GET_TABLE_SCHEMA_FAILED',
+        domain: ErrorDomain.STORAGE,
+        category: ErrorCategory.THIRD_PARTY,
+        details: { tableName },
+      },
+      error,
+    );
+  }
+}

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

@@ -0,0 +1,207 @@
+import type { Connection } from '@lancedb/lancedb';
+import type { WorkflowRun, WorkflowRunState, WorkflowRuns } from '@mastra/core';
+import { ErrorCategory, ErrorDomain, MastraError } from '@mastra/core/error';
+import { ensureDate, TABLE_WORKFLOW_SNAPSHOT, WorkflowsStorage } from '@mastra/core/storage';
+
+function parseWorkflowRun(row: any): WorkflowRun {
+  let parsedSnapshot: WorkflowRunState | string = row.snapshot;
+  if (typeof parsedSnapshot === 'string') {
+    try {
+      parsedSnapshot = JSON.parse(row.snapshot as string) as WorkflowRunState;
+    } catch (e) {
+      // If parsing fails, return the raw snapshot string
+      console.warn(`Failed to parse snapshot for workflow ${row.workflow_name}: ${e}`);
+    }
+  }
+
+  return {
+    workflowName: row.workflow_name,
+    runId: row.run_id,
+    snapshot: parsedSnapshot,
+    createdAt: ensureDate(row.createdAt)!,
+    updatedAt: ensureDate(row.updatedAt)!,
+    resourceId: row.resourceId,
+  };
+}
+
+export class StoreWorkflowsLance extends WorkflowsStorage {
+  client: Connection;
+  constructor({ client }: { client: Connection }) {
+    super();
+    this.client = client;
+  }
+
+  async persistWorkflowSnapshot({
+    workflowName,
+    runId,
+    snapshot,
+  }: {
+    workflowName: string;
+    runId: string;
+    snapshot: WorkflowRunState;
+  }): Promise<void> {
+    try {
+      const table = await this.client.openTable(TABLE_WORKFLOW_SNAPSHOT);
+
+      // Try to find the existing record
+      const query = table.query().where(`workflow_name = '${workflowName}' AND run_id = '${runId}'`);
+      const records = await query.toArray();
+      let createdAt: number;
+      const now = Date.now();
+
+      if (records.length > 0) {
+        createdAt = records[0].createdAt ?? now;
+      } else {
+        createdAt = now;
+      }
+
+      const record = {
+        workflow_name: workflowName,
+        run_id: runId,
+        snapshot: JSON.stringify(snapshot),
+        createdAt,
+        updatedAt: now,
+      };
+
+      await table
+        .mergeInsert(['workflow_name', 'run_id'])
+        .whenMatchedUpdateAll()
+        .whenNotMatchedInsertAll()
+        .execute([record]);
+    } catch (error: any) {
+      throw new MastraError(
+        {
+          id: 'LANCE_STORE_PERSIST_WORKFLOW_SNAPSHOT_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { workflowName, runId },
+        },
+        error,
+      );
+    }
+  }
+  async loadWorkflowSnapshot({
+    workflowName,
+    runId,
+  }: {
+    workflowName: string;
+    runId: string;
+  }): Promise<WorkflowRunState | null> {
+    try {
+      const table = await this.client.openTable(TABLE_WORKFLOW_SNAPSHOT);
+      const query = table.query().where(`workflow_name = '${workflowName}' AND run_id = '${runId}'`);
+      const records = await query.toArray();
+      return records.length > 0 ? JSON.parse(records[0].snapshot) : null;
+    } catch (error: any) {
+      throw new MastraError(
+        {
+          id: 'LANCE_STORE_LOAD_WORKFLOW_SNAPSHOT_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { workflowName, runId },
+        },
+        error,
+      );
+    }
+  }
+
+  async getWorkflowRunById(args: { runId: string; workflowName?: string }): Promise<{
+    workflowName: string;
+    runId: string;
+    snapshot: any;
+    createdAt: Date;
+    updatedAt: Date;
+  } | null> {
+    try {
+      const table = await this.client.openTable(TABLE_WORKFLOW_SNAPSHOT);
+      let whereClause = `run_id = '${args.runId}'`;
+      if (args.workflowName) {
+        whereClause += ` AND workflow_name = '${args.workflowName}'`;
+      }
+      const query = table.query().where(whereClause);
+      const records = await query.toArray();
+      if (records.length === 0) return null;
+      const record = records[0];
+      return parseWorkflowRun(record);
+    } catch (error: any) {
+      throw new MastraError(
+        {
+          id: 'LANCE_STORE_GET_WORKFLOW_RUN_BY_ID_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { runId: args.runId, workflowName: args.workflowName ?? '' },
+        },
+        error,
+      );
+    }
+  }
+
+  async getWorkflowRuns(args?: {
+    namespace?: string;
+    resourceId?: string;
+    workflowName?: string;
+    fromDate?: Date;
+    toDate?: Date;
+    limit?: number;
+    offset?: number;
+  }): Promise<WorkflowRuns> {
+    try {
+      const table = await this.client.openTable(TABLE_WORKFLOW_SNAPSHOT);
+
+      let query = table.query();
+
+      const conditions: string[] = [];
+
+      if (args?.workflowName) {
+        conditions.push(`workflow_name = '${args.workflowName.replace(/'/g, "''")}'`);
+      }
+
+      if (args?.resourceId) {
+        conditions.push(`\`resourceId\` = '${args.resourceId}'`);
+      }
+
+      if (args?.fromDate instanceof Date) {
+        conditions.push(`\`createdAt\` >= ${args.fromDate.getTime()}`);
+      }
+
+      if (args?.toDate instanceof Date) {
+        conditions.push(`\`createdAt\` <= ${args.toDate.getTime()}`);
+      }
+
+      let total = 0;
+
+      // Apply all conditions
+      if (conditions.length > 0) {
+        query = query.where(conditions.join(' AND '));
+        total = await table.countRows(conditions.join(' AND '));
+      } else {
+        total = await table.countRows();
+      }
+
+      if (args?.limit) {
+        query.limit(args.limit);
+      }
+
+      if (args?.offset) {
+        query.offset(args.offset);
+      }
+
+      const records = await query.toArray();
+
+      return {
+        runs: records.map(record => parseWorkflowRun(record)),
+        total: total || records.length,
+      };
+    } catch (error: any) {
+      throw new MastraError(
+        {
+          id: 'LANCE_STORE_GET_WORKFLOW_RUNS_FAILED',
+          domain: ErrorDomain.STORAGE,
+          category: ErrorCategory.THIRD_PARTY,
+          details: { namespace: args?.namespace ?? '', workflowName: args?.workflowName ?? '' },
+        },
+        error,
+      );
+    }
+  }
+}