| 10 | |
| 11 | |
| 12 | class DFS_tree_search(base_search_method): |
| 13 | |
| 14 | def __init__(self, llm, io_func, process_id=0, callbacks=None): |
| 15 | super(DFS_tree_search, self).__init__( |
| 16 | llm, io_func, process_id, callbacks) |
| 17 | """Depth-first search. |
| 18 | with_filter=True: Every time a child node is generated, choose the best multiple iterations to go. |
| 19 | with_filter=False: Do as Preorder traversal. |
| 20 | """ |
| 21 | self.io_func = io_func |
| 22 | self.llm = llm |
| 23 | self.process_id = process_id |
| 24 | self.restart() |
| 25 | |
| 26 | self.callbacks = callbacks if callbacks is not None else [] |
| 27 | |
| 28 | def restart(self): |
| 29 | self.status = 0 |
| 30 | self.terminal_node = [] |
| 31 | self.give_up_node = [] |
| 32 | self.now_expand_num = 0 |
| 33 | self.query_count = 0 |
| 34 | self.total_tokens = 0 |
| 35 | |
| 36 | def send_agent_chain_end(self, depth, agent_block_ids, chain_block_ids): |
| 37 | for i in range(len(self.callbacks)): |
| 38 | callback = self.callbacks[i] |
| 39 | callback.on_chain_end( |
| 40 | depth=depth, |
| 41 | block_id=chain_block_ids[i] |
| 42 | ) |
| 43 | if i < len(agent_block_ids): |
| 44 | callback.on_agent_end( |
| 45 | depth=depth, |
| 46 | block_id=agent_block_ids[i] |
| 47 | ) |
| 48 | |
| 49 | def to_json(self, answer=False, process=True): |
| 50 | |
| 51 | if process: |
| 52 | json_obj = { |
| 53 | "win": self.status == 1, |
| 54 | "tree": self.tree.to_json_recursive(), |
| 55 | "forward_args": self.forward_args, |
| 56 | "compare_candidates": [], |
| 57 | } |
| 58 | for node in self.terminal_node: |
| 59 | if node.pruned == False: # has answer |
| 60 | json_obj["compare_candidates"].append( |
| 61 | node.get_chain_result_from_this_node(use_messages=False)) |
| 62 | else: |
| 63 | json_obj = {} |
| 64 | |
| 65 | if answer: |
| 66 | json_obj["answer_generation"] = { |
| 67 | "valid_data": False, |
| 68 | "query_count": self.query_count, |
| 69 | "total_tokens": self.total_tokens, |