graph-agent 0.1.0

data/lib/graph_agent/channels/binary_operator_aggregate.rb

@@ -0,0 +1,56 @@
+# frozen_string_literal: true
+
+module GraphAgent
+  module Channels
+    class BinaryOperatorAggregate < BaseChannel
+      attr_reader :operator
+
+      def initialize(operator:, key: "", default: MISSING)
+        super(key: key)
+        @operator = operator
+        @value = default
+      end
+
+      def get
+        raise EmptyChannelError.new("Channel '#{key}' is empty") if @value.equal?(MISSING)
+
+        @value
+      end
+
+      def update(values)
+        return false if values.empty?
+
+        if @value.equal?(MISSING)
+          @value = values.first
+          values = values[1..]
+        end
+
+        values.each do |value|
+          @value = @operator.call(@value, value)
+        end
+
+        true
+      end
+
+      def available?
+        !@value.equal?(MISSING)
+      end
+
+      def checkpoint
+        @value
+      end
+
+      def from_checkpoint(checkpoint)
+        ch = self.class.new(operator: @operator, key: key)
+        ch.instance_variable_set(:@value, checkpoint.equal?(MISSING) ? MISSING : checkpoint)
+        ch
+      end
+
+      def copy
+        ch = self.class.new(operator: @operator, key: key)
+        ch.instance_variable_set(:@value, @value)
+        ch
+      end
+    end
+  end
+end

data/lib/graph_agent/channels/ephemeral_value.rb

@@ -0,0 +1,59 @@
+# frozen_string_literal: true
+
+module GraphAgent
+  module Channels
+    class EphemeralValue < BaseChannel
+      def initialize(key: "", guard: true)
+        super(key: key)
+        @guard = guard
+        @value = MISSING
+      end
+
+      def get
+        raise EmptyChannelError.new("Channel '#{key}' is empty") if @value.equal?(MISSING)
+
+        @value
+      end
+
+      def update(values)
+        if values.empty?
+          if @value.equal?(MISSING)
+            return false
+          else
+            @value = MISSING
+            return true
+          end
+        end
+
+        if values.length != 1 && @guard
+          raise InvalidUpdateError.new(
+            "At key '#{key}': EphemeralValue(guard=true) can receive only one value per step."
+          )
+        end
+
+        @value = values.last
+        true
+      end
+
+      def available?
+        !@value.equal?(MISSING)
+      end
+
+      def checkpoint
+        @value
+      end
+
+      def from_checkpoint(checkpoint)
+        ch = self.class.new(key: key, guard: @guard)
+        ch.instance_variable_set(:@value, checkpoint.equal?(MISSING) ? MISSING : checkpoint)
+        ch
+      end
+
+      def copy
+        ch = self.class.new(key: key, guard: @guard)
+        ch.instance_variable_set(:@value, @value)
+        ch
+      end
+    end
+  end
+end

data/lib/graph_agent/channels/last_value.rb

@@ -0,0 +1,49 @@
+# frozen_string_literal: true
+
+module GraphAgent
+  module Channels
+    class LastValue < BaseChannel
+      def initialize(key: "", default: MISSING)
+        super(key: key)
+        @value = default
+      end
+
+      def get
+        raise EmptyChannelError.new("Channel '#{key}' is empty") if @value.equal?(MISSING)
+
+        @value
+      end
+
+      def update(values)
+        return false if values.empty?
+
+        if values.length != 1
+          raise InvalidUpdateError.new("At key '#{key}': Can receive only one value per step.")
+        end
+
+        @value = values.last
+        true
+      end
+
+      def available?
+        !@value.equal?(MISSING)
+      end
+
+      def checkpoint
+        @value
+      end
+
+      def from_checkpoint(checkpoint)
+        ch = self.class.new(key: key)
+        ch.instance_variable_set(:@value, checkpoint.equal?(MISSING) ? MISSING : checkpoint)
+        ch
+      end
+
+      def copy
+        ch = self.class.new(key: key)
+        ch.instance_variable_set(:@value, @value)
+        ch
+      end
+    end
+  end
+end

data/lib/graph_agent/channels/topic.rb

