Cleanup actions

Author: Val Brodsky

libs/labelbox/src/labelbox/schema/tool_building/step_reasoning_tool.py

@@ -1,11 +1,17 @@
 import warnings
 from dataclasses import dataclass, field
+from enum import Enum
 from typing import Any, Dict, List, Optional
 
 from labelbox.schema.tool_building.tool_type import ToolType
 from labelbox.schema.tool_building.variant import Variant, VariantWithActions
 
 
+class IncorrectStepActions(Enum):
+    REGENERATE_STEPS = "regenerateSteps"
+    GENERATE_AND_RATE_ALTERNATIVE_STEPS = "generateAndRateAlternativeSteps"
+
+
 @dataclass
 class StepReasoningVariants:
     """
@@ -25,8 +31,7 @@ class StepReasoningVariants:
             id=2,
             name="Incorrect",
             _available_actions={
-                "regenerateSteps",
-                "generateAndRateAlternativeSteps",
+                action.value for action in IncorrectStepActions
             },
             actions=["regenerateSteps"],  # regenerateSteps is on by default
         )
@@ -110,20 +115,12 @@ def __post_init__(self):
             "This feature is experimental and subject to change.",
         )
 
-    def reset_regenerate_conversations_after_incorrect_step(self):
-        """
-        For live models, the default acation will invoke the model to generate alternatives if a step is marked as incorrect
-        This method will reset the action to not regenerate the conversation
-        """
-        self.definition.variants.incorrect_step.regenerate_conversations_after_incorrect_step = False
-
-    def reset_rate_alternative_responses(self):
+    def set_incorrect_step_actions(self, actions: List[IncorrectStepActions]):
         """
-        For live models, will require labelers to rate the alternatives generated by the model
+        For live models, will invoke the model to generate alternatives if a step is marked as incorrect
         """
-        self.definition.variants.incorrect_step.rate_alternative_responses = (
-            False
-        )
+        actions_values = [action.value for action in actions]
+        self.definition.variants.incorrect_step.set_actions(actions_values)
 
     def asdict(self) -> Dict[str, Any]:
         return {

libs/labelbox/src/labelbox/schema/tool_building/tool_type_mapping.py

@@ -5,12 +5,10 @@
 
 def map_tool_type_to_tool_cls(tool_type_str: str):
     if not ToolType.valid(tool_type_str):
-        raise ValueError(f"Invalid tool type {tool_type_str}")
+        return None
 
     tool_type = ToolType(tool_type_str.lower())
     if tool_type == ToolType.STEP_REASONING:
         return StepReasoningTool
     elif tool_type == ToolType.FACT_CHECKING:
         return FactCheckingTool
-
-    return None

libs/labelbox/src/labelbox/schema/tool_building/variant.py

@@ -22,7 +22,8 @@ class VariantWithActions:
     actions: List[str] = field(default_factory=list)
     _available_actions: Set[str] = field(default_factory=set)
 
-    def set_actions(self, actions: Set[str]) -> None:
+    def set_actions(self, actions: List[str]) -> None:
+        self.actions = []
         for action in actions:
             if action in self._available_actions:
                 self.actions.append(action)

libs/labelbox/tests/integration/conftest.py

@@ -21,7 +21,6 @@
 )
 from labelbox.schema.data_row import DataRowMetadataField
 from labelbox.schema.ontology_kind import OntologyKind
-from labelbox.schema.tool_building.fact_checking_tool import FactCheckingTool
 from labelbox.schema.tool_building.step_reasoning_tool import StepReasoningTool
 from labelbox.schema.user import User
 
@@ -580,7 +579,6 @@ def chat_evaluation_ontology(client, rand_gen):
                 name="model output multi ranking",
             ),
             StepReasoningTool(name="step reasoning"),
-            FactCheckingTool(name="fact checking"),
         ],
         classifications=[
             Classification(

libs/labelbox/tests/integration/test_ontology.py

@@ -339,7 +339,7 @@ def test_step_reasoning_ontology(chat_evaluation_ontology):
         {
             "id": 2,
             "name": "Incorrect",
-            "actions": [],
+            "actions": ["regenerateSteps"],
         },
     ]
     assert step_reasoning_tool["definition"]["version"] == 1
@@ -364,6 +364,6 @@ def test_step_reasoning_ontology(chat_evaluation_ontology):
         {
             "id": 2,
             "name": "Incorrect",
-            "actions": [],
+            "actions": ["regenerateSteps"],
         },
     ]

libs/labelbox/tests/unit/test_unit_step_reasoning_tool.py

@@ -1,4 +1,7 @@
-from labelbox.schema.tool_building.step_reasoning_tool import StepReasoningTool
+from labelbox.schema.tool_building.step_reasoning_tool import (
+    IncorrectStepActions,
+    StepReasoningTool,
+)
 
 
 def test_step_reasoning_as_dict_default():
@@ -16,7 +19,7 @@ def test_step_reasoning_as_dict_default():
                 {
                     "id": 2,
                     "name": "Incorrect",
-                    "actions": [],
+                    "actions": ["regenerateSteps"],
                 },
             ],
             "version": 1,
@@ -26,8 +29,12 @@ def test_step_reasoning_as_dict_default():
 
 def test_step_reasoning_as_dict_with_actions():
     tool = StepReasoningTool(name="step reasoning")
-    tool.reset_rate_alternative_responses()
-    tool.reset_regenerate_conversations_after_incorrect_step()
+    tool.set_incorrect_step_actions(
+        [
+            IncorrectStepActions.REGENERATE_STEPS,
+            IncorrectStepActions.GENERATE_AND_RATE_ALTERNATIVE_STEPS,
+        ]
+    )
     assert tool.asdict() == {
         "tool": "step-reasoning",
         "name": "step reasoning",