# Copyright 2025 - Oumi
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import base64
import json
from collections.abc import Callable, Generator
from enum import Enum
from typing import Any, NamedTuple
import pydantic
from jinja2 import Template
from transformers.utils.chat_template_utils import get_json_schema
from oumi.core.types.tool_call import ToolCall, ToolDefinition
[docs]
class Role(str, Enum):
"""Role of the entity sending the message."""
SYSTEM = "system"
"""Represents a system message in the conversation."""
USER = "user"
"""Represents a user message in the conversation."""
ASSISTANT = "assistant"
"""Represents an assistant message in the conversation."""
TOOL = "tool"
"""Represents a tool message in the conversation."""
[docs]
def __str__(self) -> str:
"""Return the string representation of the Role enum.
Returns:
str: The string value of the Role enum.
"""
return self.value
[docs]
class Type(str, Enum):
"""Type of the message."""
TEXT = "text"
"""Represents a text message."""
IMAGE_PATH = "image_path"
"""Represents an image referenced by its file path."""
IMAGE_URL = "image_url"
"""Represents an image referenced by its URL."""
IMAGE_BINARY = "image_binary"
"""Represents an image stored as binary data."""
[docs]
def __str__(self) -> str:
"""Return the string representation of the Type enum.
Returns:
str: The string value of the Type enum.
"""
return self.value
[docs]
class FinishReason(str, Enum):
"""Reason why the model stopped generating tokens."""
STOP = "stop"
"""Model hit a natural stopping point or stop sequence."""
LENGTH = "length"
"""Model reached max_tokens limit."""
TOOL_CALLS = "tool_calls"
"""Model made a tool/function call."""
CONTENT_FILTER = "content_filter"
"""Content was filtered due to safety/moderation."""
ERROR = "error"
"""Generation failed due to an error."""
UNKNOWN = "unknown"
"""Finish reason could not be determined."""
[docs]
def __str__(self) -> str:
"""Return the string representation of the FinishReason enum.
Returns:
str: The string value of the FinishReason enum.
"""
return self.value
[docs]
class ContentItemCounts(NamedTuple):
"""Contains counts of content items in a message by type."""
total_items: int
"""The total number of content items in a message."""
text_items: int
"""The number of text content items in a message."""
image_items: int
"""The number of image content items in a message."""
[docs]
class ContentItem(pydantic.BaseModel):
"""A sub-part of `Message.content`.
For example, a multimodal message from `USER` may include
two `ContentItem`-s: one for text, and another for image.
Note:
Either content or binary must be provided when creating an instance.
"""
model_config = pydantic.ConfigDict(frozen=True)
type: Type
"""The type of the content (e.g., text, image path, image URL)."""
content: str | None = None
"""Optional text content of the content item.
One of content or binary must be provided.
"""
binary: bytes | None = None
"""Optional binary data for the message content item, used for image data.
One of content or binary must be provided.
The field is required for `IMAGE_BINARY`, and can be optionally populated for
`IMAGE_URL`, `IMAGE_PATH` in which case it must be the loaded bytes of
the image specified in the `content` field.
The field must be `None` for `TEXT`.
"""
[docs]
def is_image(self) -> bool:
"""Checks if the item contains an image."""
return self.type in (Type.IMAGE_BINARY, Type.IMAGE_URL, Type.IMAGE_PATH)
[docs]
def is_text(self) -> bool:
"""Checks if the item contains text."""
return self.type == Type.TEXT
@pydantic.field_serializer("binary")
def _encode_binary(self, value: bytes | None) -> str:
"""Encode binary value as base64 ASCII string.
This is needed for compatibility with JSON.
"""
if value is None or len(value) == 0:
return ""
return base64.b64encode(value).decode("ascii")
@pydantic.field_validator("binary", mode="before")
def _decode_binary(cls, value: str | bytes | None) -> bytes | None:
if value is None:
return None
elif isinstance(value, str):
return base64.b64decode(value)
return value
[docs]
def model_post_init(self, __context) -> None:
"""Post-initialization method for the `ContentItem` model.
This method is automatically called after the model is initialized.
Performs additional validation e.g., to ensure that either content or binary
is provided for the message.
Raises:
ValueError: If fields are set to invalid or inconsistent values.
"""
if self.binary is None and self.content is None:
raise ValueError(
"Either content or binary must be provided for the message item "
f"(Item type: {self.type})."
)
if self.is_image():
if self.type == Type.IMAGE_BINARY and (
self.binary is None or len(self.binary) == 0
):
raise ValueError(
f"No image bytes in message content item (Item type: {self.type})."
)
if self.type in (Type.IMAGE_PATH, Type.IMAGE_URL) and (
self.content is None or len(self.content) == 0
):
raise ValueError(f"Content not provided for {self.type} message item.")
else:
if self.binary is not None:
raise ValueError(
f"Binary can only be provided for images (Item type: {self.type})."
)
[docs]
def __repr__(self) -> str:
"""Returns a string representation of the message item."""
return f"{self.content}" if self.is_text() else f"<{self.type.upper()}>"
[docs]
class Message(pydantic.BaseModel):
"""A message in a conversation.
This class represents a single message within a conversation, containing
various attributes such as role, content, identifier.
"""
model_config = pydantic.ConfigDict(frozen=True)
id: str | None = None
"""Optional unique identifier for the message.
This attribute can be used to assign a specific identifier to the message,
which may be useful for tracking or referencing messages within a conversation.
Returns:
Optional[str]: The unique identifier of the message, if set; otherwise None.
"""
content: str | list[ContentItem] | None = None
"""Content of the message.
For text messages, `content` can be set to a string value.
For multimodal messages, `content` should be a list of content items of
potentially different types e.g., text and image.
May be ``None`` on assistant messages that only contain ``tool_calls``
(OpenAI tool-calling wire format).
"""
role: Role
"""The role of the entity sending the message (e.g., user, assistant, system)."""
tool_calls: list[ToolCall] | None = None
"""Structured tool calls emitted by an assistant message.
Uses the OpenAI function-calling wire format. Pydantic auto-coerces
dict input (e.g., from JSONL) into ``ToolCall`` instances::
[{"id": "call_abc", "type": "function",
"function": {"name": "get_weather", "arguments": "{...}"}}]
Only set on assistant messages. The chat template renders these via
``tokenizer.apply_chat_template(messages, tools=...)``; serialization
(``to_dict()`` / ``to_json()``) emits them in the OpenAI dict shape.
"""
tool_call_id: str | None = None
"""Identifier linking a tool response message to the call it responds to.
Only set on messages with ``role == 'tool'``, matching the ``id`` of a
prior assistant ``tool_calls`` entry.
"""
reasoning_content: str | None = None
"""Model-emitted reasoning / chain-of-thought content.
Populated by inference engines that surface a separate reasoning field
in the API response (e.g., Fireworks ``reasoning_content``, Together
``reasoning``). Distinct from ``content``: this is the model's
internal deliberation, not the final answer.
Only set on assistant messages.
"""
[docs]
def model_post_init(self, __context) -> None:
"""Post-initialization method for the Message model.
This method is automatically called after the model is initialized.
It validates that the message has at least one of ``content`` or
``tool_calls`` and that ``content``, when set, is a string or list.
Raises:
ValueError: If ``content`` is ``None`` and ``tool_calls`` is
missing or empty, or if ``content`` has an unsupported type.
"""
if self.content is None and not self.tool_calls:
raise ValueError(
"Message must have at least one of `content` (string/list) "
"or a non-empty `tool_calls`. An empty list of tool calls "
"is not valid."
)
if self.content is not None and not isinstance(self.content, str | list):
raise ValueError(
f"Unexpected content type: {type(self.content)}. "
f"Must be a Python string, a list, or None."
)
[docs]
def get(self, key: str, default: Any = None) -> Any:
"""Returns a field by name, dict-style, or ``default`` if absent.
Some chat templates (e.g. Gemma 4's) read message fields via ``.get()``, so
``Message`` must support it when passed to ``apply_chat_template``. Only
declared fields are exposed, so a key that collides with a method name
(e.g. ``"get"``) returns ``default`` rather than the bound method.
"""
if key in type(self).model_fields:
return getattr(self, key)
return default
def _iter_content_items(
self, *, return_text: bool = False, return_images: bool = False
) -> Generator[ContentItem, None, None]:
"""Returns a list of content items."""
if isinstance(self.content, str):
if return_text:
yield ContentItem(type=Type.TEXT, content=self.content)
elif isinstance(self.content, list):
if return_text and return_images:
yield from self.content
else:
for item in self.content:
if (return_text and item.is_text()) or (
return_images and item.is_image()
):
yield item
def _iter_all_content_items(self) -> Generator[ContentItem, None, None]:
return self._iter_content_items(return_text=True, return_images=True)
[docs]
def count_content_items(self) -> ContentItemCounts:
"""Counts content items by type."""
total_items: int = 0
num_text_items: int = 0
num_image_items: int = 0
for item in self._iter_all_content_items():
total_items += 1
if item.is_text():
num_text_items += 1
elif item.is_image():
num_image_items += 1
return ContentItemCounts(
total_items=total_items,
text_items=num_text_items,
image_items=num_image_items,
)
@property
def content_items(self) -> list[ContentItem]:
"""Returns a list of text content items."""
return [item for item in self._iter_all_content_items()]
@property
def image_content_items(self) -> list[ContentItem]:
"""Returns a list of image content items."""
return [item for item in self._iter_content_items(return_images=True)]
@property
def text_content_items(self) -> list[ContentItem]:
"""Returns a list of text content items."""
return [item for item in self._iter_content_items(return_text=True)]
[docs]
def compute_flattened_text_content(self, separator=" ") -> str:
"""Joins contents of all text items."""
return separator.join(
[(item.content or "") for item in self.text_content_items]
)
[docs]
def contains_images(self) -> bool:
"""Checks if the message contains at least one image."""
first_image = next(self._iter_content_items(return_images=True), None)
return first_image is not None
[docs]
def contains_text(self) -> bool:
"""Checks if the message contains at least one text item."""
first_text = next(self._iter_content_items(return_text=True), None)
return first_text is not None
[docs]
def contains_image_content_items_only(self) -> bool:
"""Checks if the message contains only image items.
At least one image item is required.
"""
counts = self.count_content_items()
return counts.image_items > 0 and counts.image_items == counts.total_items
[docs]
def contains_text_content_items_only(self) -> bool:
"""Checks if the message contains only text items.
At least one text item is required.
"""
counts = self.count_content_items()
return counts.text_items > 0 and counts.text_items == counts.total_items
[docs]
def contains_single_text_content_item_only(self) -> bool:
"""Checks if the message contains exactly 1 text item, and nothing else.
These are the most common and simple messages, and may need special handling.
"""
counts = self.count_content_items()
return counts.text_items == 1 and counts.text_items == counts.total_items
[docs]
def contains_single_image_content_item_only(self) -> bool:
"""Checks if the message contains exactly 1 image item, and nothing else."""
counts = self.count_content_items()
return counts.image_items == 1 and counts.image_items == counts.total_items
[docs]
def __repr__(self) -> str:
"""Returns a string representation of the message."""
id_str = ""
if self.id:
id_str = f"{self.id} - "
return f"{id_str}{self.role.upper()}: " + " | ".join(
[repr(item) for item in self._iter_all_content_items()]
)
[docs]
class Conversation(pydantic.BaseModel):
"""Represents a conversation, which is a sequence of messages."""
conversation_id: str | None = None
"""Optional unique identifier for the conversation.
This attribute can be used to assign a specific identifier to the conversation,
which may be useful for tracking or referencing conversations in a larger context.
"""
messages: list[Message]
"""List of Message objects that make up the conversation."""
metadata: dict[str, Any] = pydantic.Field(default_factory=dict)
"""Optional metadata associated with the conversation.
This attribute allows for storing additional information about the conversation
in a key-value format. It can be used to include any relevant contextual data.
"""
tools: list[ToolDefinition] | None = None
"""Tool definitions available to the model for this conversation.
Accepts ``ToolDefinition`` instances, OpenAI-format dicts, or Python
callables on input. Callables are converted via
``transformers.utils.chat_template_utils.get_json_schema`` (which
requires a Google-style docstring and type hints on every user-facing
argument) and dicts are validated by Pydantic into ``ToolDefinition``.
After construction, every entry is a ``ToolDefinition`` instance.
Uses the OpenAI function-calling schema, e.g.::
[{"type": "function",
"function": {"name": "get_weather",
"description": "...",
"parameters": {...}}}]
When set, ``to_dict()`` emits ``tools`` as a top-level key (in dict
form, via Pydantic's ``model_dump(mode="json")``), which TRL's
``SFTTrainer`` forwards to ``tokenizer.apply_chat_template(messages, tools=...)``
so the model's chat template can render tools natively.
"""
@pydantic.field_validator("tools", mode="before")
@classmethod
def _coerce_tools_input(cls, raw_tools: list[Any] | None) -> list[Any] | None:
"""Normalize each tool input to a shape Pydantic can validate.
Runs in ``mode="before"`` so we can intercept callables (which
Pydantic doesn't know how to validate into ``ToolDefinition``)
and convert them via ``get_json_schema``. ``ToolDefinition``
instances and dicts pass through; Pydantic's main validator
coerces dicts into ``ToolDefinition`` after this hook returns.
"""
if raw_tools is None:
return None
coerced_tools: list[Any] = []
for tool in raw_tools:
if isinstance(tool, ToolDefinition) or isinstance(tool, dict):
coerced_tools.append(tool)
elif callable(tool):
coerced_tools.append(get_json_schema(tool))
else:
raise ValueError(
f"Each tool must be a ToolDefinition, dict (OpenAI "
f"schema), or a callable; got {type(tool).__name__}."
)
return coerced_tools
[docs]
def __getitem__(self, idx: int) -> Message:
"""Gets the message at the specified index.
Args:
idx (int): The index of the message to retrieve.
Returns:
Any: The message at the specified index.
"""
return self.messages[idx]
[docs]
def first_message(self, role: Role | None = None) -> Message | None:
"""Gets the first message in the conversation, optionally filtered by role.
Args:
role: The role to filter messages by.
If None, considers all messages.
Returns:
Optional[Message]: The first message matching the criteria,
or None if no messages are found.
"""
messages = self.filter_messages(role=role)
return messages[0] if len(messages) > 0 else None
[docs]
def last_message(self, role: Role | None = None) -> Message | None:
"""Gets the last message in the conversation, optionally filtered by role.
Args:
role: The role to filter messages by.
If None, considers all messages.
Returns:
Optional[Message]: The last message matching the criteria,
or None if no messages are found.
"""
messages = self.filter_messages(role=role)
return messages[-1] if len(messages) > 0 else None
[docs]
def filter_messages(
self,
*,
role: Role | None = None,
filter_fn: Callable[[Message], bool] | None = None,
) -> list[Message]:
"""Gets all messages in the conversation, optionally filtered by role.
Args:
role (Optional): The role to filter messages by. If None, no filtering
by role is applied.
filter_fn (Optional): A predicate to filter messages by. If the predicate
returns True for a message, then the message is returned.
Otherwise, the message is excluded.
Returns:
List[Message]: A list of all messages matching the criteria.
"""
if role is not None:
messages = [message for message in self.messages if role == message.role]
else:
messages = self.messages
if filter_fn is not None:
messages = [message for message in messages if filter_fn(message)]
return messages
[docs]
def to_dict(self):
"""Converts the conversation to a dictionary."""
data = self.model_dump(
mode="json", exclude_unset=True, exclude_defaults=False, exclude_none=True
)
self._restore_null_content_on_tool_call_messages(data)
return data
def _restore_null_content_on_tool_call_messages(self, data: dict) -> None:
"""Re-add ``content: null`` on assistant messages that only carry tool_calls.
Matches the OpenAI wire format. ``exclude_none=True`` strips
``None`` globally, which would drop the ``content`` key and trip
HF chat templates that don't tolerate a missing ``content``.
Mutates ``data`` in place.
"""
msg_dicts = data.get("messages", [])
assert len(msg_dicts) == len(self.messages), (
"Pydantic dump produced a different number of messages than "
f"the source ({len(msg_dicts)} vs {len(self.messages)})."
)
for msg_dict, msg in zip(msg_dicts, self.messages):
if msg.tool_calls and "content" not in msg_dict:
msg_dict["content"] = None
[docs]
def append_id_to_string(self, s: str) -> str:
"""Appends conversation ID to a string.
Can be useful for log or exception errors messages to allow users
to identify relevant conversation.
"""
if not self.conversation_id:
return s
suffix = f"Conversation id: {self.conversation_id}."
return (s.strip() + " " + suffix) if s else suffix
[docs]
@classmethod
def from_dict(cls, data: dict) -> "Conversation":
"""Converts a dictionary to a conversation."""
return cls.model_validate(data)
[docs]
def to_json(self) -> str:
"""Converts the conversation to a JSON string."""
# to_dict() keeps the null `content` HF chat templates need. The
# separator and escaping options reproduce pydantic's bytes exactly, so
# every other conversation hashes and caches as it did before.
return json.dumps(self.to_dict(), separators=(",", ":"), ensure_ascii=False)
[docs]
@classmethod
def from_json(cls, data: str) -> "Conversation":
"""Converts a JSON string to a conversation."""
return cls.model_validate_json(data)
[docs]
def __repr__(self) -> str:
"""Returns a string representation of the conversation."""
return "\n".join([repr(m) for m in self.messages])
[docs]
class TemplatedMessage(pydantic.BaseModel):
"""Represents a templated message.
This class is used to create messages with dynamic content using a template.
The template can be rendered with variables to produce the final message content.
"""
template: str
"""The template string used to generate the message content."""
role: Role
"""The role of the message sender (e.g., USER, ASSISTANT, SYSTEM)."""
@property
def content(self) -> str:
"""Renders the content of the message."""
template = Template(self.template)
fields = self.model_dump()
fields.pop("template") # remove the template from the fields
return template.render(**fields).strip()
@property
def message(self) -> Message:
"""Returns the message in oumi format."""
content = str(self.content)
return Message(content=content, role=self.role)
class PlannedTurn(pydantic.BaseModel):
"""A single turn in a planned multi-turn conversation."""
model_config = pydantic.ConfigDict(extra="forbid")
turn: int = pydantic.Field(description="1-indexed turn number.")
instruction: str = pydantic.Field(description="What should happen on this turn.")
class PlannerOutput(pydantic.BaseModel):
"""Top-level planner output: a list of planned turns wrapped in an object."""
model_config = pydantic.ConfigDict(extra="forbid")
turns: list[PlannedTurn]
PLANNER_JSON_SCHEMA: dict = PlannerOutput.model_json_schema()