@@ -0,0 +1,58 @@
+# frozen_string_literal: true
+
+module GraphAgent
+  module Channels
+    class Topic < BaseChannel
+      attr_reader :accumulate
+
+      def initialize(key: "", accumulate: false)
+        super(key: key)
+        @accumulate = accumulate
+        @values = []
+      end
+
+      def get
+        raise EmptyChannelError.new("Channel '#{key}' is empty") if @values.empty?
+
+        @values.dup
+      end
+
+      def update(values)
+        updated = false
+
+        unless @accumulate
+          updated = !@values.empty?
+          @values = []
+        end
+
+        flat = values.flat_map { |v| v.is_a?(Array) ? v : [v] }
+        unless flat.empty?
+          updated = true
+          @values.concat(flat)
+        end
+
+        updated
+      end
+
+      def available?
+        !@values.empty?
+      end
+
+      def checkpoint
+        @values.dup
+      end
+
+      def from_checkpoint(checkpoint)
+        ch = self.class.new(key: key, accumulate: @accumulate)
+        ch.instance_variable_set(:@values, checkpoint.equal?(MISSING) ? [] : Array(checkpoint))
+        ch
+      end
+
+      def copy
+        ch = self.class.new(key: key, accumulate: @accumulate)
+        ch.instance_variable_set(:@values, @values.dup)
+        ch
+      end
+    end
+  end
+end

data/lib/graph_agent/checkpoint/base_saver.rb

@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+module GraphAgent
+  module Checkpoint
+    CheckpointTuple = Data.define(:config, :checkpoint, :metadata, :parent_config, :pending_writes) do
+      def initialize(config:, checkpoint:, metadata: nil, parent_config: nil, pending_writes: [])
+        super
+      end
+    end
+
+    class BaseSaver
+      def get(config)
+        tuple = get_tuple(config)
+        tuple&.checkpoint
+      end
+
+      def get_tuple(config)
+        raise NotImplementedError.new("#{self.class}#get_tuple must be implemented")
+      end
+
+      def list(config, filter: nil, before: nil, limit: nil)
+        raise NotImplementedError.new("#{self.class}#list must be implemented")
+      end
+
+      def put(config, checkpoint, metadata, new_versions)
+        raise NotImplementedError.new("#{self.class}#put must be implemented")
+      end
+
+      def put_writes(config, writes, task_id)
+        raise NotImplementedError.new("#{self.class}#put_writes must be implemented")
+      end
+
+      def delete_thread(thread_id)
+        raise NotImplementedError.new("#{self.class}#delete_thread must be implemented")
+      end
+    end
+  end
+end

data/lib/graph_agent/checkpoint/in_memory_saver.rb

