Update to support create_ontology_from_feature_schemas

Author: Val Brodsky

libs/labelbox/src/labelbox/client.py

@@ -57,7 +57,7 @@
     FeatureSchema,
     Ontology,
     PromptResponseClassification,
-    Tool,
+    tool_type_cls_from_type,
 )
 from labelbox.schema.ontology_kind import (
     EditorTaskType,
@@ -1106,7 +1106,8 @@ def create_ontology_from_feature_schemas(
             if "tool" in feature_schema.normalized:
                 tool = feature_schema.normalized["tool"]
                 try:
-                    Tool.Type(tool)
+                    tool_type_cls = tool_type_cls_from_type(tool)
+                    tool_type_cls(tool)
                     tools.append(feature_schema.normalized)
                 except ValueError:
                     raise ValueError(

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

@@ -500,6 +500,12 @@ def tool_cls_from_type(tool_type: str):
     return Tool
 
 
+def tool_type_cls_from_type(tool_type: str):
+    if tool_type.lower() == ToolType.STEP_REASONING.value:
+        return ToolType
+    return Tool.Type
+
+
 class Ontology(DbObject):
     """An ontology specifies which tools and classifications are available
     to a project. This is read only for now.

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

@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional
 
 from labelbox.schema.tool_building.tool_type import ToolType
 
@@ -28,6 +28,19 @@ def asdict(self) -> Dict[str, Any]:
             actions.append("generateAlternatives")
         return {"id": self.id, "name": self.name, "actions": actions}
 
+    @classmethod
+    def from_dict(
+        cls, dictionary: Dict[str, Any]
+    ) -> "IncorrectStepReasoningVariant":
+        return cls(
+            id=dictionary["id"],
+            name=dictionary["name"],
+            regenerate_conversations_after_incorrect_step="regenerateSteps"
+            in dictionary.get("actions", []),
+            rate_alternative_responses="generateAlternatives"
+            in dictionary.get("actions", []),
+        )
+
 
 @dataclass
 class StepReasoningVariants:
@@ -36,10 +49,10 @@ class StepReasoningVariants:
     INCORRECT_STEP_ID = 2
 
     correct_step: StepReasoningVariant = field(
-        default=StepReasoningVariant(CORRECT_STEP_ID, "Correct"), init=False
+        default=StepReasoningVariant(CORRECT_STEP_ID, "Correct")
     )
     neutral_step: StepReasoningVariant = field(
-        default=StepReasoningVariant(NEUTRAL_STEP_ID, "Neutral"), init=False
+        default=StepReasoningVariant(NEUTRAL_STEP_ID, "Neutral")
     )
     incorrect_step: IncorrectStepReasoningVariant = field(
         default=IncorrectStepReasoningVariant(INCORRECT_STEP_ID, "Incorrect"),
@@ -52,6 +65,31 @@ def asdict(self):
             self.incorrect_step.asdict(),
         ]
 
+    @classmethod
+    def from_dict(cls, dictionary: List[Dict[str, Any]]):
+        correct_step = None
+        neutral_step = None
+        incorrect_step = None
+
+        for variant in dictionary:
+            if variant["id"] == cls.CORRECT_STEP_ID:
+                correct_step = StepReasoningVariant(**variant)
+            elif variant["id"] == cls.NEUTRAL_STEP_ID:
+                neutral_step = StepReasoningVariant(**variant)
+            elif variant["id"] == cls.INCORRECT_STEP_ID:
+                incorrect_step = IncorrectStepReasoningVariant.from_dict(
+                    variant
+                )
+
+        if not all([correct_step, neutral_step, incorrect_step]):
+            raise ValueError("Invalid step reasoning variants")
+
+        return cls(
+            correct_step=correct_step,  # type: ignore
+            neutral_step=neutral_step,  # type: ignore
+            incorrect_step=incorrect_step,  # type: ignore
+        )
+
 
 @dataclass
 class StepReasoningDefinition:
@@ -73,6 +111,14 @@ def asdict(self) -> Dict[str, Any]:
             result["color"] = self.color
         return result
 
+    @classmethod
+    def from_dict(cls, dictionary: Dict[str, Any]) -> "StepReasoningDefinition":
+        variants = StepReasoningVariants.from_dict(dictionary["variants"])
+        title = dictionary.get("title", None)
+        value = dictionary.get("value", None)
+        color = dictionary.get("color", None)
+        return cls(variants=variants, title=title, value=value, color=color)
+
 
 @dataclass
 class StepReasoningTool:
@@ -113,5 +159,7 @@ def from_dict(cls, dictionary: Dict[str, Any]) -> "StepReasoningTool":
             schema_id=dictionary.get("schemaNodeId", None),
             feature_schema_id=dictionary.get("featureSchemaId", None),
             required=dictionary.get("required", False),
-            definition=StepReasoningDefinition(**dictionary["definition"]),
+            definition=StepReasoningDefinition.from_dict(
+                dictionary["definition"]
+            ),
         )

libs/labelbox/tests/integration/conftest.py

@@ -643,14 +643,12 @@ def chat_evaluation_ontology(client, rand_gen):
             ),
         ],
     )
-
     ontology = client.create_ontology(
         ontology_name,
         ontology_builder.asdict(),
         media_type=MediaType.Conversational,
         ontology_kind=OntologyKind.ModelEvaluation,
     )
-
     yield ontology
 
     try:

libs/labelbox/tests/integration/test_chat_evaluation_ontology_project.py

@@ -43,7 +43,7 @@ def test_create_chat_evaluation_ontology_project(
 
 
 def test_create_chat_evaluation_ontology_project_existing_dataset(
-    client, chat_evaluation_ontology, chat_evaluation_project_append_to_dataset
+    chat_evaluation_ontology, chat_evaluation_project_append_to_dataset
 ):
     ontology = chat_evaluation_ontology
 
@@ -85,6 +85,26 @@ def tools_json():
             "schemaNodeId": None,
             "featureSchemaId": None,
         },
+        {
+            "tool": "step-reasoning",
+            "name": "step reasoning",
+            "required": True,
+            "schemaNodeId": None,
+            "featureSchemaId": None,
+            "color": "#0000ff",
+            "definition": {
+                "variants": [
+                    {"id": 0, "name": "Correct"},
+                    {"id": 1, "name": "Neutral"},
+                    {
+                        "id": 2,
+                        "name": "Incorrect",
+                        "actions": ["regenerateSteps", "generateAlternatives"],
+                    },
+                ],
+                "version": 1,
+            },
+        },
     ]
 
     return tools

libs/labelbox/tests/integration/test_offline_chat_evaluation_project.py

@@ -2,7 +2,6 @@
 
 
 def test_create_offline_chat_evaluation_project(
-    client,
     rand_gen,
     offline_chat_evaluation_project,
     chat_evaluation_ontology,

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.step_reasoning_tool import StepReasoningTool
 
 
 def test_feature_schema_is_not_archived(client, ontology):
@@ -322,3 +323,47 @@ def test_unarchive_feature_schema_node_for_non_existing_ontology(
         client.unarchive_feature_schema_node(
             "invalid-ontology", feature_schema_to_unarchive["featureSchemaId"]
         )
+
+
+def test_step_reasoning_ontology(chat_evaluation_ontology):
+    ontology = chat_evaluation_ontology
+    step_reasoning_tool = None
+    for tool in ontology.normalized["tools"]:
+        if tool["tool"] == "step-reasoning":
+            step_reasoning_tool = tool
+            break
+    assert step_reasoning_tool is not None
+    assert step_reasoning_tool["definition"]["variants"] == [
+        {"id": 0, "name": "Correct"},
+        {"id": 1, "name": "Neutral"},
+        {
+            "id": 2,
+            "name": "Incorrect",
+            "actions": ["regenerateSteps", "generateAlternatives"],
+        },
+    ]
+    assert step_reasoning_tool["definition"]["version"] == 1
+    assert step_reasoning_tool["schemaNodeId"] is not None
+    assert step_reasoning_tool["featureSchemaId"] is not None
+
+    step_reasoning_tool = None
+    for tool in ontology.tools():
+        if isinstance(tool, StepReasoningTool):
+            step_reasoning_tool = tool
+            break
+    assert step_reasoning_tool is not None
+    assert step_reasoning_tool.definition.variants.asdict() == [
+        {
+            "id": 0,
+            "name": "Correct",
+        },
+        {
+            "id": 1,
+            "name": "Neutral",
+        },
+        {
+            "id": 2,
+            "name": "Incorrect",
+            "actions": ["regenerateSteps", "generateAlternatives"],
+        },
+    ]

libs/labelbox/tests/unit/test_unit_step_reasoning_tool.py

@@ -9,11 +9,18 @@ def test_step_reasoning_as_dict_default():
         "required": False,
         "schemaNodeId": None,
         "featureSchemaId": None,
-        "variants": [
-            {"id": 0, "name": "Correct"},
-            {"id": 1, "name": "Neutral"},
-            {"id": 2, "name": "Incorrect", "actions": []},
-        ],
+        "definition": {
+            "variants": [
+                {"id": 0, "name": "Correct"},
+                {"id": 1, "name": "Neutral"},
+                {
+                    "id": 2,
+                    "name": "Incorrect",
+                    "actions": ["regenerateSteps", "generateAlternatives"],
+                },
+            ],
+            "version": 1,
+        },
     }
 
 
@@ -27,13 +34,16 @@ def test_step_reasoning_as_dict_with_actions():
         "required": False,
         "schemaNodeId": None,
         "featureSchemaId": None,
-        "variants": [
-            {"id": 0, "name": "Correct"},
-            {"id": 1, "name": "Neutral"},
-            {
-                "id": 2,
-                "name": "Incorrect",
-                "actions": ["regenerateSteps", "generateAlternatives"],
-            },
-        ],
+        "definition": {
+            "variants": [
+                {"id": 0, "name": "Correct"},
+                {"id": 1, "name": "Neutral"},
+                {
+                    "id": 2,
+                    "name": "Incorrect",
+                    "actions": ["regenerateSteps", "generateAlternatives"],
+                },
+            ],
+            "version": 1,
+        },
     }