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

kiln_ai/adapters/{test_langchain_adapter.py → model_adapters/test_langchain_adapter.py}

@@ -7,21 +7,31 @@ from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
 from langchain_fireworks import ChatFireworks
 from langchain_groq import ChatGroq
 from langchain_ollama import ChatOllama
-from langchain_openai import ChatOpenAI
 
-from kiln_ai.adapters.langchain_adapters import (
+from kiln_ai.adapters.ml_model_list import (
+    KilnModelProvider,
+    ModelProviderName,
+    StructuredOutputMode,
+)
+from kiln_ai.adapters.model_adapters.base_adapter import COT_FINAL_ANSWER_PROMPT
+from kiln_ai.adapters.model_adapters.langchain_adapters import (
     LangchainAdapter,
-    get_structured_output_options,
     langchain_model_from_provider,
 )
-from kiln_ai.adapters.ml_model_list import KilnModelProvider, ModelProviderName
 from kiln_ai.adapters.prompt_builders import SimpleChainOfThoughtPromptBuilder
 from kiln_ai.adapters.test_prompt_adaptors import build_test_task
 
 
-def test_langchain_adapter_munge_response(tmp_path):
-    task = build_test_task(tmp_path)
-    lca = LangchainAdapter(kiln_task=task, model_name="llama_3_1_8b", provider="ollama")
+@pytest.fixture
+def mock_adapter(tmp_path):
+    return LangchainAdapter(
+        kiln_task=build_test_task(tmp_path),
+        model_name="llama_3_1_8b",
+        provider="ollama",
+    )
+
+
+def test_langchain_adapter_munge_response(mock_adapter):
     # Mistral Large tool calling format is a bit different
     response = {
         "name": "task_response",
@@ -30,12 +40,12 @@ def test_langchain_adapter_munge_response(tmp_path):
             "punchline": "Because she wanted to be a moo-sician!",
         },
     }
-    munged = lca._munge_response(response)
+    munged = mock_adapter._munge_response(response)
     assert munged["setup"] == "Why did the cow join a band?"
     assert munged["punchline"] == "Because she wanted to be a moo-sician!"
 
     # non mistral format should continue to work
-    munged = lca._munge_response(response["arguments"])
+    munged = mock_adapter._munge_response(response["arguments"])
     assert munged["setup"] == "Why did the cow join a band?"
     assert munged["punchline"] == "Because she wanted to be a moo-sician!"
 
@@ -89,9 +99,7 @@ async def test_langchain_adapter_with_cot(tmp_path):
 
     # Patch both the langchain_model_from function and self.model()
     with (
-        patch(
-            "kiln_ai.adapters.langchain_adapters.langchain_model_from", mock_model_from
-        ),
+        patch.object(LangchainAdapter, "langchain_model_from", mock_model_from),
         patch.object(LangchainAdapter, "model", return_value=mock_model_instance),
     ):
         response = await lca._run("test input")
@@ -121,8 +129,8 @@ async def test_langchain_adapter_with_cot(tmp_path):
     invoke_args = mock_model_instance.ainvoke.call_args[0][0]
     assert isinstance(invoke_args[3], AIMessage)
     assert "Chain of thought reasoning..." in invoke_args[3].content
