OpenViking/openviking/message/message.py

# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd.
# SPDX-License-Identifier: AGPL-3.0
"""Message class definition - based on opencode Message design.

Message = role + parts, supports serialization to JSONL.
"""

import json
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import List, Literal, Optional

from openviking.core.peer_id import normalize_peer_id
from openviking.message.part import ContextPart, ImagePart, Part, TextPart, ToolPart, part_from_dict
from openviking.utils.token_estimation import estimate_text_tokens


@dataclass
class Message:
    """Message = role + parts."""

    id: str
    role: Literal["user", "assistant"]
    parts: List[Part]
    peer_id: Optional[str] = None
    created_at: str = None
    turn_id: Optional[str] = None
    message_kind: Optional[
        Literal["user_query", "assistant_step", "tool_transport", "checkpoint"]
    ] = None
    source_message_ids: Optional[List[str]] = None

    @property
    def content(self) -> str:
        """Quick access to first TextPart content."""
        for p in self.parts:
            if isinstance(p, TextPart):
                return p.text
        return ""

    @property
    def estimated_tokens(self) -> int:
        """Estimate token count from all parts using a CJK-aware fallback.

        Counts fields that actually appear in the assembled prompt:
        - TextPart.text: always emitted
        - ContextPart.abstract: injected as text (uri is not sent to the model)
        - ImagePart: not counted here; image captioning/model usage is tracked
          by the VLM call that converts images into text for extraction
        - ToolPart: tool_id (appears in toolUse.id / toolResult.toolCallId),
          tool_name, tool_input (JSON), tool_output

        Known limitation: ToolPart estimation undercounts by ~10-20 tokens per
        tool call because tool_id/toolName appear twice in the assembled transcript
        (toolUse + toolResult), and small literals like "(no output)" / "{}" are
        not counted. Under 128k budgets this is negligible; for smaller budgets
        (8k/16k) or tool-dense sessions, consider adding a conservative per-tool
        buffer instead of mirroring the full convertToAgentMessages logic.
        """
        token_text = []
        for p in self.parts:
            if isinstance(p, TextPart):
                token_text.append(p.text)
            elif isinstance(p, ContextPart):
                token_text.append(p.abstract)
            elif isinstance(p, ToolPart):
                token_text.extend([p.tool_id, p.tool_name])
                if p.tool_input:
                    token_text.append(json.dumps(p.tool_input, ensure_ascii=False))
                if p.tool_output:
                    token_text.append(p.tool_output)
        return estimate_text_tokens("".join(token_text))

    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")
            )
        data = {
            "id": self.id,
            "role": self.role,
            "parts": [self._part_to_dict(p) for p in self.parts],
            "created_at": created_at_val,
        }
        if self.peer_id is not None:
            data["peer_id"] = self.peer_id
        if self.turn_id is not None:
            data["turn_id"] = self.turn_id
        if self.message_kind is not None:
            data["message_kind"] = self.message_kind
        if self.source_message_ids is not None:
            data["source_message_ids"] = list(self.source_message_ids)
        return data

    def _part_to_dict(self, part: Part) -> dict:
        if isinstance(part, TextPart):
            return {"type": part.type, "text": part.text}
        elif isinstance(part, ContextPart):
            return {
                "type": part.type,
                "uri": part.uri,
                "context_type": part.context_type,
                "abstract": part.abstract,
            }
        elif isinstance(part, ImagePart):
            image_url = {"url": part.url}
            if part.detail is not None:
                image_url["detail"] = part.detail
            return {
                "type": part.type,
                "image_url": image_url,
            }
        elif isinstance(part, ToolPart):
            d = {
                "type": part.type,
                "tool_id": part.tool_id,
                "tool_name": part.tool_name,
                "tool_uri": part.tool_uri,
                "skill_uri": part.skill_uri,
                "tool_status": part.tool_status,
            }
            if part.tool_input:
                d["tool_input"] = part.tool_input
            if part.tool_output:
                d["tool_output"] = part.tool_output
            if part.duration_ms is not None:
                d["duration_ms"] = part.duration_ms
            if part.prompt_tokens is not None:
                d["prompt_tokens"] = part.prompt_tokens
            if part.completion_tokens is not None:
                d["completion_tokens"] = part.completion_tokens
            if part.tool_output_ref:
                d["tool_output_ref"] = part.tool_output_ref
            if part.tool_output_truncated:
                d["tool_output_truncated"] = part.tool_output_truncated
            if part.tool_output_original_chars is not None:
                d["tool_output_original_chars"] = part.tool_output_original_chars
            if part.tool_output_preview_chars is not None:
                d["tool_output_preview_chars"] = part.tool_output_preview_chars
            if part.tool_output_sha256:
                d["tool_output_sha256"] = part.tool_output_sha256
            if part.tool_output_storage_uri:
                d["tool_output_storage_uri"] = part.tool_output_storage_uri
            if part.tool_output_mime_type and part.tool_output_mime_type != "text/plain":
                d["tool_output_mime_type"] = part.tool_output_mime_type
            if part.tool_output_source_ref:
                d["tool_output_source_ref"] = part.tool_output_source_ref
            if part.tool_output_source_offset is not None:
                d["tool_output_source_offset"] = part.tool_output_source_offset
            if part.tool_output_source_limit is not None:
                d["tool_output_source_limit"] = part.tool_output_source_limit
            if part.tool_output_externalization_error:
                d["tool_output_externalization_error"] = part.tool_output_externalization_error
            if part.tool_output_group_id:
                d["tool_output_group_id"] = part.tool_output_group_id
            if part.tool_output_externalized_reason:
                d["tool_output_externalized_reason"] = part.tool_output_externalized_reason
            if part.tool_output_group_original_chars is not None:
                d["tool_output_group_original_chars"] = part.tool_output_group_original_chars
            if part.tool_output_group_budget_chars is not None:
                d["tool_output_group_budget_chars"] = part.tool_output_group_budget_chars
            return d
        return {}

    @classmethod
    def from_dict(cls, data: dict) -> "Message":
        """Deserialize from JSONL."""
        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:
            parts.append(part_from_dict(p))
        try:
            peer_id = normalize_peer_id(data.get("peer_id"))
        except ValueError:
            peer_id = data.get("peer_id")

        return cls(
            id=data["id"],
            role=data["role"],
            parts=parts,
            peer_id=peer_id,
            created_at=data.get("created_at"),
            turn_id=data.get("turn_id"),
            message_kind=data.get("message_kind"),
            source_message_ids=(
                list(data.get("source_message_ids"))
                if isinstance(data.get("source_message_ids"), list)
                else None
            ),
        )

    def to_jsonl(self) -> str:
        """Serialize to JSONL string."""
        return json.dumps(self.to_dict(), ensure_ascii=False)