graph-agent 0.1.0

data/lib/graph_agent/graph/compiled_state_graph.rb

@@ -0,0 +1,362 @@
+# frozen_string_literal: true
+
+require "securerandom"
+
+module GraphAgent
+  module Graph
+    class CompiledStateGraph # rubocop:disable Metrics/ClassLength
+      DEFAULT_RECURSION_LIMIT = 25
+
+      attr_reader :builder, :checkpointer
+
+      def initialize(builder:, checkpointer: nil, interrupt_before: [], interrupt_after: [], debug: false)
+        @builder = builder
+        @checkpointer = checkpointer
+        @interrupt_before = _normalize_interrupt(interrupt_before)
+        @interrupt_after = _normalize_interrupt(interrupt_after)
+        @debug = debug
+      end
+
+      def invoke(input, config: {}, recursion_limit: DEFAULT_RECURSION_LIMIT)
+        last_state = nil
+
+        _run_pregel(input, config: config, recursion_limit: recursion_limit) do |_event|
+          last_state = _event[:state] if _event[:type] == :values
+        end
+
+        last_state
+      end
+
+      def stream(input, config: {}, recursion_limit: DEFAULT_RECURSION_LIMIT, stream_mode: :values, &block)
+        unless block
+          return enum_for(:stream, input, config: config, recursion_limit: recursion_limit, stream_mode: stream_mode)
+        end
+
+        _run_pregel(input, config: config, recursion_limit: recursion_limit) do |event|
+          _emit_stream_event(event, stream_mode, &block)
+        end
+      end
+
+      def get_state(config)
+        return nil unless @checkpointer
+
+        tuple = @checkpointer.get_tuple(config)
+        return nil unless tuple
+
+        StateSnapshot.new(
+          values: tuple.checkpoint[:channel_values] || {},
+          config: tuple.config,
+          metadata: tuple.metadata || {},
+          parent_config: tuple.parent_config,
+          next_nodes: tuple.checkpoint[:next_nodes] || []
+        )
+      end
+
+      def update_state(config, values, as_node: nil)
+        return nil unless @checkpointer
+
+        tuple = @checkpointer.get_tuple(config)
+        return nil unless tuple
+
+        current_state = tuple.checkpoint[:channel_values] || {}
+        new_state = _apply_updates(current_state, values)
+
+        checkpoint = _build_checkpoint(new_state, [])
+        metadata = { source: :update, step: (tuple.metadata || {})[:step].to_i + 1, writes: values }
+
+        @checkpointer.put(tuple.config, checkpoint, metadata, {})
+      end
+
+      def get_graph
+        { nodes: @builder.nodes.keys, edges: @builder.edges.to_a.map { |e| [e.source, e.target] } }
+      end
+
+      private
+
+      def _run_pregel(input, config:, recursion_limit:, &event_handler)
+        state = _initialize_state(input, config)
+        step = 0
+        current_nodes = _resolve_entry_nodes(state, config)
+
+        _save_checkpoint(config, state, step, :input, current_nodes)
+
+        while current_nodes.any? { |n| !_is_terminal?(n) }
+          _check_recursion_limit!(step, recursion_limit)
+          runnable_nodes = current_nodes.reject { |n| _is_terminal?(n) }
+          _check_interrupts_before!(runnable_nodes, config, state, step)
+
+          state, next_nodes = _execute_superstep(runnable_nodes, state, config, step, &event_handler)
+
+          _check_interrupts_after!(runnable_nodes, config, state, step)
+
+          current_nodes = next_nodes.uniq
+          step += 1
+          _save_checkpoint(config, state, step, :loop, current_nodes)
+        end
+
+        _save_checkpoint(config, state, step, :exit, [])
+        event_handler.call({ type: :values, step: step, state: state })
+      end
+
+      def _execute_superstep(runnable_nodes, state, config, step, &event_handler)
+        state_snapshot = _deep_dup(state)
+        step_updates = {}
+        next_nodes_from_commands = []
+        sends = []
+
+        runnable_nodes.each do |node_name|
+          result = _execute_node(node_name, state_snapshot, config)
+          _collect_results(result, node_name, step_updates, next_nodes_from_commands, sends)
+        end
+
+        step_updates.each_value { |updates| state = _apply_updates(state, updates) }
+
+        event_handler.call({ type: :updates, step: step, updates: step_updates.transform_keys(&:to_s) })
+        event_handler.call({ type: :values, step: step, state: state })
+
+        next_nodes = _resolve_next_nodes_for_step(runnable_nodes, state, config)
+        next_nodes.concat(next_nodes_from_commands)
+
+        state = _execute_sends(sends, state, config)
+
+        [state, next_nodes]
+      end
+
+      def _execute_node(node_name, state_snapshot, config)
+        node = @builder.nodes[node_name]
+        raise InvalidGraphError.new("Node '#{node_name}' not found") unless node
+
+        node.call(state_snapshot, config)
+      rescue GraphInterrupt, GraphRecursionError
+        raise
+      rescue => e
+        raise NodeExecutionError.new(node_name, e)
+      end
+
+      def _collect_results(result, node_name, step_updates, next_nodes_from_commands, sends)
+        node_updates, node_commands, node_sends = _process_result(result)
+
+        step_updates[node_name] = node_updates if node_updates && !node_updates.empty?
+
+        node_commands.each do |cmd|
+          _collect_command(cmd, node_name, step_updates, next_nodes_from_commands, sends)
+        end
+
+        sends.concat(node_sends)
+      end
+
+      def _collect_command(cmd, node_name, step_updates, next_nodes_from_commands, sends)
+        if cmd.update
+          cmd_updates = cmd.update.is_a?(Hash) ? cmd.update : {}
+          step_updates["#{node_name}:command"] = cmd_updates unless cmd_updates.empty?
+        end
+        cmd.goto.each do |target|
+          if target.is_a?(GraphAgent::Send)
+            sends << target
+          else
+            next_nodes_from_commands << target.to_s
+          end
+        end
+      end
+
+      def _execute_sends(sends, state, config)
+        sends.each do |send_obj|
+          target = send_obj.node.to_s
+          next unless @builder.nodes.key?(target)
+
+          send_state = send_obj.arg.is_a?(Hash) ? _apply_updates(_deep_dup(state), send_obj.arg) : state
+          result = @builder.nodes[target].call(send_state, config)
+          node_updates, = _process_result(result)
+          state = _apply_updates(state, node_updates) if node_updates
+        end
+        state
+      end
+
+      def _check_recursion_limit!(step, limit)
+        return unless step >= limit
+
+        raise GraphRecursionError.new("Recursion limit of #{limit} reached without hitting END node")
+      end
+
+      def _check_interrupts_before!(runnable_nodes, config, state, step)
+        runnable_nodes.each do |node_name|
+          next unless _should_interrupt_before?(node_name)
+
+          _save_checkpoint(config, state, step, :interrupt, [node_name])
+          raise GraphInterrupt.new([Interrupt.new("Interrupted before '#{node_name}'")])
+        end
+      end
+
+      def _check_interrupts_after!(runnable_nodes, config, state, step)
+        runnable_nodes.each do |node_name|
+          next unless _should_interrupt_after?(node_name)
+
+          _save_checkpoint(config, state, step, :interrupt, [])
+          raise GraphInterrupt.new([Interrupt.new("Interrupted after '#{node_name}'")])
+        end
+      end
+
+      def _emit_stream_event(event, stream_mode)
+        case stream_mode
+        when :values
+          yield event[:state].dup if event[:type] == :values
+        when :updates
+          yield event[:updates] if event[:type] == :updates && event[:updates]
+        when :debug
+          yield event
+        end
+      end
+
+      def _process_result(result)
+        updates = {}
+        commands = []
+        sends = []
+
+        case result
+        when Hash      then updates = result
+        when Command   then commands << result
+                            updates = result.update if result.update.is_a?(Hash)
+        when Send      then sends << result
+        when Array     then _process_array_result(result, updates, commands, sends)
+        end
+
+        [updates, commands, sends]
+      end
+
+      def _process_array_result(result, updates, commands, sends)
+        result.each do |item|
+          case item
+          when Hash    then updates.merge!(item)
+          when Command then commands << item
+                            updates.merge!(item.update) if item.update.is_a?(Hash)
+          when Send    then sends << item
+          end
+        end
+      end
+
+      def _initialize_state(input, config)
+        state = _restore_from_checkpoint(input, config)
+        return state if state
+
+        if @builder.schema.is_a?(State::Schema)
+          state = @builder.schema.initial_state
+          _apply_updates(state, input) if input.is_a?(Hash)
+          state
+        elsif input.is_a?(Hash)
+          input.transform_keys(&:to_sym).dup
+        else
+          {}
+        end
+      end
+
+      def _restore_from_checkpoint(input, config)
+        return nil unless @checkpointer && config.dig(:configurable, :thread_id)
+
+        tuple = @checkpointer.get_tuple(config)
+        return nil unless tuple
+
+        state = tuple.checkpoint[:channel_values]&.dup || {}
+        input.is_a?(Hash) ? _apply_updates(state, input) : state
+      end
+
+      def _apply_updates(state, updates)
+        return state unless updates.is_a?(Hash) && state.is_a?(Hash)
+
+        updates.each do |key, value|
+          key = key.to_sym
+          field = @builder.schema.is_a?(State::Schema) ? @builder.schema.fields[key] : nil
+
+          if field&.reducer
+            state[key] = field.reducer.call(state[key], value)
+          else
+            state[key] = value
+          end
+        end
+
+        state
+      end
+
+      def _resolve_entry_nodes(state, config)
+        nodes = @builder.edges.select { |e| e.source == START.to_s }.map(&:target)
+
+        if @builder.branches.key?(START.to_s)
+          @builder.branches[START.to_s].each_value do |branch|
+            nodes.concat(Array(branch.resolve(state, config)).map(&:to_s))
+          end
+        end
+
+        nodes.uniq
+      end
+
+      def _resolve_next_nodes_for_step(executed_nodes, state, config)
+        next_nodes = []
+
+        executed_nodes.each do |node_name|
+          @builder.edges.each { |edge| next_nodes << edge.target if edge.source == node_name }
+          _resolve_branches(node_name, state, config, next_nodes)
+        end
+
+        _resolve_waiting_edges(executed_nodes, next_nodes)
+        next_nodes.uniq
+      end
+
+      def _resolve_branches(node_name, state, config, next_nodes)
+        return unless @builder.branches.key?(node_name)
+
+        @builder.branches[node_name].each_value do |branch|
+          next_nodes.concat(Array(branch.resolve(state, config)).map(&:to_s))
+        end
+      end
+
+      def _resolve_waiting_edges(executed_nodes, next_nodes)
+        @builder.waiting_edges.each do |sources, target|
+          next_nodes << target if sources.all? { |s| executed_nodes.include?(s) }
+        end
+      end
+
+      def _is_terminal?(node_name)
+        node_name == END_NODE.to_s
+      end
+
+      def _should_interrupt_before?(node_name)
+        @interrupt_before.include?("*") || @interrupt_before.include?(node_name)
+      end
+
+      def _should_interrupt_after?(node_name)
+        @interrupt_after.include?("*") || @interrupt_after.include?(node_name)
+      end
+
+      def _normalize_interrupt(value)
+        case value
+        when Array then value.map(&:to_s)
+        when String, Symbol then [value.to_s]
+        when nil then []
+        else Array(value).map(&:to_s)
+        end
+      end
+
+      def _save_checkpoint(config, state, step, source, next_nodes)
+        return unless @checkpointer && config.dig(:configurable, :thread_id)
+
+        checkpoint = _build_checkpoint(_deep_dup(state), next_nodes)
+        metadata = { source: source, step: step, parents: {} }
+        @checkpointer.put(config, checkpoint, metadata, {})
+      end
+
+      def _build_checkpoint(channel_values, next_nodes)
+        { id: SecureRandom.uuid, channel_values: channel_values, channel_versions: {},
+          versions_seen: {}, next_nodes: next_nodes }
+      end
+
+      def _deep_dup(obj)
+        case obj
+        when Hash    then obj.transform_values { |v| _deep_dup(v) }
+        when Array   then obj.map { |v| _deep_dup(v) }
+        when String  then obj.dup
+        when Integer, Float, Symbol, TrueClass, FalseClass, NilClass then obj
+        else obj.respond_to?(:dup) ? obj.dup : obj
+        end
+      end
+    end
+  end
+end

