Class: GraphAgent::Graph::CompiledStateGraph

Inherits:
Object
  • Object
show all
Defined in:
lib/graph_agent/graph/compiled_state_graph.rb

Overview

rubocop:disable Metrics/ClassLength

Constant Summary collapse

DEFAULT_RECURSION_LIMIT = 
25

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(builder:, checkpointer: nil, interrupt_before: [], interrupt_after: [], debug: false) ⇒ CompiledStateGraph

Returns a new instance of CompiledStateGraph.



12
13
14
15
16
17
18
 
# File 'lib/graph_agent/graph/compiled_state_graph.rb', line 12

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

Instance Attribute Details

#builder ⇒ Object (readonly)

Returns the value of attribute builder.



10
11
12
 
# File 'lib/graph_agent/graph/compiled_state_graph.rb', line 10

def builder
  @builder
end

#checkpointer ⇒ Object (readonly)

Returns the value of attribute checkpointer.



10
11
12
 
# File 'lib/graph_agent/graph/compiled_state_graph.rb', line 10

def checkpointer
  @checkpointer
end

Instance Method Details

#get_graph ⇒ Object 



70
71
72
 
# File 'lib/graph_agent/graph/compiled_state_graph.rb', line 70

def get_graph
  { nodes: @builder.nodes.keys, edges: @builder.edges.to_a.map { |e| [e.source, e.target] } }
end

#get_state(config) ⇒ Object 



40
41
42
43
44
45
46
47
48
49
50
51
52
53
 
# File 'lib/graph_agent/graph/compiled_state_graph.rb', line 40

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

#invoke(input, config: {}, recursion_limit: DEFAULT_RECURSION_LIMIT) ⇒ Object 



20
21
22
23
24
25
26
27
28
 
# File 'lib/graph_agent/graph/compiled_state_graph.rb', line 20

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

#stream(input, config: {}, recursion_limit: DEFAULT_RECURSION_LIMIT, stream_mode: :values, &block) ⇒ Object 



30
31
32
33
34
35
36
37
38
 
# File 'lib/graph_agent/graph/compiled_state_graph.rb', line 30

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

#update_state(config, values, as_node: nil) ⇒ Object 



55
56
57
58
59
60
61
62
63
64
65
66
67
68
 
# File 'lib/graph_agent/graph/compiled_state_graph.rb', line 55

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