Method set_node_vis

toolbench/tooleval/evaluation/dataclass.py:132–160

(gnode,node:ExecutionNode)

            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))