run black

Author: Yifu Cai

aide/agent.py

@@ -15,9 +15,11 @@
 
 logger = logging.getLogger("aide")
 
+
 def format_time(time_in_sec: int):
     return f"{time_in_sec // 3600}hrs {(time_in_sec % 3600) // 60}mins {time_in_sec % 60}secs"
-    
+
+
 ExecCallbackType = Callable[[str, bool], ExecutionResult]
 
 review_func_spec = FunctionSpec(
@@ -65,7 +67,7 @@ def __init__(
         self.current_step = 0
         if self.acfg.cost_limit:
             self.token_counter = TokenCounter(
-                max_cost=self.acfg.cost_limit,
+                cost_limit=self.acfg.cost_limit,
             )
         else:
             self.token_counter = None
@@ -121,7 +123,7 @@ def _prompt_environment(self):
         pkg_str = ", ".join([f"`{p}`" for p in pkgs])
 
         ts_pksg = [
-            "sktime", 
+            "sktime",
             "statsforecast",
             "tsfresh",
             "neuralforecast",
@@ -142,16 +144,18 @@ def _prompt_impl_guideline(self):
         exec_timeout = int(min(self.cfg.exec.timeout, tot_time_remaining))
 
         if self.acfg.remind_resource_limit:
-            impl_guideline = [f"<TOTAL_TIME_REMAINING: {format_time(tot_time_remaining)}>",
-                              f"<TOTAL_STEPS_REMAINING: {self.acfg.steps - self.current_step}>"]
-            
+            impl_guideline = [
+                f"<TOTAL_TIME_REMAINING: {format_time(tot_time_remaining)}>",
+                f"<TOTAL_STEPS_REMAINING: {self.acfg.steps - self.current_step}>",
+            ]
+
             if self.token_counter:
                 impl_guideline.append(
                     f"<OUTPUT_TOKEN_LIMIT_REMAINING: {self.token_counter.remaining_output_tokens(self.acfg.code.model)}>"
                 )
         else:
             impl_guideline = []
-        
+
         impl_guideline += [
             "The code should **implement the proposed solution** and **print the value of the evaluation metric computed on a hold-out validation set**.",
             "**AND MOST IMPORTANTLY SAVE PREDICTIONS ON THE PROVIDED UNLABELED TEST DATA IN REQUIRED FILE FORMAT IN THE ./submission/ DIRECTORY.**",
@@ -352,7 +356,9 @@ def step(self, exec_callback: ExecCallbackType) -> bool:
                 for item in submission_idr.iterdir():
                     if item.is_file():
                         shutil.copy(item, best_submission_dir / item.name)
-                        logger.info(f"Copied {item.name} to {best_submission_dir / item.name}")
+                        logger.info(
+                            f"Copied {item.name} to {best_submission_dir / item.name}"
+                        )
                 # copy solution.py and relevant node id to best_solution/
                 with open(best_solution_dir / "solution.py", "w") as f:
                     f.write(result_node.code)
@@ -366,7 +372,7 @@ def step(self, exec_callback: ExecCallbackType) -> bool:
 
         exceed_budget_limit = self.token_counter.exceed_budget_limit()
         return exceed_budget_limit
-    
+
     def parse_exec_result(self, node: Node, exec_result: ExecutionResult):
         logger.info(f"Agent is parsing execution results for node {node.id}")
 
@@ -407,9 +413,7 @@ def parse_exec_result(self, node: Node, exec_result: ExecutionResult):
         )
 
         if node.is_buggy:
-            logger.info(
-                f"Parsed results: Node {node.id} is buggy"
-            )
+            logger.info(f"Parsed results: Node {node.id} is buggy")
             node.metric = WorstMetricValue()
         else:
             logger.info(f"Parsed results: Node {node.id} is not buggy")

aide/backend/__init__.py

@@ -7,15 +7,21 @@
 
 logger = logging.getLogger("aide")
 
-#cost per input/output token for each model 
+# cost per input/output token for each model
 MODEL_COST = {
-    "gpt-4o-2024-08-06": {"input": 2.5/1000000, "output": 10/1000000},
-    "o3-mini": {"input": 1.1/1000000, "output": 4.4/1000000},
-    "o3": {"input": 10/1000000, "output": 40/1000000},
+    "gpt-4o-2024-08-06": {"input": 2.5 / 1000000, "output": 10 / 1000000},
+    "o3-mini": {"input": 1.1 / 1000000, "output": 4.4 / 1000000},
+    "o3": {"input": 10 / 1000000, "output": 40 / 1000000},
 }
 
+
 def determine_provider(model: str) -> str:
-    if model.startswith("gpt-") or model.startswith("o1-") or model.startswith("o3-") or model.startswith("o4-"):
+    if (
+        model.startswith("gpt-")
+        or model.startswith("o1-")
+        or model.startswith("o3-")
+        or model.startswith("o4-")
+    ):
         return "openai"
     elif model.startswith("claude-"):
         return "anthropic"
@@ -33,67 +39,69 @@ def determine_provider(model: str) -> str:
     "openrouter": backend_openrouter.query,
 }
 
+
 class TokenCounter:
-    def __init__(self, cost_limit:int):
+    def __init__(self, cost_limit: int):
         self.cost_limit = cost_limit
         self.total_input_tokens = defaultdict(int)
         self.total_output_tokens = defaultdict(int)
-        
+
     def cost(self) -> float:
-        '''
+        """
         compute to total cost of the tokens used
-        '''
+        """
         total_cost = 0
 
-        #compute cost for input tokens
+        # compute cost for input tokens
         for model_name, input_tokens in self.total_input_tokens.items():
             if model_name not in MODEL_COST:
                 raise ValueError(f"Model {model_name} not supported for token counting")
             total_cost += input_tokens * MODEL_COST[model_name]["input"]
-        
-        #compute cost for output tokens
+
+        # compute cost for output tokens
         for model_name, output_tokens in self.total_output_tokens.items():
             if model_name not in MODEL_COST:
                 raise ValueError(f"Model {model_name} not supported for token counting")
             total_cost += output_tokens * MODEL_COST[model_name]["output"]
         return total_cost
-    
-    def add_tokens(self, model_name:str, input_tokens=None, output_tokens=None):
-        '''
+
+    def add_tokens(self, model_name: str, input_tokens=None, output_tokens=None):
+        """
         update the token counts
-        '''
+        """
         if model_name not in MODEL_COST:
             raise ValueError(f"Model {model_name} not supported for token counting")
-        
+
         if input_tokens is not None:
             self.total_input_tokens[model_name] += input_tokens
         if output_tokens is not None:
             self.total_output_tokens[model_name] += output_tokens
 
-    def remaining_output_tokens(self, model_name:str, max_budget:int) -> int:
-        '''
+    def remaining_output_tokens(self, model_name: str, max_budget: int) -> int:
+        """
         max_budget: the maximum dollar budget for the model
         compute the remaining tokens for a model
-        '''
+        """
         if model_name not in MODEL_COST:
             raise ValueError(f"Model {model_name} not supported for token counting")
-        
+
         current_cost = self.cost
         remaining_budget = max_budget - current_cost
         if remaining_budget <= 0:
             return 0
         else:
             output_tokens_cost = MODEL_COST[model_name]["output"]
             return int(remaining_budget / output_tokens_cost)
-    
+
     def exceed_budget_limit(self) -> bool:
-        '''
+        """
         check if the budget limit is exceeded
-        '''
-        
+        """
+
         current_cost = self.cost
         return current_cost > self.cost_limit
-    
+
+
 def query(
     system_message: PromptType | None,
     user_message: PromptType | None,
@@ -150,6 +158,8 @@ def query(
     logger.info(f"response: {output}", extra={"verbose": True})
     logger.info("---Query complete---", extra={"verbose": True})
     if token_counter is not None:
-        token_counter.add_tokens(model, input_tokens=in_tok_count, output_tokens=out_tok_count)
+        token_counter.add_tokens(
+            model, input_tokens=in_tok_count, output_tokens=out_tok_count
+        )
 
     return output

aide/backend/backend_anthropic.py

@@ -73,4 +73,4 @@ def query(
         "stop_reason": message.stop_reason,
     }
 
-    return output, req_time, in_tokens, out_tokens, info
+    return output, req_time, in_tokens, out_tokens, info

aide/backend/backend_gdm.py

@@ -95,4 +95,4 @@ def query(
     info = {}  # this isnt used anywhere, but is an expected return value
 
     # only `output` is actually used by scaffolding
-    return output, req_time, in_tokens, out_tokens, info
+    return output, req_time, in_tokens, out_tokens, info

aide/backend/backend_openai.py

@@ -42,7 +42,9 @@ def query(
     _setup_openai_client()
     filtered_kwargs: dict = select_values(notnone, model_kwargs)  # type: ignore
 
-    messages = opt_messages_to_list(system_message, user_message, convert_system_to_user=convert_system_to_user)
+    messages = opt_messages_to_list(
+        system_message, user_message, convert_system_to_user=convert_system_to_user
+    )
 
     if func_spec is not None:
         filtered_kwargs["tools"] = [func_spec.as_openai_tool_dict]
@@ -86,4 +88,4 @@ def query(
         "created": completion.created,
     }
 
-    return output, req_time, in_tokens, out_tokens, info
+    return output, req_time, in_tokens, out_tokens, info

aide/backend/backend_openrouter.py

@@ -83,4 +83,4 @@ def query(
         "created": completion.created,
     }
 
-    return output, req_time, in_tokens, out_tokens, info
+    return output, req_time, in_tokens, out_tokens, info

aide/backend/utils.py

@@ -86,4 +86,4 @@ def openai_tool_choice_dict(self):
         return {
             "type": "function",
             "function": {"name": self.name},
-        }
+        }

aide/journal.py

@@ -192,6 +192,7 @@ def generate_summary(self, include_code: bool = False) -> str:
             summary.append(summary_part)
         return "\n-------------------------------\n".join(summary)
 
+
 def get_path_to_node(journal: Journal, node_id: str) -> list[str]:
     path = [node_id]
 
@@ -243,4 +244,4 @@ def filter_journal(journal: Journal) -> Journal:
     else:
         filtered_journal = filter_for_longest_path(journal)
 
-    return filtered_journal
+    return filtered_journal

aide/run.py

@@ -26,6 +26,7 @@
 from rich.tree import Tree
 from .utils.config import load_task_desc, prep_agent_workspace, save_run, load_cfg
 
+
 class VerboseFilter(logging.Filter):
     """
     Filter (remove) logs that have verbose attribute set to True
@@ -34,6 +35,7 @@ class VerboseFilter(logging.Filter):
     def filter(self, record):
         return not (hasattr(record, "verbose") and record.verbose)
 
+
 def journal_to_rich_tree(journal: Journal):
     best_node = journal.get_best_node()