kiln-ai 0.19.0__py3-none-any.whl → 0.21.0__py3-none-any.whl

kiln_ai/adapters/fine_tune/__init__.py

@@ -8,7 +8,7 @@ from . import base_finetune, dataset_formatter, finetune_registry, openai_finetu
 
 __all__ = [
     "base_finetune",
-    "openai_finetune",
     "dataset_formatter",
     "finetune_registry",
+    "openai_finetune",
 ]

kiln_ai/adapters/fine_tune/openai_finetune.py

@@ -13,9 +13,16 @@ from kiln_ai.adapters.fine_tune.dataset_formatter import DatasetFormat, DatasetF
 from kiln_ai.datamodel import DatasetSplit, StructuredOutputMode, Task
 from kiln_ai.utils.config import Config
 
-oai_client = openai.AsyncOpenAI(
-    api_key=Config.shared().open_ai_api_key or "",
-)
+
+def _get_openai_client():
+    key = Config.shared().open_ai_api_key
+    if not key:
+        raise RuntimeError(
+            "OpenAI API key not set. You must connect OpenAI in settings."
+        )
+    return openai.AsyncOpenAI(
+        api_key=key,
+    )
 
 
 class OpenAIFinetune(BaseFinetuneAdapter):
@@ -45,6 +52,7 @@ class OpenAIFinetune(BaseFinetuneAdapter):
 
         try:
             # Will raise an error if the job is not found, or for other issues
+            oai_client = _get_openai_client()
             response = await oai_client.fine_tuning.jobs.retrieve(
                 self.datamodel.provider_id
             )
@@ -79,7 +87,7 @@ class OpenAIFinetune(BaseFinetuneAdapter):
                 )
             return FineTuneStatus(
                 status=FineTuneStatusType.unknown,
-                message=f"Unknown error: [{str(e)}]",
+                message=f"Unknown error: [{e!s}]",
             )
 
         if not response or not isinstance(response, FineTuningJob):
@@ -145,6 +153,7 @@ class OpenAIFinetune(BaseFinetuneAdapter):
             if k in ["n_epochs", "learning_rate_multiplier", "batch_size"]
         }
 
+        oai_client = _get_openai_client()
         ft = await oai_client.fine_tuning.jobs.create(
             training_file=train_file_id,
             model=self.datamodel.base_model_id,
@@ -168,6 +177,7 @@ class OpenAIFinetune(BaseFinetuneAdapter):
         )
         path = formatter.dump_to_file(split_name, format, self.datamodel.data_strategy)
 
