Merge branch 'release/4.0.0-beta1'

Author: bsorrentino

README.md

@@ -23,11 +23,10 @@
 - [x] Async support 
 - [x] Streaming support 
 - [x] Subgraph support 
-- [ ] Checkpoints (_save and replay feature_)
-- [ ] Threads (_checkpointing of multiple different runs_)
-- [ ] Update state (_interact with the state directly and update it_)
-- [ ] Breakpoints (_pause and resume feature_)
-- [ ] Graph migration
+- [x] Checkpoints (_save and replay feature_)
+    - [x] Threads (_checkpointing of multiple different runs_)
+    - [x] Update state (_interact with the state directly and update it_)
+    - [x] Breakpoints (_pause and resume feature_)
 - [ ] Graph visualization
   - [ ] [PlantUML]
   - [ ] [Mermaid]
@@ -183,12 +182,83 @@ In the [LangChainDemo](LangChainDemo) project, you can find the porting of [Agen
 
     let app = try workflow.compile()
 
-    let result = try await app.invoke(inputs: [ "input": input, "chat_history": [] ])
+    let result = try await app.invoke(inputs: .args([ "input": input, "chat_history": [] ]) )
 
     print( result )
 
 ```
 
+## User interruptions
+
+LangGraph support pause and resume of execution. You must provide a `CheckpointSaver` through `CompileConfig` and a `ThreadId`(aka Session) through `RunnableConfig` to enable it. Below an example
+
+```swift
+
+// Create a memory-based checkpoint saver
+let saver = MemoryCheckpointSaver()
+
+// Build the workflow with an initial state
+let workflow = try StateGraph { BinaryOpState($0) }
+
+// Add node "agent_1" that returns "add1": 37
+.addNode("agent_1") { state in
+    print( "agent_1", state )
+    return ["add1": 37]
+}
+// Add node "agent_2" that returns "add2": 10
+.addNode("agent_2") { state in
+    print( "agent_2", state )
+    return ["add2": 10]
+}
+// Add node "sum" that sums add1 and add2
+.addNode("sum") { state in
+    print( "sum", state )
+    guard let add1 = state.add1, let add2 = state.add2 else {
+        throw CompiledGraphError.executionError("agent state is not valid! expect 'add1', 'add2'")
+    }
+    return ["result": add1 + add2 ]
+}
+// Define the edges between nodes
+.addEdge(sourceId: "agent_1", targetId: "agent_2")
+.addEdge(sourceId: "agent_2", targetId: "sum")
+.addEdge( sourceId: START, targetId: "agent_1")
+.addEdge(sourceId: "sum", targetId: END )
+
+// Compile the workflow, instructing it to interrupt before executing "sum"
+let app = try workflow.compile( config: CompileConfig(checkpointSaver: saver, interruptionsBefore: ["sum"]) )
+
+// Start a new run in a different thread 
+let runnableConfig = RunnableConfig( threadId: "T1" )
+
+let initValue:( lastState:BinaryOpState?, nodes:[String]) = ( nil, [] )
+
+let result = try await app.stream( .args([:]), config: runnableConfig ).reduce( initValue, { partialResult, output in
+    print( output )
+    return ( output.state,  partialResult.1 + [output.node ] )
+})
+
+// This run is also interrupted before "sum"
+#expect( dictionaryOfAnyEqual(  ["add1": 37, "add2": 10 ], result.lastState!.data) )
+
+// Resume the third run with updated state: change add2 from 10 to 13
+let lastCheckpoint2 = try #require( saver.last( config: runnableConfig ) )
+var runnableConfig2 = runnableConfig.with { $0.checkpointId = lastCheckpoint2.id }
+runnableConfig2 = try await app.updateState(config: runnableConfig2, values: ["add2": 13] )
+
+// Resume and complete execution with updated value
+let initValue2:( lastState:BinaryOpState?, nodes:[String]) = ( nil, [] )
+let result2 = try await app.stream( .resume, config: runnableConfig2 ).reduce( initValue2, { partialResult, output in
+    print( output )
+    return ( output.state,  partialResult.1 + [output.node ] )
+})
+
+// Verify that "result" now reflects the updated input
+#expect( dictionaryOfAnyEqual(  ["add1": 37, "add2": 13, "result": 50 ], result2.lastState!.data) )
+
+```
+
+
+
 # References
 
 * [AI Agent on iOS with LangGraph for Swift](https://dev.to/bsorrentino/ai-agent-on-ios-with-langgraph-for-swift-1740)

Sources/LangGraph/Checkpoints.swift

@@ -0,0 +1,249 @@
+import Foundation
+
+
+public enum CheckpointError: Error, LocalizedError {
+    
+    case missingThreadIdentifier(String)
+    
+    public var errorDescription: String? {
+        switch self {
+        case .missingThreadIdentifier(let message):
+            return message
+        }
+    }
+
+}
+
+/// Represents a checkpoint of an agent state.
+///
+/// The checkpoint is an immutable object that holds an agent state
+/// and a string that represents the next state.
+/// It is designed to be serializable and restorable.
+public struct Checkpoint : Equatable  {
+    
+    public static func == (lhs: Checkpoint, rhs: Checkpoint) -> Bool {
+        lhs.id == rhs.id
+    }
+    
+    let id: UUID
+    var state: [String: Any]
+    var nodeId: String
+    var nextNodeId: String
+    
+    public init( state: [String: Any], nodeId: String, nextNodeId: String) {
+        self.id = UUID()
+        self.state = state
+        self.nodeId = nodeId
+        self.nextNodeId = nextNodeId
+    }
+
+    func updateState(values: PartialAgentState, channels: Channels) throws -> Self {
+        
+        var editable = self
+        editable.state = try LangGraph.updateState(currentState: self.state , partialState: values , channels: channels)
+        return editable
+    }
+}
+
+extension Checkpoint: Codable {
+    private enum CodingKeys: String, CodingKey {
+        case id
+        case state
+        case nodeId
+        case nextNodeId
+    }
+
+    public init(from decoder: any Decoder) throws {
+        let container = try decoder.container(keyedBy: CodingKeys.self)
+        
+        id = try container.decode(UUID.self, forKey: .id)
+        state = try container.decode([String: LangGraph.AnyDecodable].self, forKey: .state).mapValues { $0.value }
+        nodeId = try container.decode(String.self, forKey: .nodeId)
+        nextNodeId = try container.decode(String.self, forKey: .nextNodeId)
+    }
+    
+    public func encode(to encoder: any Encoder) throws {
+        
+        var container = encoder.container(keyedBy: CodingKeys.self)
+        try container.encode(id, forKey: .id)
+        try container.encode( toEncodableStateData(data: state), forKey: .state)
+        try container.encode(nodeId, forKey: .nodeId)
+        try container.encode(nextNodeId, forKey: .nextNodeId)
+        
+    }
+}
+
+public struct Tag {
+    let threadId: String
+    let checkpoints: AnyCollection<Checkpoint>
+}
+
+/// A protocol that defines an interface for saving and retrieving `Checkpoint` instances.
+///
+/// Conforming types manage checkpoint data associated with a specific `RunnableConfig`,
+/// allowing retrieval, update, listing, and release of checkpoints. This protocol enables
+/// persistence strategies to be customized for different runtime environments or threading models.
+public protocol CheckpointSaver {
+    /// Returns all checkpoints associated with the provided `RunnableConfig`.
+    ///
+    /// - Parameter config: The configuration that identifies the context or thread.
+    /// - Returns: A collection of `Checkpoint` instances.
+    func list(config: RunnableConfig) -> AnyCollection<Checkpoint>;
+
+    /// Retrieves a specific checkpoint based on the provided `RunnableConfig`.
+    ///
+    /// If `checkpointId` is set in the configuration, the corresponding checkpoint is returned.
+    /// Otherwise, returns the latest checkpoint for the context.
+    ///
+    /// - Parameter config: The configuration that may include a specific checkpoint identifier.
+    /// - Returns: The requested `Checkpoint` instance, or `nil` if not found.
+    func get(config: RunnableConfig) ->  Checkpoint?;
+
+    /// Persists a new checkpoint or updates an existing one based on the configuration.
+    ///
+    /// If the `checkpointId` is set in the configuration, the checkpoint with that ID is updated.
+    /// Otherwise, the new checkpoint is added to the thread's checkpoint stack.
+    ///
+    /// - Parameters:
+    ///   - config: The configuration identifying the context or thread.
+    ///   - checkpoint: The checkpoint to be saved or updated.
+    /// - Returns: A modified `RunnableConfig` reflecting the new checkpoint ID.
+    /// - Throws: An error if the checkpoint cannot be persisted.
+    func put(config: RunnableConfig,  checkpoint: Checkpoint) throws -> RunnableConfig;
+
+    /// Releases all checkpoints associated with the provided `RunnableConfig`.
+    ///
+    /// This method is responsible for cleanup of checkpoints tied to a specific context or thread.
+    ///
+    /// - Parameter config: The configuration identifying the context or thread.
+    /// - Returns: A `Tag` representing the final state of the thread's checkpoints.
+    /// - Throws: An error if the release operation fails.
+    func release(config: RunnableConfig) throws -> Tag;
+}
+
+extension CheckpointSaver {
+    @inline(__always) func THREAD_ID_DEFAULT() -> String { "$default" };
+    
+    @inline(__always) func last( config: RunnableConfig ) ->  Checkpoint? {
+        list( config: config ).first
+    }
+}
+
+struct Stack<T>  {
+    var elements: [T] = []
+
+    var isEmpty: Bool {
+        return elements.isEmpty
+    }
+
+    var count: Int {
+        return elements.count
+    }
+
+    mutating func push(_ value: T) {
+        elements.append(value)
+    }
+
+    mutating func pop() -> T? {
+        return elements.popLast()
+    }
+
+    func peek() -> T? {
+        return elements.last
+    }
+}
+
+
+extension Stack: Sequence {
+    public func makeIterator() -> IndexingIterator<Array<T>> {
+        return elements.reversed().makeIterator()
+    }
+}
+
+extension Stack {
+    
+    subscript(id: UUID) -> T? where T == Checkpoint {
+        get {
+            return elements.first { $0.id == id }
+        }
+        set(newValue) {
+            guard let newValue else {
+                fatalError( "Cannot set checkpoint with id \(id) to nil")
+            }
+            
+            guard let index = elements.firstIndex(where: { $0.id == id }) else {
+                fatalError( "Cannot find checkpoint with id \(id)" )
+            }
+            
+            elements[index] = newValue
+        }
+    }
+}
+
+public class MemoryCheckpointSaver: CheckpointSaver {
+    var checkpointsByThread: [String: Stack<Checkpoint>] = [:];
+    
+    private func checkpoints(config: RunnableConfig ) -> Stack<Checkpoint> {
+        let threadId = config.threadId ?? THREAD_ID_DEFAULT()
+            
+        guard let result = self.checkpointsByThread[threadId] else {
+            let result = Stack<Checkpoint>();
+            self.checkpointsByThread[threadId] = result;
+            return result;
+        }
+        return result
+
+    }
+    
+    private func updateCheckpoint( config: RunnableConfig, checkpoints: Stack<Checkpoint>  ) {
+        let threadId = config.threadId ?? THREAD_ID_DEFAULT()
+        
+        self.checkpointsByThread[threadId] = checkpoints
+    }
+    
+    public func get(config: RunnableConfig) -> Checkpoint? {
+        let checkpoints = checkpoints(config: config);
+        
+        guard let checkpointId = config.checkpointId else {
+            return checkpoints.peek()
+        }
+        
+        return checkpoints.first(where: { $0.id == checkpointId })
+    }
+    
+    public func put(config: RunnableConfig, checkpoint: Checkpoint) throws -> RunnableConfig {
+        var checkpoints = checkpoints(config: config);
+        
+        if let checkpointId = config.checkpointId, checkpointId == checkpoint.id {
+            
+            checkpoints[checkpointId] = checkpoint
+        }
+        
+        checkpoints.push(checkpoint)
+        
+        updateCheckpoint( config: config, checkpoints: checkpoints )
+        
+        return config.with {
+            $0.checkpointId = checkpoint.id
+        }
+        
+    }
+    
+    @discardableResult
+    public func release(config: RunnableConfig) throws -> Tag {
+        let threadId = config.threadId ?? THREAD_ID_DEFAULT()
+        
+        guard let removedCheckpoints = self.checkpointsByThread.removeValue(forKey: threadId) else {
+            throw CheckpointError.missingThreadIdentifier("No checkpoint found for thread \(threadId)")
+        }
+        
+        return Tag( threadId: threadId, checkpoints: AnyCollection(removedCheckpoints.elements) )
+    }
+    
+    public func list(config: RunnableConfig) -> AnyCollection<Checkpoint> {
+        let checkpoints = checkpoints(config: config);
+        
+        return AnyCollection(checkpoints.elements.reversed())
+    }
+}
+