Implement logic to add and update tasks

Author: xeon27

src/capability.py

@@ -12,6 +12,7 @@
     NO_ANSWER_STR,
     NON_SEED_CAPABILITIES_SCORE_DIR,
     SEED_CAPABILITIES_SCORE_DIR,
+    TAB_W_SPACES,
 )
 from src.utils.data_utils import load_data
 from src.utils.prompts import TASK_SOLVER_SYSTEM_PROMPT
@@ -140,7 +141,10 @@ def from_dict(cls, capability_dict: Dict[str, Any], base_dir: str) -> "Capabilit
             f"capability_{c_dict['name']}", os.path.join(c_dir, "capability.py")
         )
         c_obj = c_module.Capability()
-        initial_tasks = list(c_obj.repr_tasks().values())
+        initial_tasks = [
+            {"id": k, "problem": v["problem"], "answer": v["answer"]}
+            for k, v in c_obj.repr_tasks().items()
+        ]
         template_instructions = c_obj.get_instructions({"problem": '{t["problem"]}'})
         template_instructions = f'f"""{template_instructions}"""'
 
@@ -232,6 +236,95 @@ def get_repr_tasks(self) -> List[Dict[str, Any]]:
             )
         return repr_tasks
 
+    def add_and_update_tasks(self, tasks: List[Dict[str, Any]]) -> None:
+        """
+        Add and/or update tasks for the capability.
+
+        Args
+        ----
+            tasks (List[Dict[str, Any]]): A list of dictionaries containing the tasks
+            to be added. Each task dict consists of id, problem, and answer keys.
+        """
+        if not all(
+            "id" in task and "problem" in task and "answer" in task for task in tasks
+        ):
+            raise ValueError(
+                "Each task must contain 'id', 'problem', and 'answer' keys."
+            )
+
+        existing_task_ids = [task["id"] for task in self._data]
+        new_task_ids = [task["id"] for task in tasks]
+        # Keep new task for overlapping tasks
+        # TODO: Add `overwrite` flag to update existing tasks
+        tasks_to_keep = [
+            task
+            for task in self._data
+            if task["id"]
+            not in list(set.intersection(set(existing_task_ids), set(new_task_ids)))
+        ] + tasks
+        # Sort by task id
+        tasks_to_keep.sort(key=lambda x: x["id"])
+
+        # Check if the new task list consists of representative tasks
+        # If yes, update the capability class python file
+        repr_tasks = [
+            task
+            for task in tasks
+            if task["id"] in self.capability_repr_class.repr_tasks()
+        ]
+        if len(repr_tasks) > 0:
+            partial_repr_task_ids = [task["id"] for task in repr_tasks]
+            if len(partial_repr_task_ids) < len(
+                self.capability_repr_class.repr_tasks()
+            ):
+                # Get remaining tasks from existing task list
+                for k, v in self.capability_repr_class.repr_tasks().items():
+                    if k not in partial_repr_task_ids:
+                        repr_task = {"id": k}
+                        repr_task.update(v)
+                        repr_tasks.append(repr_task)
+                repr_tasks.sort(key=lambda x: x["id"])
+            # Update the capability class python file
+            # Extract str which contains the repr_tasks dictionary
+            # TODO: Since these are hardcoded, update when the format changes
+            prefix_str = f"def repr_tasks() -> dict[str, dict]:\n{TAB_W_SPACES}{TAB_W_SPACES}return "
+            suffix_str = f"\n\n{TAB_W_SPACES}@staticmethod\n{TAB_W_SPACES}def get_instructions(t: dict) -> str:"
+            prev_repr_tasks_str = self.capability_repr_class_str.split(prefix_str)[
+                1
+            ].split(suffix_str)[0]
+            # Restructure to match the original format
+            repr_tasks_dict = {}
+            for elm in repr_tasks:
+                repr_tasks_dict[elm["id"]] = {k: v for k, v in elm.items() if k != "id"}
+            # Replace the repr_tasks dictionary in the capability class string
+            # with the updated one
+            updated_repr_tasks_str = json.dumps(repr_tasks_dict, indent=4)
+            newline = "\n"
+            capability_repr_class_str = self.capability_repr_class_str.lstrip(
+                f"```python{newline}"
+            ).rstrip(f"{newline}```")
+            capability_repr_class_str = capability_repr_class_str.replace(
+                prev_repr_tasks_str,
+                updated_repr_tasks_str,
+            )
+            with open(os.path.join(self.source_dir, "capability.py"), "w") as f:
+                f.write(capability_repr_class_str)
+
+        # Update the capability data in the capability json file
+        c_dict = {
+            "capability_name": self.name,
+            "capability_description": self.description,
+            "capability_domain": self.domain,
+            "capability_instructions": self.instructions,
+            "capability_data": tasks_to_keep,
+        }
+        with open(os.path.join(self.source_dir, "capability.json"), "w") as f:
+            json.dump(c_dict, f, indent=4)
+
+        # Reload the capability class to reflect these changes
+        self._load_capability_json()
+        self._load_capability_repr_class()
+
     def _to_dict(self) -> Dict[str, Any]:
         return {
             "name": self.name,

src/cfg/run_cfg.yaml

@@ -24,7 +24,7 @@ capabilities_cfg:
   num_seed_capabilities: -1
   num_gen_capabilities: 4
   num_gen_capabilities_per_run: 2
-  num_gen_tasks_per_capability: 2
+  num_gen_tasks_per_capability: 1
   task_gen_few_shot: true
 
 lbo_cfg:

src/generate_tasks.py

@@ -147,7 +147,7 @@ def generate_tasks_using_llm(
     # Combine with sample tasks to get the full set of tasks
     start_id = len(sample_tasks) + 1
     all_tasks = sample_tasks + [
-        {"id": (start_id + idx), "problem": new_tasks[idx]}
+        {"id": str(start_id + idx), "problem": new_tasks[idx]}
         for idx in range(len(new_tasks))
     ]
 
@@ -159,3 +159,7 @@ def generate_tasks_using_llm(
     )
     print(json.dumps(solved_tasks, indent=4))
     print(task_solver_metadata)
+
+    capability.add_and_update_tasks(
+        tasks=solved_tasks,
+    )