kiln-ai 0.8.0__py3-none-any.whl → 0.11.1__py3-none-any.whl

kiln_ai/adapters/fine_tune/test_dataset_formatter.py

@@ -8,15 +8,20 @@ import pytest
 from kiln_ai.adapters.fine_tune.dataset_formatter import (
     DatasetFormat,
     DatasetFormatter,
+    ModelTrainingData,
+    build_training_data,
     generate_chat_message_response,
     generate_chat_message_toolcall,
     generate_huggingface_chat_template,
     generate_huggingface_chat_template_toolcall,
+    generate_vertex_gemini_1_5,
 )
+from kiln_ai.adapters.model_adapters.base_adapter import COT_FINAL_ANSWER_PROMPT
 from kiln_ai.datamodel import (
     DatasetSplit,
     DataSource,
     DataSourceType,
+    FinetuneDataStrategy,
     Task,
     TaskOutput,
     TaskRun,
@@ -25,32 +30,60 @@ from kiln_ai.datamodel import (
 
 @pytest.fixture
 def mock_task():
-    task = Mock(spec=Task)
+    task = Mock(spec=Task, thinking_instruction=None)
     task_runs = [
-        TaskRun(
-            id=f"run{i}",
-            input='{"test": "input"}',
-            input_source=DataSource(
-                type=DataSourceType.human, properties={"created_by": "test"}
-            ),
-            output=TaskOutput(
-                output='{"test": "output"}',
-                source=DataSource(
-                    type=DataSourceType.synthetic,
-                    properties={
-                        "model_name": "test",
-                        "model_provider": "test",
-                        "adapter_name": "test",
+        Mock(
+            spec=TaskRun,
+            **{
+                "id": f"run{i}",
+                "input": '{"test": "input 你好"}',
+                "repaired_output": None,
+                "intermediate_outputs": {},
+                "input_source": Mock(
+                    spec=DataSource,
+                    **{
+                        "type": DataSourceType.human,
+                        "properties": {"created_by": "test"},
                     },
                 ),
-            ),
+                "output": Mock(
+                    spec=TaskOutput,
+                    **{
+                        "output": '{"test":   "output 你好"}',
+                        "source": Mock(
+                            spec=DataSource,
+                            **{
+                                "type": DataSourceType.synthetic,
+                                "properties": {
+                                    "model_name": "test",
+                                    "model_provider": "test",
+                                    "adapter_name": "test",
+                                },
+                            },
+                        ),
+                    },
+                ),
+            },
         )
         for i in range(1, 4)
     ]
+
+    # Set up parent_task reference for each TaskRun
+    for run in task_runs:
+        run.parent_task = Mock(return_value=task)
+
     task.runs.return_value = task_runs
     return task
 
 
+@pytest.fixture
+def mock_intermediate_outputs(mock_task):
+    for run in mock_task.runs():
+        run.intermediate_outputs = {"reasoning": "thinking output"}
+    mock_task.thinking_instruction = "thinking instructions"
+    return mock_task
+
+
 @pytest.fixture
 def mock_dataset(mock_task):
     dataset = Mock(spec=DatasetSplit)
@@ -61,26 +94,13 @@ def mock_dataset(mock_task):
 
 
 def test_generate_chat_message_response():
-    task_run = TaskRun(
-        id="run1",
+    thinking_data = ModelTrainingData(
         input="test input",
-        input_source=DataSource(
-            type=DataSourceType.human, properties={"created_by": "test"}
-        ),
-        output=TaskOutput(
-            output="test output",
-            source=DataSource(
-                type=DataSourceType.synthetic,
-                properties={
-                    "model_name": "test",
-                    "model_provider": "test",
-                    "adapter_name": "test",
-                },
-            ),
-        ),
+        system_message="system message",
+        final_output="test output",
     )
 
-    result = generate_chat_message_response(task_run, "system message")
+    result = generate_chat_message_response(thinking_data)
 
     assert result == {
         "messages": [
@@ -91,32 +111,80 @@ def test_generate_chat_message_response():
     }
 
 
+def test_generate_chat_message_response_thinking():
+    thinking_data = ModelTrainingData(
+        input="test input",
+        system_message="system message",
+        final_output="test output",
+        thinking="thinking output",
+        thinking_instructions="thinking instructions",
+        thinking_final_answer_prompt="thinking final answer prompt",
+    )
+
+    result = generate_chat_message_response(thinking_data)
+
+    assert result == {
+        "messages": [
+            {"role": "system", "content": "system message"},
+            {"role": "user", "content": "test input"},
+            {"role": "user", "content": "thinking instructions"},
+            {"role": "assistant", "content": "thinking output"},
+            {"role": "user", "content": "thinking final answer prompt"},
+            {"role": "assistant", "content": "test output"},
+        ]
+    }
+
+
 def test_generate_chat_message_toolcall():
-    task_run = TaskRun(
-        id="run1",
+    training_data = ModelTrainingData(
+        input="test input 你好",
+        system_message="system message 你好",
+        final_output='{"key": "value 你好"}',
+    )
+
+    result = generate_chat_message_toolcall(training_data)
+
+    assert result == {
+        "messages": [
+            {"role": "system", "content": "system message 你好"},
+            {"role": "user", "content": "test input 你好"},
+            {
+                "role": "assistant",
+                "content": None,
+                "tool_calls": [
+                    {
+                        "id": "call_1",
+                        "type": "function",
+                        "function": {
+                            "name": "task_response",
+                            "arguments": '{"key": "value 你好"}',
+                        },
+                    }
+                ],
+            },
+        ]
+    }
+
+
+def test_generate_chat_message_toolcall_thinking():
+    training_data = ModelTrainingData(
         input="test input",
-        input_source=DataSource(
-            type=DataSourceType.human, properties={"created_by": "test"}
-        ),
-        output=TaskOutput(
-            output='{"key": "value"}',
-            source=DataSource(
-                type=DataSourceType.synthetic,
-                properties={
-                    "model_name": "test",
-                    "model_provider": "test",
-                    "adapter_name": "test",
-                },
-            ),
-        ),
+        system_message="system message",
+        final_output='{"key": "value"}',
+        thinking="thinking output",
+        thinking_instructions="thinking instructions",
+        thinking_final_answer_prompt="thinking final answer prompt",
     )
 
-    result = generate_chat_message_toolcall(task_run, "system message")
+    result = generate_chat_message_toolcall(training_data)
 
     assert result == {
         "messages": [
             {"role": "system", "content": "system message"},
             {"role": "user", "content": "test input"},
+            {"role": "user", "content": "thinking instructions"},
+            {"role": "assistant", "content": "thinking output"},
+            {"role": "user", "content": "thinking final answer prompt"},
             {
                 "role": "assistant",
                 "content": None,
@@ -136,27 +204,14 @@ def test_generate_chat_message_toolcall():
 
 
 def test_generate_chat_message_toolcall_invalid_json():
-    task_run = TaskRun(
-        id="run1",
+    training_data = ModelTrainingData(
         input="test input",
-        input_source=DataSource(
-            type=DataSourceType.human, properties={"created_by": "test"}
-        ),
-        output=TaskOutput(
-            output="invalid json",
-            source=DataSource(
-                type=DataSourceType.synthetic,
-                properties={
-                    "model_name": "test",
-                    "model_provider": "test",
-                    "adapter_name": "test",
-                },
-            ),
-        ),
+        system_message="system message",
+        final_output="invalid json",
     )
 
     with pytest.raises(ValueError, match="Invalid JSON in for tool call"):
-        generate_chat_message_toolcall(task_run, "system message")
+        generate_chat_message_toolcall(training_data)
 
 
 def test_dataset_formatter_init_no_parent_task(mock_dataset):
@@ -170,14 +225,20 @@ def test_dataset_formatter_dump_invalid_format(mock_dataset):
     formatter = DatasetFormatter(mock_dataset, "system message")
 
     with pytest.raises(ValueError, match="Unsupported format"):
-        formatter.dump_to_file("train", "invalid_format")  # type: ignore
+        formatter.dump_to_file(
+            "train", "invalid_format", FinetuneDataStrategy.final_only
+        )  # type: ignore
 
 
 def test_dataset_formatter_dump_invalid_split(mock_dataset):
     formatter = DatasetFormatter(mock_dataset, "system message")
 
     with pytest.raises(ValueError, match="Split invalid_split not found in dataset"):
-        formatter.dump_to_file("invalid_split", DatasetFormat.OPENAI_CHAT_JSONL)
+        formatter.dump_to_file(
+            "invalid_split",
+            DatasetFormat.OPENAI_CHAT_JSONL,
+            FinetuneDataStrategy.final_only,
+        )
 
 
 def test_dataset_formatter_dump_to_file(mock_dataset, tmp_path):
@@ -185,7 +246,10 @@ def test_dataset_formatter_dump_to_file(mock_dataset, tmp_path):
     output_path = tmp_path / "output.jsonl"
 
     result_path = formatter.dump_to_file(
-        "train", DatasetFormat.OPENAI_CHAT_JSONL, output_path
+        "train",
+        DatasetFormat.OPENAI_CHAT_JSONL,
+        path=output_path,
+        data_strategy=FinetuneDataStrategy.final_only,
     )
 
     assert result_path == output_path
@@ -200,23 +264,38 @@ def test_dataset_formatter_dump_to_file(mock_dataset, tmp_path):
             assert "messages" in data
             assert len(data["messages"]) == 3
             assert data["messages"][0]["content"] == "system message"
-            assert data["messages"][1]["content"] == '{"test": "input"}'
-            assert data["messages"][2]["content"] == '{"test": "output"}'
+            assert data["messages"][1]["content"] == '{"test": "input 你好"}'
+            # Raw chat doesn't fix json issues, like extra spaces
+            assert data["messages"][2]["content"] == '{"test":   "output 你好"}'
 
 
 def test_dataset_formatter_dump_to_temp_file(mock_dataset):
-    formatter = DatasetFormatter(mock_dataset, "system message")
+    formatter = DatasetFormatter(mock_dataset, "system message 你好")
 
-    result_path = formatter.dump_to_file("train", DatasetFormat.OPENAI_CHAT_JSONL)
+    result_path = formatter.dump_to_file(
+        "train",
+        DatasetFormat.OPENAI_CHAT_JSONL,
+        data_strategy=FinetuneDataStrategy.final_only,
+    )
 
     assert result_path.exists()
     assert result_path.parent == Path(tempfile.gettempdir())
-    assert result_path.name.startswith("test_dataset_train_")
+    # Test our nice naming
+    assert result_path.name.startswith(
+        "test_dataset -- split-train -- format-openai_chat_jsonl -- no-cot.jsonl"
+    )
     assert result_path.name.endswith(".jsonl")
     # Verify file contents
     with open(result_path) as f:
         lines = f.readlines()
         assert len(lines) == 2
+        # check non-ascii characters are not escaped
+        assert "你好" in lines[0]
+        assert "你好" in lines[1]
+
+        # confirm didn't use COT for final_only
+        assert "thinking output" not in lines[0]
+        assert "thinking instructions" not in lines[0]
 
 
 def test_dataset_formatter_dump_to_file_tool_format(mock_dataset, tmp_path):
@@ -224,7 +303,10 @@ def test_dataset_formatter_dump_to_file_tool_format(mock_dataset, tmp_path):
     output_path = tmp_path / "output.jsonl"
 
     result_path = formatter.dump_to_file(
-        "train", DatasetFormat.OPENAI_CHAT_TOOLCALL_JSONL, output_path
+        "train",
+        DatasetFormat.OPENAI_CHAT_TOOLCALL_JSONL,
+        path=output_path,
+        data_strategy=FinetuneDataStrategy.final_only,
     )
 
     assert result_path == output_path
@@ -240,7 +322,7 @@ def test_dataset_formatter_dump_to_file_tool_format(mock_dataset, tmp_path):
             assert len(data["messages"]) == 3
             # Check system and user messages
             assert data["messages"][0]["content"] == "system message"
-            assert data["messages"][1]["content"] == '{"test": "input"}'
+            assert data["messages"][1]["content"] == '{"test": "input 你好"}'
             # Check tool call format
             assistant_msg = data["messages"][2]
             assert assistant_msg["content"] is None
@@ -249,30 +331,78 @@ def test_dataset_formatter_dump_to_file_tool_format(mock_dataset, tmp_path):
             tool_call = assistant_msg["tool_calls"][0]
             assert tool_call["type"] == "function"
             assert tool_call["function"]["name"] == "task_response"
-            assert tool_call["function"]["arguments"] == '{"test": "output"}'
+            assert tool_call["function"]["arguments"] == '{"test": "output 你好"}'
+
+
+def test_dataset_formatter_dump_with_intermediate_data(
+    mock_dataset, mock_intermediate_outputs
+):
+    formatter = DatasetFormatter(
+        mock_dataset,
+        "system message 你好",
+        thinking_instructions="thinking instructions",
+    )
+
+    result_path = formatter.dump_to_file(
+        "train",
+        DatasetFormat.OPENAI_CHAT_JSONL,
+        data_strategy=FinetuneDataStrategy.final_and_intermediate,
+    )
+
+    assert result_path.exists()
+    assert result_path.parent == Path(tempfile.gettempdir())
+    # Test our nice naming, with cot
+    assert (
+        result_path.name
+        == "test_dataset -- split-train -- format-openai_chat_jsonl -- cot.jsonl"
+    )
+    # Verify file contents
+    with open(result_path) as f:
+        lines = f.readlines()
+        assert len(lines) == 2
+        for line in lines:
+            assert "thinking output" in line
+            assert "thinking instructions" in line
+
+
+def test_dataset_formatter_dump_with_intermediate_data_custom_instructions(
+    mock_dataset, mock_intermediate_outputs
+):
+    formatter = DatasetFormatter(
+        mock_dataset, "custom system message 你好", "custom thinking instructions"
+    )
+
+    result_path = formatter.dump_to_file(
+        "train",
+        DatasetFormat.OPENAI_CHAT_JSONL,
+        data_strategy=FinetuneDataStrategy.final_and_intermediate,
+    )
+
+    assert result_path.exists()
+    assert result_path.parent == Path(tempfile.gettempdir())
+    # Test our nice naming, with cot
+    assert (
+        result_path.name
+        == "test_dataset -- split-train -- format-openai_chat_jsonl -- cot.jsonl"
+    )
+    # Verify file contents
+    with open(result_path) as f:
+        lines = f.readlines()
+        assert len(lines) == 2
+        for line in lines:
+            assert "custom system message 你好" in line
+            assert "custom thinking instructions" in line
+            assert "thinking output" in line
 
 
 def test_generate_huggingface_chat_template():
-    task_run = TaskRun(
-        id="run1",
+    training_data = ModelTrainingData(
         input="test input",
-        input_source=DataSource(
-            type=DataSourceType.human, properties={"created_by": "test"}
-        ),
-        output=TaskOutput(
-            output="test output",
-            source=DataSource(
-                type=DataSourceType.synthetic,
-                properties={
-                    "model_name": "test",
-                    "model_provider": "test",
-                    "adapter_name": "test",
-                },
-            ),
-        ),
+        system_message="system message",
+        final_output="test output",
     )
 
-    result = generate_huggingface_chat_template(task_run, "system message")
+    result = generate_huggingface_chat_template(training_data)
 
     assert result == {
         "conversations": [
@@ -283,27 +413,96 @@ def test_generate_huggingface_chat_template():
     }
 
 
+def test_generate_huggingface_chat_template_thinking():
+    training_data = ModelTrainingData(
+        input="test input",
+        system_message="system message",
+        final_output="test output",
+        thinking="thinking output",
+        thinking_instructions="thinking instructions",
+        thinking_final_answer_prompt="thinking final answer prompt",
+    )
+
+    result = generate_huggingface_chat_template(training_data)
+
+    assert result == {
+        "conversations": [
+            {"role": "system", "content": "system message"},
+            {"role": "user", "content": "test input"},
+            {"role": "user", "content": "thinking instructions"},
+            {"role": "assistant", "content": "thinking output"},
+            {"role": "user", "content": "thinking final answer prompt"},
+            {"role": "assistant", "content": "test output"},
+        ]
+    }
+
+
+def test_generate_vertex_template():
+    training_data = ModelTrainingData(
+        input="test input",
+        system_message="system message",
+        final_output="test output",
+    )
+
+    result = generate_vertex_gemini_1_5(training_data)
+
+    assert result == {
+        "systemInstruction": {
+            "role": "system",
+            "parts": [
+                {
+                    "text": "system message",
+                }
+            ],
+        },
+        "contents": [
+            {"role": "user", "parts": [{"text": "test input"}]},
+            {"role": "model", "parts": [{"text": "test output"}]},
+        ],
+    }
+
+
+def test_generate_vertex_template_thinking():
+    training_data = ModelTrainingData(
+        input="test input",
+        system_message="system message",
+        final_output="test output",
+        thinking="thinking output",
+        thinking_instructions="thinking instructions",
+        thinking_final_answer_prompt="thinking final answer prompt",
+    )
+
+    result = generate_vertex_gemini_1_5(training_data)
+
+    print(result)
+
+    assert result == {
+        "systemInstruction": {
+            "role": "system",
+            "parts": [
+                {
+                    "text": "system message",
+                }
+            ],
+        },
+        "contents": [
+            {"role": "user", "parts": [{"text": "test input"}]},
+            {"role": "user", "parts": [{"text": "thinking instructions"}]},
+            {"role": "model", "parts": [{"text": "thinking output"}]},
+            {"role": "user", "parts": [{"text": "thinking final answer prompt"}]},
+            {"role": "model", "parts": [{"text": "test output"}]},
+        ],
+    }
+
+
 def test_generate_huggingface_chat_template_toolcall():
-    task_run = TaskRun(
-        id="run1",
+    training_data = ModelTrainingData(
         input="test input",
-        input_source=DataSource(
-            type=DataSourceType.human, properties={"created_by": "test"}
-        ),
-        output=TaskOutput(
-            output='{"key": "value"}',
-            source=DataSource(
-                type=DataSourceType.synthetic,
-                properties={
-                    "model_name": "test",
-                    "model_provider": "test",
-                    "adapter_name": "test",
-                },
-            ),
-        ),
+        system_message="system message",
+        final_output='{"key": "value"}',
     )
 
-    result = generate_huggingface_chat_template_toolcall(task_run, "system message")
+    result = generate_huggingface_chat_template_toolcall(training_data)
 
     assert result["conversations"][0] == {"role": "system", "content": "system message"}
     assert result["conversations"][1] == {"role": "user", "content": "test input"}
@@ -318,25 +517,166 @@ def test_generate_huggingface_chat_template_toolcall():
     assert tool_call["function"]["arguments"] == {"key": "value"}
 
 
+def test_generate_huggingface_chat_template_toolcall_thinking():
+    training_data = ModelTrainingData(
+        input="test input",
+        system_message="system message",
+        final_output='{"key": "value"}',
+        thinking="thinking output",
+        thinking_instructions="thinking instructions",
+        thinking_final_answer_prompt="thinking final answer prompt",
+    )
+
+    result = generate_huggingface_chat_template_toolcall(training_data)
+
+    assert result["conversations"][0] == {"role": "system", "content": "system message"}
+    assert result["conversations"][1] == {"role": "user", "content": "test input"}
+    assert result["conversations"][2] == {
+        "role": "user",
+        "content": "thinking instructions",
+    }
+    assert result["conversations"][3] == {
+        "role": "assistant",
+        "content": "thinking output",
+    }
+    assert result["conversations"][4] == {
+        "role": "user",
+        "content": "thinking final answer prompt",
+    }
+
+    assistant_msg = result["conversations"][5]
+    assert assistant_msg["role"] == "assistant"
+    assert len(assistant_msg["tool_calls"]) == 1
+    tool_call = assistant_msg["tool_calls"][0]
+    assert tool_call["type"] == "function"
+    assert tool_call["function"]["name"] == "task_response"
+    assert len(tool_call["function"]["id"]) == 9  # UUID is truncated to 9 chars
+    assert tool_call["function"]["id"].isalnum()  # Check ID is alphanumeric
+    assert tool_call["function"]["arguments"] == {"key": "value"}
+
+
 def test_generate_huggingface_chat_template_toolcall_invalid_json():
-    task_run = TaskRun(
-        id="run1",
+    training_data = ModelTrainingData(
         input="test input",
-        input_source=DataSource(
-            type=DataSourceType.human, properties={"created_by": "test"}
-        ),
-        output=TaskOutput(
-            output="invalid json",
-            source=DataSource(
-                type=DataSourceType.synthetic,
-                properties={
-                    "model_name": "test",
-                    "model_provider": "test",
-                    "adapter_name": "test",
-                },
-            ),
-        ),
+        system_message="system message",
+        final_output="invalid json",
     )
 
     with pytest.raises(ValueError, match="Invalid JSON in for tool call"):
-        generate_huggingface_chat_template_toolcall(task_run, "system message")
+        generate_huggingface_chat_template_toolcall(training_data)
+
+
+def test_build_training_data(mock_task):
+    # Non repaired should use original output
+    mock_task_run = mock_task.runs()[0]
+    training_data_output = build_training_data(mock_task_run, "system message", False)
+    assert training_data_output.final_output == '{"test":   "output 你好"}'
+    assert training_data_output.thinking is None
+    assert training_data_output.thinking_instructions is None
+    assert training_data_output.thinking_final_answer_prompt is None
+    assert training_data_output.input == '{"test": "input 你好"}'
+    assert training_data_output.system_message == "system message"
+    assert not training_data_output.supports_cot()
+
+
+def test_build_training_data_with_COT(mock_task):
+    # Setup with needed fields for thinking
+    mock_task_run = mock_task.runs()[0]
+    assert mock_task_run.parent_task() == mock_task
+    mock_task_run.intermediate_outputs = {"chain_of_thought": "cot output"}
+
+    training_data_output = build_training_data(
+        mock_task_run,
+        "system message",
+        True,
+        thinking_instructions="thinking instructions",
+    )
+    assert training_data_output.final_output == '{"test":   "output 你好"}'
+    assert training_data_output.thinking == "cot output"
+    assert training_data_output.thinking_instructions == "thinking instructions"
+    assert training_data_output.thinking_final_answer_prompt == COT_FINAL_ANSWER_PROMPT
+    assert training_data_output.input == '{"test": "input 你好"}'
+    assert training_data_output.system_message == "system message"
+    assert training_data_output.supports_cot()
+
+
+def test_build_training_data_with_thinking(mock_task):
+    # Setup with needed fields for thinking
+    mock_task_run = mock_task.runs()[0]
+    assert mock_task_run.parent_task() == mock_task
+    # It should just use the reasoning output if both thinking and chain_of_thought are present
+    mock_task_run.intermediate_outputs = {
+        "reasoning": "thinking output",
+        "chain_of_thought": "cot output",
+    }
+    mock_task.thinking_instruction = "thinking instructions"
+    assert mock_task.thinking_instruction == "thinking instructions"
+
+    training_data_output = build_training_data(
+        mock_task_run,
+        "system message",
+        True,
+        thinking_instructions="thinking instructions",
+    )
+    assert training_data_output.final_output == '{"test":   "output 你好"}'
+    assert training_data_output.thinking == "thinking output"
+    assert training_data_output.thinking_instructions == "thinking instructions"
+    assert training_data_output.thinking_final_answer_prompt == COT_FINAL_ANSWER_PROMPT
+    assert training_data_output.input == '{"test": "input 你好"}'
+    assert training_data_output.system_message == "system message"
+    assert training_data_output.supports_cot()
+
+
+def test_build_training_data_with_repaired_output(mock_task):
+    # use repaired output if available
+    mock_task_run = mock_task.runs()[0]
+    mock_task_run.repair_instructions = "repair instructions"
+    mock_task_run.repaired_output = TaskOutput(
+        output='{"test": "repaired output"}',
+        source=DataSource(
+            type=DataSourceType.human,
+            properties={"created_by": "test-user"},
+        ),
+    )
+
+    training_data_output = build_training_data(mock_task_run, "system message", False)
+    assert training_data_output.final_output == '{"test": "repaired output"}'
+    assert training_data_output.thinking is None
+    assert training_data_output.thinking_instructions is None
+    assert training_data_output.thinking_final_answer_prompt is None
+    assert training_data_output.input == '{"test": "input 你好"}'
+    assert training_data_output.system_message == "system message"
+
+
+def test_dataset_formatter_dump_to_file_json_schema_format(mock_dataset, tmp_path):
+    formatter = DatasetFormatter(mock_dataset, "system message")
+    output_path = tmp_path / "output.jsonl"
+
+    result_path = formatter.dump_to_file(
+        "train",
+        DatasetFormat.OPENAI_CHAT_JSON_SCHEMA_JSONL,
+        path=output_path,
+        data_strategy=FinetuneDataStrategy.final_only,
+    )
+
+    assert result_path == output_path
+    assert output_path.exists()
+
+    # Verify file contents
+    with open(output_path) as f:
+        lines = f.readlines()
+        assert len(lines) == 2  # Should have 2 entries for train split
+        for line in lines:
+            data = json.loads(line)
+            assert "messages" in data
+            assert len(data["messages"]) == 3
+            # Check system and user messages
+            assert data["messages"][0]["content"] == "system message"
+            assert data["messages"][1]["content"] == '{"test": "input 你好"}'
+            # Check JSON format
+            assistant_msg = data["messages"][2]
+            assert assistant_msg["role"] == "assistant"
+            # Verify the content is valid JSON
+            assert assistant_msg["content"] == '{"test": "output 你好"}'
+            json_content = json.loads(assistant_msg["content"])
+            assert json_content == {"test": "output 你好"}