change print to mcts logger

metagpt/ext/sela/data/custom_task.py

@@ -6,6 +6,7 @@ from metagpt.ext.sela.runner.mle_bench.instructions import (
     INSTRUCTIONS,
     INSTRUCTIONS_OBFUSCATED,
 )
+from metagpt.ext.sela.utils import mcts_logger
 
 MLE_BENCH_FILES = ["description.md", "description_obfuscated.md"]
 
@@ -70,5 +71,5 @@ def get_mle_bench_requirements(dataset_dir, data_config, special_instruction, ob
         output_dir=output_dir,
         special_instruction=special_instruction,
     )
-    print(mle_requirement)
+    mcts_logger.info(mle_requirement)
     return mle_requirement

metagpt/ext/sela/data/dataset.py

@@ -10,7 +10,7 @@ import yaml
 from sklearn.model_selection import train_test_split
 
 from metagpt.ext.sela.insights.solution_designer import SolutionDesigner
-from metagpt.ext.sela.utils import DATA_CONFIG
+from metagpt.ext.sela.utils import DATA_CONFIG, mcts_logger
 
 BASE_USER_REQUIREMENT = """
 This is a {datasetname} dataset. Your goal is to predict the target column `{target_col}`.
@@ -191,7 +191,7 @@ def generate_task_requirement(task_name, data_config, is_di=True, special_instru
         additional_instruction=additional_instruction,
         data_info_path=data_info_path,
     )
-    print(user_requirement)
+    mcts_logger.info(user_requirement)
     return user_requirement
 
 
@@ -286,16 +286,16 @@ class ExpDataset:
     def save_dataset(self, target_col):
         df, test_df = self.get_raw_dataset()
         if not self.check_dataset_exists() or self.force_update:
-            print(f"Saving Dataset {self.name} in {self.dataset_dir}")
+            mcts_logger.info(f"Saving Dataset {self.name} in {self.dataset_dir}")
             self.split_and_save(df, target_col, test_df=test_df)
         else:
-            print(f"Dataset {self.name} already exists")
+            mcts_logger.info(f"Dataset {self.name} already exists")
         if not self.check_datasetinfo_exists() or self.force_update:
-            print(f"Saving Dataset info for {self.name}")
+            mcts_logger.info(f"Saving Dataset info for {self.name}")
             dataset_info = self.get_dataset_info()
             self.save_datasetinfo(dataset_info)
         else:
-            print(f"Dataset info for {self.name} already exists")
+            mcts_logger.info(f"Dataset info for {self.name} already exists")
 
     def save_datasetinfo(self, dataset_info):
         with open(Path(self.dataset_dir, self.name, "dataset_info.json"), "w", encoding="utf-8") as file:

metagpt/ext/sela/evaluation/visualize_mcts.py

@@ -4,6 +4,7 @@ import matplotlib.pyplot as plt
 import networkx as nx
 
 from metagpt.ext.sela.search.tree_search import Node
+from metagpt.ext.sela.utils import mcts_logger
 
 NODE_TEMPLATE = """\
 [Node {id}]
@@ -139,7 +140,7 @@ def build_tree_recursive(graph, parent_id, node, node_order, start_task_id=2):
         instruction = "\n\n".join([role.planner.plan.tasks[i].instruction for i in range(start_task_id)])
     else:
         instruction = role.planner.plan.tasks[depth + start_task_id - 1].instruction
-    print(instruction)
+    mcts_logger.info(instruction)
     # Add the current node with attributes to the graph
     dev_score = node.raw_reward.get("dev_score", 0) * 100
     avg_score = node.avg_value() * 100

metagpt/ext/sela/experimenter.py

@@ -133,9 +133,9 @@ class Experimenter(DataInterpreter):
         if self.planner.plan.goal != "":
             self.set_actions([WriteAnalysisCode])
             self._set_state(0)
-            print("Plan already exists, skipping initialization.")
+            mcts_logger.info("Plan already exists, skipping initialization.")
             return self
-        print("Initializing plan and tool...")
+        mcts_logger.info("Initializing plan and tool...")
         return super().set_plan_and_tool()
 
     async def _act_on_task(self, current_task: Task) -> TaskResult:

metagpt/ext/sela/runner/aide.py

@@ -3,6 +3,8 @@ import time
 
 import aide
 
+from metagpt.ext.sela.utils import mcts_logger
+
 os.environ["OPENAI_API_KEY"] = "sk-xxx"
 os.environ["OPENAI_BASE_URL"] = "your url"
 