-    assert isinstance(invoke_args[4], SystemMessage)
-    assert "Considering the above, return a final result." in invoke_args[4].content
+    assert isinstance(invoke_args[4], HumanMessage)
+    assert COT_FINAL_ANSWER_PROMPT in invoke_args[4].content
 
     assert (
         response.intermediate_outputs["chain_of_thought"]
@@ -131,46 +139,28 @@ async def test_langchain_adapter_with_cot(tmp_path):
     assert response.output == {"count": 1}
 
 
-async def test_get_structured_output_options():
+@pytest.mark.parametrize(
+    "structured_output_mode,expected_method",
+    [
+        (StructuredOutputMode.function_calling, "function_calling"),
+        (StructuredOutputMode.json_mode, "json_mode"),
+        (StructuredOutputMode.json_schema, "json_schema"),
+        (StructuredOutputMode.json_instruction_and_object, "json_mode"),
+        (StructuredOutputMode.default, None),
+    ],
+)
+async def test_get_structured_output_options(
+    mock_adapter, structured_output_mode, expected_method
+):
     # Mock the provider response
     mock_provider = MagicMock()
-    mock_provider.adapter_options = {
-        "langchain": {
-            "with_structured_output_options": {
-                "force_json_response": True,
-                "max_retries": 3,
-            }
-        }
-    }
+    mock_provider.structured_output_mode = structured_output_mode
 
-    # Test with provider that has options
-    with patch(
-        "kiln_ai.adapters.langchain_adapters.kiln_model_provider_from",
-        AsyncMock(return_value=mock_provider),
-    ):
-        options = await get_structured_output_options("model_name", "provider")
-        assert options == {"force_json_response": True, "max_retries": 3}
+    # Mock adapter.model_provider()
+    mock_adapter.model_provider = MagicMock(return_value=mock_provider)
 
-    # Test with provider that has no options
-    with patch(
-        "kiln_ai.adapters.langchain_adapters.kiln_model_provider_from",
-        AsyncMock(return_value=None),
-    ):
-        options = await get_structured_output_options("model_name", "provider")
-        assert options == {}
-
-
-@pytest.mark.asyncio
-async def test_langchain_model_from_provider_openai():
-    provider = KilnModelProvider(
-        name=ModelProviderName.openai, provider_options={"model": "gpt-4"}
-    )
-
-    with patch("kiln_ai.adapters.langchain_adapters.Config.shared") as mock_config:
-        mock_config.return_value.open_ai_api_key = "test_key"
-        model = await langchain_model_from_provider(provider, "gpt-4")
-        assert isinstance(model, ChatOpenAI)
-        assert model.model_name == "gpt-4"
+    options = mock_adapter.get_structured_output_options("model_name", "provider")
+    assert options.get("method") == expected_method
 
 
 @pytest.mark.asyncio
@@ -179,7 +169,9 @@ async def test_langchain_model_from_provider_groq():
         name=ModelProviderName.groq, provider_options={"model": "mixtral-8x7b"}
     )
 
-    with patch("kiln_ai.adapters.langchain_adapters.Config.shared") as mock_config:
+    with patch(
+        "kiln_ai.adapters.model_adapters.langchain_adapters.Config.shared"
+    ) as mock_config:
         mock_config.return_value.groq_api_key = "test_key"
         model = await langchain_model_from_provider(provider, "mixtral-8x7b")
         assert isinstance(model, ChatGroq)
@@ -193,7 +185,9 @@ async def test_langchain_model_from_provider_bedrock():
         provider_options={"model": "anthropic.claude-v2", "region_name": "us-east-1"},
     )
 
-    with patch("kiln_ai.adapters.langchain_adapters.Config.shared") as mock_config:
+    with patch(
+        "kiln_ai.adapters.model_adapters.langchain_adapters.Config.shared"
+    ) as mock_config:
         mock_config.return_value.bedrock_access_key = "test_access"
         mock_config.return_value.bedrock_secret_key = "test_secret"
         model = await langchain_model_from_provider(provider, "anthropic.claude-v2")
@@ -208,7 +202,9 @@ async def test_langchain_model_from_provider_fireworks():
         name=ModelProviderName.fireworks_ai, provider_options={"model": "mixtral-8x7b"}
     )
 
-    with patch("kiln_ai.adapters.langchain_adapters.Config.shared") as mock_config:
+    with patch(
+        "kiln_ai.adapters.model_adapters.langchain_adapters.Config.shared"
+    ) as mock_config:
         mock_config.return_value.fireworks_api_key = "test_key"
         model = await langchain_model_from_provider(provider, "mixtral-8x7b")
         assert isinstance(model, ChatFireworks)