+        oai_client = _get_openai_client()
         response = await oai_client.files.create(
             file=open(path, "rb"),
             purpose="fine-tune",

kiln_ai/adapters/fine_tune/test_dataset_formatter.py

@@ -206,7 +206,7 @@ def test_generate_chat_message_toolcall(mock_training_chat_two_step_json):
 
 def test_generate_chat_message_toolcall_invalid_json(mock_training_chat_two_step_json):
     mock_training_chat_two_step_json[-1].content = "invalid json"
-    with pytest.raises(ValueError, match="^Last message is not JSON"):
+    with pytest.raises(ValueError, match=r"^Last message is not JSON"):
         generate_chat_message_toolcall(mock_training_chat_two_step_json)
 
 
@@ -536,7 +536,7 @@ def test_generate_huggingface_chat_template_toolcall_invalid_json(
 ):
     mock_training_chat_two_step_json[-1].content = "invalid json"
 
-    with pytest.raises(ValueError, match="^Last message is not JSON"):
+    with pytest.raises(ValueError, match=r"^Last message is not JSON"):
         generate_huggingface_chat_template_toolcall(mock_training_chat_two_step_json)
 
 

kiln_ai/adapters/fine_tune/test_fireworks_tinetune.py

@@ -14,11 +14,7 @@ from kiln_ai.adapters.fine_tune.fireworks_finetune import (
     DeployStatus,
     FireworksFinetune,
 )
-from kiln_ai.datamodel import (
-    DatasetSplit,
-    StructuredOutputMode,
-    Task,
-)
+from kiln_ai.datamodel import DatasetSplit, StructuredOutputMode, Task
 from kiln_ai.datamodel import Finetune as FinetuneModel
 from kiln_ai.datamodel.datamodel_enums import ChatStrategy
 from kiln_ai.datamodel.dataset_split import Train80Test20SplitDefinition
@@ -1053,7 +1049,7 @@ async def test_fetch_all_deployments_invalid_json(fireworks_finetune, mock_api_k
 
         with pytest.raises(
             ValueError,
-            match="Invalid response from Fireworks. Expected list of deployments in 'deployments' key",
+            match=r"Invalid response from Fireworks. Expected list of deployments in 'deployments' key",
         ):
             await fireworks_finetune._fetch_all_deployments()
 

kiln_ai/adapters/fine_tune/test_openai_finetune.py

@@ -21,6 +21,25 @@ from kiln_ai.datamodel.dataset_split import Train80Test20SplitDefinition
 from kiln_ai.utils.config import Config
 
 
+@pytest.fixture
+def mock_openai_client():
+    """Mock the OpenAI client returned by _get_openai_client()"""
+    from unittest.mock import AsyncMock
+
+    with patch(
+        "kiln_ai.adapters.fine_tune.openai_finetune._get_openai_client"
+    ) as mock_get_client:
+        mock_client = MagicMock()
+
+        # Use AsyncMock for async methods
+        mock_client.fine_tuning.jobs.retrieve = AsyncMock()
+        mock_client.fine_tuning.jobs.create = AsyncMock()
+        mock_client.files.create = AsyncMock()
+
+        mock_get_client.return_value = mock_client
+        yield mock_client
+
+
 @pytest.fixture
 def openai_finetune(tmp_path):
     tmp_file = tmp_path / "test-finetune.kiln"
@@ -122,15 +141,12 @@ async def test_setup(openai_finetune):
     ],
 )
 async def test_status_api_errors(
-    openai_finetune, exception, expected_status, expected_message
+    openai_finetune, mock_openai_client, exception, expected_status, expected_message
 ):