@@ -0,0 +1,145 @@
+# frozen_string_literal: true
+
+require "securerandom"
+
+module GraphAgent
+  module Checkpoint
+    class InMemorySaver < BaseSaver
+      def initialize
+        super
+        @storage = Hash.new { |h, k| h[k] = Hash.new { |h2, k2| h2[k2] = {} } }
+        @writes = Hash.new { |h, k| h[k] = {} }
+        @order = Hash.new { |h, k| h[k] = Hash.new { |h2, k2| h2[k2] = [] } }
+      end
+
+      def get_tuple(config)
+        thread_id = config.dig(:configurable, :thread_id)
+        checkpoint_ns = config.dig(:configurable, :checkpoint_ns) || ""
+        checkpoint_id = config.dig(:configurable, :checkpoint_id)
+
+        if checkpoint_id
+          saved = @storage[thread_id][checkpoint_ns][checkpoint_id]
+          return nil unless saved
+
+          build_tuple(thread_id, checkpoint_ns, checkpoint_id, saved, config)
+        else
+          checkpoints = @storage[thread_id][checkpoint_ns]
+          return nil if checkpoints.empty?
+
+          checkpoint_id = @order[thread_id][checkpoint_ns].last
+          return nil unless checkpoint_id
+
+          saved = checkpoints[checkpoint_id]
+          result_config = {
+            configurable: {
+              thread_id: thread_id,
+              checkpoint_ns: checkpoint_ns,
+              checkpoint_id: checkpoint_id
+            }
+          }
+          build_tuple(thread_id, checkpoint_ns, checkpoint_id, saved, result_config)
+        end
+      end
+
+      def list(config, filter: nil, before: nil, limit: nil)
+        thread_ids = config ? [config.dig(:configurable, :thread_id)] : @storage.keys
+        results = []
+
+        thread_ids.each do |thread_id|
+          @storage[thread_id].each do |checkpoint_ns, checkpoints|
+            _list_ns(thread_id, checkpoint_ns, checkpoints, results, filter: filter, before: before, limit: limit)
+          end
+        end
+
+        results
+      end
+
+      def put(config, checkpoint, metadata, new_versions)
+        thread_id = config.dig(:configurable, :thread_id)
+        checkpoint_ns = config.dig(:configurable, :checkpoint_ns) || ""
+        checkpoint_id = checkpoint[:id] || SecureRandom.uuid
+
+        @storage[thread_id][checkpoint_ns][checkpoint_id] = {
+          checkpoint: checkpoint,
+          metadata: metadata,
+          parent_checkpoint_id: config.dig(:configurable, :checkpoint_id)
+        }
+        @order[thread_id][checkpoint_ns] << checkpoint_id
+
+        {
+          configurable: {
+            thread_id: thread_id,
+            checkpoint_ns: checkpoint_ns,
+            checkpoint_id: checkpoint_id
+          }
+        }
+      end
+
+      def put_writes(config, writes, task_id)
+        thread_id = config.dig(:configurable, :thread_id)
+        checkpoint_ns = config.dig(:configurable, :checkpoint_ns) || ""
+        checkpoint_id = config.dig(:configurable, :checkpoint_id)
+        key = [thread_id, checkpoint_ns, checkpoint_id]
+
+        writes.each_with_index do |(channel, value), idx|
+          @writes[key][[task_id, idx]] = { task_id: task_id, channel: channel, value: value }
+        end
+      end
+
+      def delete_thread(thread_id)
+        @storage.delete(thread_id)
+        @order.delete(thread_id)
+        @writes.delete_if { |key, _| key.first == thread_id }
+      end
+
+      private
+
+      def _list_ns(thread_id, checkpoint_ns, checkpoints, results, filter:, before:, limit:)
+        ordered_ids = @order[thread_id][checkpoint_ns].reverse
+        entries = ordered_ids.filter_map { |id| [id, checkpoints[id]] if checkpoints[id] }
+
+        entries.each do |checkpoint_id, saved|
+          next if before && _skip_before?(checkpoint_id, before)
+          next if filter && !_matches_filter?(saved[:metadata], filter)
+          break if limit && results.length >= limit
+
+          tuple_config = { configurable: { thread_id: thread_id, checkpoint_ns: checkpoint_ns,
+                                           checkpoint_id: checkpoint_id } }
+          results << build_tuple(thread_id, checkpoint_ns, checkpoint_id, saved, tuple_config)
+        end
+      end
+
+      def _skip_before?(checkpoint_id, before)
+        before_id = before.dig(:configurable, :checkpoint_id)
+        before_id && checkpoint_id >= before_id
+      end
+
+      def _matches_filter?(metadata, filter)
+        filter.all? { |k, v| metadata[k] == v }
+      end
+
+      def build_tuple(thread_id, checkpoint_ns, checkpoint_id, saved, config)
+        pending = @writes[[thread_id, checkpoint_ns, checkpoint_id]].values
+        parent_id = saved[:parent_checkpoint_id]
+
+        parent_config = if parent_id
+                          {
+                            configurable: {
+                              thread_id: thread_id,
+                              checkpoint_ns: checkpoint_ns,
+                              checkpoint_id: parent_id
+                            }
+                          }
+                        end
+
+        CheckpointTuple.new(
+          config: config,
+          checkpoint: saved[:checkpoint],
+          metadata: saved[:metadata],
+          parent_config: parent_config,
+          pending_writes: pending.map { |w| [w[:task_id], w[:channel], w[:value]] }
+        )
+      end
+    end
+  end
+end

data/lib/graph_agent/constants.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+module GraphAgent
+  START = :__start__
+  END_NODE = :__end__
+
+  TAG_NOSTREAM = :nostream
+  TAG_HIDDEN = :"langsmith:hidden"
+end

data/lib/graph_agent/errors.rb

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+module GraphAgent
+  class GraphError < StandardError; end
+
+  class GraphRecursionError < GraphError
+    def initialize(msg = nil)
+      super(msg || "Graph has exhausted the maximum number of steps. " \
+                   "To increase the limit, set recursion_limit in config.")
+    end
+  end
+
+  class InvalidUpdateError < GraphError; end
+
+  class EmptyChannelError < GraphError; end
+
+  class InvalidGraphError < GraphError; end
+
+  class NodeExecutionError < GraphError
+    attr_reader :node_name, :original_error
+
+    def initialize(node_name, original_error)
+      @node_name = node_name
+      @original_error = original_error
+      super("Error in node '#{node_name}': #{original_error.message}")
+    end
+  end
+
+  class GraphInterrupt < GraphError
+    attr_reader :interrupts
+
+    def initialize(interrupts = [])
+      @interrupts = interrupts
+      super("Graph interrupted with #{interrupts.length} interrupt(s)")
+    end
+  end
+
+  class EmptyInputError < GraphError; end
+
+  class TaskNotFound < GraphError; end
+end