Feature refactor LLMDatasetMessage

council/llm/llm_message.py

@@ -4,7 +4,7 @@
 import base64
 import mimetypes
 from enum import Enum
-from typing import Iterable, List, Optional, Sequence
+from typing import Dict, Iterable, List, Optional, Sequence
 
 from council.contexts import ChatMessage, ChatMessageKind
 
@@ -248,6 +248,20 @@ def normalize(self) -> str:
 
         return normalized_content
 
+    @classmethod
+    def from_dict(cls, values: Dict[str, str]) -> LLMMessage:
+        """Create an instance from OpenAI-compatible dict with role and content (name and data not supported)."""
+
+        role = values.get("role")
+        content = values.get("content")
+        if role is None or content is None:
+            raise ValueError("Both 'role' and 'content' must be defined for LLMMessage")
+        return LLMMessage(LLMMessageRole(role), content.strip())
+
+    def to_dict(self) -> Dict[str, str]:
+        """Convert an instance to OpenAI-compatible dict with role and content (name and data not supported)."""
+        return {"role": self.role, "content": self.content}
+
 
 class LLMMessageTokenCounterBase(abc.ABC):
     @abc.abstractmethod

council/prompt/__init__.py

@@ -1,5 +1,4 @@
 from .llm_dataset import (
-    LLMDatasetMessage,
     LLMDatasetConversation,
     LLMDatasetObject,
     LLMDatasetSpec,

council/prompt/llm_dataset.py

@@ -6,41 +6,16 @@
 from typing import Any, Counter, DefaultDict, Dict, List, Mapping, Optional, Sequence
 
 import yaml
-from council.llm import LLMMessage, LLMMessageRole
+from council.llm import LLMMessage
 from council.utils import DataObject, DataObjectSpecBase
 
 
-class LLMDatasetMessage:
-    """
-    Represents a single chat message in a conversation.
-    """
-
-    def __init__(self, role: LLMMessageRole, content: str):
-        self.role = role
-        self.content = content.strip()
-
-    @classmethod
-    def from_dict(cls, values: Dict[str, str]) -> LLMDatasetMessage:
-        role = values.get("role")
-        content = values.get("content")
-        if role is None or content is None:
-            raise ValueError("Both 'role' and 'content' must be defined for a message")
-        return LLMDatasetMessage(LLMMessageRole(role), content)
-
-    @classmethod
-    def from_llm_message(cls, message: LLMMessage) -> LLMDatasetMessage:
-        return LLMDatasetMessage(role=message.role, content=message.content)
-
-    def to_dict(self) -> Dict[str, str]:
-        return {"role": self.role, "content": self.content}
-
-
 class LLMDatasetConversation:
     """
     Represents a conversation between user and assistant with optional labels.
     """
 
-    def __init__(self, messages: Sequence[LLMDatasetMessage], labels: Optional[Mapping[str, str]]):
+    def __init__(self, messages: Sequence[LLMMessage], labels: Optional[Mapping[str, str]]):
         self.messages = list(messages)
         self.labels: Dict[str, str] = dict(labels) if labels is not None else {}
 
@@ -49,7 +24,7 @@ def from_dict(cls, values: Dict[str, Any]) -> LLMDatasetConversation:
         messages = values.get("messages", [])
         if not messages:
             raise ValueError("Conversation must contain at least one message")
-        llm_dataset_messages = [LLMDatasetMessage.from_dict(message) for message in messages]
+        llm_dataset_messages = [LLMMessage.from_dict(message) for message in messages]
         labels = values.get("labels")
         return LLMDatasetConversation(llm_dataset_messages, labels)
 

docs/source/reference/llm/llm_dataset.md

@@ -1,9 +1,3 @@
-# LLMDatasetMessage
-
-```{eval-rst}
-.. autoclass:: council.prompt.LLMDatasetMessage
-```
-
 # LLMDatasetConversation
 
 ```{eval-rst}