data/lib/graph_agent/graph/conditional_edge.rb

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+module GraphAgent
+  module Graph
+    class ConditionalEdge
+      attr_reader :source, :path, :path_map
+
+      def initialize(source, path, path_map: nil)
+        @source = source.to_s
+        @path = path
+        @path_map = path_map
+      end
+
+      def resolve(state, config = {})
+        result = invoke_path(state, config)
+        _resolve_result(result)
+      end
+
+      private
+
+      def invoke_path(state, config)
+        unless @path.is_a?(Proc) || @path.respond_to?(:call)
+          raise GraphError.new("Conditional edge path is not callable")
+        end
+
+        case @path.arity
+        when 0 then @path.call
+        when 1, -1 then @path.call(state)
+        else @path.call(state, config)
+        end
+      end
+
+      def _resolve_result(result)
+        case result
+        when GraphAgent::Send then [result]
+        when Array            then _resolve_array(result)
+        else @path_map ? _map_result(result) : result
+        end
+      end
+
+      def _resolve_array(result)
+        return result if result.all? { |r| r.is_a?(GraphAgent::Send) }
+        return result.map { |r| _map_result(r) } if @path_map
+
+        result
+      end
+
+      def _map_result(result)
+        mapped = @path_map[result]
+        mapped = @path_map[:default] if mapped.nil? && @path_map.key?(:default)
+        raise InvalidGraphError.new("Unknown path result: #{result.inspect}") if mapped.nil?
+
+        mapped
+      end
+    end
+  end
+end

