Update FactCheckingTool to use _BaseStepReasoningTool

Author: vbrodsky

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

@@ -17,6 +17,10 @@ def set_actions(self, actions: List[str]) -> None:
         for action in actions:
             if action in self._available_actions:
                 self.actions.append(action)
+            else:
+                warnings.warn(
+                    f"Variant ID {self.id} {action} is an invalid action, skipping"
+                )
 
     def reset_actions(self) -> None:
         self.actions = []
@@ -28,6 +32,10 @@ def asdict(self) -> Dict[str, Any]:
             "actions": self.actions,
         }
 
+    def _post_init(self):
+        # Call set_actions to remove any invalid actions
+        self.set_actions(self.actions)
+
 
 @dataclass
 class _Definition:

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

@@ -1,233 +1,73 @@
-import warnings
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import Any, Dict, List, Optional
 
-from labelbox.schema.tool_building.tool_type import ToolType
-from labelbox.schema.tool_building.variant import (
-    VariantWithActions,
+from labelbox.schema.tool_building.base_step_reasoning_tool import (
+    _BaseStepReasoningTool,
+    _Definition,
+    _Variant,
 )
+from labelbox.schema.tool_building.tool_type import ToolType
 
 
 class FactCheckingActions(Enum):
     WRITE_JUSTIFICATION = "justification"
 
 
-class UnsupportedStepActions(Enum):
-    WRITE_JUSTIFICATION = "justification"
-
-
-class CanConfidentlyAssessStepActions(Enum):
-    WRITE_JUSTIFICATION = "justification"
-
-
-class NoFactualInformationStepActions(Enum):
-    WRITE_JUSTIFICATION = "justification"
-
-
-@dataclass
-class FactCheckingVariants:
-    """
-    This class is used to define the possible options for fact-checking a step
-    NOTE do not change the variants directly
-    """
-
-    accurate_step: VariantWithActions = field(
-        default_factory=lambda: VariantWithActions(
-            id=0,
-            name="Accurate",
-            _available_actions={action.value for action in FactCheckingActions},
-            actions=[action.value for action in FactCheckingActions],
-        )
+def build_fact_checking_definition():
+    accurate_step = _Variant(
+        id=0,
+        name="Accurate",
+        _available_actions={action.value for action in FactCheckingActions},
+        actions=[action.value for action in FactCheckingActions],
     )
-
-    inaccurate_step: VariantWithActions = field(
-        default_factory=lambda: VariantWithActions(
-            id=1,
-            name="Inaccurate",
-            _available_actions={action.value for action in FactCheckingActions},
-            actions=[action.value for action in FactCheckingActions],
-        )
+    inaccurate_step = _Variant(
+        id=1,
+        name="Inaccurate",
+        _available_actions={action.value for action in FactCheckingActions},
+        actions=[action.value for action in FactCheckingActions],
     )
-    disputed_step: VariantWithActions = field(
-        default_factory=lambda: VariantWithActions(
-            id=2,
-            name="Disputed",
-            _available_actions={action.value for action in FactCheckingActions},
-            actions=[action.value for action in FactCheckingActions],
-        )
+    disputed_step = _Variant(
+        id=2,
+        name="Disputed",
+        _available_actions={action.value for action in FactCheckingActions},
+        actions=[action.value for action in FactCheckingActions],
     )
-    unsupported_step: VariantWithActions = field(
-        default_factory=lambda: VariantWithActions(
-            id=3,
-            name="Unsupported",
-            _available_actions=set(),
-            actions=[],
-        )
+    unsupported_step = _Variant(
+        id=3,
+        name="Unsupported",
+        _available_actions=set(),
+        actions=[],
     )
-    cant_confidently_assess_step: VariantWithActions = field(
-        default_factory=lambda: VariantWithActions(
-            id=4,
-            name="Can't confidently assess",
-            _available_actions=set(),
-            actions=[],
-        )
+    cant_confidently_assess_step = _Variant(
+        id=4,
+        name="Can't confidently assess",
+        _available_actions=set(),
+        actions=[],
     )
-    no_factual_information_step: VariantWithActions = field(
-        default_factory=lambda: VariantWithActions(
-            id=5,
-            name="No factual information",
-            _available_actions=set(),
-            actions=[],
-        )
+    no_factual_information_step = _Variant(
+        id=5,
+        name="No factual information",
+        _available_actions=set(),
+        actions=[],
     )
-
-    def asdict(self):
-        return [
-            self.accurate_step.asdict(),
-            self.inaccurate_step.asdict(),
-            self.disputed_step.asdict(),
-            self.unsupported_step.asdict(),
-            self.cant_confidently_assess_step.asdict(),
-            self.no_factual_information_step.asdict(),
-        ]
-
-    @classmethod
-    def from_dict(cls, dictionary: List[Dict[str, Any]]):
-        accurate_step = None
-        inaccurate_step = None
-        disputed_step = None
-        unsupported_step = None
-        cant_confidently_assess_step = None
-        no_factual_information_step = None
-
-        for variant in dictionary:
-            if variant["id"] == 0:
-                accurate_step = VariantWithActions(**variant)
-            elif variant["id"] == 1:
-                inaccurate_step = VariantWithActions(**variant)
-            elif variant["id"] == 2:
-                disputed_step = VariantWithActions(**variant)
-            elif variant["id"] == 3:
-                unsupported_step = VariantWithActions(**variant)
-            elif variant["id"] == 4:
-                cant_confidently_assess_step = VariantWithActions(**variant)
-            elif variant["id"] == 5:
-                no_factual_information_step = VariantWithActions(**variant)
-
-        if not all(
-            [
-                accurate_step,
-                inaccurate_step,
-                disputed_step,
-                unsupported_step,
-                cant_confidently_assess_step,
-                no_factual_information_step,
-            ]
-        ):
-            raise ValueError("Missing variant")
-
-        return cls(
-            accurate_step=accurate_step,  # type: ignore
-            inaccurate_step=inaccurate_step,  # type: ignore
-            disputed_step=disputed_step,  # type: ignore
-            unsupported_step=unsupported_step,  # type: ignore
-            cant_confidently_assess_step=cant_confidently_assess_step,  # type: ignore
-            no_factual_information_step=no_factual_information_step,  # type: ignore
-        )
+    variants = [
+        accurate_step,
+        inaccurate_step,
+        disputed_step,
+        unsupported_step,
+        cant_confidently_assess_step,
+        no_factual_information_step,
+    ]
+    return _Definition(variants=variants)
 
 
 @dataclass
-class FactCheckingDefinition:
-    variants: FactCheckingVariants = field(default_factory=FactCheckingVariants)
-    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]) -> "FactCheckingDefinition":
-        variants = FactCheckingVariants.from_dict(dictionary["variants"])
-        title = dictionary.get("title", None)
-        value = dictionary.get("value", None)
-        return cls(variants=variants, title=title, value=value)
-
-
-@dataclass
-class FactCheckingTool:
+class FactCheckingTool(_BaseStepReasoningTool):
     """
     Use this class in OntologyBuilder to create a tool for fact checking
     """
 
-    name: str
     type: ToolType = field(default=ToolType.FACT_CHECKING, init=False)
-    required: bool = False
-    schema_id: Optional[str] = None
-    feature_schema_id: Optional[str] = None
-    color: Optional[str] = None
-    definition: FactCheckingDefinition = field(
-        default_factory=FactCheckingDefinition
+    definition: _Definition = field(
+        default_factory=build_fact_checking_definition
     )
-
-    def __post_init__(self):
-        warnings.warn(
-            "This feature is experimental and subject to change.",
-        )
-
-        if self.name.strip() == "":
-            raise ValueError("Name cannot be empty")
-
-    def set_unsupported_step_actions(
-        self, actions: List[UnsupportedStepActions]
-    ) -> None:
-        actions_values = [action.value for action in actions]
-        self.definition.variants.unsupported_step.set_actions(actions_values)
-
-    def set_cant_confidently_assess_step_actions(
-        self, actions: List[CanConfidentlyAssessStepActions]
-    ) -> None:
-        actions_values = [action.value for action in actions]
-        self.definition.variants.cant_confidently_assess_step.set_actions(
-            actions_values
-        )
-
-    def set_no_factual_information_step_actions(
-        self, actions: List[NoFactualInformationStepActions]
-    ) -> None:
-        actions_values = [action.value for action in actions]
-        self.definition.variants.no_factual_information_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]) -> "FactCheckingTool":
-        return cls(
-            name=dictionary["name"],
-            schema_id=dictionary.get("schemaNodeId", None),
-            feature_schema_id=dictionary.get("featureSchemaId", None),
-            required=dictionary.get("required", False),
-            definition=FactCheckingDefinition.from_dict(
-                dictionary["definition"]
-            ),
-        )

libs/labelbox/tests/integration/test_ontology.py

@@ -422,19 +422,24 @@ def test_fact_checking_ontology(chat_evaluation_ontology):
             break
     assert fact_checking is not None
 
-    assert fact_checking.definition.variants.asdict() == [
-        {"id": 0, "name": "Accurate", "actions": ["justification"]},
-        {"id": 1, "name": "Inaccurate", "actions": ["justification"]},
-        {"id": 2, "name": "Disputed", "actions": ["justification"]},
-        {"id": 3, "name": "Unsupported", "actions": []},
-        {
-            "id": 4,
-            "name": "Can't confidently assess",
-            "actions": [],
-        },
-        {
-            "id": 5,
-            "name": "No factual information",
-            "actions": [],
-        },
-    ]
+    assert fact_checking.definition.asdict() == {
+        "title": "fact checking",
+        "value": "fact_checking",
+        "variants": [
+            {"id": 0, "name": "Accurate", "actions": ["justification"]},
+            {"id": 1, "name": "Inaccurate", "actions": ["justification"]},
+            {"id": 2, "name": "Disputed", "actions": ["justification"]},
+            {"id": 3, "name": "Unsupported", "actions": []},
+            {
+                "id": 4,
+                "name": "Can't confidently assess",
+                "actions": [],
+            },
+            {
+                "id": 5,
+                "name": "No factual information",
+                "actions": [],
+            },
+        ],
+        "version": 1,
+    }

libs/labelbox/tests/unit/test_unit_fact_checking_tool.py

@@ -14,50 +14,7 @@ def test_fact_checking_as_dict_default():
         "required": False,
         "schemaNodeId": None,
         "featureSchemaId": None,
-        "definition": {
-            "variants": [
-                {"id": 0, "name": "Accurate", "actions": ["justification"]},
-                {"id": 1, "name": "Inaccurate", "actions": ["justification"]},
-                {"id": 2, "name": "Disputed", "actions": ["justification"]},
-                {
-                    "id": 3,
-                    "name": "Unsupported",
-                    "actions": [],
-                },
-                {
-                    "id": 4,
-                    "name": "Can't confidently assess",
-                    "actions": [],
-                },
-                {
-                    "id": 5,
-                    "name": "No factual information",
-                    "actions": [],
-                },
-            ],
-            "version": 1,
-        },
-    }
-
-    assert tool_dict == expected_dict
-
-
-def test_step_reasoning_as_dict_with_actions():
-    tool = FactCheckingTool(name="Fact Checking Tool")
-    tool.set_unsupported_step_actions([])
-    tool.set_cant_confidently_assess_step_actions([])
-    tool.set_no_factual_information_step_actions([])
-
-    # Get the dictionary representation
-    tool_dict = tool.asdict()
-
-    # Expected dictionary structure
-    expected_dict = {
-        "tool": "fact-checking",
-        "name": "Fact Checking Tool",
-        "required": False,
-        "schemaNodeId": None,
-        "featureSchemaId": None,
+        "color": None,
         "definition": {
             "variants": [
                 {"id": 0, "name": "Accurate", "actions": ["justification"]},