Integrate with Ontology

Author: Val Brodsky

libs/labelbox/src/labelbox/schema/ontology.py

@@ -12,8 +12,12 @@
 
 from labelbox.orm.db_object import DbObject
 from labelbox.orm.model import Field, Relationship
+from labelbox.schema.tool_building.fact_checking_tool import FactCheckingTool
 from labelbox.schema.tool_building.step_reasoning_tool import StepReasoningTool
 from labelbox.schema.tool_building.tool_type import ToolType
+from labelbox.schema.tool_building.tool_type_mapping import (
+    map_tool_type_to_tool_cls,
+)
 
 FeatureSchemaId: Type[str] = Annotated[
     str, StringConstraints(min_length=25, max_length=25)
@@ -491,13 +495,14 @@ def add_classification(self, classification: Classification) -> None:
 
 
 def tool_cls_from_type(tool_type: str):
-    if tool_type.lower() == ToolType.STEP_REASONING.value:
-        return StepReasoningTool
+    tool_cls = map_tool_type_to_tool_cls(tool_type)
+    if tool_cls is not None:
+        return tool_cls
     return Tool
 
 
 def tool_type_cls_from_type(tool_type: str):
-    if tool_type.lower() == ToolType.STEP_REASONING.value:
+    if ToolType.valid(tool_type):
         return ToolType
     return Tool.Type
 
@@ -596,7 +601,9 @@ class OntologyBuilder:
 
     """
 
-    tools: List[Union[Tool, StepReasoningTool]] = field(default_factory=list)
+    tools: List[Union[Tool, StepReasoningTool, FactCheckingTool]] = field(
+        default_factory=list
+    )
     classifications: List[
         Union[Classification, PromptResponseClassification]
     ] = field(default_factory=list)

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

@@ -2,3 +2,4 @@
 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/tool_type.py

@@ -4,3 +4,11 @@
 class ToolType(Enum):
     STEP_REASONING = "step-reasoning"
     FACT_CHECKING = "fact-checking"
+
+    @classmethod
+    def valid(cls, tool_type: str) -> bool:
+        try:
+            ToolType(tool_type.lower())
+            return True
+        except ValueError:
+            return False

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

@@ -0,0 +1,16 @@
+from labelbox.schema.tool_building.fact_checking_tool import FactCheckingTool
+from labelbox.schema.tool_building.step_reasoning_tool import StepReasoningTool
+from labelbox.schema.tool_building.tool_type import ToolType
+
+
+def map_tool_type_to_tool_cls(tool_type_str: str):
+    if not ToolType.valid(tool_type_str):
+        raise ValueError(f"Invalid tool type {tool_type_str}")
+
+    tool_type = ToolType(tool_type_str.lower())
+    if tool_type == ToolType.STEP_REASONING:
+        return StepReasoningTool
+    elif tool_type == ToolType.FACT_CHECKING:
+        return FactCheckingTool
+
+    return None

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

@@ -31,4 +31,4 @@ def reset_actions(self) -> None:
         self.actions = set()
 
     def asdict(self) -> Dict[str, Any]:
-        return {"id": self.id, "name": self.name, "actions": self.actions}
+        return {"id": self.id, "name": self.name, "actions": list(self.actions)}