Class: GraphAgent::Graph::StateGraph

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

Direct Known Subclasses

MessageGraph

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(schema = nil, input_schema: nil, output_schema: nil) ⇒ StateGraph

Returns a new instance of StateGraph.



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

def initialize(schema = nil, input_schema: nil, output_schema: nil)
  @schema = _normalize_schema(schema)
  @input_schema = input_schema ? _normalize_schema(input_schema) : nil
  @output_schema = output_schema ? _normalize_schema(output_schema) : nil
  @nodes = {}
  @edges = Set.new
  @branches = Hash.new { |h, k| h[k] = {} }
  @waiting_edges = Set.new
end

Instance Attribute Details

#branches ⇒ Object (readonly)

Returns the value of attribute branches.



8
9
10
 
# File 'lib/graph_agent/graph/state_graph.rb', line 8

def branches
  @branches
end

#edges ⇒ Object (readonly)

Returns the value of attribute edges.



8
9
10
 
# File 'lib/graph_agent/graph/state_graph.rb', line 8

def edges
  @edges
end

#nodes ⇒ Object (readonly)

Returns the value of attribute nodes.



8
9
10
 
# File 'lib/graph_agent/graph/state_graph.rb', line 8

def nodes
  @nodes
end

#schema ⇒ Object (readonly)

Returns the value of attribute schema.



8
9
10
 
# File 'lib/graph_agent/graph/state_graph.rb', line 8

def schema
  @schema
end

#waiting_edges ⇒ Object (readonly)

Returns the value of attribute waiting_edges.



8
9
10
 
# File 'lib/graph_agent/graph/state_graph.rb', line 8

def waiting_edges
  @waiting_edges
end

Instance Method Details

#add_conditional_edges(source, path, path_map = nil) ⇒ Object 



54
55
56
57
58
59
60
61
62
63
64
 
# File 'lib/graph_agent/graph/state_graph.rb', line 54

def add_conditional_edges(source, path, path_map = nil)
  source = source.to_s
  name = _branch_name(path)

  if @branches[source].key?(name)
    raise InvalidGraphError.new("Branch '#{name}' already exists for node '#{source}'")
  end

  @branches[source][name] = ConditionalEdge.new(source, path, path_map: path_map)
  self
end

#add_edge(start_key, end_key) ⇒ Object

Raises:



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

def add_edge(start_key, end_key)
  if start_key.is_a?(Array)
    targets = start_key.map(&:to_s)
    @waiting_edges.add([targets, end_key.to_s])
    return self
  end

  start_key = start_key.to_s
  end_key = end_key.to_s

  raise InvalidGraphError.new("END cannot be a start node") if start_key == END_NODE.to_s
  raise InvalidGraphError.new("START cannot be an end node") if end_key == START.to_s

  @edges.add(Edge.new(start_key, end_key))
  self
end

#add_node(name, action = nil, metadata: nil, retry_policy: nil, cache_policy: nil, &block) ⇒ Object

Raises:



20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
 
# File 'lib/graph_agent/graph/state_graph.rb', line 20

def add_node(name, action = nil, metadata: nil, retry_policy: nil, cache_policy: nil, &block)
  action = block if block && action.nil?
  name = _get_node_name(name, action)

  raise InvalidGraphError.new("Node action must be provided") if action.nil?
  raise InvalidGraphError.new("Node '#{name}' already exists") if @nodes.key?(name)
  raise InvalidGraphError.new("Node name '#{name}' is reserved") if [END_NODE.to_s, START.to_s].include?(name)

  @nodes[name] = Node.new(
    name, action,
    metadata: metadata,
    retry_policy: retry_policy,
    cache_policy: cache_policy
  )
  self
end

#add_sequence(nodes) ⇒ Object 



66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
 
# File 'lib/graph_agent/graph/state_graph.rb', line 66

def add_sequence(nodes)
  node_names = nodes.map do |node|
    if node.is_a?(Array)
      name, action = node
      add_node(name, action)
      name.to_s
    elsif node.respond_to?(:call)
      name = _get_node_name(nil, node)
      add_node(name, node)
      name
    else
      node.to_s
    end
  end

  node_names.each_cons(2) { |a, b| add_edge(a, b) }
  self
end

#compile(checkpointer: nil, interrupt_before: nil, interrupt_after: nil, debug: false) ⇒ Object 



97
98
99
100
101
102
103
104
105
106
107
 
# File 'lib/graph_agent/graph/state_graph.rb', line 97

def compile(checkpointer: nil, interrupt_before: nil, interrupt_after: nil, debug: false)
  validate!

  CompiledStateGraph.new(
    builder: self,
    checkpointer: checkpointer,
    interrupt_before: interrupt_before || [],
    interrupt_after: interrupt_after || [],
    debug: debug
  )
end

#print_mermaid(options = {}) ⇒ Object

Print Mermaid diagram to stdout



115
116
117
 
# File 'lib/graph_agent/graph/state_graph.rb', line 115

def print_mermaid(options = {})
  puts to_mermaid(options)
end

#set_conditional_entry_point(path, path_map = nil) ⇒ Object 



93
94
95
 
# File 'lib/graph_agent/graph/state_graph.rb', line 93

def set_conditional_entry_point(path, path_map = nil)
  add_conditional_edges(START, path, path_map)
end

#set_entry_point(node_name) ⇒ Object 



85
86
87
 
# File 'lib/graph_agent/graph/state_graph.rb', line 85

def set_entry_point(node_name)
  add_edge(START, node_name)
end

#set_finish_point(node_name) ⇒ Object 



89
90
91
 
# File 'lib/graph_agent/graph/state_graph.rb', line 89

def set_finish_point(node_name)
  add_edge(node_name, END_NODE)
end

#to_mermaid(options = {}) ⇒ Object

Generate a Mermaid diagram representation of the graph



110
111
112
 
# File 'lib/graph_agent/graph/state_graph.rb', line 110

def to_mermaid(options = {})
  MermaidVisualizer.render(self, options)
end