def get_state_graph(self):
cache_key = "state_graph"
if cache_key not in self.cache:
graph = Graph()
for agent_name, agent_def in self.agents.iteritems():
g = graph_builder.generate_state_graph(agent_def)
# store agent_name in each node before merging graphs
for node in g.get_nodes():
g.set_node_attribute(node, "agent_name", agent_name)
graph.extend(g)
self.cache[cache_key] = graph
return self.cache[cache_key]
def state_to_function_graph(sg):
assert type(sg) == Graph
fg = Graph()
for node in sg.traverse():
node_type = sg.get_node_attribute(node, "type")
parents = sg.get_parents(node)
children = sg.get_children(node)
_validate_sg_node(sg, node, node_type, parents, children)
if node_type == "state":
if not parents or not children:
continue
elif sg.get_node_attribute(children[0], "type") == "state":
f1 = sg.get_edge_attribute(node, children[0], "label")
attr = NODE_ATTR["task"]
else: # child_type == "branch"
f1 = children[0]
attr = NODE_ATTR["task_branch"]
for p in parents:
f0 = sg.get_edge_attribute(p, node, "label")
fg.add_edge(f0, f1, EDGE_ATTR["task_edge"])
fg.set_node_attributes(f0, NODE_ATTR["task"])
fg.set_node_attributes(f1, attr)
elif node_type == "branch_cond":
f1 = sg.get_edge_attribute(node, children[0], "label")
fg.add_edge(parents[0], f1, EDGE_ATTR["task_edge"])
fg.set_node_attributes(f1, NODE_ATTR["task"])
# nothing to do for node_type == "branch"
return fg
def generate_state_graph(agent_def):
assert type(agent_def) == AgentDefinition
graph = Graph()
agent_name = agent_def.name
for idx, st in enumerate(agent_def.klass.state_transitions):
st_from, tf, st_to = st
if type(st_to) == str: # basic state transition
s0 = "[%s] %s" % (agent_name, st_from)
s1 = "[%s] %s" % (agent_name, st_to)
tf_name = "%s.%s" % (agent_name, tf)
graph.add_edge(s0, s1, EDGE_ATTR["st_func"])
graph.set_edge_attribute(s0, s1, "label", tf_name)
graph.set_edge_attribute(s0, s1, "st_func", tf)
graph.set_node_attributes(s0, NODE_ATTR["state"])
graph.set_node_attributes(s1, NODE_ATTR["state"])
else:
assert type(st_to) in (list, tuple)
branch = tf
end_states = st_to
# connect state_from to branch
s0 = "[%s] %s" % (agent_name, st_from)
sb = "%s.%s" % (agent_name, branch)
graph.add_edge(s0, sb, EDGE_ATTR["branch_edge"])
graph.set_node_attributes(s0, NODE_ATTR["state"])
graph.set_node_attributes(sb, NODE_ATTR["branch"])
graph.set_node_attribute(sb, "st_func", branch)
# for each branching condition
for tf, st_to, condition in end_states:
# connect state_from to branch condition
sc = "[%s] %s" % (agent_name, condition)
graph.add_edge(sb, sc, EDGE_ATTR["branch_edge"])
graph.set_node_attributes(sc, NODE_ATTR["branch_cond"])
# connect branch condition to state_to
s1 = "[%s] %s" % (agent_name, st_to)
tf_name = "%s.%s" % (agent_name, tf)
graph.add_edge(sc, s1, EDGE_ATTR["st_func"])
graph.set_edge_attribute(sc, s1, "label", tf_name)
graph.set_edge_attribute(sc, s1, "st_func", tf)
graph.set_node_attributes(s1, NODE_ATTR["state"])
return graph
