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

@@ -353,7 +353,7 @@ class Type(Enum):
     @classmethod
     def from_dict(cls, dictionary: Dict[str, Any]) -> Dict[str, Any]:
         return cls(
-            class_type=Type(dictionary["type"]),
+            class_type=PromptResponseClassification.Type(dictionary["type"]),
             name=dictionary["name"],
             instructions=dictionary["instructions"],
             required=True,  # always required
@@ -492,14 +492,16 @@ def add_classification(self, classification: Classification) -> None:
 
 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
 
 
+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.
@@ -533,12 +535,6 @@ 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:

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

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

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
 
@@ -25,9 +25,40 @@ def asdict(self) -> Dict[str, Any]:
         if self.regenerate_conversations_after_incorrect_step:
             actions.append("regenerateSteps")
         if self.rate_alternative_responses:
-            actions.append("generateAlternatives")
+            actions.append("generateAndRateAlternativeSteps")
         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="generateAndRateAlternativeSteps"
+            in dictionary.get("actions", []),
+        )
+
+
+def _create_correct_step() -> StepReasoningVariant:
+    return StepReasoningVariant(
+        id=StepReasoningVariants.CORRECT_STEP_ID, name="Correct"
+    )
+
+
+def _create_neutral_step() -> StepReasoningVariant:
+    return StepReasoningVariant(
+        id=StepReasoningVariants.NEUTRAL_STEP_ID, name="Neutral"
+    )
+
+
+def _create_incorrect_step() -> IncorrectStepReasoningVariant:
+    return IncorrectStepReasoningVariant(
+        id=StepReasoningVariants.INCORRECT_STEP_ID, name="Incorrect"
+    )
+
 
 @dataclass
 class StepReasoningVariants:
@@ -36,13 +67,13 @@ class StepReasoningVariants:
     INCORRECT_STEP_ID = 2
 
     correct_step: StepReasoningVariant = field(
-        default=StepReasoningVariant(CORRECT_STEP_ID, "Correct"), init=False
+        default_factory=_create_correct_step
     )
     neutral_step: StepReasoningVariant = field(
-        default=StepReasoningVariant(NEUTRAL_STEP_ID, "Neutral"), init=False
+        default_factory=_create_neutral_step
     )
     incorrect_step: IncorrectStepReasoningVariant = field(
-        default=IncorrectStepReasoningVariant(INCORRECT_STEP_ID, "Incorrect"),
+        default_factory=_create_incorrect_step
     )
 
     def asdict(self):
@@ -52,6 +83,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:
@@ -61,18 +117,22 @@ class StepReasoningDefinition:
     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
 
+    @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)
+        return cls(variants=variants, title=title, value=value)
+
 
 @dataclass
 class StepReasoningTool:
@@ -113,5 +173,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,29 @@ 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",
+                            "generateAndRateAlternativeSteps",
+                        ],
+                    },
+                ],
+                "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", "generateAndRateAlternativeSteps"],
+        },
+    ]
+    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", "generateAndRateAlternativeSteps"],
+        },
+    ]

libs/labelbox/tests/unit/test_unit_step_reasoning_tool.py

@@ -9,11 +9,21 @@ 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",
+                        "generateAndRateAlternativeSteps",
+                    ],
+                },
+            ],
+            "version": 1,
+        },
     }
 
 
@@ -27,13 +37,19 @@ 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",
+                        "generateAndRateAlternativeSteps",
+                    ],
+                },
+            ],
+            "version": 1,
+        },
     }