@@ -27,9 +29,9 @@ exp = aide.Experiment(
 
 best_solution = exp.run(steps=10)
 
-print(f"Best solution has validation metric: {best_solution.valid_metric}")
-print(f"Best solution code: {best_solution.code}")
+mcts_logger.info(f"Best solution has validation metric: {best_solution.valid_metric}")
+mcts_logger.info(f"Best solution code: {best_solution.code}")
 end_time = time.time()
 execution_time = end_time - start_time
 
-print(f"run time : {execution_time} seconds")
+mcts_logger.info(f"run time : {execution_time} seconds")

metagpt/ext/sela/runner/mcts.py

@@ -7,6 +7,7 @@ from metagpt.ext.sela.evaluation.evaluation import (
 from metagpt.ext.sela.evaluation.visualize_mcts import get_tree_text
 from metagpt.ext.sela.runner.runner import Runner
 from metagpt.ext.sela.search.search_algorithm import MCTS, Greedy, Random
+from metagpt.ext.sela.utils import mcts_logger
 
 
 class MCTSRunner(Runner):
@@ -46,7 +47,7 @@ class MCTSRunner(Runner):
         text += f"Best node: {best_node.id}, score: {best_node.raw_reward}\n"
         text += f"Dev best node: {dev_best_node.id}, score: {dev_best_node.raw_reward}\n"
         text += f"Grader score: {additional_scores['grader']}\n"
-        print(text)
+        mcts_logger.info(text)
         results = [
             {
                 "best_node": best_node.id,

metagpt/ext/sela/runner/random_search.py

@@ -1,7 +1,7 @@
 from metagpt.ext.sela.experimenter import Experimenter
 from metagpt.ext.sela.insights.instruction_generator import InstructionGenerator
 from metagpt.ext.sela.runner.runner import Runner
-from metagpt.ext.sela.utils import get_exp_pool_path
+from metagpt.ext.sela.utils import get_exp_pool_path, mcts_logger
 
 EXPS_PROMPT = """
 When doing the tasks, you can refer to the insights below:
@@ -37,7 +37,7 @@ class RandomSearchRunner(Runner):
             di = Experimenter(node_id=str(i), use_reflection=self.args.reflection, role_timeout=self.args.role_timeout)
             di.role_dir = f"{di.role_dir}_{self.args.task}"
             requirement = user_requirement + EXPS_PROMPT.format(experience=exps[i])
-            print(requirement)
+            mcts_logger.info(requirement)
             score_dict = await self.run_di(di, requirement, run_idx=i)
             results.append(
                 {

metagpt/ext/sela/runner/runner.py

@@ -8,7 +8,7 @@ import pandas as pd
 from metagpt.ext.sela.evaluation.evaluation import evaluate_score
 from metagpt.ext.sela.experimenter import Experimenter
 from metagpt.ext.sela.search.tree_search import create_initial_state
-from metagpt.ext.sela.utils import DATA_CONFIG, save_notebook
+from metagpt.ext.sela.utils import DATA_CONFIG, mcts_logger, save_notebook
 
 
 class Runner:
@@ -38,7 +38,7 @@ class Runner:
                 score_dict = self.evaluate(score_dict, self.state)
                 run_finished = True
             except Exception as e:
-                print(f"Error: {e}")
+                mcts_logger.info(f"Error: {e}")
                 num_runs += 1
         # save_notebook(role=di, save_dir=self.result_path, name=f"{self.args.task}_{self.start_time}_{run_idx}")
         save_name = self.get_save_name()

metagpt/ext/sela/search/tree_search.py

@@ -374,7 +374,7 @@ class BaseTreeSearch:
                 return best_score, best_child
             for child in self.children[node]:
                 score = child.normalized_reward[split]
-                print(child.id, split, score)
+                mcts_logger.info(f"{child.id} {split} {score}")
                 if score > best_score:
                     best_score = score
                     best_child = child

metagpt/ext/sela/utils.py

@@ -57,7 +57,7 @@ def get_exp_pool_path(task_name, data_config, pool_name="analysis_pool"):
 
 
 def change_plan(role, plan):
-    print(f"Change next plan to: {plan}")
+    mcts_logger.info(f"Change next plan to: {plan}")
     tasks = role.planner.plan.tasks
     finished = True
     for i, task in enumerate(tasks):
@@ -115,8 +115,8 @@ async def load_execute_notebook(role):
     # await executor.build()
     for code in codes:
         outputs, success = await executor.run(code)
-        print(f"Execution success: {success}, Output: {outputs}")
-    print("Finish executing the loaded notebook")
+        mcts_logger.info(f"Execution success: {success}, Output: {outputs}")
+    mcts_logger.info("Finish executing the loaded notebook")
     return executor