Update StepReasoningTool to use _BaseStepReasoningTool

Author: vbrodsky

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

@@ -0,0 +1,100 @@
+import warnings
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Set
+
+from labelbox.schema.tool_building.tool_type import ToolType
+
+
+@dataclass
+class _Variant:
+    id: int
+    name: str
+    actions: List[str] = field(default_factory=list)
+    _available_actions: Set[str] = field(default_factory=set)
+
+    def set_actions(self, actions: List[str]) -> None:
+        self.actions = []
+        for action in actions:
+            if action in self._available_actions:
+                self.actions.append(action)
+
+    def reset_actions(self) -> None:
+        self.actions = []
+
+    def asdict(self) -> Dict[str, Any]:
+        return {
+            "id": self.id,
+            "name": self.name,
+            "actions": self.actions,
+        }
+
+
+@dataclass
+class _Definition:
+    variants: List[_Variant]
+    version: int = field(default=1)
+    title: Optional[str] = None
+    value: Optional[str] = None
+
+    def __post_init__(self):
+        if self.version != 1:
+            raise ValueError("Invalid version")
+
+    def asdict(self) -> Dict[str, Any]:
+        result = {
+            "variants": [variant.asdict() for variant in self.variants],
+            "version": self.version,
+        }
+        if self.title is not None:
+            result["title"] = self.title
+        if self.value is not None:
+            result["value"] = self.value
+        return result
+
+    @classmethod
+    def from_dict(cls, dictionary: Dict[str, Any]) -> "_Definition":
+        variants = [_Variant(**variant) for variant in dictionary["variants"]]
+        title = dictionary.get("title", None)
+        value = dictionary.get("value", None)
+        return cls(variants=variants, title=title, value=value)
+
+
+@dataclass
+class _BaseStepReasoningTool:
+    name: str
+    type: ToolType
+    definition: _Definition
+    schema_id: Optional[str] = None
+    feature_schema_id: Optional[str] = None
+    color: Optional[str] = None
+    required: bool = False
+
+    def __post_init__(self):
+        warnings.warn(
+            "This feature is experimental and subject to change.",
+        )
+
+        if not self.name.strip():
+            raise ValueError("Name is required")
+
+    def asdict(self) -> Dict[str, Any]:
+        return {
+            "tool": self.type.value,
+            "name": self.name,
+            "required": self.required,
+            "schemaNodeId": self.schema_id,
+            "featureSchemaId": self.feature_schema_id,
+            "definition": self.definition.asdict(),
+            "color": self.color,
+        }
+
+    @classmethod
+    def from_dict(cls, dictionary: Dict[str, Any]) -> "_BaseStepReasoningTool":
+        return cls(
+            name=dictionary["name"],
+            schema_id=dictionary.get("schemaNodeId", None),
+            feature_schema_id=dictionary.get("featureSchemaId", None),
+            required=dictionary.get("required", False),
+            definition=_Definition.from_dict(dictionary["definition"]),
+            color=dictionary.get("color", None),
+        )

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

@@ -1,10 +1,12 @@
-import warnings
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import Any, Dict, List, Optional
 
+from labelbox.schema.tool_building.base_step_reasoning_tool import (
+    _BaseStepReasoningTool,
+    _Definition,
+    _Variant,
+)
 from labelbox.schema.tool_building.tool_type import ToolType
-from labelbox.schema.tool_building.variant import VariantWithActions
 
 
 class IncorrectStepActions(Enum):
@@ -14,156 +16,29 @@ class IncorrectStepActions(Enum):
     JUSTIFICATION = "justification"
 
 