data/lib/graph_agent/graph/edge.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+module GraphAgent
+  module Graph
+    class Edge
+      attr_reader :source, :target
+
+      def initialize(source, target)
+        @source = source.to_s
+        @target = target.to_s
+      end
+
+      def ==(other)
+        other.is_a?(Edge) && source == other.source && target == other.target
+      end
+      alias eql? ==
+
+      def hash
+        [source, target].hash
+      end
+    end
+  end
+end

data/lib/graph_agent/graph/mermaid_visualizer.rb

@@ -0,0 +1,154 @@
+# frozen_string_literal: true
+
+require "set"
+
+module GraphAgent
+  module Graph
+    class MermaidVisualizer
+      START_LABEL = "START"
+      END_LABEL = "END"
+      CONDITION_PREFIX = "cond_"
+
+      # Generate Mermaid diagram for a StateGraph
+      def self.render(state_graph, options = {})
+        new(state_graph, options).render
+      end
+
+      def initialize(state_graph, options = {})
+        @graph = state_graph
+        @options = options
+      end
+
+      def render
+        lines = ["graph TD"]
+        lines << _style_definitions
+        lines << ""
+
+        # Render entry point (START)
+        lines << _node_definition(START.to_s, START_LABEL, :start)
+
+        # Render all graph nodes
+        @graph.nodes.each do |name, node|
+          lines << _node_definition(name, _node_label(name, node), :node)
+        end
+
+        # Render exit point (END)
+        lines << _node_definition(END_NODE.to_s, END_LABEL, :end)
+
+        lines << ""
+
+        # Render regular edges
+        @graph.edges.each do |edge|
+          lines << _edge_definition(edge)
+        end
+
+        # Render conditional edges (branches)
+        @graph.branches.each do |source, branches|
+          branches.each_with_index do |(name, branch), idx|
+            cond_id = "#{source}_#{CONDITION_PREFIX}#{idx}"
+            lines << _node_definition(cond_id, _condition_label(branch), :condition)
+
+            # Edge from source to condition node
+            lines << "  #{_safe_id(source)} --> #{_safe_id(cond_id)}"
+
+            # Edges from condition to targets
+            if branch.path_map
+              _render_path_map_edges(branch, cond_id, lines)
+            else
+              # Simple condition - render with note
+              lines << "  #{_safe_id(cond_id)} -.->|condition| #{_safe_id(source)}_next"
+              lines << "  #{_safe_id(source)}_next[\"?\"]"
+            end
+          end
+        end
+
+        # Render waiting edges (multi-source edges)
+        @graph.waiting_edges.each do |(sources, target)|
+          sources.each do |source|
+            lines << "  #{_safe_id(source)} --> #{_safe_id(target)}"
+          end
+        end
+
+        lines.join("\n")
+      end
+
+      private
+
+      def _style_definitions
+        <<~STYLES
+          classDef start fill:#e1f5e1,stroke:#4caf50,stroke-width:2px
+          classDef end fill:#ffebee,stroke:#f44336,stroke-width:2px
+          classDef node fill:#e3f2fd,stroke:#2196f3,stroke-width:2px,rx:5px
+          classDef condition fill:#fff9c4,stroke:#ffc107,stroke-width:2px
+
+          class #{_safe_id(START.to_s)} start
+          class #{_safe_id(END_NODE.to_s)} end
+        STYLES
+      end
+
+      def _node_definition(id, label, type)
+        safe_id = _safe_id(id)
+        case type
+        when :condition
+          "  #{safe_id}{#{label.inspect}}"
+        else
+          "  #{safe_id}[#{label.inspect}]"
+        end
+      end
+
+      def _edge_definition(edge)
+        "  #{_safe_id(edge.source)} --> #{_safe_id(edge.target)}"
+      end
+
+      def _render_path_map_edges(branch, cond_id, lines)
+        branch.path_map.each do |condition, target|
+          next if condition == :default
+
+          lines << "  #{_safe_id(cond_id)} -.->|#{condition}| #{_safe_id(target)}"
+        end
+        # Handle default path
+        if branch.path_map.key?(:default)
+          lines << "  #{_safe_id(cond_id)} -.->|default| #{_safe_id(branch.path_map[:default])}"
+        end
+      end
+
+      def _safe_id(id)
+        # Escape special characters for Mermaid
+        id.to_s.gsub(/[^a-zA-Z0-9_]/, "_")
+      end
+
+      def _node_label(name, node)
+        # Use node name or try to extract a readable name from action
+        if @options[:show_node_names]
+          name
+        else
+          _extract_readable_name(name, node)
+        end
+      end
+
+      def _extract_readable_name(name, node)
+        # Try to get a readable name from the action
+        action = node.action
+        if action.respond_to?(:name) && !action.name.nil? && !action.name.to_s.empty?
+          action.name.to_s.split("::").last
+        else
+          name
+        end
+      end
+
+      def _condition_label(branch)
+        # Try to extract a meaningful label from the condition
+        path = branch.path
+        if path.respond_to?(:name) && !path.name.nil? && !path.name.empty?
+          path.name
+        elsif branch.path_map
+          # Show the possible conditions
+          conditions = branch.path_map.keys - [:default]
+          conditions.join(" / ")
+        else
+          "?"
+        end
+      end
+    end
+  end
+end

