mistral-vibe/tests/test_message_id.py

from __future__ import annotations

import json
from uuid import UUID

from vibe.core.types import (
    AssistantEvent,
    LLMMessage,
    ReasoningEvent,
    Role,
    UserMessageEvent,
)


class TestLLMMessageId:
    def test_user_message_gets_message_id(self) -> None:
        msg = LLMMessage(role=Role.user, content="Hello")
        assert msg.message_id is not None
        UUID(msg.message_id)  # Validates it's a valid UUID

    def test_assistant_message_gets_message_id(self) -> None:
        msg = LLMMessage(role=Role.assistant, content="Hi there")
        assert msg.message_id is not None
        UUID(msg.message_id)

    def test_system_message_gets_message_id(self) -> None:
        msg = LLMMessage(role=Role.system, content="You are helpful")
        assert msg.message_id is not None
        UUID(msg.message_id)

    def test_tool_message_does_not_get_message_id(self) -> None:
        msg = LLMMessage(role=Role.tool, content="result", tool_call_id="tc_123")
        assert msg.message_id is None

    def test_each_message_gets_unique_id(self) -> None:
        msg1 = LLMMessage(role=Role.user, content="First")
        msg2 = LLMMessage(role=Role.user, content="Second")
        assert msg1.message_id != msg2.message_id

    def test_message_id_preserved_from_dict(self) -> None:
        expected_id = "custom-message-id-123"
        msg = LLMMessage.model_validate({
            "role": "user",
            "content": "Hello",
            "message_id": expected_id,
        })
        assert msg.message_id == expected_id

    def test_message_id_preserved_for_tool_from_dict(self) -> None:
        expected_id = "tool-message-id"
        msg = LLMMessage.model_validate({
            "role": "tool",
            "content": "result",
            "tool_call_id": "tc_123",
            "message_id": expected_id,
        })
        assert msg.message_id == expected_id

    def test_tool_message_no_id_from_dict_without_id(self) -> None:
        msg = LLMMessage.model_validate({
            "role": "tool",
            "content": "result",
            "tool_call_id": "tc_123",
        })
        assert msg.message_id is None


class TestLLMMessageAccumulation:
    def test_message_id_preserved_on_add(self) -> None:
        msg1 = LLMMessage(role=Role.assistant, content="Hello")
        msg2 = LLMMessage(role=Role.assistant, content=" world")

        result = msg1 + msg2

        assert result.message_id == msg1.message_id
        assert result.content == "Hello world"

    def test_message_id_preserved_after_multiple_adds(self) -> None:
        msg1 = LLMMessage(role=Role.assistant, content="A")
        msg2 = LLMMessage(role=Role.assistant, content="B")
        msg3 = LLMMessage(role=Role.assistant, content="C")

        result = msg1 + msg2 + msg3

        assert result.message_id == msg1.message_id
        assert result.content == "ABC"


class TestEventMessageId:
    def test_user_message_event_has_message_id(self) -> None:
        event = UserMessageEvent(content="Hello", message_id="user-msg-id")
        assert event.message_id == "user-msg-id"
        assert event.content == "Hello"

    def test_assistant_event_has_message_id(self) -> None:
        event = AssistantEvent(content="test", message_id="test-id")
        assert event.message_id == "test-id"

    def test_assistant_event_message_id_optional(self) -> None:
        event = AssistantEvent(content="test")
        assert event.message_id is None

    def test_reasoning_event_has_message_id(self) -> None:
        event = ReasoningEvent(content="thinking...", message_id="reason-id")
        assert event.message_id == "reason-id"

    def test_reasoning_event_message_id_optional(self) -> None:
        event = ReasoningEvent(content="thinking...")
        assert event.message_id is None

    def test_assistant_event_add_preserves_message_id(self) -> None:
        event1 = AssistantEvent(content="Hello", message_id="first-id")
        event2 = AssistantEvent(content=" world", message_id="second-id")

        result = event1 + event2

        assert result.message_id == "first-id"
        assert result.content == "Hello world"


class TestMessageIdExcludedFromAPI:
    def test_message_id_excluded_with_exclude_param(self) -> None:
        msg = LLMMessage(role=Role.user, content="Hello")
        dumped = msg.model_dump(exclude_none=True, exclude={"message_id"})

        assert "message_id" not in dumped
        assert dumped["role"] == "user"
        assert dumped["content"] == "Hello"

    def test_message_id_included_in_normal_dump(self) -> None:
        msg = LLMMessage(role=Role.user, content="Hello")
        dumped = msg.model_dump(exclude_none=True)

        assert "message_id" in dumped
        assert dumped["message_id"] == msg.message_id


class TestMessageIdInLogs:
    def test_message_id_in_json_dump(self) -> None:
        msg = LLMMessage(role=Role.assistant, content="Response")
        dumped = msg.model_dump(exclude_none=True)

        json_str = json.dumps(dumped)
        loaded = json.loads(json_str)

        assert "message_id" in loaded
        assert loaded["message_id"] == msg.message_id

    def test_message_id_roundtrip(self) -> None:
        original = LLMMessage(role=Role.user, content="Test")
        original_id = original.message_id

        dumped = original.model_dump(exclude_none=True)
        restored = LLMMessage.model_validate(dumped)

        assert restored.message_id == original_id

    def test_tool_message_id_none_in_json(self) -> None:
        msg = LLMMessage(role=Role.tool, content="result", tool_call_id="tc_1")
        dumped = msg.model_dump(exclude_none=True)

        assert "message_id" not in dumped