kiln-ai 0.6.1__py3-none-any.whl → 0.7.1__py3-none-any.whl

kiln_ai/adapters/fine_tune/test_base_finetune.py

@@ -0,0 +1,290 @@
+from unittest.mock import Mock
+
+import pytest
+
+from kiln_ai.adapters.fine_tune.base_finetune import (
+    BaseFinetuneAdapter,
+    FineTuneParameter,
+    FineTuneStatus,
+    FineTuneStatusType,
+)
+from kiln_ai.datamodel import DatasetSplit, Task
+from kiln_ai.datamodel import Finetune as FinetuneModel
+
+
+class MockFinetune(BaseFinetuneAdapter):
+    """Mock implementation of BaseFinetune for testing"""
+
+    async def _start(self, dataset: DatasetSplit) -> None:
+        pass
+
+    async def status(self) -> FineTuneStatus:
+        return FineTuneStatus(status=FineTuneStatusType.pending, message="loading...")
+
+    @classmethod
+    def available_parameters(cls) -> list[FineTuneParameter]:
+        return [
+            FineTuneParameter(
+                name="learning_rate",
+                type="float",
+                description="Learning rate for training",
+            ),
+            FineTuneParameter(
+                name="epochs",
+                type="int",
+                description="Number of training epochs",
+                optional=False,
+            ),
+        ]
+
+
+@pytest.fixture
+def sample_task(tmp_path):
+    task_path = tmp_path / "task.kiln"
+    task = Task(
+        name="Test Task",
+        path=task_path,
+        description="Test task for fine-tuning",
+        instruction="Test instruction",
+    )
+    task.save_to_file()
+    return task
+
+
+@pytest.fixture
+def basic_finetune(sample_task):
+    return MockFinetune(
+        datamodel=FinetuneModel(
+            parent=sample_task,
+            name="test_finetune",
+            provider="test_provider",
+            provider_id="model_1234",
+            base_model_id="test_model",
+            train_split_name="train",
+            dataset_split_id="dataset-123",
+            system_message="Test system message",
+        ),
+    )
+
+
+async def test_finetune_status(basic_finetune):
+    status = await basic_finetune.status()
+    assert status.status == FineTuneStatusType.pending
+    assert status.message == "loading..."
+    assert isinstance(status, FineTuneStatus)
+
+
+def test_available_parameters():
+    params = MockFinetune.available_parameters()
+    assert len(params) == 2
+
+    learning_rate_param = params[0]
+    assert learning_rate_param.name == "learning_rate"
+    assert learning_rate_param.type == "float"
+    assert learning_rate_param.optional is True
+
+    epochs_param = params[1]
+    assert epochs_param.name == "epochs"
+    assert epochs_param.type == "int"
+    assert epochs_param.optional is False
+
+
+def test_validate_parameters_valid():
+    # Test valid parameters
+    valid_params = {
+        "learning_rate": 0.001,
+        "epochs": 10,
+    }
+    MockFinetune.validate_parameters(valid_params)  # Should not raise
+
+
+def test_validate_parameters_missing_required():
+    # Test missing required parameter
+    invalid_params = {
+        "learning_rate": 0.001,
+        # missing required 'epochs'
+    }
+    with pytest.raises(ValueError, match="Parameter epochs is required"):
+        MockFinetune.validate_parameters(invalid_params)
+
+
+def test_validate_parameters_wrong_type():
+    # Test wrong parameter types
+    invalid_params = {
+        "learning_rate": "0.001",  # string instead of float
+        "epochs": 10,
+    }
+    with pytest.raises(ValueError, match="Parameter learning_rate must be a float"):
+        MockFinetune.validate_parameters(invalid_params)
+
+    invalid_params = {
+        "learning_rate": 0.001,
+        "epochs": 10.5,  # float instead of int
+    }
+    with pytest.raises(ValueError, match="Parameter epochs must be an integer"):
+        MockFinetune.validate_parameters(invalid_params)
+
+
+def test_validate_parameters_unknown_parameter():
+    # Test unknown parameter
+    invalid_params = {
+        "learning_rate": 0.001,
+        "epochs": 10,
+        "unknown_param": "value",
+    }
+    with pytest.raises(ValueError, match="Parameter unknown_param is not available"):
+        MockFinetune.validate_parameters(invalid_params)
+
+
+@pytest.fixture
+def mock_dataset(sample_task):
+    dataset = Mock(spec=DatasetSplit)
+    dataset.id = "dataset_123"
+    dataset.parent_task.return_value = sample_task
+    dataset.split_contents = {"train": [], "validation": [], "test": []}
+    return dataset
+
+
+async def test_create_and_start_success(mock_dataset):
+    # Test successful creation with minimal parameters
+    adapter, datamodel = await MockFinetune.create_and_start(
+        dataset=mock_dataset,
+        provider_id="openai",
+        provider_base_model_id="gpt-4o-mini-2024-07-18",
+        train_split_name="train",
+        parameters={"epochs": 10},  # Required parameter
+        system_message="Test system message",
+    )
+
+    assert isinstance(adapter, MockFinetune)
+    assert isinstance(datamodel, FinetuneModel)
+    assert len(datamodel.name.split()) == 2  # 2 word memorable name
+    assert datamodel.provider == "openai"
+    assert datamodel.base_model_id == "gpt-4o-mini-2024-07-18"
+    assert datamodel.dataset_split_id == mock_dataset.id
+    assert datamodel.train_split_name == "train"
+    assert datamodel.parameters == {"epochs": 10}
+    assert datamodel.system_message == "Test system message"
+    assert datamodel.path.exists()
+
+
+async def test_create_and_start_with_all_params(mock_dataset):
+    # Test creation with all optional parameters
+    adapter, datamodel = await MockFinetune.create_and_start(
+        dataset=mock_dataset,
+        provider_id="openai",
+        provider_base_model_id="gpt-4o-mini-2024-07-18",
+        train_split_name="train",
+        parameters={"epochs": 10, "learning_rate": 0.001},
+        name="Custom Name",
+        description="Custom Description",
+        validation_split_name="test",
+        system_message="Test system message",
+    )
+
+    assert datamodel.name == "Custom Name"
+    assert datamodel.description == "Custom Description"
+    assert datamodel.validation_split_name == "test"
+    assert datamodel.parameters == {"epochs": 10, "learning_rate": 0.001}
+    assert datamodel.system_message == "Test system message"
+    assert adapter.datamodel == datamodel
+
+    # load the datamodel from the file, confirm it's saved
+    loaded_datamodel = FinetuneModel.load_from_file(datamodel.path)
+    assert loaded_datamodel.model_dump_json() == datamodel.model_dump_json()
+
+
+async def test_create_and_start_invalid_parameters(mock_dataset):
+    # Test with invalid parameters
+    with pytest.raises(ValueError, match="Parameter epochs is required"):
+        await MockFinetune.create_and_start(
+            dataset=mock_dataset,
+            provider_id="openai",
+            provider_base_model_id="gpt-4o-mini-2024-07-18",
+            train_split_name="train",
+            parameters={"learning_rate": 0.001},  # Missing required 'epochs'
+            system_message="Test system message",
+        )
+
+
+async def test_create_and_start_no_parent_task():
+    # Test with dataset that has no parent task
+    dataset = Mock(spec=DatasetSplit)
+    dataset.id = "dataset_123"
+    dataset.parent_task.return_value = None
+    dataset.split_contents = {"train": [], "validation": [], "test": []}
+
+    with pytest.raises(ValueError, match="Dataset must have a parent task with a path"):
+        await MockFinetune.create_and_start(
+            dataset=dataset,
+            provider_id="openai",
+            provider_base_model_id="gpt-4o-mini-2024-07-18",
+            train_split_name="train",
+            parameters={"epochs": 10},
+            system_message="Test system message",
+        )
+
+
+async def test_create_and_start_no_parent_task_path():
+    # Test with dataset that has parent task but no path
+    task = Mock(spec=Task)
+    task.path = None
+
+    dataset = Mock(spec=DatasetSplit)
+    dataset.id = "dataset_123"
+    dataset.parent_task.return_value = task
+    dataset.split_contents = {"train": [], "validation": [], "test": []}
+
+    with pytest.raises(ValueError, match="Dataset must have a parent task with a path"):
+        await MockFinetune.create_and_start(
+            dataset=dataset,
+            provider_id="openai",
+            provider_base_model_id="gpt-4o-mini-2024-07-18",
+            train_split_name="train",
+            parameters={"epochs": 10},
+            system_message="Test system message",
+        )
+
+
+def test_check_valid_provider_model():
+    MockFinetune.check_valid_provider_model("openai", "gpt-4o-mini-2024-07-18")
+
+    with pytest.raises(
+        ValueError, match="Provider openai with base model gpt-99 is not available"
+    ):
+        MockFinetune.check_valid_provider_model("openai", "gpt-99")
+
+
+async def test_create_and_start_invalid_train_split(mock_dataset):
+    # Test with an invalid train split name
+    mock_dataset.split_contents = {"valid_train": [], "valid_test": []}
+
+    with pytest.raises(
+        ValueError, match="Train split invalid_train not found in dataset"
+    ):
+        await MockFinetune.create_and_start(
+            dataset=mock_dataset,
+            provider_id="openai",
+            provider_base_model_id="gpt-4o-mini-2024-07-18",
+            train_split_name="invalid_train",  # Invalid train split
+            parameters={"epochs": 10},
+            system_message="Test system message",
+        )
+
+
+async def test_create_and_start_invalid_validation_split(mock_dataset):
+    # Test with an invalid validation split name
+    mock_dataset.split_contents = {"valid_train": [], "valid_test": []}
+
+    with pytest.raises(
+        ValueError, match="Validation split invalid_test not found in dataset"
+    ):
+        await MockFinetune.create_and_start(
+            dataset=mock_dataset,
+            provider_id="openai",
+            provider_base_model_id="gpt-4o-mini-2024-07-18",
+            train_split_name="valid_train",
+            validation_split_name="invalid_test",  # Invalid validation split
+            parameters={"epochs": 10},
+            system_message="Test system message",
+        )