data/lib/graph_agent/graph/message_graph.rb

@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+
+module GraphAgent
+  module Graph
+    class MessagesState < State::Schema
+      def initialize
+        super
+        field :messages, type: Array, reducer: GraphAgent::Reducers.method(:add_messages), default: []
+      end
+    end
+
+    class MessageGraph < StateGraph
+      def initialize
+        super(MessagesState.new)
+      end
+    end
+  end
+end

data/lib/graph_agent/graph/node.rb

@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+module GraphAgent
+  module Graph
+    class Node
+      attr_reader :name, :action, :metadata, :retry_policy, :cache_policy
+
+      def initialize(name, action, metadata: nil, retry_policy: nil, cache_policy: nil)
+        @name = name.to_s
+        @action = action
+        @metadata = metadata || {}
+        @retry_policy = retry_policy
+        @cache_policy = cache_policy
+      end
+
+      def call(state, config = {})
+        result = execute_with_retry(state, config)
+        normalize_result(result)
+      end
+
+      private
+
+      def execute_with_retry(state, config)
+        if @retry_policy
+          attempt = 0
+          begin
+            attempt += 1
+            invoke_action(state, config)
+          rescue => e
+            if attempt < @retry_policy.max_attempts && @retry_policy.should_retry?(e)
+              sleep(@retry_policy.interval_for(attempt))
+              retry
+            end
+            raise
+          end
+        else
+          invoke_action(state, config)
+        end
+      end
+
+      def invoke_action(state, config)
+        if @action.is_a?(Proc) || @action.respond_to?(:call)
+          case @action.arity
+          when 0
+            @action.call
+          when 1, -1
+            @action.call(state)
+          else
+            @action.call(state, config)
+          end
+        else
+          raise GraphError.new("Node '#{name}' action is not callable")
+        end
+      end
+
+      def normalize_result(result)
+        result
+      end
+    end
+  end
+end