diff --git a/council/llm/llm_message.py b/council/llm/llm_message.py index 86060f22..722aec48 100644 --- a/council/llm/llm_message.py +++ b/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 diff --git a/council/prompt/__init__.py b/council/prompt/__init__.py index 03068db5..aa9a89cf 100644 --- a/council/prompt/__init__.py +++ b/council/prompt/__init__.py @@ -1,5 +1,4 @@ from .llm_dataset import ( - LLMDatasetMessage, LLMDatasetConversation, LLMDatasetObject, LLMDatasetSpec, diff --git a/council/prompt/llm_dataset.py b/council/prompt/llm_dataset.py index 77c1f79d..7a5c3a87 100644 --- a/council/prompt/llm_dataset.py +++ b/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) diff --git a/docs/source/reference/llm/llm_dataset.md b/docs/source/reference/llm/llm_dataset.md index 083cd275..e27cebed 100644 --- a/docs/source/reference/llm/llm_dataset.md +++ b/docs/source/reference/llm/llm_dataset.md @@ -1,9 +1,3 @@ -# LLMDatasetMessage - -```{eval-rst} -.. autoclass:: council.prompt.LLMDatasetMessage -``` - # LLMDatasetConversation ```{eval-rst}