Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 17 additions & 4 deletions openviking/message/message.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,6 @@
from typing import List, Literal, Optional

from openviking.message.part import ContextPart, Part, TextPart, ToolPart
from openviking.utils.time_utils import format_iso8601, parse_iso_datetime


@dataclass
Expand Down Expand Up @@ -65,6 +64,12 @@ def estimated_tokens(self) -> int:
def to_dict(self) -> dict:
"""Serialize to JSONL."""
created_at_val = self.created_at or datetime.now(timezone.utc).isoformat()
if isinstance(created_at_val, datetime):
created_at_val = (
created_at_val.astimezone(timezone.utc)
.isoformat(timespec="milliseconds")
.replace("+00:00", "Z")
)
return {
"id": self.id,
"role": self.role,
Expand Down Expand Up @@ -108,7 +113,15 @@ def _part_to_dict(self, part: Part) -> dict:
def from_dict(cls, data: dict) -> "Message":
"""Deserialize from JSONL."""
parts = []
for p in data.get("parts", []):
raw_parts = data.get("parts")
if raw_parts is None:
legacy_content = data.get("content")
if legacy_content is not None:
raw_parts = [{"type": "text", "text": legacy_content}]
else:
raw_parts = []

for p in raw_parts:
if p["type"] == "text":
parts.append(TextPart(text=p.get("text", "")))
elif p["type"] == "context":
Expand Down Expand Up @@ -138,7 +151,7 @@ def from_dict(cls, data: dict) -> "Message":
id=data["id"],
role=data["role"],
parts=parts,
created_at=data["created_at"],
created_at=data.get("created_at"),
)

@classmethod
Expand Down Expand Up @@ -193,7 +206,7 @@ def create_assistant(
id=msg_id or f"msg_{uuid4().hex}",
role="assistant",
parts=parts,
created_at=datetime.now(timezone.utc).isoformat()
created_at=datetime.now(timezone.utc).isoformat(),
)

def get_context_parts(self) -> List[ContextPart]:
Expand Down
41 changes: 41 additions & 0 deletions tests/unit/test_message.py
Original file line number Diff line number Diff line change
Expand Up @@ -514,6 +514,24 @@ def test_from_dict_with_tool_part(self):
assert isinstance(msg.parts[0], ToolPart)
assert msg.parts[0].tool_id == "call-1"

def test_from_dict_supports_legacy_content_only_messages(self):
"""Legacy messages with only content should load as a TextPart."""
d = {
"id": "msg-legacy",
"role": "user",
"content": "Hello from legacy storage",
}

msg = Message.from_dict(d)

assert msg.id == "msg-legacy"
assert msg.role == "user"
assert len(msg.parts) == 1
assert isinstance(msg.parts[0], TextPart)
assert msg.parts[0].text == "Hello from legacy storage"
assert msg.content == "Hello from legacy storage"
assert msg.created_at is None

def test_roundtrip(self):
"""Test to_dict -> from_dict roundtrip."""
original = Message(
Expand All @@ -534,6 +552,29 @@ def test_roundtrip(self):
assert isinstance(restored.parts[0], TextPart)
assert isinstance(restored.parts[1], ContextPart)

def test_legacy_message_can_be_reloaded_and_extended(self):
"""Legacy content-only rows should survive reload before appending new messages."""
legacy_row = {
"id": "msg-legacy",
"role": "user",
"content": "Legacy message",
"created_at": "2026-03-26T10:30:00Z",
}
fresh = Message.create_user("Fresh message", msg_id="msg-fresh")

reloaded_messages = [Message.from_dict(legacy_row), Message.from_dict(fresh.to_dict())]

assert [message.content for message in reloaded_messages] == [
"Legacy message",
"Fresh message",
]
assert [
json.loads(message.to_jsonl())["parts"][0]["text"] for message in reloaded_messages
] == [
"Legacy message",
"Fresh message",
]


class TestMessageFactoryMethods:
"""Test Message factory methods."""
Expand Down
Loading