@@ -224,15 +220,15 @@ async def test_langchain_model_from_provider_ollama():
     mock_connection = MagicMock()
     with (
         patch(
-            "kiln_ai.adapters.langchain_adapters.get_ollama_connection",
+            "kiln_ai.adapters.model_adapters.langchain_adapters.get_ollama_connection",
             return_value=AsyncMock(return_value=mock_connection),
         ),
         patch(
-            "kiln_ai.adapters.langchain_adapters.ollama_model_installed",
+            "kiln_ai.adapters.model_adapters.langchain_adapters.ollama_model_installed",
             return_value=True,
         ),
         patch(
-            "kiln_ai.adapters.langchain_adapters.ollama_base_url",
+            "kiln_ai.adapters.model_adapters.langchain_adapters.ollama_base_url",
             return_value="http://localhost:11434",
         ),
     ):
@@ -283,33 +279,27 @@ async def test_langchain_adapter_model_structured_output(tmp_path):
     mock_model.with_structured_output = MagicMock(return_value="structured_model")
 
     adapter = LangchainAdapter(
-        kiln_task=task, model_name="test_model", provider="test_provider"
+        kiln_task=task, model_name="test_model", provider="ollama"
     )
+    adapter.get_structured_output_options = MagicMock(
+        return_value={"option1": "value1"}
+    )
+    adapter.langchain_model_from = AsyncMock(return_value=mock_model)
 
-    with (
-        patch(
-            "kiln_ai.adapters.langchain_adapters.langchain_model_from",
-            AsyncMock(return_value=mock_model),
-        ),
-        patch(
-            "kiln_ai.adapters.langchain_adapters.get_structured_output_options",
-            AsyncMock(return_value={"option1": "value1"}),
-        ),
-    ):
-        model = await adapter.model()
-
-        # Verify the model was configured with structured output
-        mock_model.with_structured_output.assert_called_once_with(
-            {
-                "type": "object",
-                "properties": {"count": {"type": "integer"}},
-                "title": "task_response",
-                "description": "A response from the task",
-            },
-            include_raw=True,
-            option1="value1",
-        )
-        assert model == "structured_model"
+    model = await adapter.model()
+
+    # Verify the model was configured with structured output
+    mock_model.with_structured_output.assert_called_once_with(
+        {
+            "type": "object",
+            "properties": {"count": {"type": "integer"}},
+            "title": "task_response",
+            "description": "A response from the task",
+        },
+        include_raw=True,
+        option1="value1",
+    )
+    assert model == "structured_model"
 
 
 @pytest.mark.asyncio
@@ -324,12 +314,37 @@ async def test_langchain_adapter_model_no_structured_output_support(tmp_path):
     del mock_model.with_structured_output
 
     adapter = LangchainAdapter(
-        kiln_task=task, model_name="test_model", provider="test_provider"
+        kiln_task=task, model_name="test_model", provider="ollama"
     )
+    adapter.langchain_model_from = AsyncMock(return_value=mock_model)
 
-    with patch(
-        "kiln_ai.adapters.langchain_adapters.langchain_model_from",
-        AsyncMock(return_value=mock_model),
-    ):
-        with pytest.raises(ValueError, match="does not support structured output"):
-            await adapter.model()
+    with pytest.raises(ValueError, match="does not support structured output"):
+        await adapter.model()
+
+
+import pytest
+
+from kiln_ai.adapters.ml_model_list import KilnModelProvider, ModelProviderName
+from kiln_ai.adapters.model_adapters.langchain_adapters import (
+    langchain_model_from_provider,
+)
+
+
+@pytest.mark.parametrize(
+    "provider_name",
+    [
+        (ModelProviderName.openai),
+        (ModelProviderName.openai_compatible),
+        (ModelProviderName.openrouter),
+    ],
+)
+@pytest.mark.asyncio
+async def test_langchain_model_from_provider_unsupported_providers(provider_name):
+    # Arrange
+    provider = KilnModelProvider(
+        name=provider_name, provider_options={}, structured_output_mode="default"
+    )
+
+    # Assert unsupported providers raise an error
+    with pytest.raises(ValueError):
+        await langchain_model_from_provider(provider, "test-model")

