Tests for FactChecking

Author: Val Brodsky

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

@@ -1,5 +1,6 @@
 from dataclasses import dataclass, field
-from typing import Any, Dict, List, Optional, Set
+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 (
@@ -8,6 +9,18 @@
 )
 
 
+class UnsupportedStepActions(Enum):
+    WRITE_JUSTIFICATION = "writeJustification"
+
+
+class CanConfidentlyAssessStepActions(Enum):
+    WRITE_JUSTIFICATION = "writeJustification"
+
+
+class NoFactualInformationStepActions(Enum):
+    WRITE_JUSTIFICATION = "writeJustification"
+
+
 @dataclass
 class FactCheckingVariants:
     """
@@ -26,21 +39,32 @@ class FactCheckingVariants:
     )
     unsupported_step: VariantWithActions = field(
         default_factory=lambda: VariantWithActions(
-            id=3, name="Unsupported", _available_actions={"writeJustification"}
+            id=3,
+            name="Unsupported",
+            _available_actions={
+                action.value for action in UnsupportedStepActions
+            },
+            actions=[UnsupportedStepActions.WRITE_JUSTIFICATION.value],
         )
     )
     cant_confidently_assess_step: VariantWithActions = field(
         default_factory=lambda: VariantWithActions(
             id=4,
             name="Can't confidently assess",
-            _available_actions={"writeJustification"},
+            _available_actions={
+                action.value for action in CanConfidentlyAssessStepActions
+            },
+            actions=[CanConfidentlyAssessStepActions.WRITE_JUSTIFICATION.value],
         )
     )
     no_factual_information_step: VariantWithActions = field(
         default_factory=lambda: VariantWithActions(
             id=5,
             name="No factual information",
-            _available_actions={"writeJustification"},
+            _available_actions={
+                action.value for action in NoFactualInformationStepActions
+            },
+            actions=[NoFactualInformationStepActions.WRITE_JUSTIFICATION.value],
         )
     )
 
@@ -139,22 +163,25 @@ class FactCheckingTool:
     )
 
     def set_unsupported_step_actions(
-        self, actions: Set[str] = {"writeJustification"}
+        self, actions: List[UnsupportedStepActions]
     ) -> None:
-        self.definition.variants.unsupported_step.set_actions(actions)
+        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: Set[str] = {"writeJustification"}
+        self, actions: List[CanConfidentlyAssessStepActions]
     ) -> None:
+        actions_values = [action.value for action in actions]
         self.definition.variants.cant_confidently_assess_step.set_actions(
-            actions
+            actions_values
         )
 
     def set_no_factual_information_step_actions(
-        self, actions: Set[str] = {"writeJustification"}
+        self, actions: List[NoFactualInformationStepActions]
     ) -> None:
+        actions_values = [action.value for action in actions]
         self.definition.variants.no_factual_information_step.set_actions(
-            actions
+            actions_values
         )
 
     def asdict(self) -> Dict[str, Any]:

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

@@ -32,8 +32,11 @@ def reset_actions(self) -> None:
         self.actions = []
 
     def asdict(self) -> Dict[str, Any]:
-        return {
+        data = {
             "id": self.id,
             "name": self.name,
-            "actions": list(set(self.actions)),
         }
+        if len(self.actions) > 0:
+            data["actions"] = self.actions
+
+        return data

libs/labelbox/tests/integration/conftest.py

@@ -21,6 +21,7 @@
 )
 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
 
@@ -579,6 +580,7 @@ def chat_evaluation_ontology(client, rand_gen):
                 name="model output multi ranking",
             ),
             StepReasoningTool(name="step reasoning"),
+            FactCheckingTool(name="fact checking"),
         ],
         classifications=[
             Classification(

libs/labelbox/tests/integration/test_ontology.py

@@ -5,6 +5,7 @@
 
 from labelbox import MediaType, OntologyBuilder, Tool
 from labelbox.orm.model import Entity
+from labelbox.schema.tool_building.fact_checking_tool import FactCheckingTool
 from labelbox.schema.tool_building.step_reasoning_tool import StepReasoningTool
 
 
@@ -367,3 +368,56 @@ def test_step_reasoning_ontology(chat_evaluation_ontology):
             "actions": ["regenerateSteps"],
         },
     ]
+
+
+def test_fact_checking_ontology(chat_evaluation_ontology):
+    ontology = chat_evaluation_ontology
+    fact_checking = None
+    for tool in ontology.normalized["tools"]:
+        if tool["tool"] == "fact-checking":
+            fact_checking = tool
+            break
+    assert fact_checking is not None
+    assert fact_checking["definition"]["variants"] == [
+        {"id": 0, "name": "Accurate"},
+        {"id": 1, "name": "Inaccurate"},
+        {"id": 2, "name": "Disputed"},
+        {"id": 3, "name": "Unsupported", "actions": ["writeJustification"]},
+        {
+            "id": 4,
+            "name": "Can't confidently assess",
+            "actions": ["writeJustification"],
+        },
+        {
+            "id": 5,
+            "name": "No factual information",
+            "actions": ["writeJustification"],
+        },
+    ]
+    assert fact_checking["definition"]["version"] == 1
+    assert fact_checking["schemaNodeId"] is not None
+    assert fact_checking["featureSchemaId"] is not None
+
+    fact_checking = None
+    for tool in ontology.tools():
+        if isinstance(tool, FactCheckingTool):
+            fact_checking = tool
+            break
+    assert fact_checking is not None
+
+    assert fact_checking.definition.variants.asdict() == [
+        {"id": 0, "name": "Accurate"},
+        {"id": 1, "name": "Inaccurate"},
+        {"id": 2, "name": "Disputed"},
+        {"id": 3, "name": "Unsupported", "actions": ["writeJustification"]},
+        {
+            "id": 4,
+            "name": "Can't confidently assess",
+            "actions": ["writeJustification"],
+        },
+        {
+            "id": 5,
+            "name": "No factual information",
+            "actions": ["writeJustification"],
+        },
+    ]