-    with patch(
-        "kiln_ai.adapters.fine_tune.openai_finetune.oai_client.fine_tuning.jobs.retrieve",
-        side_effect=exception,
-    ):
-        status = await openai_finetune.status()
-        assert status.status == expected_status
-        assert expected_message in status.message
+    mock_openai_client.fine_tuning.jobs.retrieve.side_effect = exception
+    status = await openai_finetune.status()
+    assert status.status == expected_status
+    assert expected_message in status.message
 
 
 @pytest.mark.parametrize(
@@ -151,63 +167,57 @@ async def test_status_api_errors(
 )
 async def test_status_job_states(
     openai_finetune,
+    mock_openai_client,
     mock_response,
     job_status,
     expected_status,
     message_contains,
 ):
     mock_response.status = job_status
+    mock_openai_client.fine_tuning.jobs.retrieve.return_value = mock_response
 
-    with patch(
-        "kiln_ai.adapters.fine_tune.openai_finetune.oai_client.fine_tuning.jobs.retrieve",
-        return_value=mock_response,
-    ):
-        status = await openai_finetune.status()
-        assert status.status == expected_status
-        assert message_contains in status.message
+    status = await openai_finetune.status()
+    assert status.status == expected_status
+    assert message_contains in status.message
 
 
-async def test_status_with_error_response(openai_finetune, mock_response):
+async def test_status_with_error_response(
+    openai_finetune, mock_openai_client, mock_response
+):
     mock_response.error = MagicMock()
     mock_response.error.message = "Something went wrong"
+    mock_openai_client.fine_tuning.jobs.retrieve.return_value = mock_response
 
-    with patch(
-        "kiln_ai.adapters.fine_tune.openai_finetune.oai_client.fine_tuning.jobs.retrieve",
-        return_value=mock_response,
-    ):
-        status = await openai_finetune.status()
-        assert status.status == FineTuneStatusType.failed
-        assert status.message.startswith("Something went wrong [Code:")
+    status = await openai_finetune.status()
+    assert status.status == FineTuneStatusType.failed
+    assert status.message.startswith("Something went wrong [Code:")
 
 
-async def test_status_with_estimated_finish_time(openai_finetune, mock_response):
+async def test_status_with_estimated_finish_time(
+    openai_finetune, mock_openai_client, mock_response
+):
     current_time = time.time()
     mock_response.status = "running"
     mock_response.estimated_finish = current_time + 300  # 5 minutes from now
+    mock_openai_client.fine_tuning.jobs.retrieve.return_value = mock_response
 
-    with patch(
-        "kiln_ai.adapters.fine_tune.openai_finetune.oai_client.fine_tuning.jobs.retrieve",
-        return_value=mock_response,
-    ):
-        status = await openai_finetune.status()
-        assert status.status == FineTuneStatusType.running
-        assert (
-            "Estimated finish time: 299 seconds" in status.message
-        )  # non zero time passes
+    status = await openai_finetune.status()
+    assert status.status == FineTuneStatusType.running
+    assert (
+        "Estimated finish time: 299 seconds" in status.message
+    )  # non zero time passes
 
 
-async def test_status_empty_response(openai_finetune):
-    with patch(
-        "kiln_ai.adapters.fine_tune.openai_finetune.oai_client.fine_tuning.jobs.retrieve",
-        return_value=mock_response,
-    ):
-        status = await openai_finetune.status()
-        assert status.status == FineTuneStatusType.unknown
-        assert "Invalid response from OpenAI" in status.message
+async def test_status_empty_response(openai_finetune, mock_openai_client):
+    mock_openai_client.fine_tuning.jobs.retrieve.return_value = None
+
+    status = await openai_finetune.status()
+    assert status.status == FineTuneStatusType.unknown
+    assert "Invalid response from OpenAI" in status.message
 
 
 async def test_generate_and_upload_jsonl_success(
-    openai_finetune, mock_dataset, mock_task
+    openai_finetune, mock_openai_client, mock_dataset, mock_task
 ):
     mock_path = Path("mock_path.jsonl")
     mock_file_id = "file-123"
@@ -219,16 +229,13 @@ async def test_generate_and_upload_jsonl_success(
     # Mock the file response
     mock_file_response = MagicMock()
     mock_file_response.id = mock_file_id
+    mock_openai_client.files.create.return_value = mock_file_response
 
     with (
         patch(
             "kiln_ai.adapters.fine_tune.openai_finetune.DatasetFormatter",
             return_value=mock_formatter,
         ) as mock_formatter_class,
-        patch(
-            "kiln_ai.adapters.fine_tune.openai_finetune.oai_client.files.create",
-            return_value=mock_file_response,
-        ) as mock_create,
         patch("builtins.open") as mock_open,
     ):
         result = await openai_finetune.generate_and_upload_jsonl(
@@ -252,13 +259,13 @@ async def test_generate_and_upload_jsonl_success(
 
         # Verify file was opened and uploaded
         mock_open.assert_called_once_with(mock_path, "rb")
-        mock_create.assert_called_once()
+        mock_openai_client.files.create.assert_called_once()
 
         assert result == mock_file_id
 
 
 async def test_generate_and_upload_jsonl_schema_success(
-    openai_finetune, mock_dataset, mock_task
+    openai_finetune, mock_openai_client, mock_dataset, mock_task
 ):
     mock_path = Path("mock_path.jsonl")
     mock_file_id = "file-123"
@@ -271,16 +278,13 @@ async def test_generate_and_upload_jsonl_schema_success(
     # Mock the file response
     mock_file_response = MagicMock()
     mock_file_response.id = mock_file_id
+    mock_openai_client.files.create.return_value = mock_file_response
 
     with (
         patch(
             "kiln_ai.adapters.fine_tune.openai_finetune.DatasetFormatter",
             return_value=mock_formatter,
         ) as mock_formatter_class,
-        patch(
-            "kiln_ai.adapters.fine_tune.openai_finetune.oai_client.files.create",
-            return_value=mock_file_response,
-        ) as mock_create,
         patch("builtins.open") as mock_open,
     ):
         result = await openai_finetune.generate_and_upload_jsonl(
@@ -304,13 +308,13 @@ async def test_generate_and_upload_jsonl_schema_success(
 
         # Verify file was opened and uploaded
         mock_open.assert_called_once_with(mock_path, "rb")
-        mock_create.assert_called_once()
+        mock_openai_client.files.create.assert_called_once()
 
         assert result == mock_file_id
 
 
 async def test_generate_and_upload_jsonl_upload_failure(
-    openai_finetune, mock_dataset, mock_task
+    openai_finetune, mock_openai_client, mock_dataset, mock_task
 ):
     mock_path = Path("mock_path.jsonl")
 
@@ -320,16 +324,13 @@ async def test_generate_and_upload_jsonl_upload_failure(
     # Mock response with no ID
     mock_file_response = MagicMock()
     mock_file_response.id = None
+    mock_openai_client.files.create.return_value = mock_file_response
 
     with (
         patch(
             "kiln_ai.adapters.fine_tune.openai_finetune.DatasetFormatter",
             return_value=mock_formatter,
         ),
-        patch(
-            "kiln_ai.adapters.fine_tune.openai_finetune.oai_client.files.create",
-            return_value=mock_file_response,
-        ),
         patch("builtins.open"),
     ):
         with pytest.raises(ValueError, match="Failed to upload file to OpenAI"):
@@ -339,24 +340,21 @@ async def test_generate_and_upload_jsonl_upload_failure(
 
 
 async def test_generate_and_upload_jsonl_api_error(
-    openai_finetune, mock_dataset, mock_task
+    openai_finetune, mock_openai_client, mock_dataset, mock_task
 ):
     mock_path = Path("mock_path.jsonl")
 
     mock_formatter = MagicMock(spec=DatasetFormatter)
     mock_formatter.dump_to_file.return_value = mock_path
+    mock_openai_client.files.create.side_effect = openai.APIError(
+        message="API error", request=MagicMock(), body={}
+    )
 
     with (
         patch(
             "kiln_ai.adapters.fine_tune.openai_finetune.DatasetFormatter",
             return_value=mock_formatter,
         ),
-        patch(
-            "kiln_ai.adapters.fine_tune.openai_finetune.oai_client.files.create",
-            side_effect=openai.APIError(
-                message="API error", request=MagicMock(), body={}
-            ),
-        ),
         patch("builtins.open"),
     ):
         with pytest.raises(openai.APIError):
@@ -378,6 +376,7 @@ async def test_generate_and_upload_jsonl_api_error(
 )
 async def test_start_success(
     openai_finetune,
+    mock_openai_client,
     mock_dataset,
     mock_task,
     output_schema,
@@ -401,6 +400,7 @@ async def test_start_success(
     mock_ft_response.id = "ft-123"
     mock_ft_response.fine_tuned_model = None
     mock_ft_response.model = "gpt-4o-mini-2024-07-18"
+    mock_openai_client.fine_tuning.jobs.create.return_value = mock_ft_response
 
     with (
         patch.object(
@@ -408,10 +408,6 @@ async def test_start_success(
             "generate_and_upload_jsonl",
             side_effect=["train-file-123", "val-file-123"],
         ) as mock_upload,
-        patch(
-            "kiln_ai.adapters.fine_tune.openai_finetune.oai_client.fine_tuning.jobs.create",
-            return_value=mock_ft_response,
-        ) as mock_create,
     ):
         await openai_finetune._start(mock_dataset)
 
@@ -425,7 +421,7 @@ async def test_start_success(
         )
 
         # Verify fine-tune creation
-        mock_create.assert_called_once_with(
+        mock_openai_client.fine_tuning.jobs.create.assert_called_once_with(
             training_file="train-file-123",
             model="gpt-4o",
             validation_file=None,
@@ -444,7 +440,9 @@ async def test_start_success(
         assert openai_finetune.datamodel.structured_output_mode == expected_mode
 
 
-async def test_start_with_validation(openai_finetune, mock_dataset, mock_task):
+async def test_start_with_validation(
+    openai_finetune, mock_openai_client, mock_dataset, mock_task
+):
     openai_finetune.datamodel.parent = mock_task
     openai_finetune.datamodel.validation_split_name = "validation"
 
@@ -452,6 +450,7 @@ async def test_start_with_validation(openai_finetune, mock_dataset, mock_task):
     mock_ft_response.id = "ft-123"
     mock_ft_response.fine_tuned_model = None
     mock_ft_response.model = "gpt-4o-mini-2024-07-18"
+    mock_openai_client.fine_tuning.jobs.create.return_value = mock_ft_response
 
     with (
         patch.object(
@@ -459,10 +458,6 @@ async def test_start_with_validation(openai_finetune, mock_dataset, mock_task):
             "generate_and_upload_jsonl",
             side_effect=["train-file-123", "val-file-123"],
         ) as mock_upload,
-        patch(
-            "kiln_ai.adapters.fine_tune.openai_finetune.oai_client.fine_tuning.jobs.create",
-            return_value=mock_ft_response,
-        ) as mock_create,
     ):
         await openai_finetune._start(mock_dataset)
 
@@ -486,8 +481,11 @@ async def test_start_with_validation(openai_finetune, mock_dataset, mock_task):
         )
 
         # Verify validation file was included
-        mock_create.assert_called_once()
-        assert mock_create.call_args[1]["validation_file"] == "val-file-123"
+        mock_openai_client.fine_tuning.jobs.create.assert_called_once()
+        assert (
+            mock_openai_client.fine_tuning.jobs.create.call_args[1]["validation_file"]
+            == "val-file-123"
+        )
 
 
 async def test_start_no_task(openai_finetune, mock_dataset):
@@ -498,7 +496,9 @@ async def test_start_no_task(openai_finetune, mock_dataset):
         await openai_finetune._start(mock_dataset)
 
 
-async def test_status_updates_model_ids(openai_finetune, mock_response):
+async def test_status_updates_model_ids(
+    openai_finetune, mock_openai_client, mock_response
+):
     # Set up initial model IDs
     openai_finetune.datamodel.fine_tune_model_id = "old-ft-model"
     openai_finetune.datamodel.base_model_id = "old-base-model"
@@ -507,49 +507,41 @@ async def test_status_updates_model_ids(openai_finetune, mock_response):
     mock_response.fine_tuned_model = "new-ft-model"
     mock_response.model = "new-base-model"
     mock_response.status = "succeeded"
+    mock_openai_client.fine_tuning.jobs.retrieve.return_value = mock_response
 
-    with (
-        patch(
-            "kiln_ai.adapters.fine_tune.openai_finetune.oai_client.fine_tuning.jobs.retrieve",
-            return_value=mock_response,
-        ),
-    ):
-        status = await openai_finetune.status()
+    status = await openai_finetune.status()
 
-        # Verify model IDs were updated
-        assert openai_finetune.datamodel.fine_tune_model_id == "new-ft-model"
-        assert openai_finetune.datamodel.base_model_id == "new-base-model"
+    # Verify model IDs were updated
+    assert openai_finetune.datamodel.fine_tune_model_id == "new-ft-model"
+    assert openai_finetune.datamodel.base_model_id == "new-base-model"
 
-        # Verify save was called
-        # This isn't properly mocked, so not checking
-        # assert openai_finetune.datamodel.save.called
+    # Verify save was called
+    # This isn't properly mocked, so not checking
+    # assert openai_finetune.datamodel.save.called
 
-        # Verify status is still returned correctly
-        assert status.status == FineTuneStatusType.completed
-        assert status.message == "Training job completed"
+    # Verify status is still returned correctly
+    assert status.status == FineTuneStatusType.completed
+    assert status.message == "Training job completed"
 
 
-async def test_status_updates_latest_status(openai_finetune, mock_response):
+async def test_status_updates_latest_status(
+    openai_finetune, mock_openai_client, mock_response
+):
     # Set initial status
     openai_finetune.datamodel.latest_status = FineTuneStatusType.running
     assert openai_finetune.datamodel.latest_status == FineTuneStatusType.running
     mock_response.status = "succeeded"
+    mock_openai_client.fine_tuning.jobs.retrieve.return_value = mock_response
 
-    with (
-        patch(
-            "kiln_ai.adapters.fine_tune.openai_finetune.oai_client.fine_tuning.jobs.retrieve",
-            return_value=mock_response,
-        ),
-    ):
-        status = await openai_finetune.status()
+    status = await openai_finetune.status()
 
-        # Verify status was updated in datamodel
-        assert openai_finetune.datamodel.latest_status == FineTuneStatusType.completed
-        assert status.status == FineTuneStatusType.completed
-        assert status.message == "Training job completed"
+    # Verify status was updated in datamodel
+    assert openai_finetune.datamodel.latest_status == FineTuneStatusType.completed
+    assert status.status == FineTuneStatusType.completed
+    assert status.message == "Training job completed"
 
-        # Verify file was saved
-        assert openai_finetune.datamodel.path.exists()
+    # Verify file was saved
+    assert openai_finetune.datamodel.path.exists()
 
 
 @pytest.mark.parametrize(
@@ -595,11 +587,16 @@ async def test_generate_and_upload_jsonl_with_data_strategy(
             return_value=mock_formatter,
         ),
         patch(
-            "kiln_ai.adapters.fine_tune.openai_finetune.oai_client.files.create",
-            return_value=mock_file_response,
-        ),
+            "kiln_ai.adapters.fine_tune.openai_finetune._get_openai_client"
+        ) as mock_get_client,
         patch("builtins.open"),
     ):
+        from unittest.mock import AsyncMock
+
+        mock_client = MagicMock()
+        mock_client.files.create = AsyncMock(return_value=mock_file_response)
+        mock_get_client.return_value = mock_client
+
         result = await openai_finetune.generate_and_upload_jsonl(
             mock_dataset, "train", mock_task, DatasetFormat.OPENAI_CHAT_JSONL
         )

kiln_ai/adapters/fine_tune/test_together_finetune.py

@@ -17,11 +17,7 @@ from kiln_ai.adapters.fine_tune.together_finetune import (
     _pending_statuses,
     _running_statuses,
 )
-from kiln_ai.datamodel import (
-    DatasetSplit,
-    StructuredOutputMode,
-    Task,
-)
+from kiln_ai.datamodel import DatasetSplit, StructuredOutputMode, Task
 from kiln_ai.datamodel import Finetune as FinetuneModel
 from kiln_ai.datamodel.dataset_split import Train80Test20SplitDefinition
 from kiln_ai.utils.config import Config
@@ -105,7 +101,7 @@ def mock_api_key():
 def test_init_missing_api_key(finetune):
     with patch.object(Config, "shared") as mock_config:
         mock_config.return_value.together_api_key = None
-        with pytest.raises(ValueError, match="Together.ai API key not set"):
+        with pytest.raises(ValueError, match=r"Together.ai API key not set"):
             TogetherFinetune(datamodel=finetune)