Method draw

toolbench/tooleval/evaluation/dataclass.py:109–169

(self)

        return eg
    
    def draw(self):
        import pygraphviz as pgv
        G = pgv.AGraph(directed=True)
        G.add_nodes_from([str(node) for node in self.nodes.values()])
        VIS_CONFIG={
            'system':{'shape':'plaintext'},
            'user': {'fillcolor':'yellow','style':'filled','shape':'circle'},
            'tool': {
                # 'fillcolor':'red','style':'filled',
                'shape':'diamond'},
            'assistant': {
                # 'fillcolor':'green','style':'filled',
                'shape':'box'}
        }
        def wrap_text(text:str, width=20):
            wrapped_text = ''
            for i in range(0, min(width*5,len(text)), width):
                wrapped_text += text[i:i+width] + '\n'
            escaped_chars = re.findall(r'\\[nrt\'"\\]', wrapped_text)
            for escaped_char in escaped_chars:
                wrapped_text = wrapped_text.replace(escaped_char, '')
            return wrapped_text
        
        def set_node_vis(gnode,node:ExecutionNode):
            for k,v in VIS_CONFIG[node.role].items():
                gnode.attr[k] = v
                
            if node.role == 'system':
                gnode.attr['label']='SystemPrompt'
            elif node.role == 'tool':
                if node.message['name'] == 'Finish':
                    # args = json.loads(node.message['arguments'])
                    args = str(node.message['arguments'])
                    idx = args.find('return_type')
                    
                    if 'give_answer' in args[idx:idx+30]:
                        gnode.attr['fillcolor'] = 'green'
                        # gnode.attr['xlabel'] = f"{wrap_text(args.get('final_answer',''))}"
                        gnode.attr['label'] = wrap_text(args[args.find('final_answer'):])
                    else:
                        gnode.attr['fillcolor'] = 'red'
                        gnode.attr['label'] = 'restart'
                    gnode.attr['style'] = 'filled'
                    gnode.attr['shape'] = 'ellipse'
                else:
                    gnode.attr['label'] = f"tool: {wrap_text(node.message['name'])}"
                    gnode.attr['xlabel'] = f"{wrap_text(node.message['response'])}"
            elif node.role =='assistant':
                gnode.attr['label'] = node.role.upper() +'\n'+ wrap_text(str(node.message))
            else:
                gnode.attr['xlabel'] = wrap_text(str(node.message))
                gnode.attr['label'] = node.role.upper()
                
        for node in self.nodes.values():
            gnode = G.get_node(str(node))
            set_node_vis(gnode,node)
            to_nodes = list(self.edges.get(node.node_id,{}).keys())
            G.add_edges_from([(str(node),str(to_node)) for to_node in to_nodes])