@farukada/aws-langgraph-dynamodb-ts 1.0.0-rc.1 → 1.0.0

package/dist/esm/checkpointer/internal/read.js

@@ -0,0 +1,125 @@
+import { maxChannelVersion, TASKS, } from '@langchain/langgraph-checkpoint';
+import { nowSeconds } from '../../shared/clock.js';
+import { LIST_SCAN_WARN_THRESHOLD, paginateQuery } from '../../shared/dynamodb/paginate.js';
+import { withDynamoDBRetry, retryFor } from '../../shared/dynamodb/retry.js';
+import { isExpiredRow, withoutExpired, assertReadableRow, } from '../../shared/dynamodb/table-schema.js';
+import { truncateForLog } from '../../shared/logging/truncate.js';
+import { beginsWithQuery, metaRowKey, metaSortKeyPrefix, parseHeadRow, parseWriteRow, partitionKey, payloadRowKey, readCheckpoint, readMetadata, toPendingWrites, writeSortKeyPrefix, } from './rows.js';
+const LATEST_META_PAGE_SIZE = 50;
+export async function fetchTargetMeta(context, address, signal, reader = 'getTuple') {
+    const { threadId, checkpointNs, checkpointId } = address;
+    const now = nowSeconds();
+    if (checkpointId !== undefined) {
+        const result = await withDynamoDBRetry((request) => context.client.get({
+            TableName: context.tableName,
+            Key: metaRowKey({ threadId, checkpointNs, checkpointId }),
+            ConsistentRead: true,
+        }, request), retryFor(context, signal));
+        const meta = parseHeadRow(context, result.Item, reader);
+        return meta && !isExpiredRow(meta, now) ? meta : undefined;
+    }
+    const params = beginsWithQuery(context.tableName, partitionKey(threadId), metaSortKeyPrefix(checkpointNs), {
+        limit: context.ttl === undefined ? 1 : LATEST_META_PAGE_SIZE,
+        consistent: true,
+    });
+    const rows = paginateQuery({
+        retry: retryFor(context, signal),
+        signal,
+        client: context.client,
+        params: withoutExpired(params, now),
+        maxItems: Number.POSITIVE_INFINITY,
+        maxIterations: Number.POSITIVE_INFINITY,
+    });
+    for await (const raw of rows) {
+        const meta = parseHeadRow(context, raw, reader);
+        if (meta && !isExpiredRow(meta, now))
+            return meta;
+    }
+    return undefined;
+}
+export async function fetchPayload(context, at, read = {}) {
+    const result = await withDynamoDBRetry((request) => context.client.get({
+        TableName: context.tableName,
+        Key: payloadRowKey(at),
+        ConsistentRead: read.consistent ?? true,
+    }, request), retryFor(context, read.signal));
+    const item = result.Item;
+    if (item !== undefined)
+        assertReadableRow(item, 'checkpoint payload');
+    return item;
+}
+export async function fetchPendingWrites(context, at, read = {}) {
+    const params = beginsWithQuery(context.tableName, partitionKey(at.threadId), writeSortKeyPrefix(at.checkpointNs, at.checkpointId), { ascending: true, consistent: read.consistent ?? true });
+    const items = [];
+    for await (const item of paginateQuery({
+        retry: retryFor(context, read.signal),
+        signal: read.signal,
+        client: context.client,
+        params,
+        maxItems: Number.POSITIVE_INFINITY,
+        maxIterations: Number.POSITIVE_INFINITY,
+    })) {
+        const row = parseWriteRow(item);
+        if (row === undefined) {
+            context.logger.warn('getTuple: skipped a row that is not a pending-write item', {
+                sortKey: truncateForLog(item.SK),
+            });
+            continue;
+        }
+        items.push(row);
+    }
+    if (items.length >= LIST_SCAN_WARN_THRESHOLD) {
+        context.logger.warn('getTuple: a checkpoint carries very many pending-write rows; the read is complete but slow', {
+            threadId: at.threadId,
+            checkpointId: at.checkpointId,
+            rows: items.length,
+        });
+    }
+    return toPendingWrites(context, items, at.threadId, read.signal);
+}
+function configFor(threadId, checkpointNs, checkpointId) {
+    return {
+        configurable: { thread_id: threadId, checkpoint_ns: checkpointNs, checkpoint_id: checkpointId },
+    };
+}
+export async function assembleTuple(context, thread, meta, options) {
+    const read = { signal: options.signal, consistent: options.consistent };
+    const at = { ...thread, checkpointId: meta.checkpointId };
+    const payload = await fetchPayload(context, at, read);
+    if (!payload)
+        return undefined;
+    const [checkpoint, metadata, pendingWrites] = await Promise.all([
+        readCheckpoint(context, payload, thread.threadId, options.signal).then((stored) => migratePendingSends(context, stored, { ...thread, parentCheckpointId: meta.parentCheckpointId }, read)),
+        options.metadata ?? readMetadata(context, meta, thread.threadId, options.signal),
+        fetchPendingWrites(context, at, read),
+    ]);
+    const tuple = {
+        config: configFor(thread.threadId, thread.checkpointNs, meta.checkpointId),
+        checkpoint,
+        metadata,
+        pendingWrites,
+    };
+    if (meta.parentCheckpointId !== undefined) {
+        tuple.parentConfig = configFor(thread.threadId, thread.checkpointNs, meta.parentCheckpointId);
+    }
+    return tuple;
+}
+export async function migratePendingSends(context, checkpoint, parent, read) {
+    if (checkpoint.v >= 4 || parent.parentCheckpointId === undefined)
+        return checkpoint;
+    const writes = await fetchPendingWrites(context, {
+        threadId: parent.threadId,
+        checkpointNs: parent.checkpointNs,
+        checkpointId: parent.parentCheckpointId,
+    }, read);
+    const sends = writes.filter(([, channel]) => channel === TASKS).map(([, , value]) => value);
+    const versions = Object.values(checkpoint.channel_versions);
+    return {
+        ...checkpoint,
+        channel_values: { ...checkpoint.channel_values, [TASKS]: sends },
+        channel_versions: {
+            ...checkpoint.channel_versions,
+            [TASKS]: versions.length > 0 ? maxChannelVersion(...versions) : 1,
+        },
+    };
+}