kiln_ai/adapters/fine_tune/test_dataset_formatter.py

@@ -0,0 +1,342 @@
+import json
+import tempfile
+from pathlib import Path
+from unittest.mock import Mock
+
+import pytest
+
+from kiln_ai.adapters.fine_tune.dataset_formatter import (
+    DatasetFormat,
+    DatasetFormatter,
+    generate_chat_message_response,
+    generate_chat_message_toolcall,
+    generate_huggingface_chat_template,
+    generate_huggingface_chat_template_toolcall,
+)
+from kiln_ai.datamodel import (
+    DatasetSplit,
+    DataSource,
+    DataSourceType,
+    Task,
+    TaskOutput,
+    TaskRun,
+)
+
+
+@pytest.fixture
+def mock_task():
+    task = Mock(spec=Task)
+    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",
+                    },
+                ),
+            ),
+        )
+        for i in range(1, 4)
+    ]
+    task.runs.return_value = task_runs
+    return task
+
+
+@pytest.fixture
+def mock_dataset(mock_task):
+    dataset = Mock(spec=DatasetSplit)
+    dataset.name = "test_dataset"
+    dataset.parent_task.return_value = mock_task
+    dataset.split_contents = {"train": ["run1", "run2"], "test": ["run3"]}
+    return dataset
+
+
+def test_generate_chat_message_response():
+    task_run = TaskRun(
+        id="run1",
+        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",
+                },
+            ),
+        ),
+    )
+
+    result = generate_chat_message_response(task_run, "system message")
+
+    assert result == {
+        "messages": [
+            {"role": "system", "content": "system message"},
+            {"role": "user", "content": "test input"},
+            {"role": "assistant", "content": "test output"},
+        ]
+    }
+
+
+def test_generate_chat_message_toolcall():
+    task_run = TaskRun(
+        id="run1",
+        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",
+                },
+            ),
+        ),
+    )
+
+    result = generate_chat_message_toolcall(task_run, "system message")
+
+    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_invalid_json():
+    task_run = TaskRun(
+        id="run1",
+        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",
+                },
+            ),
+        ),
+    )
+
+    with pytest.raises(ValueError, match="Invalid JSON in for tool call"):
+        generate_chat_message_toolcall(task_run, "system message")
+
+
+def test_dataset_formatter_init_no_parent_task(mock_dataset):
+    mock_dataset.parent_task.return_value = None
+
+    with pytest.raises(ValueError, match="Dataset has no parent task"):
+        DatasetFormatter(mock_dataset, "system message")
+
+
+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
+
+
+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)
+
+
+def test_dataset_formatter_dump_to_file(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_JSONL, output_path
+    )
+
+    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
+            assert data["messages"][0]["content"] == "system message"
+            assert data["messages"][1]["content"] == '{"test": "input"}'
+            assert data["messages"][2]["content"] == '{"test": "output"}'
+
+
+def test_dataset_formatter_dump_to_temp_file(mock_dataset):
+    formatter = DatasetFormatter(mock_dataset, "system message")
+
+    result_path = formatter.dump_to_file("train", DatasetFormat.OPENAI_CHAT_JSONL)
+
+    assert result_path.exists()
+    assert result_path.parent == Path(tempfile.gettempdir())
+    assert result_path.name.startswith("test_dataset_train_")
+    assert result_path.name.endswith(".jsonl")
+    # Verify file contents
+    with open(result_path) as f:
+        lines = f.readlines()
+        assert len(lines) == 2
+
+
+def test_dataset_formatter_dump_to_file_tool_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_TOOLCALL_JSONL, output_path
+    )
+
+    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 tool call format
+            assistant_msg = data["messages"][2]
+            assert assistant_msg["content"] is None
+            assert "tool_calls" in assistant_msg
+            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 tool_call["function"]["arguments"] == '{"test": "output"}'
+
+
+def test_generate_huggingface_chat_template():
+    task_run = TaskRun(
+        id="run1",
+        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",
+                },
+            ),
+        ),
+    )
+
+    result = generate_huggingface_chat_template(task_run, "system message")
+
+    assert result == {
+        "conversations": [
+            {"role": "system", "content": "system message"},
+            {"role": "user", "content": "test input"},
+            {"role": "assistant", "content": "test output"},
+        ]
+    }
+
+
+def test_generate_huggingface_chat_template_toolcall():
+    task_run = TaskRun(
+        id="run1",
+        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",
+                },
+            ),
+        ),
+    )
+
+    result = generate_huggingface_chat_template_toolcall(task_run, "system message")
+
+    assert result["conversations"][0] == {"role": "system", "content": "system message"}
+    assert result["conversations"][1] == {"role": "user", "content": "test input"}
+    assistant_msg = result["conversations"][2]
+    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",
+        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",
+                },
+            ),
+        ),
+    )
+
+    with pytest.raises(ValueError, match="Invalid JSON in for tool call"):
+        generate_huggingface_chat_template_toolcall(task_run, "system message")