[PLT-1866][SDK][STEP REASONING] Actions Required for All Variants (#1894)

Author: vbrodsky

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

@@ -9,24 +9,31 @@
 class StepReasoningVariant:
     id: int
     name: str
+    actions: List[str] = field(default_factory=list)
 
     def asdict(self) -> Dict[str, Any]:
-        return {"id": self.id, "name": self.name}
+        return {"id": self.id, "name": self.name, "actions": self.actions}
 
 
 @dataclass
 class IncorrectStepReasoningVariant:
     id: int
     name: str
-    regenerate_conversations_after_incorrect_step: Optional[bool] = True
-    rate_alternative_responses: Optional[bool] = True
+    regenerate_steps: Optional[bool] = True
+    generate_and_rate_alternative_steps: Optional[bool] = True
+    rewrite_steps: Optional[bool] = True
+    justification: Optional[bool] = True
 
     def asdict(self) -> Dict[str, Any]:
         actions = []
-        if self.regenerate_conversations_after_incorrect_step:
+        if self.regenerate_steps:
             actions.append("regenerateSteps")
-        if self.rate_alternative_responses:
+        if self.generate_and_rate_alternative_steps:
             actions.append("generateAndRateAlternativeSteps")
+        if self.rewrite_steps:
+            actions.append("rewriteSteps")
+        if self.justification:
+            actions.append("justification")
         return {"id": self.id, "name": self.name, "actions": actions}
 
     @classmethod
@@ -36,10 +43,11 @@ def from_dict(
         return cls(
             id=dictionary["id"],
             name=dictionary["name"],
-            regenerate_conversations_after_incorrect_step="regenerateSteps"
-            in dictionary.get("actions", []),
-            rate_alternative_responses="generateAndRateAlternativeSteps"
+            regenerate_steps="regenerateSteps" in dictionary.get("actions", []),
+            generate_and_rate_alternative_steps="generateAndRateAlternativeSteps"
             in dictionary.get("actions", []),
+            rewrite_steps="rewriteSteps" in dictionary.get("actions", []),
+            justification="justification" in dictionary.get("actions", []),
         )
 
 
@@ -162,20 +170,30 @@ def __post_init__(self):
             "This feature is experimental and subject to change.",
         )
 
-    def reset_regenerate_conversations_after_incorrect_step(self):
+    def reset_regenerate_steps(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
+        self.definition.variants.incorrect_step.regenerate_steps = False
 
-    def reset_rate_alternative_responses(self):
+    def reset_generate_and_rate_alternative_steps(self):
         """
         For live models, will require labelers to rate the alternatives generated by the model
         """
-        self.definition.variants.incorrect_step.rate_alternative_responses = (
-            False
-        )
+        self.definition.variants.incorrect_step.generate_and_rate_alternative_steps = False
+
+    def reset_rewrite_steps(self):
+        """
+        For live models, will require labelers to rewrite the conversation
+        """
+        self.definition.variants.incorrect_step.rewrite_steps = False
+
+    def reset_justification(self):
+        """
+        For live models, will require labelers to provide a justification for their evaluation
+        """
+        self.definition.variants.incorrect_step.justification = False
 
     def asdict(self) -> Dict[str, Any]:
         return {

libs/labelbox/tests/unit/test_unit_step_reasoning_tool.py

@@ -11,14 +11,16 @@ def test_step_reasoning_as_dict_default():
         "featureSchemaId": None,
         "definition": {
             "variants": [
-                {"id": 0, "name": "Correct"},
-                {"id": 1, "name": "Neutral"},
+                {"id": 0, "name": "Correct", "actions": []},
+                {"id": 1, "name": "Neutral", "actions": []},
                 {
                     "id": 2,
                     "name": "Incorrect",
                     "actions": [
                         "regenerateSteps",
                         "generateAndRateAlternativeSteps",
+                        "rewriteSteps",
+                        "justification",
                     ],
                 },
             ],
@@ -29,8 +31,10 @@ 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.reset_generate_and_rate_alternative_steps()
+    tool.reset_regenerate_steps()
+    tool.reset_rewrite_steps()
+    tool.reset_justification()
     assert tool.asdict() == {
         "tool": "step-reasoning",
         "name": "step reasoning",
@@ -39,8 +43,8 @@ def test_step_reasoning_as_dict_with_actions():
         "featureSchemaId": None,
         "definition": {
             "variants": [
-                {"id": 0, "name": "Correct"},
-                {"id": 1, "name": "Neutral"},
+                {"id": 0, "name": "Correct", "actions": []},
+                {"id": 1, "name": "Neutral", "actions": []},
                 {
                     "id": 2,
                     "name": "Incorrect",