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

kiln_ai/adapters/fine_tune/test_openai_finetune.py

@@ -0,0 +1,503 @@
+import time
+from pathlib import Path
+from unittest import mock
+from unittest.mock import MagicMock, patch
+
+import openai
+import pytest
+from openai.types.fine_tuning import FineTuningJob
+
+from kiln_ai.adapters.fine_tune.base_finetune import FineTuneStatusType
+from kiln_ai.adapters.fine_tune.dataset_formatter import DatasetFormat, DatasetFormatter
+from kiln_ai.adapters.fine_tune.openai_finetune import OpenAIFinetune
+from kiln_ai.datamodel import DatasetSplit, Task, Train80Test20SplitDefinition
+from kiln_ai.datamodel import Finetune as FinetuneModel
+from kiln_ai.utils.config import Config
+
+
+@pytest.fixture
+def openai_finetune(tmp_path):
+    tmp_file = tmp_path / "test-finetune.kiln"
+    finetune = OpenAIFinetune(
+        datamodel=FinetuneModel(
+            name="test-finetune",
+            provider="openai",
+            provider_id="openai-123",
+            base_model_id="gpt-4o",
+            train_split_name="train",
+            dataset_split_id="dataset-123",
+            system_message="Test system message",
+            fine_tune_model_id="ft-123",
+            path=tmp_file,
+        ),
+    )
+    return finetune
+
+
+@pytest.fixture
+def mock_response():
+    response = MagicMock(spec=FineTuningJob)
+    response.error = None
+    response.status = "succeeded"
+    response.finished_at = time.time()
+    response.estimated_finish = None
+    response.fine_tuned_model = "ft-123"
+    response.model = "gpt-4o"
+    return response
+
+
+@pytest.fixture
+def mock_dataset():
+    return DatasetSplit(
+        id="test-dataset-123",
+        name="Test Dataset",
+        splits=Train80Test20SplitDefinition,
+        split_contents={"train": [], "test": []},
+    )
+
+
+@pytest.fixture
+def mock_task():
+    return Task(
+        id="test-task-123",
+        name="Test Task",
+        output_json_schema=None,  # Can be modified in specific tests
+        instruction="Test instruction",
+    )
+
+
+async def test_setup(openai_finetune):
+    if not Config.shared().open_ai_api_key:
+        pytest.skip("OpenAI API key not set")
+    openai_finetune.provider_id = "openai-123"
+    openai_finetune.provider = "openai"
+
+    # Real API call, with fake ID
+    status = await openai_finetune.status()
+    # fake id fails
+    assert status.status == FineTuneStatusType.unknown
+    assert "Job with this ID not found. It may have been deleted." == status.message
+
+
+@pytest.mark.parametrize(
+    "exception,expected_status,expected_message",
+    [
+        (
+            openai.APIConnectionError(request=MagicMock()),
+            FineTuneStatusType.unknown,
+            "Server connection error",
+        ),
+        (
+            openai.RateLimitError(
+                message="Rate limit exceeded", body={}, response=MagicMock()
+            ),
+            FineTuneStatusType.unknown,
+            "Rate limit exceeded",
+        ),
+        (
+            openai.APIStatusError(
+                "Not found",
+                response=MagicMock(status_code=404),
+                body={},
+            ),
+            FineTuneStatusType.unknown,
+            "Job with this ID not found",
+        ),
+        (
+            openai.APIStatusError(
+                "Server error",
+                response=MagicMock(status_code=500),
+                body={},
+            ),
+            FineTuneStatusType.unknown,
+            "Unknown error",
+        ),
+    ],
+)
+async def test_status_api_errors(
+    openai_finetune, 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
+
+
+@pytest.mark.parametrize(
+    "job_status,expected_status,message_contains",
+    [
+        ("failed", FineTuneStatusType.failed, "Job failed"),
+        ("cancelled", FineTuneStatusType.failed, "Job cancelled"),
+        ("succeeded", FineTuneStatusType.completed, "Training job completed"),
+        ("running", FineTuneStatusType.running, "Fine tune job is running"),
+        ("queued", FineTuneStatusType.running, "Fine tune job is running"),
+        (
+            "validating_files",
+            FineTuneStatusType.running,
+            "Fine tune job is running",
+        ),
+        ("unknown_status", FineTuneStatusType.unknown, "Unknown status"),
+    ],
+)
+async def test_status_job_states(
+    openai_finetune,
+    mock_response,
+    job_status,
+    expected_status,
+    message_contains,
+):
+    mock_response.status = job_status
+
+    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
+
+
+async def test_status_with_error_response(openai_finetune, mock_response):
+    mock_response.error = MagicMock()
+    mock_response.error.message = "Something went wrong"
+
+    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:")
+
+
+async def test_status_with_estimated_finish_time(openai_finetune, mock_response):
+    current_time = time.time()
+    mock_response.status = "running"
+    mock_response.estimated_finish = current_time + 300  # 5 minutes from now
+
+    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
+
+
+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_generate_and_upload_jsonl_success(
+    openai_finetune, mock_dataset, mock_task
+):
+    mock_path = Path("mock_path.jsonl")
+    mock_file_id = "file-123"
+
+    # Mock the formatter
+    mock_formatter = MagicMock(spec=DatasetFormatter)
+    mock_formatter.dump_to_file.return_value = mock_path
+
+    # Mock the file response
+    mock_file_response = MagicMock()
+    mock_file_response.id = mock_file_id
+
+    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(
+            mock_dataset, "train", mock_task
+        )
+
+        # Verify formatter was created with correct parameters
+        mock_formatter_class.assert_called_once_with(
+            mock_dataset, openai_finetune.datamodel.system_message
+        )
+
+        # Verify correct format was used
+        mock_formatter.dump_to_file.assert_called_once_with(
+            "train", DatasetFormat.OPENAI_CHAT_JSONL
+        )
+
+        # Verify file was opened and uploaded
+        mock_open.assert_called_once_with(mock_path, "rb")
+        mock_create.assert_called_once()
+
+        assert result == mock_file_id
+
+
+async def test_generate_and_upload_jsonl_toolcall_success(
+    openai_finetune, mock_dataset, mock_task
+):
+    mock_path = Path("mock_path.jsonl")
+    mock_file_id = "file-123"
+    mock_task.output_json_schema = '{"type": "object", "properties": {"key": {"type": "string"}}}'  # Add JSON schema
+
+    # Mock the formatter
+    mock_formatter = MagicMock(spec=DatasetFormatter)
+    mock_formatter.dump_to_file.return_value = mock_path
+
+    # Mock the file response
+    mock_file_response = MagicMock()
+    mock_file_response.id = mock_file_id
+
+    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(
+            mock_dataset, "train", mock_task
+        )
+
+        # Verify formatter was created with correct parameters
+        mock_formatter_class.assert_called_once_with(
+            mock_dataset, openai_finetune.datamodel.system_message
+        )
+
+        # Verify correct format was used
+        mock_formatter.dump_to_file.assert_called_once_with(
+            "train", DatasetFormat.OPENAI_CHAT_TOOLCALL_JSONL
+        )
+
+        # Verify file was opened and uploaded
+        mock_open.assert_called_once_with(mock_path, "rb")
+        mock_create.assert_called_once()
+
+        assert result == mock_file_id
+
+
+async def test_generate_and_upload_jsonl_upload_failure(
+    openai_finetune, 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 response with no ID
+    mock_file_response = MagicMock()
+    mock_file_response.id = None
+
+    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"):
+            await openai_finetune.generate_and_upload_jsonl(
+                mock_dataset, "train", mock_task
+            )
+
+
+async def test_generate_and_upload_jsonl_api_error(
+    openai_finetune, mock_dataset, mock_task
+):
+    mock_path = Path("mock_path.jsonl")
+
+    mock_formatter = MagicMock(spec=DatasetFormatter)
+    mock_formatter.dump_to_file.return_value = mock_path
+
+    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):
+            await openai_finetune.generate_and_upload_jsonl(
+                mock_dataset, "train", mock_task
+            )
+
+
+async def test_start_success(openai_finetune, mock_dataset, mock_task):
+    openai_finetune.datamodel.parent = mock_task
+
+    # Mock parameters
+    openai_finetune.datamodel.parameters = {
+        "n_epochs": 3,
+        "learning_rate_multiplier": 0.1,
+        "batch_size": 4,
+        "ignored_param": "value",
+    }
+
+    # Mock the fine-tuning response
+    mock_ft_response = MagicMock()
+    mock_ft_response.id = "ft-123"
+    mock_ft_response.fine_tuned_model = None
+    mock_ft_response.model = "gpt-4o-mini-2024-07-18"
+
+    with (
+        patch.object(
+            openai_finetune,
+            "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)
+
+        # Verify file uploads
+        assert mock_upload.call_count == 1  # Only training file
+        mock_upload.assert_called_with(
+            mock_dataset, openai_finetune.datamodel.train_split_name, mock_task
+        )
+
+        # Verify fine-tune creation
+        mock_create.assert_called_once_with(
+            training_file="train-file-123",
+            model="gpt-4o",
+            validation_file=None,
+            seed=None,
+            hyperparameters={
+                "n_epochs": 3,
+                "learning_rate_multiplier": 0.1,
+                "batch_size": 4,
+            },
+            suffix=f"kiln_ai.{openai_finetune.datamodel.id}",
+        )
+
+        # Verify model updates
+        assert openai_finetune.datamodel.provider_id == "ft-123"
+        assert openai_finetune.datamodel.base_model_id == "gpt-4o-mini-2024-07-18"
+
+
+async def test_start_with_validation(openai_finetune, mock_dataset, mock_task):
+    openai_finetune.datamodel.parent = mock_task
+    openai_finetune.datamodel.validation_split_name = "validation"
+
+    mock_ft_response = MagicMock()
+    mock_ft_response.id = "ft-123"
+    mock_ft_response.fine_tuned_model = None
+    mock_ft_response.model = "gpt-4o-mini-2024-07-18"
+
+    with (
+        patch.object(
+            openai_finetune,
+            "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)
+
+        # Verify both files were uploaded
+        assert mock_upload.call_count == 2
+        mock_upload.assert_has_calls(
+            [
+                mock.call(
+                    mock_dataset, openai_finetune.datamodel.train_split_name, mock_task
+                ),
+                mock.call(mock_dataset, "validation", mock_task),
+            ]
+        )
+
+        # Verify validation file was included
+        mock_create.assert_called_once()
+        assert mock_create.call_args[1]["validation_file"] == "val-file-123"
+
+
+async def test_start_no_task(openai_finetune, mock_dataset):
+    openai_finetune.datamodel.parent = None
+    openai_finetune.datamodel.path = None
+
+    with pytest.raises(ValueError, match="Task is required to start a fine-tune"):
+        await openai_finetune._start(mock_dataset)
+
+
+async def test_status_updates_model_ids(openai_finetune, 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"
+
+    # Configure mock response with different model IDs
+    mock_response.fine_tuned_model = "new-ft-model"
+    mock_response.model = "new-base-model"
+    mock_response.status = "succeeded"
+
+    with (
+        patch(
+            "kiln_ai.adapters.fine_tune.openai_finetune.oai_client.fine_tuning.jobs.retrieve",
+            return_value=mock_response,
+        ),
+    ):
+        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 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"
+
+
+async def test_status_updates_latest_status(openai_finetune, mock_response):
+    # Set initial status
+    openai_finetune.datamodel.latest_status = FineTuneStatusType.running
+    assert openai_finetune.datamodel.latest_status == FineTuneStatusType.running
+    mock_response.status = "succeeded"
+
+    with (
+        patch(
+            "kiln_ai.adapters.fine_tune.openai_finetune.oai_client.fine_tuning.jobs.retrieve",
+            return_value=mock_response,
+        ),
+    ):
+        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 file was saved
+        assert openai_finetune.datamodel.path.exists()

kiln_ai/adapters/langchain_adapters.py

@@ -1,21 +1,35 @@
-from typing import Dict
+import os
+from os import getenv
+from typing import Any, Dict
 
+from langchain_aws import ChatBedrockConverse
 from langchain_core.language_models import LanguageModelInput
 from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
 from langchain_core.messages.base import BaseMessage
 from langchain_core.runnables import Runnable
+from langchain_fireworks import ChatFireworks
+from langchain_groq import ChatGroq
+from langchain_ollama import ChatOllama
+from langchain_openai import ChatOpenAI
 from pydantic import BaseModel
 
 import kiln_ai.datamodel as datamodel
+from kiln_ai.adapters.ollama_tools import (
+    get_ollama_connection,
+    ollama_base_url,
+    ollama_model_installed,
+)
+from kiln_ai.utils.config import Config
 
 from .base_adapter import AdapterInfo, BaseAdapter, BasePromptBuilder, RunOutput
-from .ml_model_list import langchain_model_from
+from .ml_model_list import KilnModelProvider, ModelProviderName
+from .provider_tools import kiln_model_provider_from
 
 LangChainModelType = BaseChatModel | Runnable[LanguageModelInput, Dict | BaseModel]
 
 
-class LangChainPromptAdapter(BaseAdapter):
+class LangchainAdapter(BaseAdapter):
     _model: LangChainModelType | None = None
 
     def __init__(
@@ -51,12 +65,6 @@ class LangChainPromptAdapter(BaseAdapter):
                 "model_name and provider must be provided if custom_model is not provided"
             )
 
-    def adapter_specific_instructions(self) -> str | None:
-        # TODO: would be better to explicitly use bind_tools:tool_choice="task_response" here
-        if self.has_structured_output():
-            return "Always respond with a tool call. Never respond with a human readable message."
-        return None
-
     async def model(self) -> LangChainModelType:
         # cached model
         if self._model:
@@ -79,8 +87,13 @@ class LangChainPromptAdapter(BaseAdapter):
                 )
             output_schema["title"] = "task_response"
             output_schema["description"] = "A response from the task"
+            with_structured_output_options = await get_structured_output_options(
+                self.model_name, self.model_provider
+            )
             self._model = self._model.with_structured_output(
-                output_schema, include_raw=True
+                output_schema,
+                include_raw=True,
+                **with_structured_output_options,
             )
         return self._model
 
@@ -108,17 +121,16 @@ class LangChainPromptAdapter(BaseAdapter):
             )
 
             cot_messages = [*messages]
-            cot_response = base_model.invoke(cot_messages)
+            cot_response = await base_model.ainvoke(cot_messages)
             intermediate_outputs["chain_of_thought"] = cot_response.content
             messages.append(AIMessage(content=cot_response.content))
             messages.append(
                 SystemMessage(content="Considering the above, return a final result.")
             )
         elif cot_prompt:
-            # for plaintext output, we just add COT instructions. We still only make one call.
             messages.append(SystemMessage(content=cot_prompt))
 
-        response = chain.invoke(messages)
+        response = await chain.ainvoke(messages)
 
         if self.has_structured_output():
             if (
@@ -160,3 +172,81 @@ class LangChainPromptAdapter(BaseAdapter):
         ):
             return response["arguments"]
         return response
+
+
+async def get_structured_output_options(
+    model_name: str, model_provider: str
+) -> Dict[str, Any]:
+    finetune_provider = await kiln_model_provider_from(model_name, model_provider)
+    if finetune_provider and finetune_provider.adapter_options.get("langchain"):
+        return finetune_provider.adapter_options["langchain"].get(
+            "with_structured_output_options", {}
+        )
+    return {}
+
+
+async def langchain_model_from(
+    name: str, provider_name: str | None = None
+) -> BaseChatModel:
+    provider = await kiln_model_provider_from(name, provider_name)
+    return await langchain_model_from_provider(provider, name)
+
+
+async def langchain_model_from_provider(
+    provider: KilnModelProvider, model_name: str
+) -> BaseChatModel:
+    if provider.name == ModelProviderName.openai:
+        api_key = Config.shared().open_ai_api_key
+        return ChatOpenAI(**provider.provider_options, openai_api_key=api_key)  # type: ignore[arg-type]
+    elif provider.name == ModelProviderName.groq:
+        api_key = Config.shared().groq_api_key
+        if api_key is None:
+            raise ValueError(
+                "Attempted to use Groq without an API key set. "
+                "Get your API key from https://console.groq.com/keys"
+            )
+        return ChatGroq(**provider.provider_options, groq_api_key=api_key)  # type: ignore[arg-type]
+    elif provider.name == ModelProviderName.amazon_bedrock:
+        api_key = Config.shared().bedrock_access_key
+        secret_key = Config.shared().bedrock_secret_key
+        # langchain doesn't allow passing these, so ugly hack to set env vars
+        os.environ["AWS_ACCESS_KEY_ID"] = api_key
+        os.environ["AWS_SECRET_ACCESS_KEY"] = secret_key
+        return ChatBedrockConverse(
+            **provider.provider_options,
+        )
+    elif provider.name == ModelProviderName.fireworks_ai:
+        api_key = Config.shared().fireworks_api_key
+        return ChatFireworks(**provider.provider_options, api_key=api_key)
+    elif provider.name == ModelProviderName.ollama:
+        # Ollama model naming is pretty flexible. We try a few versions of the model name
+        potential_model_names = []
+        if "model" in provider.provider_options:
+            potential_model_names.append(provider.provider_options["model"])
+        if "model_aliases" in provider.provider_options:
+            potential_model_names.extend(provider.provider_options["model_aliases"])
+
+        # Get the list of models Ollama supports
+        ollama_connection = await get_ollama_connection()
+        if ollama_connection is None:
+            raise ValueError("Failed to connect to Ollama. Ensure Ollama is running.")
+
+        for model_name in potential_model_names:
+            if ollama_model_installed(ollama_connection, model_name):
+                return ChatOllama(model=model_name, base_url=ollama_base_url())
+
+        raise ValueError(f"Model {model_name} not installed on Ollama")
+    elif provider.name == ModelProviderName.openrouter:
+        api_key = Config.shared().open_router_api_key
+        base_url = getenv("OPENROUTER_BASE_URL") or "https://openrouter.ai/api/v1"
+        return ChatOpenAI(
+            **provider.provider_options,
+            openai_api_key=api_key,  # type: ignore[arg-type]
+            openai_api_base=base_url,  # type: ignore[arg-type]
+            default_headers={
+                "HTTP-Referer": "https://getkiln.ai/openrouter",
+                "X-Title": "KilnAI",
+            },
+        )
+    else:
+        raise ValueError(f"Invalid model or provider: {model_name} - {provider.name}")