@mastra/upstash 0.2.2-alpha.1 → 0.2.2-alpha.3

package/src/storage/index.ts

@@ -1,5 +1,12 @@
+import type { MetricResult, TestInfo } from '@mastra/core/eval';
 import type { StorageThreadType, MessageType } from '@mastra/core/memory';
-import { MastraStorage } from '@mastra/core/storage';
+import {
+  MastraStorage,
+  TABLE_MESSAGES,
+  TABLE_THREADS,
+  TABLE_WORKFLOW_SNAPSHOT,
+  TABLE_EVALS,
+} from '@mastra/core/storage';
 import type { TABLE_NAMES, StorageColumn, StorageGetMessagesArg, EvalRow } from '@mastra/core/storage';
 import type { WorkflowRunState } from '@mastra/core/workflows';
 import { Redis } from '@upstash/redis';
@@ -13,18 +20,124 @@ export class UpstashStore extends MastraStorage {
   batchInsert(_input: { tableName: TABLE_NAMES; records: Record<string, any>[] }): Promise<void> {
     throw new Error('Method not implemented.');
   }
-  getEvalsByAgentName(_agentName: string, _type?: 'test' | 'live'): Promise<EvalRow[]> {
-    throw new Error('Method not implemented.');
+
+  async getEvalsByAgentName(agentName: string, type?: 'test' | 'live'): Promise<EvalRow[]> {
+    try {
+      // Get all keys that match the evals table pattern
+      const pattern = `${TABLE_EVALS}:*`;
+      const keys = await this.redis.keys(pattern);
+
+      // Fetch all eval records
+      const evalRecords = await Promise.all(
+        keys.map(async key => {
+          const data = await this.redis.get<Record<string, any>>(key);
+          return data;
+        }),
+      );
+
+      // Filter by agent name and remove nulls
+      const nonNullRecords = evalRecords.filter(
+        (record): record is Record<string, any> =>
+          record !== null && typeof record === 'object' && 'agent_name' in record && record.agent_name === agentName,
+      );
+
+      // Apply additional filtering based on type
+      let filteredEvals = nonNullRecords;
+
+      if (type === 'test') {
+        filteredEvals = filteredEvals.filter(record => {
+          if (!record.test_info) return false;
+
+          // Handle test_info as a JSON string
+          try {
+            if (typeof record.test_info === 'string') {
+              const parsedTestInfo = JSON.parse(record.test_info);
+              return parsedTestInfo && typeof parsedTestInfo === 'object' && 'testPath' in parsedTestInfo;
+            }
+
+            // Handle test_info as an object
+            return typeof record.test_info === 'object' && 'testPath' in record.test_info;
+          } catch (_e) {
+            return false;
+          }
+        });
+      } else if (type === 'live') {
+        filteredEvals = filteredEvals.filter(record => {
+          if (!record.test_info) return true;
+
+          // Handle test_info as a JSON string
+          try {
+            if (typeof record.test_info === 'string') {
+              const parsedTestInfo = JSON.parse(record.test_info);
+              return !(parsedTestInfo && typeof parsedTestInfo === 'object' && 'testPath' in parsedTestInfo);
+            }
+
+            // Handle test_info as an object
+            return !(typeof record.test_info === 'object' && 'testPath' in record.test_info);
+          } catch (_e) {
+            return true;
+          }
+        });
+      }
+
+      // Transform to EvalRow format
+      return filteredEvals.map(record => this.transformEvalRecord(record));
+    } catch (error) {
+      console.error('Failed to get evals for the specified agent:', error);
+      return [];
+    }
   }
+
+  private transformEvalRecord(record: Record<string, any>): EvalRow {
+    // Parse JSON strings if needed
+    let result = record.result;
+    if (typeof result === 'string') {
+      try {
+        result = JSON.parse(result);
+      } catch (_e) {
+        console.warn('Failed to parse result JSON:');
+      }
+    }
+
+    let testInfo = record.test_info;
+    if (typeof testInfo === 'string') {
+      try {
+        testInfo = JSON.parse(testInfo);
+      } catch (_e) {
+        console.warn('Failed to parse test_info JSON:');
+      }
+    }
+
+    return {
+      agentName: record.agent_name,
+      input: record.input,
+      output: record.output,
+      result: result as MetricResult,
+      metricName: record.metric_name,
+      instructions: record.instructions,
+      testInfo: testInfo as TestInfo | undefined,
+      globalRunId: record.global_run_id,
+      runId: record.run_id,
+      createdAt:
+        typeof record.created_at === 'string'
+          ? record.created_at
+          : record.created_at instanceof Date
+            ? record.created_at.toISOString()
+            : new Date().toISOString(),
+    };
+  }
+
   getTraces(_input: {
     name?: string;
     scope?: string;
     page: number;
     perPage: number;
     attributes?: Record<string, string>;
+    filters?: Record<string, any>;
   }): Promise<any[]> {
     throw new Error('Method not implemented.');
   }
+
   private redis: Redis;
 
   constructor(config: UpstashConfig) {
@@ -76,9 +189,17 @@ export class UpstashStore extends MastraStorage {
   async insert({ tableName, record }: { tableName: TABLE_NAMES; record: Record<string, any> }): Promise<void> {
     let key: string;
 
-    if (tableName === MastraStorage.TABLE_MESSAGES) {
+    if (tableName === TABLE_MESSAGES) {
       // For messages, use threadId as the primary key component
       key = this.getKey(tableName, { threadId: record.threadId, id: record.id });
+    } else if (tableName === TABLE_WORKFLOW_SNAPSHOT) {
+      key = this.getKey(tableName, {
+        namespace: record.namespace || 'workflows',
+        workflow_name: record.workflow_name,
+        run_id: record.run_id,
+      });
+    } else if (tableName === TABLE_EVALS) {
+      key = this.getKey(tableName, { id: record.run_id });
     } else {
       key = this.getKey(tableName, { id: record.id });
     }
@@ -101,7 +222,7 @@ export class UpstashStore extends MastraStorage {
 
   async getThreadById({ threadId }: { threadId: string }): Promise<StorageThreadType | null> {
     const thread = await this.load<StorageThreadType>({
-      tableName: MastraStorage.TABLE_THREADS,
+      tableName: TABLE_THREADS,
       keys: { id: threadId },
     });
 
@@ -116,7 +237,7 @@ export class UpstashStore extends MastraStorage {
   }
 
   async getThreadsByResourceId({ resourceId }: { resourceId: string }): Promise<StorageThreadType[]> {
-    const pattern = `${MastraStorage.TABLE_THREADS}:*`;
+    const pattern = `${TABLE_THREADS}:*`;
     const keys = await this.redis.keys(pattern);
     const threads = await Promise.all(
       keys.map(async key => {
@@ -137,7 +258,7 @@ export class UpstashStore extends MastraStorage {
 
   async saveThread({ thread }: { thread: StorageThreadType }): Promise<StorageThreadType> {
     await this.insert({
-      tableName: MastraStorage.TABLE_THREADS,
+      tableName: TABLE_THREADS,
       record: thread,
     });
     return thread;
@@ -171,12 +292,12 @@ export class UpstashStore extends MastraStorage {
   }
 
   async deleteThread({ threadId }: { threadId: string }): Promise<void> {
-    const key = this.getKey(MastraStorage.TABLE_THREADS, { id: threadId });
+    const key = this.getKey(TABLE_THREADS, { id: threadId });
     await this.redis.del(key);
   }
 
   private getMessageKey(threadId: string, messageId: string): string {
-    return this.getKey(MastraStorage.TABLE_MESSAGES, { threadId, id: messageId });
+    return this.getKey(TABLE_MESSAGES, { threadId, id: messageId });
   }
 
   private getThreadMessagesKey(threadId: string): string {
@@ -275,12 +396,17 @@ export class UpstashStore extends MastraStorage {
     snapshot: WorkflowRunState;
   }): Promise<void> {
     const { namespace = 'workflows', workflowName, runId, snapshot } = params;
-    const key = this.getKey(MastraStorage.TABLE_WORKFLOW_SNAPSHOT, {
-      namespace,
-      workflow_name: workflowName,
-      run_id: runId,
+    await this.insert({
+      tableName: TABLE_WORKFLOW_SNAPSHOT,
+      record: {
+        namespace,
+        workflow_name: workflowName,
+        run_id: runId,
+        snapshot,
+        createdAt: new Date(),
+        updatedAt: new Date(),
+      },
     });
-    await this.redis.set(key, snapshot); // Store snapshot directly without wrapping
   }
 
   async loadWorkflowSnapshot(params: {
@@ -289,13 +415,104 @@ export class UpstashStore extends MastraStorage {
     runId: string;
   }): Promise<WorkflowRunState | null> {
     const { namespace = 'workflows', workflowName, runId } = params;
-    const key = this.getKey(MastraStorage.TABLE_WORKFLOW_SNAPSHOT, {
+    const key = this.getKey(TABLE_WORKFLOW_SNAPSHOT, {
       namespace,
       workflow_name: workflowName,
       run_id: runId,
     });
-    const data = await this.redis.get<WorkflowRunState>(key);
-    return data || null;
+    const data = await this.redis.get<{
+      namespace: string;
+      workflow_name: string;
+      run_id: string;
+      snapshot: WorkflowRunState;
+    }>(key);
+    if (!data) return null;
+    return data.snapshot;
+  }
+
+  async getWorkflowRuns(
+    {
+      namespace,
+      workflowName,
+      fromDate,
+      toDate,
+      limit,
+      offset,
+    }: {
+      namespace: string;
+      workflowName?: string;
+      fromDate?: Date;
+      toDate?: Date;
+      limit?: number;
+      offset?: number;
+    } = { namespace: 'workflows' },
+  ): Promise<{
+    runs: Array<{
+      workflowName: string;
+      runId: string;
+      snapshot: WorkflowRunState | string;
+      createdAt: Date;
+      updatedAt: Date;
+    }>;
+    total: number;
+  }> {
+    // Get all workflow keys
+    const pattern = workflowName
+      ? this.getKey(TABLE_WORKFLOW_SNAPSHOT, { namespace, workflow_name: workflowName }) + ':*'
+      : this.getKey(TABLE_WORKFLOW_SNAPSHOT, { namespace }) + ':*';
+
+    const keys = await this.redis.keys(pattern);
+
+    // Get all workflow data
+    const workflows = await Promise.all(
+      keys.map(async key => {
+        const data = await this.redis.get<{
+          workflow_name: string;
+          run_id: string;
+          snapshot: WorkflowRunState | string;
+          createdAt: string | Date;
+          updatedAt: string | Date;
+        }>(key);
+        return data;
+      }),
+    );
+
+    // Filter and transform results
+    let runs = workflows
+      .filter(w => w !== null)
+      .map(w => {
+        let parsedSnapshot: WorkflowRunState | string = w!.snapshot as string;
+        if (typeof parsedSnapshot === 'string') {
+          try {
+            parsedSnapshot = JSON.parse(w!.snapshot as string) as WorkflowRunState;
+          } catch (_e) {
+            // If parsing fails, return the raw snapshot string
+            console.warn(`Failed to parse snapshot for workflow ${w!.workflow_name}:`);
+          }
+        }
+        return {
+          workflowName: w!.workflow_name,
+          runId: w!.run_id,
+          snapshot: parsedSnapshot,
+          createdAt: this.ensureDate(w!.createdAt)!,
+          updatedAt: this.ensureDate(w!.updatedAt)!,
+        };
+      })
+      .filter(w => {
+        if (fromDate && w.createdAt < fromDate) return false;
+        if (toDate && w.createdAt > toDate) return false;
+        return true;
+      })
+      .sort((a, b) => b.createdAt.getTime() - a.createdAt.getTime());
+
+    const total = runs.length;
+
+    // Apply pagination if requested
+    if (limit !== undefined && offset !== undefined) {
+      runs = runs.slice(offset, offset + limit);
+    }
+
+    return { runs, total };
   }
 
   async close(): Promise<void> {