Refactor StepReasoning to also reuse Variants

Author: Val Brodsky

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

@@ -3,62 +3,7 @@
 from typing import Any, Dict, List, Optional
 
 from labelbox.schema.tool_building.tool_type import ToolType
-
-
-@dataclass
-class StepReasoningVariant:
-    id: int
-    name: str
-
-    def asdict(self) -> Dict[str, Any]:
-        return {"id": self.id, "name": self.name}
-
-
-@dataclass
-class IncorrectStepReasoningVariant:
-    id: int
-    name: str
-    regenerate_conversations_after_incorrect_step: Optional[bool] = True
-    rate_alternative_responses: Optional[bool] = True
-
-    def asdict(self) -> Dict[str, Any]:
-        actions = []
-        if self.regenerate_conversations_after_incorrect_step:
-            actions.append("regenerateSteps")
-        if self.rate_alternative_responses:
-            actions.append("generateAndRateAlternativeSteps")
-        return {"id": self.id, "name": self.name, "actions": actions}
-
-    @classmethod
-    def from_dict(
-        cls, dictionary: Dict[str, Any]
-    ) -> "IncorrectStepReasoningVariant":
-        return cls(
-            id=dictionary["id"],
-            name=dictionary["name"],
-            regenerate_conversations_after_incorrect_step="regenerateSteps"
-            in dictionary.get("actions", []),
-            rate_alternative_responses="generateAndRateAlternativeSteps"
-            in dictionary.get("actions", []),
-        )
-
-
-def _create_correct_step() -> StepReasoningVariant:
-    return StepReasoningVariant(
-        id=StepReasoningVariants.CORRECT_STEP_ID, name="Correct"
-    )
-
-
-def _create_neutral_step() -> StepReasoningVariant:
-    return StepReasoningVariant(
-        id=StepReasoningVariants.NEUTRAL_STEP_ID, name="Neutral"
-    )
-
-
-def _create_incorrect_step() -> IncorrectStepReasoningVariant:
-    return IncorrectStepReasoningVariant(
-        id=StepReasoningVariants.INCORRECT_STEP_ID, name="Incorrect"
-    )
+from labelbox.schema.tool_building.variant import Variant, VariantWithActions
 
 
 @dataclass
@@ -68,18 +13,23 @@ class StepReasoningVariants:
     Currently the options are correct, neutral, and incorrect
     """
 
-    CORRECT_STEP_ID = 0
-    NEUTRAL_STEP_ID = 1
-    INCORRECT_STEP_ID = 2
-
-    correct_step: StepReasoningVariant = field(
-        default_factory=_create_correct_step
+    correct_step: Variant = field(
+        default_factory=lambda: Variant(id=0, name="Correct")
     )
-    neutral_step: StepReasoningVariant = field(
-        default_factory=_create_neutral_step
+    neutral_step: Variant = field(
+        default_factory=lambda: Variant(id=1, name="Neutral")
     )
-    incorrect_step: IncorrectStepReasoningVariant = field(
-        default_factory=_create_incorrect_step
+
+    incorrect_step: VariantWithActions = field(
+        default_factory=lambda: VariantWithActions(
+            id=2,
+            name="Incorrect",
+            _available_actions={
+                "regenerateSteps",
+                "generateAndRateAlternativeSteps",
+            },
+            actions=["regenerateSteps"],  # regenerateSteps is on by default
+        )
     )
 
     def asdict(self):
@@ -96,14 +46,12 @@ def from_dict(cls, dictionary: List[Dict[str, Any]]):
         incorrect_step = None
 
         for variant in dictionary:
-            if variant["id"] == cls.CORRECT_STEP_ID:
-                correct_step = StepReasoningVariant(**variant)
-            elif variant["id"] == cls.NEUTRAL_STEP_ID:
-                neutral_step = StepReasoningVariant(**variant)
-            elif variant["id"] == cls.INCORRECT_STEP_ID:
-                incorrect_step = IncorrectStepReasoningVariant.from_dict(
-                    variant
-                )
+            if variant["id"] == 0:
+                correct_step = Variant(**variant)
+            elif variant["id"] == 1:
+                neutral_step = Variant(**variant)
+            elif variant["id"] == 2:
+                incorrect_step = VariantWithActions(**variant)
 
         if not all([correct_step, neutral_step, incorrect_step]):
             raise ValueError("Invalid step reasoning variants")