-@dataclass
-class StepReasoningVariants:
-    """
-    This class is used to define the possible options for evaluating a step
-    Currently the options are correct, neutral, and incorrect
-    NOTE: do not change the variant values
-    """
-
-    correct_step: VariantWithActions = field(
-        default_factory=lambda: VariantWithActions(
-            id=0,
-            name="Correct",
-            actions=[],
-        )
-    )
-    neutral_step: VariantWithActions = field(
-        default_factory=lambda: VariantWithActions(
-            id=1,
-            name="Neutral",
-            actions=[],
-        )
-    )
-
-    incorrect_step: VariantWithActions = field(
-        default_factory=lambda: VariantWithActions(
-            id=2,
-            name="Incorrect",
-            _available_actions={
-                action.value for action in IncorrectStepActions
-            },
-            actions=[
-                action.value for action in IncorrectStepActions
-            ],  # initialize to all IncorrectStepActions by default
-        )
-    )
-
-    def asdict(self):
-        return [
-            self.correct_step.asdict(),
-            self.neutral_step.asdict(),
-            self.incorrect_step.asdict(),
-        ]
-
-    @classmethod
-    def from_dict(cls, dictionary: List[Dict[str, Any]]):
-        correct_step = None
-        neutral_step = None
-        incorrect_step = None
-
-        for variant in dictionary:
-            if variant["id"] == 0:
-                correct_step = VariantWithActions(**variant)
-            elif variant["id"] == 1:
-                neutral_step = VariantWithActions(**variant)
-            elif variant["id"] == 2:
-                incorrect_step = VariantWithActions(**variant)
-
-        if not all([correct_step, neutral_step, incorrect_step]):
-            raise ValueError("Invalid step reasoning variants")
-
-        return cls(
-            correct_step=correct_step,  # type: ignore
-            neutral_step=neutral_step,  # type: ignore
-            incorrect_step=incorrect_step,  # type: ignore
-        )
-
-
-@dataclass
-class StepReasoningDefinition:
-    variants: StepReasoningVariants = field(
-        default_factory=StepReasoningVariants
+def build_step_reasoning_definition():
+    correct_step = _Variant(id=0, name="Correct", actions=[])
+    neutral_step = _Variant(id=1, name="Neutral", actions=[])
+    incorrect_step = _Variant(
+        id=2,
+        name="Incorrect",
+        _available_actions={action.value for action in IncorrectStepActions},
+        actions=[action.value for action in IncorrectStepActions],
     )
-    version: int = field(default=1)
-    title: Optional[str] = None
-    value: Optional[str] = None
-
-    def __post_init__(self):
-        if self.version != 1:
-            raise ValueError("Invalid version")
-
-    def asdict(self) -> Dict[str, Any]:
-        result = {"variants": self.variants.asdict(), "version": self.version}
-        if self.title is not None:
-            result["title"] = self.title
-        if self.value is not None:
-            result["value"] = self.value
-        return result
-
-    @classmethod
-    def from_dict(cls, dictionary: Dict[str, Any]) -> "StepReasoningDefinition":
-        variants = StepReasoningVariants.from_dict(dictionary["variants"])
-        title = dictionary.get("title", None)
-        value = dictionary.get("value", None)
-        return cls(variants=variants, title=title, value=value)
+    variants = [correct_step, neutral_step, incorrect_step]
+    return _Definition(variants=variants)
 
 
 @dataclass
-class StepReasoningTool:
+class StepReasoningTool(_BaseStepReasoningTool):
     """
     Use this class in OntologyBuilder to create a tool for step reasoning
     The definition field lists the possible options to evaulate a step
 
     NOTE: color attribute is for backward compatibility only and should not be set directly
     """
 
-    name: str
     type: ToolType = field(default=ToolType.STEP_REASONING, init=False)
-    required: bool = False
-    schema_id: Optional[str] = None
-    feature_schema_id: Optional[str] = None
-    color: Optional[str] = None
-    definition: StepReasoningDefinition = field(
-        default_factory=StepReasoningDefinition
+    definition: _Definition = field(
+        default_factory=build_step_reasoning_definition
     )
-
-    def __post_init__(self):
-        warnings.warn(
-            "This feature is experimental and subject to change.",
-        )
-
-        if not self.name.strip():
-            raise ValueError("Name is required")
-
-    def set_incorrect_step_actions(self, actions: List[IncorrectStepActions]):
-        """
-        For live models, will invoke the model to generate alternatives if a step is marked as incorrect
-        NOTE by default all actions are set to True
-        Pass empty list to reset to 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 {
-            "tool": self.type.value,
-            "name": self.name,
-            "required": self.required,
-            "schemaNodeId": self.schema_id,
-            "featureSchemaId": self.feature_schema_id,
-            "definition": self.definition.asdict(),
-        }
-
-    @classmethod
-    def from_dict(cls, dictionary: Dict[str, Any]) -> "StepReasoningTool":
-        return cls(
-            name=dictionary["name"],
-            schema_id=dictionary.get("schemaNodeId", None),
-            feature_schema_id=dictionary.get("featureSchemaId", None),
-            required=dictionary.get("required", False),
-            definition=StepReasoningDefinition.from_dict(
-                dictionary["definition"]
-            ),
-        )

libs/labelbox/tests/integration/test_ontology.py

@@ -358,28 +358,33 @@ def test_step_reasoning_ontology(chat_evaluation_ontology):
             step_reasoning_tool = tool
             break
     assert step_reasoning_tool is not None
-    assert step_reasoning_tool.definition.variants.asdict() == [
-        {
-            "id": 0,
-            "name": "Correct",
-            "actions": [],
-        },
-        {
-            "id": 1,
-            "name": "Neutral",
-            "actions": [],
-        },
-        {
-            "id": 2,
-            "name": "Incorrect",
-            "actions": [
-                "regenerateSteps",
-                "generateAndRateAlternativeSteps",
-                "rewriteStep",
-                "justification",
-            ],
-        },
-    ]
+    assert step_reasoning_tool.definition.asdict() == {
+        "title": "step reasoning",
+        "value": "step_reasoning",
+        "variants": [
+            {
+                "id": 0,
+                "name": "Correct",
+                "actions": [],
+            },
+            {
+                "id": 1,
+                "name": "Neutral",
+                "actions": [],
+            },
+            {
+                "id": 2,
+                "name": "Incorrect",
+                "actions": [
+                    "regenerateSteps",
+                    "generateAndRateAlternativeSteps",
+                    "rewriteStep",
+                    "justification",
+                ],
+            },
+        ],
+        "version": 1,
+    }
 
 
 def test_fact_checking_ontology(chat_evaluation_ontology):

libs/labelbox/tests/unit/test_unit_step_ontology_variants.py

@@ -1,8 +1,8 @@
-from labelbox.schema.tool_building.variant import VariantWithActions
+from labelbox.schema.tool_building.base_step_reasoning_tool import _Variant
 
 
-def test_variant_with_actions_as_dict():
-    variant = VariantWithActions(
+def test_variant():
+    variant = _Variant(
         id=0, name="Correct", _available_actions={"regenerateSteps"}
     )
     variant.set_actions(["regenerateSteps"])

libs/labelbox/tests/unit/test_unit_step_reasoning_tool.py

@@ -18,6 +18,7 @@ def test_step_reasoning_as_dict_default():
         "required": False,
         "schemaNodeId": None,
         "featureSchemaId": None,
+        "color": None,
         "definition": {
             "variants": [
                 {"id": 0, "name": "Correct", "actions": []},