Port latest updates to step reasoning

Author: Val Brodsky

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

@@ -4,12 +4,14 @@
 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
+from labelbox.schema.tool_building.variant import VariantWithActions
 
 
 class IncorrectStepActions(Enum):
     REGENERATE_STEPS = "regenerateSteps"
     GENERATE_AND_RATE_ALTERNATIVE_STEPS = "generateAndRateAlternativeSteps"
+    REWRITE_STEP = "rewriteStep"
+    JUSTIFICATION = "justification"
 
 
 @dataclass
@@ -20,11 +22,19 @@ class StepReasoningVariants:
     NOTE: do not change the variant values
     """
 
-    correct_step: Variant = field(
-        default_factory=lambda: Variant(id=0, name="Correct")
+    correct_step: VariantWithActions = field(
+        default_factory=lambda: VariantWithActions(
+            id=0,
+            name="Correct",
+            actions=[],
+        )
     )
-    neutral_step: Variant = field(
-        default_factory=lambda: Variant(id=1, name="Neutral")
+    neutral_step: VariantWithActions = field(
+        default_factory=lambda: VariantWithActions(
+            id=1,
+            name="Neutral",
+            actions=[],
+        )
     )
 
     incorrect_step: VariantWithActions = field(
@@ -35,9 +45,8 @@ class StepReasoningVariants:
                 action.value for action in IncorrectStepActions
             },
             actions=[
-                IncorrectStepActions.REGENERATE_STEPS.value,
-                IncorrectStepActions.GENERATE_AND_RATE_ALTERNATIVE_STEPS.value,
-            ],  # action On by default
+                action.value for action in IncorrectStepActions
+            ],  # initialize to all IncorrectStepActions by default
         )
     )
 
@@ -56,9 +65,9 @@ def from_dict(cls, dictionary: List[Dict[str, Any]]):
 
         for variant in dictionary:
             if variant["id"] == 0:
-                correct_step = Variant(**variant)
+                correct_step = VariantWithActions(**variant)
             elif variant["id"] == 1:
-                neutral_step = Variant(**variant)
+                neutral_step = VariantWithActions(**variant)
             elif variant["id"] == 2:
                 incorrect_step = VariantWithActions(**variant)
 

libs/labelbox/tests/integration/test_ontology.py

@@ -335,12 +335,17 @@ def test_step_reasoning_ontology(chat_evaluation_ontology):
             break
     assert step_reasoning_tool is not None
     assert step_reasoning_tool["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"],
+            "actions": [
+                "regenerateSteps",
+                "generateAndRateAlternativeSteps",
+                "rewriteStep",
+                "justification",
+            ],
         },
     ]
     assert step_reasoning_tool["definition"]["version"] == 1
@@ -357,15 +362,22 @@ def test_step_reasoning_ontology(chat_evaluation_ontology):
         {
             "id": 0,
             "name": "Correct",
+            "actions": [],
         },
         {
             "id": 1,
             "name": "Neutral",
+            "actions": [],
         },
         {
             "id": 2,
             "name": "Incorrect",
-            "actions": ["regenerateSteps", "generateAndRateAlternativeSteps"],
+            "actions": [
+                "regenerateSteps",
+                "generateAndRateAlternativeSteps",
+                "rewriteStep",
+                "justification",
+            ],
         },
     ]
 

libs/labelbox/tests/unit/test_unit_step_reasoning_tool.py

@@ -20,14 +20,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",
+                        "rewriteStep",
+                        "justification",
                     ],
                 },
             ],
@@ -47,8 +49,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",