Make sure _available_actions are always set

Author: Val Brodsky

libs/labelbox/src/labelbox/__init__.py

@@ -55,8 +55,8 @@
     ResponseOption,
     Tool,
 )
-from labelbox.schema.ontology import PromptResponseClassification
-from labelbox.schema.ontology import ResponseOption
+from labelbox.schema.tool_building.fact_checking_tool import FactCheckingTool
+from labelbox.schema.tool_building.step_reasoning_tool import StepReasoningTool
 from labelbox.schema.role import Role, ProjectRole
 from labelbox.schema.invite import Invite, InviteLimit
 from labelbox.schema.data_row_metadata import (

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

@@ -1,5 +1,4 @@
 import labelbox.schema.tool_building.tool_type
-import labelbox.schema.tool_building.variant
 import labelbox.schema.tool_building.step_reasoning_tool
 import labelbox.schema.tool_building.fact_checking_tool
 import labelbox.schema.tool_building.tool_type_mapping

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

@@ -71,3 +71,14 @@ class FactCheckingTool(_BaseStepReasoningTool):
     definition: _Definition = field(
         default_factory=build_fact_checking_definition
     )
+
+    def __post_init__(self):
+        super().__post_init__()
+        # Set available actions for variants 0, 1, 2 'out of band' since they are not passed in the definition
+        self._set_variant_available_actions()
+
+    def _set_variant_available_actions(self):
+        for variant in self.definition.variants:
+            if variant.id in [0, 1, 2]:
+                for action in FactCheckingActions:
+                    variant._available_actions.add(action.value)

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

@@ -42,3 +42,14 @@ class StepReasoningTool(_BaseStepReasoningTool):
     definition: _Definition = field(
         default_factory=build_step_reasoning_definition
     )
+
+    def __post_init__(self):
+        super().__post_init__()
+        # Set available actions for variants 0, 1, 2 'out of band' since they are not passed in the definition
+        self._set_variant_available_actions()
+
+    def _set_variant_available_actions(self):
+        for variant in self.definition.variants:
+            if variant.id == 2:
+                for action in IncorrectStepActions:
+                    variant._available_actions.add(action.value)

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

@@ -12,17 +12,17 @@
     Classification,
     Client,
     Dataset,
+    FactCheckingTool,
     MediaType,
     OntologyBuilder,
     Option,
     PromptResponseClassification,
     ResponseOption,
+    StepReasoningTool,
     Tool,
 )
 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
 
 

libs/labelbox/tests/integration/conftest.py

@@ -358,6 +358,7 @@ 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.asdict() == {
         "title": "step reasoning",
         "value": "step_reasoning",

libs/labelbox/tests/integration/test_ontology.py

@@ -11,6 +11,8 @@ def test_variant():
         "name": "Correct",
         "actions": ["regenerateSteps"],
     }
+
+    assert variant._available_actions == {"regenerateSteps"}
     variant.reset_actions()
     assert variant.asdict() == {
         "id": 0,

libs/labelbox/tests/unit/test_unit_step_ontology_variants.py

@@ -12,6 +12,13 @@ def test_validations():
 
 def test_step_reasoning_as_dict_default():
     tool = StepReasoningTool(name="step reasoning")
+    assert tool.definition.variants[2]._available_actions == {
+        "regenerateSteps",
+        "generateAndRateAlternativeSteps",
+        "rewriteStep",
+        "justification",
+    }
+
     assert tool.asdict() == {
         "tool": "step-reasoning",
         "name": "step reasoning",
@@ -37,3 +44,44 @@ def test_step_reasoning_as_dict_default():
             "version": 1,
         },
     }
+
+
+def test_from_dict():
+    dict = {
+        "schemaNodeId": "cm3pdkupv0ah8070h2ujo74th",
+        "featureSchemaId": "cm3pdkupv0ah7070hg7svdeeo",
+        "required": False,
+        "name": "step reasoning",
+        "tool": "step-reasoning",
+        "definition": {
+            "variants": [
+                {"id": 0, "name": "Correct", "actions": []},
+                {"id": 1, "name": "Neutral", "actions": []},
+                {
+                    "id": 2,
+                    "name": "Incorrect",
+                    "actions": [
+                        "regenerateSteps",
+                        "generateAndRateAlternativeSteps",
+                        "rewriteStep",
+                        "justification",
+                    ],
+                },
+            ],
+            "version": 1,
+            "title": "step reasoning",
+            "value": "step_reasoning",
+            "color": "#ff0000",
+        },
+        "color": "#ff0000",
+        "archived": 0,
+        "classifications": [],
+        "kind": "StepReasoning",
+    }
+    tool = StepReasoningTool.from_dict(dict)
+    assert tool.definition.variants[2]._available_actions == {
+        "generateAndRateAlternativeSteps",
+        "justification",
+        "rewriteStep",
+        "regenerateSteps",
+    }