Add step by step reasoning ontology tool

Author: Val Brodsky

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

@@ -12,6 +12,8 @@
 
 from labelbox.orm.db_object import DbObject
 from labelbox.orm.model import Field, Relationship
+from labelbox.schema.tool_building.step_reasoning_tool import StepReasoningTool
+from labelbox.schema.tool_building.tool_type import ToolType
 
 FeatureSchemaId: Type[str] = Annotated[
     str, StringConstraints(min_length=25, max_length=25)
@@ -187,7 +189,7 @@ def __post_init__(self):
     @classmethod
     def from_dict(cls, dictionary: Dict[str, Any]) -> Dict[str, Any]:
         return cls(
-            class_type=cls.Type(dictionary["type"]),
+            class_type=Classification.Type(dictionary["type"]),
             name=dictionary["name"],
             instructions=dictionary["instructions"],
             required=dictionary.get("required", False),
@@ -351,7 +353,7 @@ class Type(Enum):
     @classmethod
     def from_dict(cls, dictionary: Dict[str, Any]) -> Dict[str, Any]:
         return cls(
-            class_type=cls.Type(dictionary["type"]),
+            class_type=Type(dictionary["type"]),
             name=dictionary["name"],
             instructions=dictionary["instructions"],
             required=True,  # always required
@@ -458,7 +460,7 @@ def from_dict(cls, dictionary: Dict[str, Any]) -> Dict[str, Any]:
             schema_id=dictionary.get("schemaNodeId", None),
             feature_schema_id=dictionary.get("featureSchemaId", None),
             required=dictionary.get("required", False),
-            tool=cls.Type(dictionary["tool"]),
+            tool=Tool.Type(dictionary["tool"]),
             classifications=[
                 Classification.from_dict(c)
                 for c in dictionary["classifications"]
@@ -488,6 +490,16 @@ def add_classification(self, classification: Classification) -> None:
         self.classifications.append(classification)
 
 
+def tool_cls_from_type(tool_type: str):
+    if tool_type.lower() == ToolType.STEP_REASONING.value:
+        from labelbox.schema.tool_building.step_reasoning_tool import (
+            StepReasoningTool,
+        )
+
+        return StepReasoningTool
+    return Tool
+
+
 class Ontology(DbObject):
     """An ontology specifies which tools and classifications are available
     to a project. This is read only for now.
@@ -521,11 +533,18 @@ def __init__(self, *args, **kwargs) -> None:
             Union[List[Classification], List[PromptResponseClassification]]
         ] = None
 
+    def _tool_deserializer_cls(self, tool: Dict[str, Any]) -> Tool:
+        import pdb
+
+        pdb.set_trace()
+        return Tool
+
     def tools(self) -> List[Tool]:
         """Get list of tools (AKA objects) in an Ontology."""
         if self._tools is None:
             self._tools = [
-                Tool.from_dict(tool) for tool in self.normalized["tools"]
+                tool_cls_from_type(tool["tool"]).from_dict(tool)
+                for tool in self.normalized["tools"]
             ]
         return self._tools
 
@@ -581,7 +600,7 @@ class OntologyBuilder:
 
     """
 
-    tools: List[Tool] = field(default_factory=list)
+    tools: List[Union[Tool, StepReasoningTool]] = field(default_factory=list)
     classifications: List[
         Union[Classification, PromptResponseClassification]
     ] = field(default_factory=list)

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

@@ -0,0 +1,2 @@
+import labelbox.schema.tool_building.tool_type
+import labelbox.schema.tool_building.step_reasoning_tool

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

@@ -0,0 +1,117 @@
+from dataclasses import dataclass, field
+from typing import Any, Dict, Optional
+
+from labelbox.schema.tool_building.tool_type import ToolType
+
+
+@dataclass
+class StepReasoningVariant:
+    id: int
+    name: str
+
+    def asdict(self) -> Dict[str, Any]:
+        return {"id": self.id, "name": self.name}
+
+
+@dataclass
+class IncorrectStepReasoningVariant:
+    id: int
+    name: str
+    regenerate_conversations_after_incorrect_step: Optional[bool] = False
+    rate_alternative_responses: Optional[bool] = False
+
+    def asdict(self) -> Dict[str, Any]:
+        actions = []
+        if self.regenerate_conversations_after_incorrect_step:
+            actions.append("regenerateSteps")
+        if self.rate_alternative_responses:
+            actions.append("generateAlternatives")
+        return {"id": self.id, "name": self.name, "actions": actions}
+
+
+@dataclass
+class StepReasoningVariants:
+    CORRECT_STEP_ID = 0
+    NEUTRAL_STEP_ID = 1
+    INCORRECT_STEP_ID = 2
+
+    correct_step: StepReasoningVariant = field(
+        default=StepReasoningVariant(CORRECT_STEP_ID, "Correct"), init=False
+    )
+    neutral_step: StepReasoningVariant = field(
+        default=StepReasoningVariant(NEUTRAL_STEP_ID, "Neutral"), init=False
+    )
+    incorrect_step: IncorrectStepReasoningVariant = field(
+        default=IncorrectStepReasoningVariant(INCORRECT_STEP_ID, "Incorrect"),
+    )
+
+    def asdict(self):
+        return [
+            self.correct_step.asdict(),
+            self.neutral_step.asdict(),
+            self.incorrect_step.asdict(),
+        ]
+
+
+@dataclass
+class StepReasoningDefinition:
+    variants: StepReasoningVariants = field(
+        default_factory=StepReasoningVariants
+    )
+    version: int = field(default=1)
+    title: Optional[str] = None
+    value: Optional[str] = None
+    color: Optional[str] = None
+
+    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
+        if self.color is not None:
+            result["color"] = self.color
+        return result
+
+
+@dataclass
+class StepReasoningTool:
+    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
+    )
+
+    def set_regenerate_conversations_after_incorrect_step(self):
+        self.definition.variants.incorrect_step.regenerate_conversations_after_incorrect_step = True
+
+    def set_rate_alternative_responses(self):
+        self.definition.variants.incorrect_step.rate_alternative_responses = (
+            True
+        )
+
+    def asdict(self) -> Dict[str, Any]:
+        self.set_rate_alternative_responses()
+        self.set_regenerate_conversations_after_incorrect_step()
+        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(**dictionary["definition"]),
+        )

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

@@ -0,0 +1,5 @@
+from enum import Enum
+
+
+class ToolType(Enum):
+    STEP_REASONING = "step-reasoning"

libs/labelbox/tests/integration/conftest.py

@@ -34,6 +34,8 @@
 from labelbox.schema.enums import AnnotationImportState
 from labelbox.schema.invite import Invite
 from labelbox.schema.ontology_kind import OntologyKind
+from labelbox.schema.tool_building.step_reasoning_tool import StepReasoningTool
+from labelbox.schema.tool_building.tool_type import ToolType
 from labelbox.schema.user import User
 
 
@@ -575,6 +577,7 @@ def feature_schema(client, point):
 @pytest.fixture
 def chat_evaluation_ontology(client, rand_gen):
     ontology_name = f"test-chat-evaluation-ontology-{rand_gen(str)}"
+
     ontology_builder = OntologyBuilder(
         tools=[
             Tool(
@@ -589,6 +592,7 @@ def chat_evaluation_ontology(client, rand_gen):
                 tool=Tool.Type.MESSAGE_RANKING,
                 name="model output multi ranking",
             ),
+            StepReasoningTool(name="step reasoning"),
         ],
         classifications=[
             Classification(

libs/labelbox/tests/integration/test_chat_evaluation_ontology_project.py

@@ -17,7 +17,7 @@ def test_create_chat_evaluation_ontology_project(
     # here we are essentially testing the ontology creation which is a fixture
     assert ontology
     assert ontology.name
-    assert len(ontology.tools()) == 3
+    assert len(ontology.tools()) == 4
     for tool in ontology.tools():
         assert tool.schema_id
         assert tool.feature_schema_id

libs/labelbox/tests/unit/test_unit_step_reasoning_tool.py

@@ -0,0 +1,39 @@
+from labelbox.schema.tool_building.step_reasoning_tool import StepReasoningTool
+
+
+def test_step_reasoning_as_dict_default():
+    tool = StepReasoningTool(name="step reasoning")
+    assert tool.asdict() == {
+        "tool": "step-reasoning",
+        "name": "step reasoning",
+        "required": False,
+        "schemaNodeId": None,
+        "featureSchemaId": None,
+        "variants": [
+            {"id": 0, "name": "Correct"},
+            {"id": 1, "name": "Neutral"},
+            {"id": 2, "name": "Incorrect", "actions": []},
+        ],
+    }
+
+
+def test_step_reasoning_as_dict_with_actions():
+    tool = StepReasoningTool(name="step reasoning")
+    tool.set_rate_alternative_responses()
+    tool.set_regenerate_conversations_after_incorrect_step()
+    assert tool.asdict() == {
+        "tool": "step-reasoning",
+        "name": "step reasoning",
+        "required": False,
+        "schemaNodeId": None,
+        "featureSchemaId": None,
+        "variants": [
+            {"id": 0, "name": "Correct"},
+            {"id": 1, "name": "Neutral"},
+            {
+                "id": 2,
+                "name": "Incorrect",
+                "actions": ["regenerateSteps", "generateAlternatives"],
+            },
+        ],
+    }