Spaces:
Sleeping
Sleeping
File size: 4,961 Bytes
720c6ed | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 | import {
BaseCheckpointSaver,
type Checkpoint,
type CheckpointListOptions,
type CheckpointTuple,
type PendingWrite,
type CheckpointMetadata,
type ChannelVersions
} from '@langchain/langgraph-checkpoint';
import type { RunnableConfig } from '@langchain/core/runnables';
import { getSupabaseClient } from '../supabase.js';
export class SupabaseSaver extends BaseCheckpointSaver {
async getTuple(config: RunnableConfig): Promise<CheckpointTuple | undefined> {
const threadId = config.configurable?.thread_id;
if (!threadId) return undefined;
const client = getSupabaseClient();
const query = client.from('langgraph_checkpoints')
.select('*')
.eq('thread_id', threadId)
.order('checkpoint_id', { ascending: false })
.limit(1);
if (config.configurable?.checkpoint_id) {
query.eq('checkpoint_id', config.configurable.checkpoint_id);
}
const { data, error } = await query.single();
if (error || !data) return undefined;
const checkpoint = await this.serde.loadsTyped('json', data.checkpoint);
const metadata = await this.serde.loadsTyped('json', data.metadata);
return {
config: {
configurable: {
thread_id: threadId,
checkpoint_ns: '',
checkpoint_id: data.checkpoint_id
}
},
checkpoint: checkpoint as Checkpoint,
metadata: metadata as CheckpointMetadata,
parentConfig: data.parent_checkpoint_id ? {
configurable: {
thread_id: threadId,
checkpoint_ns: '',
checkpoint_id: data.parent_checkpoint_id
}
} : undefined
};
}
async *list(config: RunnableConfig, options?: CheckpointListOptions): AsyncGenerator<CheckpointTuple> {
const threadId = config.configurable?.thread_id;
if (!threadId) return;
const client = getSupabaseClient();
let query = client.from('langgraph_checkpoints')
.select('*')
.eq('thread_id', threadId)
.order('checkpoint_id', { ascending: false });
if (options?.limit) {
query = query.limit(options.limit);
}
const { data, error } = await query;
if (error || !data) return;
for (const row of data) {
const checkpoint = await this.serde.loadsTyped('json', row.checkpoint);
const metadata = await this.serde.loadsTyped('json', row.metadata);
yield {
config: { configurable: { thread_id: threadId, checkpoint_ns: '', checkpoint_id: row.checkpoint_id } },
checkpoint: checkpoint as Checkpoint,
metadata: metadata as CheckpointMetadata,
parentConfig: row.parent_checkpoint_id ? { configurable: { thread_id: threadId, checkpoint_ns: '', checkpoint_id: row.parent_checkpoint_id } } : undefined
};
}
}
async put(
config: RunnableConfig,
checkpoint: Checkpoint,
metadata: CheckpointMetadata,
newVersions: ChannelVersions
): Promise<RunnableConfig> {
const threadId = config.configurable?.thread_id;
if (!threadId) throw new Error('missing thread_id');
const [ckptStr, metaStr] = await Promise.all([
this.serde.dumpsTyped(checkpoint),
this.serde.dumpsTyped(metadata)
]);
const client = getSupabaseClient();
const { error } = await client.from('langgraph_checkpoints').upsert({
thread_id: threadId,
checkpoint_id: checkpoint.id,
parent_checkpoint_id: config.configurable?.checkpoint_id || null,
checkpoint: ckptStr[1],
metadata: metaStr[1]
});
if (error) {
throw new Error(`Failed to put checkpoint: ${error.message}`);
}
return {
configurable: {
thread_id: threadId,
checkpoint_ns: '',
checkpoint_id: checkpoint.id
}
};
}
async putWrites(config: RunnableConfig, writes: PendingWrite[], taskId: string): Promise<void> {
const threadId = config.configurable?.thread_id;
const checkpointId = config.configurable?.checkpoint_id;
if (!threadId || !checkpointId) throw new Error('missing thread_id or checkpoint_id');
const client = getSupabaseClient();
const rows = await Promise.all(writes.map(async (w, idx) => {
const [type, value] = await this.serde.dumpsTyped(w[1]);
return {
thread_id: threadId,
checkpoint_id: checkpointId,
task_id: taskId,
idx,
channel: w[0],
type,
value
};
}));
if (rows.length > 0) {
const { error } = await client.from('langgraph_checkpoint_writes').upsert(rows);
if (error) {
throw new Error(`Failed to put writes: ${error.message}`);
}
}
}
async deleteThread(threadId: string): Promise<void> {
const client = getSupabaseClient();
await client.from('langgraph_checkpoint_writes').delete().eq('thread_id', threadId);
await client.from('langgraph_checkpoints').delete().eq('thread_id', threadId);
}
}
|