Repository: ErikBjare/gptme
File: tests/test_message.py (Python)

Code before:
    assert len(msg.file_hashes) == 2
    assert msg.file_hashes.get("/src/utils/test.py") == "hash_for_src_utils"
    assert msg.file_hashes.get("/tests/test.py") == "hash_for_tests"

    # Round-trip through TOML should preserve both
    toml_str = msg.to_toml()
    loaded = Message.from_toml(toml_str)
    assert loaded.file_hashes == msg.file_hashes


def test_message_metadata():
    """Test that metadata survives JSONL and TOML round-trips."""
    import json

    from gptme.message import Message, MessageMetadata

    # Create message with metadata
    meta: MessageMetadata = {
        "input_tokens": 100,
        "output_tokens": 50,
        "cost": 0.005,
        "model": "claude-sonnet",
        "cache_read_tokens": 80,
    }
    msg = Message(role="assistant", content="Hello world", metadata=meta)

    # Verify metadata stored correctly
    assert msg.metadata is not None
    assert msg.metadata["input_tokens"] == 100
    assert msg.metadata["output_tokens"] == 50
    assert msg.metadata["cost"] == 0.005
    assert msg.metadata["model"] == "claude-sonnet"
    assert msg.metadata["cache_read_tokens"] == 80

    # Test JSON/JSONL roundtrip
    d = msg.to_dict()
    assert "metadata" in d
    json_str = json.dumps(d)
    json_data = json.loads(json_str)
    from dateutil.parser import isoparse

    json_data["timestamp"] = isoparse(json_data["timestamp"])
    msg2 = Message(**json_data)
    assert msg2.metadata == meta

    # Test TOML roundtrip
    toml_str = msg.to_toml()
    msg3 = Message.from_toml(toml_str)
    assert msg3.metadata == meta



Review comment (performance):
I'm thinking it should maybe look like:

```
{
    "model": "...",
    "tokens": {"input": {"base": 2, "cache_read": 10, "cache_write": 6}, "output": 2} 
         # or {"input": 2, "cache_read": 10, "cache_write": 6, "output": 2},
    "cost": 0.05,
}

Code after:
    assert len(msg.file_hashes) == 2
    assert msg.file_hashes.get("/src/utils/test.py") == "hash_for_src_utils"
    assert msg.file_hashes.get("/tests/test.py") == "hash_for_tests"

    # Round-trip through TOML should preserve both
    toml_str = msg.to_toml()
    loaded = Message.from_toml(toml_str)
    assert loaded.file_hashes == msg.file_hashes


def test_message_metadata():
    """Test that metadata survives JSONL and TOML round-trips."""
    import json

    from gptme.message import Message, MessageMetadata

    # Create message with metadata using flat token format
    # Per Erik's review: https://github.com/gptme/gptme/pull/943#issuecomment-3633137716
    meta: MessageMetadata = {
        "model": "claude-sonnet",
        "input_tokens": 100,
        "output_tokens": 50,
        "cache_read_tokens": 80,
        "cache_creation_tokens": 10,
        "cost": 0.005,
    }
    msg = Message(role="assistant", content="Hello world", metadata=meta)

    # Verify metadata stored correctly
    assert msg.metadata is not None
    assert msg.metadata["model"] == "claude-sonnet"
    assert msg.metadata["input_tokens"] == 100
    assert msg.metadata["output_tokens"] == 50
    assert msg.metadata["cache_read_tokens"] == 80
    assert msg.metadata["cache_creation_tokens"] == 10
    assert msg.metadata["cost"] == 0.005

    # Test JSON/JSONL roundtrip
    d = msg.to_dict()
    assert "metadata" in d
    json_str = json.dumps(d)
    json_data = json.loads(json_str)
    from dateutil.parser import isoparse

    json_data["timestamp"] = isoparse(json_data["timestamp"])
    msg2 = Message(**json_data)
    assert msg2.metadata == meta

    # Test TOML roundtrip
    toml_str = msg.to_toml()
    msg3 = Message.from_toml(toml_str)