kiln-ai 0.8.1__py3-none-any.whl → 0.12.0__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
+from kiln_ai.datamodel.task import RunConfig
 
 
-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!"
 
@@ -46,9 +56,8 @@ def test_langchain_adapter_infer_model_name(tmp_path):
 
     lca = LangchainAdapter(kiln_task=task, custom_model=custom)
 
-    model_info = lca.adapter_info()
-    assert model_info.model_name == "custom.langchain:llama-3.1-8b-instant"
-    assert model_info.model_provider == "custom.langchain:ChatGroq"
+    assert lca.run_config.model_name == "custom.langchain:llama-3.1-8b-instant"
+    assert lca.run_config.model_provider_name == "custom.langchain:ChatGroq"
 
 
 def test_langchain_adapter_info(tmp_path):
@@ -56,10 +65,9 @@ def test_langchain_adapter_info(tmp_path):
 
     lca = LangchainAdapter(kiln_task=task, model_name="llama_3_1_8b", provider="ollama")
 
-    model_info = lca.adapter_info()
-    assert model_info.adapter_name == "kiln_langchain_adapter"
-    assert model_info.model_name == "llama_3_1_8b"
-    assert model_info.model_provider == "ollama"
+    assert lca.adapter_name() == "kiln_langchain_adapter"
+    assert lca.run_config.model_name == "llama_3_1_8b"
+    assert lca.run_config.model_provider_name == "ollama"
 
 
 async def test_langchain_adapter_with_cot(tmp_path):
@@ -71,7 +79,7 @@ async def test_langchain_adapter_with_cot(tmp_path):
         kiln_task=task,
         model_name="llama_3_1_8b",
         provider="ollama",
-        prompt_builder=SimpleChainOfThoughtPromptBuilder(task),
+        prompt_id="simple_chain_of_thought_prompt_builder",
     )
 
     # Mock the base model and its invoke method
@@ -89,9 +97,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 +127,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 +137,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,
-            }
-        }
-    }
-
-    # 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}
-
-    # 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 == {}
+    mock_provider.structured_output_mode = structured_output_mode
 
+    # Mock adapter.model_provider()
+    mock_adapter.model_provider = MagicMock(return_value=mock_provider)
 
-@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 +167,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 +183,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 +200,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 +218,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 +277,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 +312,32 @@ 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
+
+
+@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,216 @@
+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 AdapterConfig
+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 config():
+    return OpenAICompatibleConfig(
+        api_key="test_key",
+        base_url="https://api.test.com",
+        model_name="test-model",
+        provider_name="openrouter",
+        default_headers={"X-Test": "test"},
+    )
+
+
+def test_initialization(config, mock_task):
+    adapter = OpenAICompatibleAdapter(
+        config=config,
+        kiln_task=mock_task,
+        prompt_id="simple_prompt_builder",
+        base_adapter_config=AdapterConfig(default_tags=["test-tag"]),
+    )
+
+    assert isinstance(adapter.client, AsyncOpenAI)
+    assert adapter.config == config
+    assert adapter.run_config.task == mock_task
+    assert adapter.run_config.prompt_id == "simple_prompt_builder"
+    assert adapter.base_adapter_config.default_tags == ["test-tag"]
+    assert adapter.run_config.model_name == config.model_name
+    assert adapter.run_config.model_provider_name == config.provider_name
+
+
+def test_adapter_info(config, mock_task):
+    adapter = OpenAICompatibleAdapter(config=config, kiln_task=mock_task)
+
+    assert adapter.adapter_name() == "kiln_openai_compatible_adapter"
+
+    assert adapter.run_config.model_name == config.model_name
+    assert adapter.run_config.model_provider_name == config.provider_name
+    assert adapter.run_config.prompt_id == "simple_prompt_builder"
+
+
+@pytest.mark.asyncio
+async def test_response_format_options_unstructured(config, mock_task):
+    adapter = OpenAICompatibleAdapter(config=config, kiln_task=mock_task)
+
+    # 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, mode):
+    adapter = OpenAICompatibleAdapter(config=config, kiln_task=mock_task)
+
+    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, mode):
+    adapter = OpenAICompatibleAdapter(config=config, kiln_task=mock_task)
+
+    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):
+    adapter = OpenAICompatibleAdapter(config=config, kiln_task=mock_task)
+
+    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):
+    adapter = OpenAICompatibleAdapter(config=config, kiln_task=mock_task)
+
+    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_weak(config, mock_task):
+    adapter = OpenAICompatibleAdapter(config=config, kiln_task=mock_task)
+
+    params = adapter.tool_call_params(strict=False)
+    expected_schema = mock_task.output_schema()
+    expected_schema["additionalProperties"] = False
+
+    assert params == {
+        "tools": [
+            {
+                "type": "function",
+                "function": {
+                    "name": "task_response",
+                    "parameters": expected_schema,
+                },
+            }
+        ],
+        "tool_choice": {
+            "type": "function",
+            "function": {"name": "task_response"},
+        },
+    }
+
+
+def test_tool_call_params_strict(config, mock_task):
+    config.provider_name = "openai"
+    adapter = OpenAICompatibleAdapter(config=config, kiln_task=mock_task)
+
+    params = adapter.tool_call_params(strict=True)
+    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"},
+        },
+    }