kiln_ai/adapters/model_adapters/test_openai_model_adapter.py

@@ -0,0 +1,225 @@
+import json
+from unittest.mock import Mock, patch
+
+import pytest
+from openai import AsyncOpenAI
+
+from kiln_ai.adapters.ml_model_list import StructuredOutputMode
+from kiln_ai.adapters.model_adapters.base_adapter import AdapterInfo, BasePromptBuilder
+from kiln_ai.adapters.model_adapters.openai_compatible_config import (
+    OpenAICompatibleConfig,
+)
+from kiln_ai.adapters.model_adapters.openai_model_adapter import OpenAICompatibleAdapter
+from kiln_ai.datamodel import Project, Task
+
+
+@pytest.fixture
+def mock_task(tmp_path):
+    # Create a project first since Task requires a parent
+    project_path = tmp_path / "test_project" / "project.kiln"
+    project_path.parent.mkdir()
+
+    project = Project(name="Test Project", path=str(project_path))
+    project.save_to_file()
+
+    schema = {
+        "type": "object",
+        "properties": {"test": {"type": "string"}},
+    }
+
+    task = Task(
+        name="Test Task",
+        instruction="Test instruction",
+        parent=project,
+        output_json_schema=json.dumps(schema),
+    )
+    task.save_to_file()
+    return task
+
+
+@pytest.fixture
+def mock_prompt_builder():
+    builder = Mock(spec=BasePromptBuilder)
+    type(builder).prompt_builder_name = Mock(return_value="test_prompt_builder")
+    builder.prompt_id = Mock(return_value="test_prompt_id")
+    return builder
+
+
+@pytest.fixture
+def config():
+    return OpenAICompatibleConfig(
+        api_key="test_key",
+        base_url="https://api.test.com",
+        model_name="test-model",
+        provider_name="test-provider",
+        default_headers={"X-Test": "test"},
+    )
+
+
+def test_initialization(config, mock_task, mock_prompt_builder):
+    adapter = OpenAICompatibleAdapter(
+        config=config,
+        kiln_task=mock_task,
+        prompt_builder=mock_prompt_builder,
+        tags=["test-tag"],
+    )
+
+    assert isinstance(adapter.client, AsyncOpenAI)
+    assert adapter.config == config
+    assert adapter.kiln_task == mock_task
+    assert adapter.prompt_builder == mock_prompt_builder
+    assert adapter.default_tags == ["test-tag"]
+    assert adapter.model_name == config.model_name
+    assert adapter.model_provider_name == config.provider_name
+
+
+def test_adapter_info(config, mock_task, mock_prompt_builder):
+    adapter = OpenAICompatibleAdapter(
+        config=config, kiln_task=mock_task, prompt_builder=mock_prompt_builder
+    )
+
+    info = adapter.adapter_info()
+    assert isinstance(info, AdapterInfo)
+    assert info.model_name == config.model_name
+    assert info.model_provider == config.provider_name
+    assert info.adapter_name == "kiln_openai_compatible_adapter"
+    assert info.prompt_builder_name == "base_prompt_builder"
+    assert info.prompt_id == "test_prompt_id"
+
+
+@pytest.mark.asyncio
+async def test_response_format_options_unstructured(
+    config, mock_task, mock_prompt_builder
+):
+    adapter = OpenAICompatibleAdapter(
+        config=config, kiln_task=mock_task, prompt_builder=mock_prompt_builder
+    )
+
+    # Mock has_structured_output to return False
+    with patch.object(adapter, "has_structured_output", return_value=False):
+        options = await adapter.response_format_options()
+        assert options == {}
+
+
+@pytest.mark.parametrize(
+    "mode",
+    [
+        StructuredOutputMode.json_mode,
+        StructuredOutputMode.json_instruction_and_object,
+    ],
+)
+@pytest.mark.asyncio
+async def test_response_format_options_json_mode(
+    config, mock_task, mock_prompt_builder, mode
+):
+    adapter = OpenAICompatibleAdapter(
+        config=config, kiln_task=mock_task, prompt_builder=mock_prompt_builder
+    )
+
+    with (
+        patch.object(adapter, "has_structured_output", return_value=True),
+        patch.object(adapter, "model_provider") as mock_provider,
+    ):
+        mock_provider.return_value.structured_output_mode = mode
+
+        options = await adapter.response_format_options()
+        assert options == {"response_format": {"type": "json_object"}}
+
+
+@pytest.mark.parametrize(
+    "mode",
+    [
+        StructuredOutputMode.default,
+        StructuredOutputMode.function_calling,
+    ],
+)
+@pytest.mark.asyncio
+async def test_response_format_options_function_calling(
+    config, mock_task, mock_prompt_builder, mode
+):
+    adapter = OpenAICompatibleAdapter(
+        config=config, kiln_task=mock_task, prompt_builder=mock_prompt_builder
+    )
+
+    with (
+        patch.object(adapter, "has_structured_output", return_value=True),
+        patch.object(adapter, "model_provider") as mock_provider,
+    ):
+        mock_provider.return_value.structured_output_mode = mode
+
+        options = await adapter.response_format_options()
+        assert "tools" in options
+        # full tool structure validated below
+
+
+@pytest.mark.asyncio
+async def test_response_format_options_json_instructions(
+    config, mock_task, mock_prompt_builder
+):
+    adapter = OpenAICompatibleAdapter(
+        config=config, kiln_task=mock_task, prompt_builder=mock_prompt_builder
+    )
+
+    with (
+        patch.object(adapter, "has_structured_output", return_value=True),
+        patch.object(adapter, "model_provider") as mock_provider,
+    ):
+        mock_provider.return_value.structured_output_mode = (
+            StructuredOutputMode.json_instructions
+        )
+        options = await adapter.response_format_options()
+        assert options == {}
+
+
+@pytest.mark.asyncio
+async def test_response_format_options_json_schema(
+    config, mock_task, mock_prompt_builder
+):
+    adapter = OpenAICompatibleAdapter(
+        config=config, kiln_task=mock_task, prompt_builder=mock_prompt_builder
+    )
+
+    with (
+        patch.object(adapter, "has_structured_output", return_value=True),
+        patch.object(adapter, "model_provider") as mock_provider,
+    ):
+        mock_provider.return_value.structured_output_mode = (
+            StructuredOutputMode.json_schema
+        )
+        options = await adapter.response_format_options()
+        assert options == {
+            "response_format": {
+                "type": "json_schema",
+                "json_schema": {
+                    "name": "task_response",
+                    "schema": mock_task.output_schema(),
+                },
+            }
+        }
+
+
+def test_tool_call_params(config, mock_task, mock_prompt_builder):
+    adapter = OpenAICompatibleAdapter(
+        config=config, kiln_task=mock_task, prompt_builder=mock_prompt_builder
+    )
+
+    params = adapter.tool_call_params()
+    expected_schema = mock_task.output_schema()
+    expected_schema["additionalProperties"] = False
+
+    assert params == {
+        "tools": [
+            {
+                "type": "function",
+                "function": {
+                    "name": "task_response",
+                    "parameters": expected_schema,
+                    "strict": True,
+                },
+            }
+        ],
+        "tool_choice": {
+            "type": "function",
+            "function": {"name": "task_response"},
+        },
+    }

