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

kiln_ai/adapters/model_adapters/test_base_adapter.py

@@ -6,7 +6,11 @@ from kiln_ai.adapters.ml_model_list import KilnModelProvider, StructuredOutputMo
 from kiln_ai.adapters.model_adapters.base_adapter import BaseAdapter, RunOutput
 from kiln_ai.datamodel import Task
 from kiln_ai.datamodel.datamodel_enums import ChatStrategy
-from kiln_ai.datamodel.task import RunConfig, RunConfigProperties
+from kiln_ai.datamodel.project import Project
+from kiln_ai.datamodel.run_config import ToolsRunConfig
+from kiln_ai.datamodel.task import RunConfigProperties
+from kiln_ai.datamodel.tool_id import KilnBuiltInToolId
+from kiln_ai.tools.base_tool import KilnToolInterface
 
 
 class MockAdapter(BaseAdapter):
@@ -27,15 +31,21 @@ def mock_provider():
 
 
 @pytest.fixture
-def base_task():
-    return Task(name="test_task", instruction="test_instruction")
+def base_project():
+    return Project(name="test_project", description="test project description")
+
+
+@pytest.fixture
+def base_task(base_project):
+    task = Task(name="test_task", instruction="test_instruction", parent=base_project)
+    return task
 
 
 @pytest.fixture
 def adapter(base_task):
     return MockAdapter(
-        run_config=RunConfig(
-            task=base_task,
+        task=base_task,
+        run_config=RunConfigProperties(
             model_name="test_model",
             model_provider_name="openai",
             prompt_id="simple_prompt_builder",
@@ -98,13 +108,16 @@ async def test_model_provider_loads_and_caches(adapter, mock_provider):
         mock_loader.assert_not_called()
 
 
-async def test_model_provider_invalid_provider_model_name(base_task):
+async def test_model_provider_invalid_provider_model_name(base_project):
     """Test error when model or provider name is missing"""
+    # Create a task with a parent project
+    task = Task(name="test_task", instruction="test_instruction", parent=base_project)
+
     # Test with missing model name
     with pytest.raises(ValueError, match="Input should be"):
         MockAdapter(
-            run_config=RunConfig(
-                task=base_task,
+            task=task,
+            run_config=RunConfigProperties(
                 model_name="test_model",
                 model_provider_name="invalid",
                 prompt_id="simple_prompt_builder",
@@ -112,12 +125,15 @@ async def test_model_provider_invalid_provider_model_name(base_task):
         )
 
 
-async def test_model_provider_missing_model_names(base_task):
+async def test_model_provider_missing_model_names(base_project):
     """Test error when model or provider name is missing"""
+    # Create a task with a parent project
+    task = Task(name="test_task", instruction="test_instruction", parent=base_project)
+
     # Test with missing model name
     adapter = MockAdapter(
-        run_config=RunConfig(
-            task=base_task,
+        task=task,
+        run_config=RunConfigProperties(
             model_name="",
             model_provider_name="openai",
             prompt_id="simple_prompt_builder",
@@ -252,12 +268,13 @@ async def test_properties_for_task_output_includes_all_run_config_properties(ada
         "temperature": "temperature",
         "top_p": "top_p",
         "structured_output_mode": "structured_output_mode",
+        "tools_config": None,
     }
 
     missing_properties = []
     for field_name in run_config_properties_fields:
         expected_key = expected_mappings.get(field_name, field_name)
-        if expected_key not in saved_property_keys:
+        if expected_key is not None and expected_key not in saved_property_keys:
             missing_properties.append(
                 f"RunConfigProperties.{field_name} -> {expected_key}"
             )
@@ -297,12 +314,13 @@ async def test_properties_for_task_output_catches_missing_new_property(adapter):
             "temperature": "temperature",
             "top_p": "top_p",
             "structured_output_mode": "structured_output_mode",
+            "tools_config": None,
         }
 
         missing_properties = []
         for field_name in run_config_properties_fields:
             expected_key = expected_mappings.get(field_name, field_name)
-            if expected_key not in saved_property_keys:
+            if expected_key is not None and expected_key not in saved_property_keys:
                 missing_properties.append(
                     f"RunConfigProperties.{field_name} -> {expected_key}"
                 )
@@ -408,12 +426,14 @@ def test_build_chat_formatter(
     ],
 )
 async def test_update_run_config_unknown_structured_output_mode(
-    base_task, initial_mode, expected_mode
+    base_project, initial_mode, expected_mode
 ):
     """Test that unknown structured output mode is updated to the default for the model provider"""
+    # Create a task with a parent project
+    task = Task(name="test_task", instruction="test_instruction", parent=base_project)
+
     # Create a run config with the initial mode
-    run_config = RunConfig(
-        task=base_task,
+    run_config = RunConfigProperties(
         model_name="test_model",
         model_provider_name="openai",
         prompt_id="simple_prompt_builder",
@@ -429,7 +449,7 @@ async def test_update_run_config_unknown_structured_output_mode(
         mock_default.return_value = StructuredOutputMode.json_mode
 
         # Create the adapter
-        adapter = MockAdapter(run_config=run_config)
+        adapter = MockAdapter(task=task, run_config=run_config)
 
         # Verify the mode was updated correctly
         assert adapter.run_config.structured_output_mode == expected_mode
@@ -443,3 +463,159 @@ async def test_update_run_config_unknown_structured_output_mode(
             mock_default.assert_called_once_with("test_model", "openai")
         else:
             mock_default.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    "tools_config,expected_tool_count,expected_tool_ids",
+    [
+        # No tools config
+        (None, 0, []),
+        # Empty tools config with None tools
+        (ToolsRunConfig(tools=[]), 0, []),
+        # Single tool
+        ([KilnBuiltInToolId.ADD_NUMBERS], 1, [KilnBuiltInToolId.ADD_NUMBERS]),
+        # Multiple tools
+        (
+            [KilnBuiltInToolId.ADD_NUMBERS, KilnBuiltInToolId.SUBTRACT_NUMBERS],
+            2,
+            [KilnBuiltInToolId.ADD_NUMBERS, KilnBuiltInToolId.SUBTRACT_NUMBERS],
+        ),
+        # All available built-in tools
+        (
+            [
+                KilnBuiltInToolId.ADD_NUMBERS,
+                KilnBuiltInToolId.SUBTRACT_NUMBERS,
+                KilnBuiltInToolId.MULTIPLY_NUMBERS,
+                KilnBuiltInToolId.DIVIDE_NUMBERS,
+            ],
+            4,
+            [
+                KilnBuiltInToolId.ADD_NUMBERS,
+                KilnBuiltInToolId.SUBTRACT_NUMBERS,
+                KilnBuiltInToolId.MULTIPLY_NUMBERS,
+                KilnBuiltInToolId.DIVIDE_NUMBERS,
+            ],
+        ),
+    ],
+)
+async def test_available_tools(
+    base_project, tools_config, expected_tool_count, expected_tool_ids
+):
+    """Test that available_tools returns correct tools based on tools_config"""
+    # Create a task with a parent project
+    task = Task(name="test_task", instruction="test_instruction", parent=base_project)
+
+    # Create tools config if we have tool IDs
+    if tools_config is None:
+        final_tools_config = None
+    elif isinstance(tools_config, list):
+        final_tools_config = ToolsRunConfig(tools=tools_config)
+    else:
+        final_tools_config = tools_config
+
+    # Create adapter with tools config
+    adapter = MockAdapter(
+        task=task,
+        run_config=RunConfigProperties(
+            model_name="test_model",
+            model_provider_name="openai",
+            prompt_id="simple_prompt_builder",
+            structured_output_mode="json_schema",
+            tools_config=final_tools_config,
+        ),
+    )
+
+    # Get available tools
+    tools = await adapter.available_tools()
+
+    # Verify tool count
+    assert len(tools) == expected_tool_count
+
+    # Verify all tools implement KilnToolInterface
+    for tool in tools:
+        assert isinstance(tool, KilnToolInterface)
+
+    # Verify tool IDs match expected
+    if expected_tool_ids:
+        actual_tool_ids = [await tool.id() for tool in tools]
+        assert actual_tool_ids == expected_tool_ids
+
+
+async def test_available_tools_with_invalid_tool_id(base_project):
+    """Test that available_tools raises ValueError for invalid tool ID"""
+    # Create a task with a parent project
+    task = Task(name="test_task", instruction="test_instruction", parent=base_project)
+
+    # Create tools config with valid tool ID
+    tools_config = ToolsRunConfig(tools=[KilnBuiltInToolId.ADD_NUMBERS])
+
+    # Create adapter
+    adapter = MockAdapter(
+        task=task,
+        run_config=RunConfigProperties(
+            model_name="test_model",
+            model_provider_name="openai",
+            prompt_id="simple_prompt_builder",
+            structured_output_mode="json_schema",
+            tools_config=tools_config,
+        ),
+    )
+
+    # Mock tool_from_id to raise ValueError for any tool ID
+    with patch(
+        "kiln_ai.adapters.model_adapters.base_adapter.tool_from_id"
+    ) as mock_tool_from_id:
+        mock_tool_from_id.side_effect = ValueError(
+            "Tool ID test_id not found in tool registry"
+        )
+
+        # Should raise ValueError when trying to get tools
+        with pytest.raises(
+            ValueError, match="Tool ID test_id not found in tool registry"
+        ):
+            await adapter.available_tools()
+
+
+async def test_available_tools_duplicate_names_raises_error(base_project):
+    """Test that available_tools raises ValueError when tools have duplicate names"""
+    # Create a task with a parent project
+    task = Task(name="test_task", instruction="test_instruction", parent=base_project)
+
+    # Create tools config with two different tool IDs
+    tools_config = ToolsRunConfig(
+        tools=[KilnBuiltInToolId.ADD_NUMBERS, KilnBuiltInToolId.SUBTRACT_NUMBERS]
+    )
+
+    # Create adapter
+    adapter = MockAdapter(
+        task=task,
+        run_config=RunConfigProperties(
+            model_name="test_model",
+            model_provider_name="openai",
+            prompt_id="simple_prompt_builder",
+            structured_output_mode="json_schema",
+            tools_config=tools_config,
+        ),
+    )
+
+    # Create mock tools with duplicate names
+    async def mock_name1():
+        return "duplicate_name"
+
+    async def mock_name2():
+        return "duplicate_name"
+
+    mock_tool1 = MagicMock(spec=KilnToolInterface)
+    mock_tool1.name = mock_name1
+    mock_tool2 = MagicMock(spec=KilnToolInterface)
+    mock_tool2.name = mock_name2  # Same name as tool1
+
+    # Mock tool_from_id to return our mock tools with duplicate names
+    with patch(
+        "kiln_ai.adapters.model_adapters.base_adapter.tool_from_id"
+    ) as mock_tool_from_id:
+        mock_tool_from_id.side_effect = [mock_tool1, mock_tool2]
+
+        # Should raise ValueError when tools have duplicate names
+        with pytest.raises(ValueError, match="Each tool must have a unique name"):
+            await adapter.available_tools()