fix: BROS-392: Add Chat tag to label interface + fix prediction valid… (#584)

Co-authored-by: hakan458 <[email protected]>

Author: hakan458

src/label_studio_sdk/label_interface/control_tags.py

@@ -39,6 +39,7 @@
     "taxonomy": "TaxonomyTag",
     "textarea": "TextAreaTag",
     "timeserieslabels": "TimeSeriesLabelsTag",
+    "chatmessage": "ChatMessageTag",
 }
 
 
@@ -950,6 +951,43 @@ def to_json_schema(self):
         }
 
 
+class ChatMessageContent(BaseModel):
+    role: str
+    content: str
+    createdAt: Optional[int] = None
+
+
+class ChatMessageValue(BaseModel):
+    chatmessage: ChatMessageContent
+
+
+class ChatMessageTag(ControlTag):
+    """Control tag for chat messages targeting a `<Chat>` object.
+
+    This tag is a hybrid where `from_name == to_name` and `type == 'chatmessage'`.
+    """
+    tag: str = "ChatMessage"
+    _value_class: Type[ChatMessageValue] = ChatMessageValue
+
+    def to_json_schema(self):
+        return {
+            "type": "object",
+            "required": ["chatmessage"],
+            "properties": {
+                "chatmessage": {
+                    "type": "object",
+                    "required": ["role", "content"],
+                    "properties": {
+                        "role": {"type": "string"},
+                        "content": {"type": "string"},
+                        "createdAt": {"type": "number"}
+                    }
+                }
+            },
+            "description": f"Chat message for {self.to_name[0]}"
+        }
+
+
 class RelationsTag(ControlTag):
     """ """
     tag: str = "Relations"

src/label_studio_sdk/label_interface/data_examples.json

@@ -10,6 +10,12 @@
     "Header": "Task header",
     "Paragraphs":  [{"author": "Alice", "text": "Hi, Bob."}, {"author": "Bob", "text": "Hello, Alice!"}, {"author": "Alice", "text": "What's up?"}, {"author": "Bob", "text": "Good. Ciao!"}, {"author": "Alice", "text": "Bye, Bob."}],
     "ParagraphsUrl": "<HOSTNAME>/samples/paragraphs.json?",
+    "Chat": [
+      {"role": "user", "content": "Hi there!", "createdAt": 1710000000000},
+      {"role": "assistant", "content": "Hello! How can I help you today?", "createdAt": 1710000005000},
+      {"role": "user", "content": "Can you summarize our onboarding checklist?", "createdAt": 1710000010000},
+      {"role": "assistant", "content": "Sure. It includes account setup, roles, labeling instructions, GT tasks, overlap and review.", "createdAt": 1710000015000}
+    ],
     "Table": {"Card number": 18799210, "First name": "Max", "Last name":  "Nobel"},
     "$videoHack": "<video src='<HOSTNAME>/static/samples/opossum_snow.mp4' width=100% controls>",
     "Video": "<HOSTNAME>/static/samples/opossum_snow.mp4",
@@ -85,6 +91,10 @@
     "AudioPlus": "<HOSTNAME>/static/samples/game.wav",
     "Header": "Task header",
     "Paragraphs":  [{"author": "Alice", "text": "Hi, Bob."}, {"author": "Bob", "text": "Hello, Alice!"}, {"author": "Alice", "text": "What's up?"}, {"author": "Bob", "text": "Good. Ciao!"}, {"author": "Alice", "text": "Bye, Bob."}],
+    "Chat": [
+      {"role": "user", "content": "Hello!", "createdAt": 1710000000000},
+      {"role": "assistant", "content": "Hi! What can I do for you?", "createdAt": 1710000004000}
+    ],
     "Table": {"Card number": 18799210, "First name": "Max", "Last name":  "Nobel"},
     "$videoHack": "<video src='static/samples/opossum_snow.mp4' width=100% controls>",
     "Video": "<HOSTNAME>/static/samples/opossum_snow.mp4",

src/label_studio_sdk/label_interface/interface.py

@@ -27,6 +27,7 @@
 
 from .base import LabelStudioTag
 from .control_tags import (
+    ChatMessageTag,
     ControlTag,
     ChoicesTag,
     LabelsTag,
@@ -623,6 +624,18 @@ def parse(self, config_string: str) -> Tuple[Dict, Dict, Dict, etree._Element]:
                 if lb:
                     labels[lb.parent_name][lb.value] = lb
 
+        # Special handling: auto-create ChatMessage control for each Chat object
+        chat_object_names = [name for name, obj in objects.items() if getattr(obj, 'tag', '').lower() == 'chat']
+        for name in chat_object_names:
+            if name not in controls:
+                controls[name] = ChatMessageTag(
+                    tag='ChatMessage',
+                    name=name,
+                    to_name=[name],
+                    attr={"name": name, "toName": name}
+                )
+
+
         return controls, objects, labels, xml_tree
 
     @classmethod

src/label_studio_sdk/label_interface/object_tags.py

@@ -20,6 +20,7 @@
     "list": "ListTag",
     "paragraphs": "ParagraphsTag",
     "timeseries": "TimeSeriesTag",
+    "chat": "ChatTag",
 }
 
 _DATA_EXAMPLES = None
@@ -306,3 +307,11 @@ def _generate_example(self, examples, only_urls=False):
         else:
             # data is JSON
             return generate_time_series_json(time_column, value_columns, time_format)
+
+class ChatTag(ObjectTag):
+    """ """
+    tag: str = "Chat"
+
+    def _generate_example(self, examples, only_urls=False):
+        """ """
+        return examples.get("Chat")