@mastra/upstash 0.10.3-alpha.1 → 0.10.3-alpha.3
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/.turbo/turbo-build.log +10 -10
- package/CHANGELOG.md +17 -0
- package/dist/_tsup-dts-rollup.d.cts +25 -51
- package/dist/_tsup-dts-rollup.d.ts +25 -51
- package/dist/index.cjs +384 -384
- package/dist/index.js +384 -384
- package/package.json +3 -3
- package/src/storage/index.ts +313 -365
- package/src/storage/upstash.test.ts +190 -102
package/src/storage/index.ts
CHANGED
|
@@ -15,6 +15,9 @@ import type {
|
|
|
15
15
|
EvalRow,
|
|
16
16
|
WorkflowRuns,
|
|
17
17
|
WorkflowRun,
|
|
18
|
+
PaginationInfo,
|
|
19
|
+
PaginationArgs,
|
|
20
|
+
StorageGetTracesArg,
|
|
18
21
|
} from '@mastra/core/storage';
|
|
19
22
|
import type { WorkflowRunState } from '@mastra/core/workflows';
|
|
20
23
|
import { Redis } from '@upstash/redis';
|
|
@@ -36,6 +39,14 @@ export class UpstashStore extends MastraStorage {
|
|
|
36
39
|
});
|
|
37
40
|
}
|
|
38
41
|
|
|
42
|
+
public get supports(): {
|
|
43
|
+
selectByIncludeResourceScope: boolean;
|
|
44
|
+
} {
|
|
45
|
+
return {
|
|
46
|
+
selectByIncludeResourceScope: true,
|
|
47
|
+
};
|
|
48
|
+
}
|
|
49
|
+
|
|
39
50
|
private transformEvalRecord(record: Record<string, any>): EvalRow {
|
|
40
51
|
// Parse JSON strings if needed
|
|
41
52
|
let result = record.result;
|
|
@@ -136,7 +147,8 @@ export class UpstashStore extends MastraStorage {
|
|
|
136
147
|
}
|
|
137
148
|
|
|
138
149
|
private getMessageKey(threadId: string, messageId: string): string {
|
|
139
|
-
|
|
150
|
+
const key = this.getKey(TABLE_MESSAGES, { threadId, id: messageId });
|
|
151
|
+
return key;
|
|
140
152
|
}
|
|
141
153
|
|
|
142
154
|
private getThreadMessagesKey(threadId: string): string {
|
|
@@ -263,82 +275,44 @@ export class UpstashStore extends MastraStorage {
|
|
|
263
275
|
}
|
|
264
276
|
}
|
|
265
277
|
|
|
266
|
-
|
|
267
|
-
|
|
268
|
-
|
|
269
|
-
|
|
270
|
-
|
|
271
|
-
|
|
272
|
-
|
|
273
|
-
|
|
274
|
-
|
|
275
|
-
|
|
276
|
-
|
|
277
|
-
|
|
278
|
-
|
|
279
|
-
page: number;
|
|
280
|
-
perPage?: number;
|
|
281
|
-
attributes?: Record<string, string>;
|
|
282
|
-
filters?: Record<string, any>;
|
|
283
|
-
fromDate?: Date;
|
|
284
|
-
toDate?: Date;
|
|
285
|
-
returnPaginationResults: true;
|
|
286
|
-
}): Promise<{
|
|
287
|
-
traces: any[];
|
|
288
|
-
total: number;
|
|
289
|
-
page: number;
|
|
290
|
-
perPage: number;
|
|
291
|
-
hasMore: boolean;
|
|
292
|
-
}>;
|
|
293
|
-
public async getTraces(args: {
|
|
294
|
-
name?: string;
|
|
295
|
-
scope?: string;
|
|
296
|
-
page: number;
|
|
297
|
-
perPage?: number;
|
|
298
|
-
attributes?: Record<string, string>;
|
|
299
|
-
filters?: Record<string, any>;
|
|
300
|
-
fromDate?: Date;
|
|
301
|
-
toDate?: Date;
|
|
302
|
-
returnPaginationResults?: boolean;
|
|
303
|
-
}): Promise<
|
|
304
|
-
| any[]
|
|
305
|
-
| {
|
|
306
|
-
traces: any[];
|
|
307
|
-
total: number;
|
|
308
|
-
page: number;
|
|
309
|
-
perPage: number;
|
|
310
|
-
hasMore: boolean;
|
|
311
|
-
}
|
|
312
|
-
> {
|
|
313
|
-
const {
|
|
314
|
-
name,
|
|
315
|
-
scope,
|
|
316
|
-
page,
|
|
317
|
-
perPage: perPageInput,
|
|
318
|
-
attributes,
|
|
319
|
-
filters,
|
|
320
|
-
fromDate,
|
|
321
|
-
toDate,
|
|
322
|
-
returnPaginationResults,
|
|
323
|
-
} = args;
|
|
278
|
+
/**
|
|
279
|
+
* @deprecated use getTracesPaginated instead
|
|
280
|
+
*/
|
|
281
|
+
public async getTraces(args: StorageGetTracesArg): Promise<any[]> {
|
|
282
|
+
if (args.fromDate || args.toDate) {
|
|
283
|
+
(args as any).dateRange = {
|
|
284
|
+
start: args.fromDate,
|
|
285
|
+
end: args.toDate,
|
|
286
|
+
};
|
|
287
|
+
}
|
|
288
|
+
const { traces } = await this.getTracesPaginated(args);
|
|
289
|
+
return traces;
|
|
290
|
+
}
|
|
324
291
|
|
|
325
|
-
|
|
292
|
+
public async getTracesPaginated(
|
|
293
|
+
args: {
|
|
294
|
+
name?: string;
|
|
295
|
+
scope?: string;
|
|
296
|
+
attributes?: Record<string, string>;
|
|
297
|
+
filters?: Record<string, any>;
|
|
298
|
+
} & PaginationArgs,
|
|
299
|
+
): Promise<PaginationInfo & { traces: any[] }> {
|
|
300
|
+
const { name, scope, page = 0, perPage = 100, attributes, filters, dateRange } = args;
|
|
301
|
+
const fromDate = dateRange?.start;
|
|
302
|
+
const toDate = dateRange?.end;
|
|
326
303
|
|
|
327
304
|
try {
|
|
328
305
|
const pattern = `${TABLE_TRACES}:*`;
|
|
329
306
|
const keys = await this.scanKeys(pattern);
|
|
330
307
|
|
|
331
308
|
if (keys.length === 0) {
|
|
332
|
-
|
|
333
|
-
|
|
334
|
-
|
|
335
|
-
|
|
336
|
-
|
|
337
|
-
|
|
338
|
-
|
|
339
|
-
};
|
|
340
|
-
}
|
|
341
|
-
return [];
|
|
309
|
+
return {
|
|
310
|
+
traces: [],
|
|
311
|
+
total: 0,
|
|
312
|
+
page,
|
|
313
|
+
perPage: perPage || 100,
|
|
314
|
+
hasMore: false,
|
|
315
|
+
};
|
|
342
316
|
}
|
|
343
317
|
|
|
344
318
|
const pipeline = this.redis.pipeline();
|
|
@@ -405,29 +379,23 @@ export class UpstashStore extends MastraStorage {
|
|
|
405
379
|
const end = start + resolvedPerPage;
|
|
406
380
|
const paginatedTraces = transformedTraces.slice(start, end);
|
|
407
381
|
const hasMore = end < total;
|
|
408
|
-
|
|
409
|
-
|
|
410
|
-
|
|
411
|
-
|
|
412
|
-
|
|
413
|
-
|
|
414
|
-
|
|
415
|
-
|
|
416
|
-
} else {
|
|
417
|
-
return paginatedTraces;
|
|
418
|
-
}
|
|
382
|
+
|
|
383
|
+
return {
|
|
384
|
+
traces: paginatedTraces,
|
|
385
|
+
total,
|
|
386
|
+
page,
|
|
387
|
+
perPage: resolvedPerPage,
|
|
388
|
+
hasMore,
|
|
389
|
+
};
|
|
419
390
|
} catch (error) {
|
|
420
391
|
console.error('Failed to get traces:', error);
|
|
421
|
-
|
|
422
|
-
|
|
423
|
-
|
|
424
|
-
|
|
425
|
-
|
|
426
|
-
|
|
427
|
-
|
|
428
|
-
};
|
|
429
|
-
}
|
|
430
|
-
return [];
|
|
392
|
+
return {
|
|
393
|
+
traces: [],
|
|
394
|
+
total: 0,
|
|
395
|
+
page,
|
|
396
|
+
perPage: perPage || 100,
|
|
397
|
+
hasMore: false,
|
|
398
|
+
};
|
|
431
399
|
}
|
|
432
400
|
}
|
|
433
401
|
|
|
@@ -506,44 +474,15 @@ export class UpstashStore extends MastraStorage {
|
|
|
506
474
|
};
|
|
507
475
|
}
|
|
508
476
|
|
|
509
|
-
|
|
510
|
-
|
|
511
|
-
|
|
512
|
-
|
|
513
|
-
page: number;
|
|
514
|
-
perPage: number;
|
|
515
|
-
hasMore: boolean;
|
|
516
|
-
}>;
|
|
517
|
-
async getThreadsByResourceId(args: { resourceId: string; page?: number; perPage?: number }): Promise<
|
|
518
|
-
| StorageThreadType[]
|
|
519
|
-
| {
|
|
520
|
-
threads: StorageThreadType[];
|
|
521
|
-
total: number;
|
|
522
|
-
page: number;
|
|
523
|
-
perPage: number;
|
|
524
|
-
hasMore: boolean;
|
|
525
|
-
}
|
|
526
|
-
> {
|
|
527
|
-
const resourceId: string = args.resourceId;
|
|
528
|
-
const page: number | undefined = args.page;
|
|
529
|
-
// Determine perPage only if page is actually provided. Otherwise, its value is not critical for the non-paginated path.
|
|
530
|
-
// If page is provided, perPage defaults to 100 if not specified.
|
|
531
|
-
const perPage: number = page !== undefined ? (args.perPage !== undefined ? args.perPage : 100) : 100;
|
|
532
|
-
|
|
477
|
+
/**
|
|
478
|
+
* @deprecated use getThreadsByResourceIdPaginated instead
|
|
479
|
+
*/
|
|
480
|
+
async getThreadsByResourceId({ resourceId }: { resourceId: string }): Promise<StorageThreadType[]> {
|
|
533
481
|
try {
|
|
534
482
|
const pattern = `${TABLE_THREADS}:*`;
|
|
535
483
|
const keys = await this.scanKeys(pattern);
|
|
536
484
|
|
|
537
485
|
if (keys.length === 0) {
|
|
538
|
-
if (page !== undefined) {
|
|
539
|
-
return {
|
|
540
|
-
threads: [],
|
|
541
|
-
total: 0,
|
|
542
|
-
page,
|
|
543
|
-
perPage, // perPage is number here
|
|
544
|
-
hasMore: false,
|
|
545
|
-
};
|
|
546
|
-
}
|
|
547
486
|
return [];
|
|
548
487
|
}
|
|
549
488
|
|
|
@@ -565,40 +504,48 @@ export class UpstashStore extends MastraStorage {
|
|
|
565
504
|
}
|
|
566
505
|
|
|
567
506
|
allThreads.sort((a, b) => b.createdAt.getTime() - a.createdAt.getTime());
|
|
568
|
-
|
|
569
|
-
if (page !== undefined) {
|
|
570
|
-
// If page is defined, perPage is also a number (due to the defaulting logic above)
|
|
571
|
-
const total = allThreads.length;
|
|
572
|
-
const start = page * perPage;
|
|
573
|
-
const end = start + perPage;
|
|
574
|
-
const paginatedThreads = allThreads.slice(start, end);
|
|
575
|
-
const hasMore = end < total;
|
|
576
|
-
return {
|
|
577
|
-
threads: paginatedThreads,
|
|
578
|
-
total,
|
|
579
|
-
page,
|
|
580
|
-
perPage,
|
|
581
|
-
hasMore,
|
|
582
|
-
};
|
|
583
|
-
} else {
|
|
584
|
-
// page is undefined, return all threads
|
|
585
|
-
return allThreads;
|
|
586
|
-
}
|
|
507
|
+
return allThreads;
|
|
587
508
|
} catch (error) {
|
|
588
509
|
console.error('Error in getThreadsByResourceId:', error);
|
|
589
|
-
if (page !== undefined) {
|
|
590
|
-
return {
|
|
591
|
-
threads: [],
|
|
592
|
-
total: 0,
|
|
593
|
-
page,
|
|
594
|
-
perPage, // perPage is number here
|
|
595
|
-
hasMore: false,
|
|
596
|
-
};
|
|
597
|
-
}
|
|
598
510
|
return [];
|
|
599
511
|
}
|
|
600
512
|
}
|
|
601
513
|
|
|
514
|
+
public async getThreadsByResourceIdPaginated(
|
|
515
|
+
args: {
|
|
516
|
+
resourceId: string;
|
|
517
|
+
} & PaginationArgs,
|
|
518
|
+
): Promise<PaginationInfo & { threads: StorageThreadType[] }> {
|
|
519
|
+
const { resourceId, page = 0, perPage = 100 } = args;
|
|
520
|
+
|
|
521
|
+
try {
|
|
522
|
+
const allThreads = await this.getThreadsByResourceId({ resourceId });
|
|
523
|
+
|
|
524
|
+
const total = allThreads.length;
|
|
525
|
+
const start = page * perPage;
|
|
526
|
+
const end = start + perPage;
|
|
527
|
+
const paginatedThreads = allThreads.slice(start, end);
|
|
528
|
+
const hasMore = end < total;
|
|
529
|
+
|
|
530
|
+
return {
|
|
531
|
+
threads: paginatedThreads,
|
|
532
|
+
total,
|
|
533
|
+
page,
|
|
534
|
+
perPage,
|
|
535
|
+
hasMore,
|
|
536
|
+
};
|
|
537
|
+
} catch (error) {
|
|
538
|
+
console.error('Error in getThreadsByResourceIdPaginated:', error);
|
|
539
|
+
return {
|
|
540
|
+
threads: [],
|
|
541
|
+
total: 0,
|
|
542
|
+
page,
|
|
543
|
+
perPage,
|
|
544
|
+
hasMore: false,
|
|
545
|
+
};
|
|
546
|
+
}
|
|
547
|
+
}
|
|
548
|
+
|
|
602
549
|
async saveThread({ thread }: { thread: StorageThreadType }): Promise<StorageThreadType> {
|
|
603
550
|
await this.insert({
|
|
604
551
|
tableName: TABLE_THREADS,
|
|
@@ -664,6 +611,17 @@ export class UpstashStore extends MastraStorage {
|
|
|
664
611
|
const { messages, format = 'v1' } = args;
|
|
665
612
|
if (messages.length === 0) return [];
|
|
666
613
|
|
|
614
|
+
const threadId = messages[0]?.threadId;
|
|
615
|
+
if (!threadId) {
|
|
616
|
+
throw new Error('Thread ID is required');
|
|
617
|
+
}
|
|
618
|
+
|
|
619
|
+
// Check if thread exists
|
|
620
|
+
const thread = await this.getThreadById({ threadId });
|
|
621
|
+
if (!thread) {
|
|
622
|
+
throw new Error(`Thread ${threadId} not found`);
|
|
623
|
+
}
|
|
624
|
+
|
|
667
625
|
// Add an index to each message to maintain order
|
|
668
626
|
const messagesWithIndex = messages.map((message, index) => ({
|
|
669
627
|
...message,
|
|
@@ -676,7 +634,8 @@ export class UpstashStore extends MastraStorage {
|
|
|
676
634
|
const pipeline = this.redis.pipeline();
|
|
677
635
|
for (const message of batch) {
|
|
678
636
|
const key = this.getMessageKey(message.threadId!, message.id);
|
|
679
|
-
const
|
|
637
|
+
const createdAtScore = new Date(message.createdAt).getTime();
|
|
638
|
+
const score = message._index !== undefined ? message._index : createdAtScore;
|
|
680
639
|
|
|
681
640
|
// Store the message data
|
|
682
641
|
pipeline.set(key, message);
|
|
@@ -696,142 +655,71 @@ export class UpstashStore extends MastraStorage {
|
|
|
696
655
|
return list.get.all.v1();
|
|
697
656
|
}
|
|
698
657
|
|
|
699
|
-
|
|
700
|
-
|
|
701
|
-
|
|
702
|
-
|
|
703
|
-
|
|
704
|
-
|
|
705
|
-
page: number;
|
|
706
|
-
perPage?: number;
|
|
707
|
-
fromDate?: Date;
|
|
708
|
-
toDate?: Date;
|
|
709
|
-
},
|
|
710
|
-
): Promise<{
|
|
711
|
-
messages: MastraMessageV1[] | MastraMessageV2[];
|
|
712
|
-
total: number;
|
|
713
|
-
page: number;
|
|
714
|
-
perPage: number;
|
|
715
|
-
hasMore: boolean;
|
|
716
|
-
}>;
|
|
717
|
-
public async getMessages({
|
|
718
|
-
threadId,
|
|
719
|
-
selectBy,
|
|
720
|
-
format,
|
|
721
|
-
page,
|
|
722
|
-
perPage = 40,
|
|
723
|
-
fromDate,
|
|
724
|
-
toDate,
|
|
725
|
-
}: StorageGetMessagesArg & {
|
|
726
|
-
format?: 'v1' | 'v2';
|
|
727
|
-
page?: number;
|
|
728
|
-
perPage?: number;
|
|
729
|
-
fromDate?: Date;
|
|
730
|
-
toDate?: Date;
|
|
731
|
-
}): Promise<
|
|
732
|
-
| MastraMessageV1[]
|
|
733
|
-
| MastraMessageV2[]
|
|
734
|
-
| {
|
|
735
|
-
messages: MastraMessageV1[] | MastraMessageV2[];
|
|
736
|
-
total: number;
|
|
737
|
-
page: number;
|
|
738
|
-
perPage: number;
|
|
739
|
-
hasMore: boolean;
|
|
740
|
-
}
|
|
741
|
-
> {
|
|
742
|
-
const threadMessagesKey = this.getThreadMessagesKey(threadId);
|
|
743
|
-
|
|
744
|
-
const allMessageIds = await this.redis.zrange(threadMessagesKey, 0, -1);
|
|
745
|
-
// If pagination is requested, use the new pagination logic
|
|
746
|
-
if (page !== undefined) {
|
|
747
|
-
try {
|
|
748
|
-
// Get all message IDs from the sorted set
|
|
749
|
-
|
|
750
|
-
if (allMessageIds.length === 0) {
|
|
751
|
-
return {
|
|
752
|
-
messages: [],
|
|
753
|
-
total: 0,
|
|
754
|
-
page,
|
|
755
|
-
perPage,
|
|
756
|
-
hasMore: false,
|
|
757
|
-
};
|
|
758
|
-
}
|
|
759
|
-
|
|
760
|
-
// Use pipeline to fetch all messages efficiently
|
|
761
|
-
const pipeline = this.redis.pipeline();
|
|
762
|
-
allMessageIds.forEach(id => pipeline.get(this.getMessageKey(threadId, id as string)));
|
|
763
|
-
const results = await pipeline.exec();
|
|
764
|
-
|
|
765
|
-
// Process messages and apply filters - handle undefined results from pipeline
|
|
766
|
-
let messages = results
|
|
767
|
-
.map((result: any) => result as MastraMessageV2 | null)
|
|
768
|
-
.filter((msg): msg is MastraMessageV2 => msg !== null) as (MastraMessageV2 & { _index?: number })[];
|
|
658
|
+
private async _getIncludedMessages(
|
|
659
|
+
threadId: string,
|
|
660
|
+
selectBy: StorageGetMessagesArg['selectBy'],
|
|
661
|
+
): Promise<MastraMessageV2[] | MastraMessageV1[]> {
|
|
662
|
+
const messageIds = new Set<string>();
|
|
663
|
+
const messageIdToThreadIds: Record<string, string> = {};
|
|
769
664
|
|
|
770
|
-
|
|
771
|
-
|
|
772
|
-
|
|
773
|
-
|
|
665
|
+
// First, get specifically included messages and their context
|
|
666
|
+
if (selectBy?.include?.length) {
|
|
667
|
+
for (const item of selectBy.include) {
|
|
668
|
+
messageIds.add(item.id);
|
|
774
669
|
|
|
775
|
-
if
|
|
776
|
-
|
|
670
|
+
// Use per-include threadId if present, else fallback to main threadId
|
|
671
|
+
const itemThreadId = item.threadId || threadId;
|
|
672
|
+
messageIdToThreadIds[item.id] = itemThreadId;
|
|
673
|
+
const itemThreadMessagesKey = this.getThreadMessagesKey(itemThreadId);
|
|
674
|
+
|
|
675
|
+
// Get the rank of this message in the sorted set
|
|
676
|
+
const rank = await this.redis.zrank(itemThreadMessagesKey, item.id);
|
|
677
|
+
if (rank === null) continue;
|
|
678
|
+
|
|
679
|
+
// Get previous messages if requested
|
|
680
|
+
if (item.withPreviousMessages) {
|
|
681
|
+
const start = Math.max(0, rank - item.withPreviousMessages);
|
|
682
|
+
const prevIds = rank === 0 ? [] : await this.redis.zrange(itemThreadMessagesKey, start, rank - 1);
|
|
683
|
+
prevIds.forEach(id => {
|
|
684
|
+
messageIds.add(id as string);
|
|
685
|
+
messageIdToThreadIds[id as string] = itemThreadId;
|
|
686
|
+
});
|
|
777
687
|
}
|
|
778
688
|
|
|
779
|
-
//
|
|
780
|
-
|
|
781
|
-
|
|
782
|
-
|
|
783
|
-
|
|
784
|
-
|
|
785
|
-
const start = page * perPage;
|
|
786
|
-
const end = start + perPage;
|
|
787
|
-
const hasMore = end < total;
|
|
788
|
-
const paginatedMessages = messages.slice(start, end);
|
|
789
|
-
|
|
790
|
-
// Remove _index before returning and handle format conversion properly
|
|
791
|
-
const prepared = paginatedMessages
|
|
792
|
-
.filter(message => message !== null && message !== undefined)
|
|
793
|
-
.map(message => {
|
|
794
|
-
const { _index, ...messageWithoutIndex } = message as MastraMessageV2 & { _index?: number };
|
|
795
|
-
return messageWithoutIndex as unknown as MastraMessageV1;
|
|
689
|
+
// Get next messages if requested
|
|
690
|
+
if (item.withNextMessages) {
|
|
691
|
+
const nextIds = await this.redis.zrange(itemThreadMessagesKey, rank + 1, rank + item.withNextMessages);
|
|
692
|
+
nextIds.forEach(id => {
|
|
693
|
+
messageIds.add(id as string);
|
|
694
|
+
messageIdToThreadIds[id as string] = itemThreadId;
|
|
796
695
|
});
|
|
797
|
-
|
|
798
|
-
// Return pagination object with correct format
|
|
799
|
-
if (format === 'v2') {
|
|
800
|
-
// Convert V1 format back to V2 format
|
|
801
|
-
const v2Messages = prepared.map(msg => ({
|
|
802
|
-
...msg,
|
|
803
|
-
content: msg.content || { format: 2, parts: [{ type: 'text', text: '' }] },
|
|
804
|
-
})) as MastraMessageV2[];
|
|
805
|
-
|
|
806
|
-
return {
|
|
807
|
-
messages: v2Messages,
|
|
808
|
-
total,
|
|
809
|
-
page,
|
|
810
|
-
perPage,
|
|
811
|
-
hasMore,
|
|
812
|
-
};
|
|
813
696
|
}
|
|
814
|
-
|
|
815
|
-
return {
|
|
816
|
-
messages: prepared,
|
|
817
|
-
total,
|
|
818
|
-
page,
|
|
819
|
-
perPage,
|
|
820
|
-
hasMore,
|
|
821
|
-
};
|
|
822
|
-
} catch (error) {
|
|
823
|
-
console.error('Failed to get paginated messages:', error);
|
|
824
|
-
return {
|
|
825
|
-
messages: [],
|
|
826
|
-
total: 0,
|
|
827
|
-
page,
|
|
828
|
-
perPage,
|
|
829
|
-
hasMore: false,
|
|
830
|
-
};
|
|
831
697
|
}
|
|
698
|
+
|
|
699
|
+
const pipeline = this.redis.pipeline();
|
|
700
|
+
Array.from(messageIds).forEach(id => {
|
|
701
|
+
const tId = messageIdToThreadIds[id] || threadId;
|
|
702
|
+
pipeline.get(this.getMessageKey(tId, id as string));
|
|
703
|
+
});
|
|
704
|
+
const results = await pipeline.exec();
|
|
705
|
+
return results.filter(result => result !== null) as MastraMessageV2[] | MastraMessageV1[];
|
|
832
706
|
}
|
|
833
707
|
|
|
834
|
-
|
|
708
|
+
return [];
|
|
709
|
+
}
|
|
710
|
+
|
|
711
|
+
/**
|
|
712
|
+
* @deprecated use getMessagesPaginated instead
|
|
713
|
+
*/
|
|
714
|
+
public async getMessages(args: StorageGetMessagesArg & { format?: 'v1' }): Promise<MastraMessageV1[]>;
|
|
715
|
+
public async getMessages(args: StorageGetMessagesArg & { format: 'v2' }): Promise<MastraMessageV2[]>;
|
|
716
|
+
public async getMessages({
|
|
717
|
+
threadId,
|
|
718
|
+
selectBy,
|
|
719
|
+
format,
|
|
720
|
+
}: StorageGetMessagesArg & { format?: 'v1' | 'v2' }): Promise<MastraMessageV1[] | MastraMessageV2[]> {
|
|
721
|
+
const threadMessagesKey = this.getThreadMessagesKey(threadId);
|
|
722
|
+
const allMessageIds = await this.redis.zrange(threadMessagesKey, 0, -1);
|
|
835
723
|
// When selectBy is undefined or selectBy.last is undefined, get ALL messages (not just 40)
|
|
836
724
|
let limit: number;
|
|
837
725
|
if (typeof selectBy?.last === 'number') {
|
|
@@ -844,62 +732,59 @@ export class UpstashStore extends MastraStorage {
|
|
|
844
732
|
}
|
|
845
733
|
|
|
846
734
|
const messageIds = new Set<string>();
|
|
735
|
+
const messageIdToThreadIds: Record<string, string> = {};
|
|
847
736
|
|
|
848
737
|
if (limit === 0 && !selectBy?.include) {
|
|
849
738
|
return [];
|
|
850
739
|
}
|
|
851
740
|
|
|
852
|
-
// First, get specifically included messages and their context
|
|
853
|
-
if (selectBy?.include?.length) {
|
|
854
|
-
for (const item of selectBy.include) {
|
|
855
|
-
messageIds.add(item.id);
|
|
856
|
-
|
|
857
|
-
if (item.withPreviousMessages || item.withNextMessages) {
|
|
858
|
-
// Get the rank of this message in the sorted set
|
|
859
|
-
const rank = await this.redis.zrank(threadMessagesKey, item.id);
|
|
860
|
-
if (rank === null) continue;
|
|
861
|
-
|
|
862
|
-
// Get previous messages if requested
|
|
863
|
-
if (item.withPreviousMessages) {
|
|
864
|
-
const start = Math.max(0, rank - item.withPreviousMessages);
|
|
865
|
-
const prevIds = rank === 0 ? [] : await this.redis.zrange(threadMessagesKey, start, rank - 1);
|
|
866
|
-
prevIds.forEach(id => messageIds.add(id as string));
|
|
867
|
-
}
|
|
868
|
-
|
|
869
|
-
// Get next messages if requested
|
|
870
|
-
if (item.withNextMessages) {
|
|
871
|
-
const nextIds = await this.redis.zrange(threadMessagesKey, rank + 1, rank + item.withNextMessages);
|
|
872
|
-
nextIds.forEach(id => messageIds.add(id as string));
|
|
873
|
-
}
|
|
874
|
-
}
|
|
875
|
-
}
|
|
876
|
-
}
|
|
877
|
-
|
|
878
741
|
// Then get the most recent messages (or all if no limit)
|
|
879
742
|
if (limit === Number.MAX_SAFE_INTEGER) {
|
|
880
743
|
// Get all messages
|
|
881
744
|
const allIds = await this.redis.zrange(threadMessagesKey, 0, -1);
|
|
882
|
-
allIds.forEach(id =>
|
|
745
|
+
allIds.forEach(id => {
|
|
746
|
+
messageIds.add(id as string);
|
|
747
|
+
messageIdToThreadIds[id as string] = threadId;
|
|
748
|
+
});
|
|
883
749
|
} else if (limit > 0) {
|
|
884
750
|
// Get limited number of recent messages
|
|
885
751
|
const latestIds = await this.redis.zrange(threadMessagesKey, -limit, -1);
|
|
886
|
-
latestIds.forEach(id =>
|
|
752
|
+
latestIds.forEach(id => {
|
|
753
|
+
messageIds.add(id as string);
|
|
754
|
+
messageIdToThreadIds[id as string] = threadId;
|
|
755
|
+
});
|
|
887
756
|
}
|
|
888
757
|
|
|
758
|
+
const includedMessages = await this._getIncludedMessages(threadId, selectBy);
|
|
759
|
+
|
|
889
760
|
// Fetch all needed messages in parallel
|
|
890
|
-
const messages =
|
|
891
|
-
|
|
892
|
-
|
|
893
|
-
|
|
894
|
-
|
|
895
|
-
|
|
896
|
-
|
|
761
|
+
const messages = [
|
|
762
|
+
...includedMessages,
|
|
763
|
+
...((
|
|
764
|
+
await Promise.all(
|
|
765
|
+
Array.from(messageIds).map(async id => {
|
|
766
|
+
const tId = messageIdToThreadIds[id] || threadId;
|
|
767
|
+
const byThreadId = await this.redis.get<MastraMessageV2 & { _index?: number }>(this.getMessageKey(tId, id));
|
|
768
|
+
if (byThreadId) return byThreadId;
|
|
769
|
+
|
|
770
|
+
return null;
|
|
771
|
+
}),
|
|
772
|
+
)
|
|
773
|
+
).filter(msg => msg !== null) as (MastraMessageV2 & { _index?: number })[]),
|
|
774
|
+
];
|
|
897
775
|
|
|
898
776
|
// Sort messages by their position in the sorted set
|
|
899
777
|
messages.sort((a, b) => allMessageIds.indexOf(a!.id) - allMessageIds.indexOf(b!.id));
|
|
900
778
|
|
|
779
|
+
const seen = new Set<string>();
|
|
780
|
+
const dedupedMessages = messages.filter(row => {
|
|
781
|
+
if (seen.has(row.id)) return false;
|
|
782
|
+
seen.add(row.id);
|
|
783
|
+
return true;
|
|
784
|
+
});
|
|
785
|
+
|
|
901
786
|
// Remove _index before returning and handle format conversion properly
|
|
902
|
-
const prepared =
|
|
787
|
+
const prepared = dedupedMessages
|
|
903
788
|
.filter(message => message !== null && message !== undefined)
|
|
904
789
|
.map(message => {
|
|
905
790
|
const { _index, ...messageWithoutIndex } = message as MastraMessageV2 & { _index?: number };
|
|
@@ -919,6 +804,89 @@ export class UpstashStore extends MastraStorage {
|
|
|
919
804
|
return prepared;
|
|
920
805
|
}
|
|
921
806
|
|
|
807
|
+
public async getMessagesPaginated(
|
|
808
|
+
args: StorageGetMessagesArg & {
|
|
809
|
+
format?: 'v1' | 'v2';
|
|
810
|
+
},
|
|
811
|
+
): Promise<PaginationInfo & { messages: MastraMessageV1[] | MastraMessageV2[] }> {
|
|
812
|
+
const { threadId, selectBy, format } = args;
|
|
813
|
+
const { page = 0, perPage = 40, dateRange } = selectBy?.pagination || {};
|
|
814
|
+
const fromDate = dateRange?.start;
|
|
815
|
+
const toDate = dateRange?.end;
|
|
816
|
+
const threadMessagesKey = this.getThreadMessagesKey(threadId);
|
|
817
|
+
const messages: (MastraMessageV2 | MastraMessageV1)[] = [];
|
|
818
|
+
|
|
819
|
+
const includedMessages = await this._getIncludedMessages(threadId, selectBy);
|
|
820
|
+
messages.push(...includedMessages);
|
|
821
|
+
|
|
822
|
+
try {
|
|
823
|
+
const allMessageIds = await this.redis.zrange(threadMessagesKey, 0, -1);
|
|
824
|
+
if (allMessageIds.length === 0) {
|
|
825
|
+
return {
|
|
826
|
+
messages: [],
|
|
827
|
+
total: 0,
|
|
828
|
+
page,
|
|
829
|
+
perPage,
|
|
830
|
+
hasMore: false,
|
|
831
|
+
};
|
|
832
|
+
}
|
|
833
|
+
|
|
834
|
+
// Use pipeline to fetch all messages efficiently
|
|
835
|
+
const pipeline = this.redis.pipeline();
|
|
836
|
+
allMessageIds.forEach(id => pipeline.get(this.getMessageKey(threadId, id as string)));
|
|
837
|
+
const results = await pipeline.exec();
|
|
838
|
+
|
|
839
|
+
// Process messages and apply filters - handle undefined results from pipeline
|
|
840
|
+
let messagesData = results.filter((msg): msg is MastraMessageV2 | MastraMessageV1 => msg !== null) as (
|
|
841
|
+
| MastraMessageV2
|
|
842
|
+
| MastraMessageV1
|
|
843
|
+
)[];
|
|
844
|
+
|
|
845
|
+
// Apply date filters if provided
|
|
846
|
+
if (fromDate) {
|
|
847
|
+
messagesData = messagesData.filter(msg => msg && new Date(msg.createdAt).getTime() >= fromDate.getTime());
|
|
848
|
+
}
|
|
849
|
+
|
|
850
|
+
if (toDate) {
|
|
851
|
+
messagesData = messagesData.filter(msg => msg && new Date(msg.createdAt).getTime() <= toDate.getTime());
|
|
852
|
+
}
|
|
853
|
+
|
|
854
|
+
// Sort messages by their position in the sorted set
|
|
855
|
+
messagesData.sort((a, b) => allMessageIds.indexOf(a!.id) - allMessageIds.indexOf(b!.id));
|
|
856
|
+
|
|
857
|
+
const total = messagesData.length;
|
|
858
|
+
|
|
859
|
+
const start = page * perPage;
|
|
860
|
+
const end = start + perPage;
|
|
861
|
+
const hasMore = end < total;
|
|
862
|
+
const paginatedMessages = messagesData.slice(start, end);
|
|
863
|
+
|
|
864
|
+
messages.push(...paginatedMessages);
|
|
865
|
+
|
|
866
|
+
const list = new MessageList().add(messages, 'memory');
|
|
867
|
+
const finalMessages = (format === `v2` ? list.get.all.v2() : list.get.all.v1()) as
|
|
868
|
+
| MastraMessageV1[]
|
|
869
|
+
| MastraMessageV2[];
|
|
870
|
+
|
|
871
|
+
return {
|
|
872
|
+
messages: finalMessages,
|
|
873
|
+
total,
|
|
874
|
+
page,
|
|
875
|
+
perPage,
|
|
876
|
+
hasMore,
|
|
877
|
+
};
|
|
878
|
+
} catch (error) {
|
|
879
|
+
console.error('Failed to get paginated messages:', error);
|
|
880
|
+
return {
|
|
881
|
+
messages: [],
|
|
882
|
+
total: 0,
|
|
883
|
+
page,
|
|
884
|
+
perPage,
|
|
885
|
+
hasMore: false,
|
|
886
|
+
};
|
|
887
|
+
}
|
|
888
|
+
}
|
|
889
|
+
|
|
922
890
|
async persistWorkflowSnapshot(params: {
|
|
923
891
|
namespace: string;
|
|
924
892
|
workflowName: string;
|
|
@@ -965,28 +933,17 @@ export class UpstashStore extends MastraStorage {
|
|
|
965
933
|
* @param options Pagination and filtering options
|
|
966
934
|
* @returns Object with evals array and total count
|
|
967
935
|
*/
|
|
968
|
-
async getEvals(
|
|
969
|
-
|
|
970
|
-
|
|
971
|
-
|
|
972
|
-
|
|
973
|
-
|
|
974
|
-
offset?: number;
|
|
975
|
-
fromDate?: Date;
|
|
976
|
-
toDate?: Date;
|
|
977
|
-
}): Promise<{
|
|
978
|
-
evals: EvalRow[];
|
|
979
|
-
total: number;
|
|
980
|
-
page?: number;
|
|
981
|
-
perPage?: number;
|
|
982
|
-
hasMore?: boolean;
|
|
983
|
-
}> {
|
|
936
|
+
async getEvals(
|
|
937
|
+
options?: {
|
|
938
|
+
agentName?: string;
|
|
939
|
+
type?: 'test' | 'live';
|
|
940
|
+
} & PaginationArgs,
|
|
941
|
+
): Promise<PaginationInfo & { evals: EvalRow[] }> {
|
|
984
942
|
try {
|
|
985
943
|
// Default pagination parameters
|
|
986
|
-
const page = options
|
|
987
|
-
const
|
|
988
|
-
const
|
|
989
|
-
const offset = options?.offset;
|
|
944
|
+
const { agentName, type, page = 0, perPage = 100, dateRange } = options || {};
|
|
945
|
+
const fromDate = dateRange?.start;
|
|
946
|
+
const toDate = dateRange?.end;
|
|
990
947
|
|
|
991
948
|
// Get all keys that match the evals table pattern using cursor-based scanning
|
|
992
949
|
const pattern = `${TABLE_EVALS}:*`;
|
|
@@ -997,8 +954,8 @@ export class UpstashStore extends MastraStorage {
|
|
|
997
954
|
return {
|
|
998
955
|
evals: [],
|
|
999
956
|
total: 0,
|
|
1000
|
-
page
|
|
1001
|
-
perPage
|
|
957
|
+
page,
|
|
958
|
+
perPage,
|
|
1002
959
|
hasMore: false,
|
|
1003
960
|
};
|
|
1004
961
|
}
|
|
@@ -1014,12 +971,12 @@ export class UpstashStore extends MastraStorage {
|
|
|
1014
971
|
.filter((record): record is Record<string, any> => record !== null && typeof record === 'object');
|
|
1015
972
|
|
|
1016
973
|
// Apply agent name filter if provided
|
|
1017
|
-
if (
|
|
1018
|
-
filteredEvals = filteredEvals.filter(record => record.agent_name ===
|
|
974
|
+
if (agentName) {
|
|
975
|
+
filteredEvals = filteredEvals.filter(record => record.agent_name === agentName);
|
|
1019
976
|
}
|
|
1020
977
|
|
|
1021
978
|
// Apply type filter if provided
|
|
1022
|
-
if (
|
|
979
|
+
if (type === 'test') {
|
|
1023
980
|
filteredEvals = filteredEvals.filter(record => {
|
|
1024
981
|
if (!record.test_info) return false;
|
|
1025
982
|
|
|
@@ -1033,7 +990,7 @@ export class UpstashStore extends MastraStorage {
|
|
|
1033
990
|
return false;
|
|
1034
991
|
}
|
|
1035
992
|
});
|
|
1036
|
-
} else if (
|
|
993
|
+
} else if (type === 'live') {
|
|
1037
994
|
filteredEvals = filteredEvals.filter(record => {
|
|
1038
995
|
if (!record.test_info) return true;
|
|
1039
996
|
|
|
@@ -1050,17 +1007,17 @@ export class UpstashStore extends MastraStorage {
|
|
|
1050
1007
|
}
|
|
1051
1008
|
|
|
1052
1009
|
// Apply date filters if provided
|
|
1053
|
-
if (
|
|
1010
|
+
if (fromDate) {
|
|
1054
1011
|
filteredEvals = filteredEvals.filter(record => {
|
|
1055
1012
|
const createdAt = new Date(record.created_at || record.createdAt || 0);
|
|
1056
|
-
return createdAt.getTime() >=
|
|
1013
|
+
return createdAt.getTime() >= fromDate.getTime();
|
|
1057
1014
|
});
|
|
1058
1015
|
}
|
|
1059
1016
|
|
|
1060
|
-
if (
|
|
1017
|
+
if (toDate) {
|
|
1061
1018
|
filteredEvals = filteredEvals.filter(record => {
|
|
1062
1019
|
const createdAt = new Date(record.created_at || record.createdAt || 0);
|
|
1063
|
-
return createdAt.getTime() <=
|
|
1020
|
+
return createdAt.getTime() <= toDate.getTime();
|
|
1064
1021
|
});
|
|
1065
1022
|
}
|
|
1066
1023
|
|
|
@@ -1073,21 +1030,11 @@ export class UpstashStore extends MastraStorage {
|
|
|
1073
1030
|
|
|
1074
1031
|
const total = filteredEvals.length;
|
|
1075
1032
|
|
|
1076
|
-
// Apply pagination
|
|
1077
|
-
|
|
1078
|
-
|
|
1079
|
-
|
|
1080
|
-
|
|
1081
|
-
// Offset-based pagination
|
|
1082
|
-
paginatedEvals = filteredEvals.slice(offset, offset + limit);
|
|
1083
|
-
hasMore = offset + limit < total;
|
|
1084
|
-
} else {
|
|
1085
|
-
// Page-based pagination
|
|
1086
|
-
const start = page * perPage;
|
|
1087
|
-
const end = start + perPage;
|
|
1088
|
-
paginatedEvals = filteredEvals.slice(start, end);
|
|
1089
|
-
hasMore = end < total;
|
|
1090
|
-
}
|
|
1033
|
+
// Apply pagination
|
|
1034
|
+
const start = page * perPage;
|
|
1035
|
+
const end = start + perPage;
|
|
1036
|
+
const paginatedEvals = filteredEvals.slice(start, end);
|
|
1037
|
+
const hasMore = end < total;
|
|
1091
1038
|
|
|
1092
1039
|
// Transform to EvalRow format
|
|
1093
1040
|
const evals = paginatedEvals.map(record => this.transformEvalRecord(record));
|
|
@@ -1095,17 +1042,18 @@ export class UpstashStore extends MastraStorage {
|
|
|
1095
1042
|
return {
|
|
1096
1043
|
evals,
|
|
1097
1044
|
total,
|
|
1098
|
-
page
|
|
1099
|
-
perPage
|
|
1045
|
+
page,
|
|
1046
|
+
perPage,
|
|
1100
1047
|
hasMore,
|
|
1101
1048
|
};
|
|
1102
1049
|
} catch (error) {
|
|
1050
|
+
const { page = 0, perPage = 100 } = options || {};
|
|
1103
1051
|
console.error('Failed to get evals:', error);
|
|
1104
1052
|
return {
|
|
1105
1053
|
evals: [],
|
|
1106
1054
|
total: 0,
|
|
1107
|
-
page
|
|
1108
|
-
perPage
|
|
1055
|
+
page,
|
|
1056
|
+
perPage,
|
|
1109
1057
|
hasMore: false,
|
|
1110
1058
|
};
|
|
1111
1059
|
}
|