feat(BaseCheckpointSaver): add support for ThreadId

work on #20

Author: bsorrentino

core-jdk8/src/main/java/org/bsc/langgraph4j/checkpoint/BaseCheckpointSaver.java

@@ -1,13 +1,13 @@
 package org.bsc.langgraph4j.checkpoint;
 
-import java.io.Externalizable;
+import org.bsc.langgraph4j.RunnableConfig;
+
 import java.util.Collection;
 import java.util.Optional;
 
 public interface BaseCheckpointSaver {
 
-
-    Collection<Checkpoint> list();
-    Optional<Checkpoint> getLast();
-    void put( Checkpoint checkpoint ) throws Exception;
+    Collection<Checkpoint> list( RunnableConfig config );
+    Optional<Checkpoint> get( RunnableConfig config );
+    RunnableConfig put( RunnableConfig config, Checkpoint checkpoint ) throws Exception;
 }

core-jdk8/src/main/java/org/bsc/langgraph4j/checkpoint/MemorySaver.java

@@ -1,39 +1,72 @@
 package org.bsc.langgraph4j.checkpoint;
 
+import org.bsc.langgraph4j.RunnableConfig;
 import org.bsc.langgraph4j.serializer.CheckpointSerializer;
 
 import java.util.*;
+import java.util.stream.IntStream;
 
+import static java.lang.String.format;
 import static java.util.Collections.unmodifiableCollection;
-import static java.util.Collections.unmodifiableSet;
+import static java.util.Optional.ofNullable;
 
 public class MemorySaver implements BaseCheckpointSaver {
+    private final Map<String, LinkedList<Checkpoint>> _checkpointsByThread = new HashMap<>();
+    private final LinkedList<Checkpoint> _defaultCheckpoints = new LinkedList<>();
 
-    private final Stack<Checkpoint> checkpoints = new Stack<>();
+    public MemorySaver() {
+    }
 
+    private LinkedList<Checkpoint> getCheckpoints( RunnableConfig config ) {
+        return config.threadId()
+                    .map( threadId -> _checkpointsByThread.computeIfAbsent(threadId, k -> new LinkedList<>()) )
+                    .orElse( _defaultCheckpoints );
+    }
 
     @Override
-    public Collection<Checkpoint> list() {
+    public Collection<Checkpoint> list( RunnableConfig config ) {
+        final LinkedList<Checkpoint> checkpoints = getCheckpoints(config);
         return unmodifiableCollection(checkpoints); // immutable checkpoints;
     }
 
-    @Override
-    public Optional<Checkpoint> getLast() {
-        if( checkpoints.isEmpty() ) {
-            return Optional.empty();
-        }
-        return Optional.ofNullable( checkpoints.peek() );
+    private Optional<Checkpoint> getLast( RunnableConfig config ) {
+        final LinkedList<Checkpoint> checkpoints = getCheckpoints(config);
+        return (checkpoints.isEmpty() ) ? Optional.empty() : ofNullable(checkpoints.peek());
     }
 
-    public Optional<Checkpoint> get(String id) {
-        return checkpoints.stream()
-                    .filter( checkpoint -> checkpoint.getId().equals(id) )
-                    .findFirst();
+    @Override
+    public Optional<Checkpoint> get(RunnableConfig config) {
+        final LinkedList<Checkpoint> checkpoints = getCheckpoints(config);
+        if( config.checkPointId().isPresent() ) {
+            return config.checkPointId()
+                    .flatMap( id -> checkpoints.stream()
+                            .filter( checkpoint -> checkpoint.getId().equals(id) )
+                            .findFirst());
+        }
+        return getLast(config);
     }
 
     @Override
-    public void put(Checkpoint checkpoint) throws Exception {
-        checkpoints.add( CheckpointSerializer.INSTANCE.cloneObject(checkpoint) );
+    public RunnableConfig put(RunnableConfig config, Checkpoint checkpoint) throws Exception {
+        final LinkedList<Checkpoint> checkpoints = getCheckpoints(config);
+
+        final Checkpoint clonedCheckpoint = CheckpointSerializer.INSTANCE.cloneObject(checkpoint);
+
+        if( config.checkPointId().isPresent() ) { // Replace Checkpoint
+            String checkPointId = config.checkPointId().get();
+            int index = IntStream.range(0, checkpoints.size())
+                    .filter(i -> checkpoints.get(i).getId().equals(checkPointId))
+                    .findFirst()
+                    .orElseThrow( () -> (new NoSuchElementException( format("Checkpoint with id %s not found!", checkPointId))) );
+            checkpoints.set( index,  clonedCheckpoint);
+            return config;
+        }
+
+        checkpoints.push( clonedCheckpoint ); // Add Checkpoint
+
+        return RunnableConfig.builder(config)
+                .checkPointId( clonedCheckpoint.getId() )
+                .build();
     }
 
 }