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

kiln_ai/datamodel/test_task.py

@@ -3,22 +3,18 @@ from pydantic import ValidationError
 
 from kiln_ai.datamodel.datamodel_enums import StructuredOutputMode, TaskOutputRatingType
 from kiln_ai.datamodel.prompt_id import PromptGenerators
-from kiln_ai.datamodel.task import RunConfig, RunConfigProperties, Task, TaskRunConfig
+from kiln_ai.datamodel.task import RunConfigProperties, Task, TaskRunConfig
 from kiln_ai.datamodel.task_output import normalize_rating
 
 
 def test_runconfig_valid_creation():
-    task = Task(id="task1", name="Test Task", instruction="Do something")
-
-    config = RunConfig(
-        task=task,
+    config = RunConfigProperties(
         model_name="gpt-4",
         model_provider_name="openai",
         prompt_id=PromptGenerators.SIMPLE,
         structured_output_mode="json_schema",
     )
 
-    assert config.task == task
     assert config.model_name == "gpt-4"
     assert config.model_provider_name == "openai"
     assert config.prompt_id == PromptGenerators.SIMPLE  # Check default value
@@ -26,13 +22,12 @@ def test_runconfig_valid_creation():
 
 def test_runconfig_missing_required_fields():
     with pytest.raises(ValidationError) as exc_info:
-        RunConfig()
+        RunConfigProperties()  # type: ignore
 
     errors = exc_info.value.errors()
     assert (
-        len(errors) == 5
+        len(errors) == 4
     )  # task, model_name, model_provider_name, and prompt_id are required
-    assert any(error["loc"][0] == "task" for error in errors)
     assert any(error["loc"][0] == "model_name" for error in errors)
     assert any(error["loc"][0] == "model_provider_name" for error in errors)
     assert any(error["loc"][0] == "prompt_id" for error in errors)
@@ -40,10 +35,7 @@ def test_runconfig_missing_required_fields():
 
 
 def test_runconfig_custom_prompt_id():
-    task = Task(id="task1", name="Test Task", instruction="Do something")
-
-    config = RunConfig(
-        task=task,
+    config = RunConfigProperties(
         model_name="gpt-4",
         model_provider_name="openai",
         prompt_id=PromptGenerators.SIMPLE_CHAIN_OF_THOUGHT,
@@ -100,30 +92,18 @@ def test_task_run_config_missing_required_fields(sample_task):
     with pytest.raises(ValidationError) as exc_info:
         TaskRunConfig(
             run_config_properties=RunConfigProperties(
-                task=sample_task, model_name="gpt-4", model_provider_name="openai"
-            ),
+                model_name="gpt-4", model_provider_name="openai"
+            ),  # type: ignore
             parent=sample_task,
-        )
+        )  # type: ignore
     assert "Field required" in str(exc_info.value)
 
     # Test missing run_config
     with pytest.raises(ValidationError) as exc_info:
-        TaskRunConfig(name="Test Config", parent=sample_task)
+        TaskRunConfig(name="Test Config", parent=sample_task)  # type: ignore
     assert "Field required" in str(exc_info.value)
 
 
-def test_task_run_config_missing_task_in_run_config(sample_task):
-    with pytest.raises(
-        ValidationError, match="Input should be a valid dictionary or instance of Task"
-    ):
-        # Create a run config without a task
-        RunConfig(
-            model_name="gpt-4",
-            model_provider_name="openai",
-            task=None,  # type: ignore
-        )
-
-
 @pytest.mark.parametrize(
     "rating_type,rating,expected",
     [
@@ -165,10 +145,8 @@ def test_normalize_rating_errors(rating_type, rating):
 
 def test_run_config_defaults():
     """RunConfig should require top_p, temperature, and structured_output_mode to be set."""
-    task = Task(id="task1", name="Test Task", instruction="Do something")
 
-    config = RunConfig(
-        task=task,
+    config = RunConfigProperties(
         model_name="gpt-4",
         model_provider_name="openai",
         prompt_id=PromptGenerators.SIMPLE,
@@ -180,11 +158,9 @@ def test_run_config_defaults():
 
 def test_run_config_valid_ranges():
     """RunConfig should accept valid ranges for top_p and temperature."""
-    task = Task(id="task1", name="Test Task", instruction="Do something")
 
     # Test valid values
-    config = RunConfig(
-        task=task,
+    config = RunConfigProperties(
         model_name="gpt-4",
         model_provider_name="openai",
         prompt_id=PromptGenerators.SIMPLE,
@@ -201,10 +177,8 @@ def test_run_config_valid_ranges():
 @pytest.mark.parametrize("top_p", [0.0, 0.5, 1.0])
 def test_run_config_valid_top_p(top_p):
     """Test that RunConfig accepts valid top_p values (0-1)."""
-    task = Task(id="task1", name="Test Task", instruction="Do something")
 
-    config = RunConfig(
-        task=task,
+    config = RunConfigProperties(
         model_name="gpt-4",
         model_provider_name="openai",
         prompt_id=PromptGenerators.SIMPLE,
@@ -219,11 +193,9 @@ def test_run_config_valid_top_p(top_p):
 @pytest.mark.parametrize("top_p", [-0.1, 1.1, 2.0])
 def test_run_config_invalid_top_p(top_p):
     """Test that RunConfig rejects invalid top_p values."""
-    task = Task(id="task1", name="Test Task", instruction="Do something")
 
     with pytest.raises(ValueError, match="top_p must be between 0 and 1"):
-        RunConfig(
-            task=task,
+        RunConfigProperties(
             model_name="gpt-4",
             model_provider_name="openai",
             prompt_id=PromptGenerators.SIMPLE,
@@ -236,10 +208,8 @@ def test_run_config_invalid_top_p(top_p):
 @pytest.mark.parametrize("temperature", [0.0, 1.0, 2.0])
 def test_run_config_valid_temperature(temperature):
     """Test that RunConfig accepts valid temperature values (0-2)."""
-    task = Task(id="task1", name="Test Task", instruction="Do something")
 
-    config = RunConfig(
-        task=task,
+    config = RunConfigProperties(
         model_name="gpt-4",
         model_provider_name="openai",
         prompt_id=PromptGenerators.SIMPLE,
@@ -254,11 +224,9 @@ def test_run_config_valid_temperature(temperature):
 @pytest.mark.parametrize("temperature", [-0.1, 2.1, 3.0])
 def test_run_config_invalid_temperature(temperature):
     """Test that RunConfig rejects invalid temperature values."""
-    task = Task(id="task1", name="Test Task", instruction="Do something")
 
     with pytest.raises(ValueError, match="temperature must be between 0 and 2"):
-        RunConfig(
-            task=task,
+        RunConfigProperties(
             model_name="gpt-4",
             model_provider_name="openai",
             prompt_id=PromptGenerators.SIMPLE,

kiln_ai/datamodel/test_tool_id.py

@@ -0,0 +1,320 @@
+import pytest
+from pydantic import BaseModel, ValidationError
+
+from kiln_ai.datamodel.tool_id import (
+    MCP_LOCAL_TOOL_ID_PREFIX,
+    MCP_REMOTE_TOOL_ID_PREFIX,
+    RAG_TOOL_ID_PREFIX,
+    KilnBuiltInToolId,
+    ToolId,
+    _check_tool_id,
+    mcp_server_and_tool_name_from_id,
+    rag_config_id_from_id,
+)
+
+
+class TestKilnBuiltInToolId:
+    """Test the KilnBuiltInToolId enum."""
+
+    def test_enum_values(self):
+        """Test that enum has expected values."""
+        assert KilnBuiltInToolId.ADD_NUMBERS == "kiln_tool::add_numbers"
+        assert KilnBuiltInToolId.SUBTRACT_NUMBERS == "kiln_tool::subtract_numbers"
+        assert KilnBuiltInToolId.MULTIPLY_NUMBERS == "kiln_tool::multiply_numbers"
+        assert KilnBuiltInToolId.DIVIDE_NUMBERS == "kiln_tool::divide_numbers"
+        for enum_value in KilnBuiltInToolId.__members__.values():
+            assert _check_tool_id(enum_value) == enum_value
+
+    def test_enum_membership(self):
+        """Test enum membership checks."""
+        assert "kiln_tool::add_numbers" in KilnBuiltInToolId.__members__.values()
+        assert "invalid_tool" not in KilnBuiltInToolId.__members__.values()
+
+
+class TestCheckToolId:
+    """Test the _check_tool_id validation function."""
+
+    def test_valid_builtin_tools(self):
+        """Test validation of valid built-in tools."""
+        for tool_id in KilnBuiltInToolId:
+            result = _check_tool_id(tool_id.value)
+            assert result == tool_id.value
+
+    def test_valid_mcp_remote_tools(self):
+        """Test validation of valid MCP remote tools."""
+        valid_ids = [
+            "mcp::remote::server1::tool1",
+            "mcp::remote::my_server::my_tool",
+            "mcp::remote::test::function_name",
+        ]
+        for tool_id in valid_ids:
+            result = _check_tool_id(tool_id)
+            assert result == tool_id
+
+    def test_valid_mcp_local_tools(self):
+        """Test validation of valid MCP local tools."""
+        valid_ids = [
+            "mcp::local::server1::tool1",
+            "mcp::local::my_server::my_tool",
+            "mcp::local::test::function_name",
+        ]
+        for tool_id in valid_ids:
+            result = _check_tool_id(tool_id)
+            assert result == tool_id
+
+    def test_invalid_empty_or_none(self):
+        """Test validation fails for empty or None values."""
+        with pytest.raises(ValueError, match="Invalid tool ID"):
+            _check_tool_id("")
+
+        with pytest.raises(ValueError, match="Invalid tool ID"):
+            _check_tool_id(None)  # type: ignore
+
+    def test_invalid_non_string(self):
+        """Test validation fails for non-string values."""
+        with pytest.raises(ValueError, match="Invalid tool ID"):
+            _check_tool_id(123)  # type: ignore
+
+        with pytest.raises(ValueError, match="Invalid tool ID"):
+            _check_tool_id(["tool"])  # type: ignore
+
+    def test_invalid_unknown_tool(self):
+        """Test validation fails for unknown tool IDs."""
+        with pytest.raises(ValueError, match="Invalid tool ID: unknown_tool"):
+            _check_tool_id("unknown_tool")
+
+    def test_invalid_mcp_format(self):
+        """Test validation fails for invalid MCP tool formats."""
+        # These IDs start with the MCP remote prefix but have invalid formats
+        mcp_remote_invalid_ids = [
+            "mcp::remote::",  # Missing server and tool
+            "mcp::remote::server",  # Missing tool
+            "mcp::remote::server::",  # Empty tool name
+            "mcp::remote::::tool",  # Empty server name
+            "mcp::remote::server::tool::extra",  # Too many parts
+        ]
+
+        for invalid_id in mcp_remote_invalid_ids:
+            with pytest.raises(ValueError, match="Invalid remote MCP tool ID"):
+                _check_tool_id(invalid_id)
+
+        # These IDs start with the MCP local prefix but have invalid formats
+        mcp_local_invalid_ids = [
+            "mcp::local::",  # Missing server and tool
+            "mcp::local::server",  # Missing tool
+            "mcp::local::server::",  # Empty tool name
+            "mcp::local::::tool",  # Empty server name
+            "mcp::local::server::tool::extra",  # Too many parts
+        ]
+
+        for invalid_id in mcp_local_invalid_ids:
+            with pytest.raises(ValueError, match="Invalid local MCP tool ID"):
+                _check_tool_id(invalid_id)
+
+        # This ID doesn't start with MCP prefix so gets generic error
+        with pytest.raises(ValueError, match="Invalid tool ID"):
+            _check_tool_id("mcp::wrong::server::tool")
+
+    def test_valid_rag_tools(self):
+        """Test validation of valid RAG tools."""
+        valid_ids = [
+            "kiln_tool::rag::config1",
+            "kiln_tool::rag::my_rag_config",
+            "kiln_tool::rag::test_config_123",
+        ]
+        for tool_id in valid_ids:
+            result = _check_tool_id(tool_id)
+            assert result == tool_id
+
+    def test_invalid_rag_format(self):
+        """Test validation fails for invalid RAG tool formats."""
+        # These IDs start with the RAG prefix but have invalid formats
+        rag_invalid_ids = [
+            "kiln_tool::rag::",  # Missing config ID
+            "kiln_tool::rag::config::extra",  # Too many parts
+        ]
+
+        for invalid_id in rag_invalid_ids:
+            with pytest.raises(ValueError, match="Invalid RAG tool ID"):
+                _check_tool_id(invalid_id)
+
+    def test_rag_tool_empty_config_id(self):
+        """Test that RAG tool with empty config ID is handled properly."""
+        # This tests the case where rag_config_id_from_id returns empty string
+        # which should trigger line 66 in the source
+        with pytest.raises(ValueError, match="Invalid RAG tool ID"):
+            _check_tool_id("kiln_tool::rag::")
+
+
+class TestMcpServerAndToolNameFromId:
+    """Test the mcp_server_and_tool_name_from_id function."""
+
+    def test_valid_mcp_ids(self):
+        """Test parsing valid MCP tool IDs."""
+        test_cases = [
+            # Remote MCP tools
+            ("mcp::remote::server1::tool1", ("server1", "tool1")),
+            ("mcp::remote::my_server::my_tool", ("my_server", "my_tool")),
+            ("mcp::remote::test::function_name", ("test", "function_name")),
+            # Local MCP tools
+            ("mcp::local::server1::tool1", ("server1", "tool1")),
+            ("mcp::local::my_server::my_tool", ("my_server", "my_tool")),
+            ("mcp::local::test::function_name", ("test", "function_name")),
+        ]
+
+        for tool_id, expected in test_cases:
+            result = mcp_server_and_tool_name_from_id(tool_id)
+            assert result == expected
+
+    def test_invalid_mcp_ids(self):
+        """Test parsing fails for invalid MCP tool IDs."""
+        # Test remote MCP tool ID errors
+        remote_invalid_ids = [
+            "mcp::remote::",  # Only 3 parts
+            "mcp::remote::server",  # Only 3 parts
+            "mcp::remote::server::tool::extra",  # 5 parts
+        ]
+
+        for invalid_id in remote_invalid_ids:
+            with pytest.raises(ValueError, match="Invalid remote MCP tool ID"):
+                mcp_server_and_tool_name_from_id(invalid_id)
+
+        # Test local MCP tool ID errors
+        local_invalid_ids = [
+            "mcp::local::",  # Only 3 parts
+            "mcp::local::server",  # Only 3 parts
+            "mcp::local::server::tool::extra",  # 5 parts
+        ]
+
+        for invalid_id in local_invalid_ids:
+            with pytest.raises(ValueError, match="Invalid local MCP tool ID"):
+                mcp_server_and_tool_name_from_id(invalid_id)
+
+        # Test generic MCP tool ID errors (not remote or local)
+        generic_invalid_ids = [
+            "not_mcp_format",  # Only 1 part
+            "single_part",  # Only 1 part
+            "",  # Empty string
+        ]
+
+        for invalid_id in generic_invalid_ids:
+            with pytest.raises(ValueError, match="Invalid MCP tool ID"):
+                mcp_server_and_tool_name_from_id(invalid_id)
+
+    def test_mcp_ids_with_wrong_prefix_still_parse(self):
+        """Test that IDs with wrong prefix but correct structure still parse (validation happens elsewhere)."""
+        # This function only checks structure (4 parts), not content
+        result = mcp_server_and_tool_name_from_id("mcp::wrong::server::tool")
+        assert result == ("server", "tool")
+
+
+class TestToolIdPydanticType:
+    """Test the ToolId pydantic type annotation."""
+
+    class _ModelWithToolId(BaseModel):
+        tool_id: ToolId
+
+    def test_valid_builtin_tools(self):
+        """Test ToolId validates built-in tools."""
+        for tool_id in KilnBuiltInToolId:
+            model = self._ModelWithToolId(tool_id=tool_id.value)
+            assert model.tool_id == tool_id.value
+
+    def test_valid_mcp_tools(self):
+        """Test ToolId validates MCP remote and local tools."""
+        valid_ids = [
+            # Remote MCP tools
+            "mcp::remote::server1::tool1",
+            "mcp::remote::my_server::my_tool",
+            # Local MCP tools
+            "mcp::local::server1::tool1",
+            "mcp::local::my_server::my_tool",
+            # RAG tools
+            "kiln_tool::rag::config1",
+            "kiln_tool::rag::my_rag_config",
+        ]
+
+        for tool_id in valid_ids:
+            model = self._ModelWithToolId(tool_id=tool_id)
+            assert model.tool_id == tool_id
+
+    def test_invalid_tools_raise_validation_error(self):
+        """Test ToolId raises ValidationError for invalid tools."""
+        invalid_ids = [
+            "",
+            "unknown_tool",
+            "mcp::remote::",
+            "mcp::remote::server",
+            "mcp::local::",
+            "mcp::local::server",
+            "kiln_tool::rag::",
+            "kiln_tool::rag::config::extra",
+        ]
+
+        for invalid_id in invalid_ids:
+            with pytest.raises(ValidationError):
+                self._ModelWithToolId(tool_id=invalid_id)
+
+    def test_non_string_raises_validation_error(self):
+        """Test ToolId raises ValidationError for non-string values."""
+        with pytest.raises(ValidationError):
+            self._ModelWithToolId(tool_id=123)  # type: ignore
+
+        with pytest.raises(ValidationError):
+            self._ModelWithToolId(tool_id=None)  # type: ignore
+
+
+class TestConstants:
+    """Test module constants."""
+
+    def test_mcp_remote_tool_id_prefix(self):
+        """Test the MCP remote tool ID prefix constant."""
+        assert MCP_REMOTE_TOOL_ID_PREFIX == "mcp::remote::"
+
+    def test_mcp_local_tool_id_prefix(self):
+        """Test the MCP local tool ID prefix constant."""
+        assert MCP_LOCAL_TOOL_ID_PREFIX == "mcp::local::"
+
+    def test_rag_tool_id_prefix(self):
+        """Test the RAG tool ID prefix constant."""
+        assert RAG_TOOL_ID_PREFIX == "kiln_tool::rag::"
+
+
+class TestRagConfigIdFromId:
+    """Test the rag_config_id_from_id function."""
+
+    def test_valid_rag_ids(self):
+        """Test parsing valid RAG tool IDs."""
+        test_cases = [
+            ("kiln_tool::rag::config1", "config1"),
+            ("kiln_tool::rag::my_rag_config", "my_rag_config"),
+            ("kiln_tool::rag::test_config_123", "test_config_123"),
+            ("kiln_tool::rag::a", "a"),  # Minimal valid case
+        ]
+
+        for tool_id, expected in test_cases:
+            result = rag_config_id_from_id(tool_id)
+            assert result == expected
+
+    def test_invalid_rag_ids(self):
+        """Test parsing fails for invalid RAG tool IDs."""
+        # Test various invalid formats that should trigger line 104
+        invalid_ids = [
+            "kiln_tool::rag::config::extra",  # Too many parts (4 parts)
+            "wrong::rag::config",  # Wrong prefix
+            "kiln_tool::wrong::config",  # Wrong middle part
+            "rag::config",  # Too few parts (2 parts)
+            "",  # Empty string
+            "single_part",  # Only 1 part
+        ]
+
+        for invalid_id in invalid_ids:
+            with pytest.raises(ValueError, match="Invalid RAG tool ID"):
+                rag_config_id_from_id(invalid_id)
+
+    def test_rag_id_with_empty_config_id(self):
+        """Test that RAG tool ID with empty config ID returns empty string."""
+        # This is actually valid according to the parser - it returns empty string
+        # The validation for empty config ID happens in _check_tool_id
+        result = rag_config_id_from_id("kiln_tool::rag::")
+        assert result == ""