kiln_ai/adapters/{test_saving_adapter_results.py → model_adapters/test_saving_adapter_results.py}

@@ -2,7 +2,11 @@ from unittest.mock import patch
 
 import pytest
 
-from kiln_ai.adapters.base_adapter import AdapterInfo, BaseAdapter, RunOutput
+from kiln_ai.adapters.model_adapters.base_adapter import (
+    AdapterInfo,
+    BaseAdapter,
+    RunOutput,
+)
 from kiln_ai.datamodel import (
     DataSource,
     DataSourceType,
@@ -22,6 +26,7 @@ class MockAdapter(BaseAdapter):
             model_name="mock_model",
             model_provider="mock_provider",
             prompt_builder_name="mock_prompt_builder",
+            prompt_id="mock_prompt_id",
         )
 
 
@@ -38,8 +43,12 @@ def test_task(tmp_path):
     return task
 
 
-def test_save_run_isolation(test_task):
-    adapter = MockAdapter(test_task)
+@pytest.fixture
+def adapter(test_task):
+    return MockAdapter(test_task, model_name="phi_3_5", model_provider_name="ollama")
+
+
+def test_save_run_isolation(test_task, adapter):
     input_data = "Test input"
     output_data = "Test output"
     run_output = RunOutput(
@@ -91,7 +100,7 @@ def test_save_run_isolation(test_task):
         reloaded_output.source.properties["prompt_builder_name"]
         == "mock_prompt_builder"
     )
-
+    assert reloaded_output.source.properties["prompt_id"] == "mock_prompt_id"
     # Run again, with same input and different output. Should create a new TaskRun.
     different_run_output = RunOutput(
         output="Different output", intermediate_outputs=None
@@ -101,13 +110,6 @@ def test_save_run_isolation(test_task):
     assert len(test_task.runs()) == 2
     assert "Different output" in set(run.output.output for run in test_task.runs())
 
-    # run again with same input and same output. Should not create a new TaskRun.
-    task_output = adapter.generate_run(input_data, None, run_output)
-    task_output.save_to_file()
-    assert len(test_task.runs()) == 2
-    assert "Different output" in set(run.output.output for run in test_task.runs())
-    assert output_data in set(run.output.output for run in test_task.runs())
-
     # run again with input of different type. Should create a new TaskRun and TaskOutput.
     task_output = adapter.generate_run(
         input_data,
@@ -130,14 +132,41 @@ def test_save_run_isolation(test_task):
     assert output_data in set(run.output.output for run in test_task.runs())
 
 
+def test_generate_run_non_ascii(test_task, adapter):
+    input_data = {"key": "input with non-ascii character: 你好"}
+    output_data = {"key": "output with non-ascii character: 你好"}
+    run_output = RunOutput(
+        output=output_data,
+        intermediate_outputs=None,
+    )
+
+    task_run = adapter.generate_run(
+        input=input_data, input_source=None, run_output=run_output
+    )
+    task_run.save_to_file()
+
+    # as these values are saved as strings, they should properly represent the non-ascii characters
+    assert task_run.input == '{"key": "input with non-ascii character: 你好"}'
+    assert task_run.output.output == '{"key": "output with non-ascii character: 你好"}'
+
+    # check that the stringified unicode strings can be read back from the file
+    reloaded_task = Task.load_from_file(test_task.path)
+    reloaded_runs = reloaded_task.runs()
+    assert len(reloaded_runs) == 1
+    reloaded_run = reloaded_runs[0]
+    assert reloaded_run.input == '{"key": "input with non-ascii character: 你好"}'
+    assert (
+        reloaded_run.output.output == '{"key": "output with non-ascii character: 你好"}'
+    )
+
+
 @pytest.mark.asyncio
-async def test_autosave_false(test_task):
+async def test_autosave_false(test_task, adapter):
     with patch("kiln_ai.utils.config.Config.shared") as mock_shared:
         mock_config = mock_shared.return_value
         mock_config.autosave_runs = False
         mock_config.user_id = "test_user"
 
-        adapter = MockAdapter(test_task)
         input_data = "Test input"
 
         run = await adapter.invoke(input_data)
@@ -150,13 +179,12 @@ async def test_autosave_false(test_task):
 
 
 @pytest.mark.asyncio
-async def test_autosave_true(test_task):
+async def test_autosave_true(test_task, adapter):
     with patch("kiln_ai.utils.config.Config.shared") as mock_shared:
         mock_config = mock_shared.return_value
         mock_config.autosave_runs = True
         mock_config.user_id = "test_user"
 
-        adapter = MockAdapter(test_task)
         input_data = "Test input"
 
         run = await adapter.invoke(input_data)