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

kiln_ai/datamodel/test_registry.py

@@ -3,14 +3,19 @@ from unittest.mock import Mock, patch
 import pytest
 
 from kiln_ai.datamodel import Project
-from kiln_ai.datamodel.registry import all_projects, project_from_id
+from kiln_ai.datamodel.registry import all_projects
+from kiln_ai.utils.project_utils import project_from_id
 
 
 @pytest.fixture
 def mock_config():
-    with patch("kiln_ai.datamodel.registry.Config") as mock:
+    with (
+        patch("kiln_ai.datamodel.registry.Config") as mock_registry,
+        patch("kiln_ai.utils.project_utils.Config") as mock_utils,
+    ):
         config_instance = Mock()
-        mock.shared.return_value = config_instance
+        mock_registry.shared.return_value = config_instance
+        mock_utils.shared.return_value = config_instance
         yield config_instance
 
 

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,239 @@
+import pytest
+from pydantic import BaseModel, ValidationError
+
+from kiln_ai.datamodel.tool_id import (
+    MCP_LOCAL_TOOL_ID_PREFIX,
+    MCP_REMOTE_TOOL_ID_PREFIX,
+    KilnBuiltInToolId,
+    ToolId,
+    _check_tool_id,
+    mcp_server_and_tool_name_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")
+
+
+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",
+        ]
+
+        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",
+        ]
+
+        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::"

kiln_ai/datamodel/tool_id.py

@@ -0,0 +1,83 @@
+from enum import Enum
+from typing import Annotated
+
+from pydantic import AfterValidator
+
+ToolId = Annotated[
+    str,
+    AfterValidator(lambda v: _check_tool_id(v)),
+]
+"""
+A pydantic type that validates strings containing a valid tool ID.
+
+Tool IDs can be one of:
+- A kiln built-in tool name: kiln_tool::add_numbers
+- A remote MCP tool: mcp::remote::<server_id>::<tool_name>
+- A local MCP tool: mcp::local::<server_id>::<tool_name>
+- More coming soon like kiln_project_tool::rag::RAG_CONFIG_ID
+"""
+
+
+class KilnBuiltInToolId(str, Enum):
+    ADD_NUMBERS = "kiln_tool::add_numbers"
+    SUBTRACT_NUMBERS = "kiln_tool::subtract_numbers"
+    MULTIPLY_NUMBERS = "kiln_tool::multiply_numbers"
+    DIVIDE_NUMBERS = "kiln_tool::divide_numbers"
+
+
+MCP_REMOTE_TOOL_ID_PREFIX = "mcp::remote::"
+MCP_LOCAL_TOOL_ID_PREFIX = "mcp::local::"
+
+
+def _check_tool_id(id: str) -> str:
+    """
+    Check that the tool ID is valid.
+    """
+    if not id or not isinstance(id, str):
+        raise ValueError(f"Invalid tool ID: {id}")
+
+    # Build in tools
+    if id in KilnBuiltInToolId.__members__.values():
+        return id
+
+    # MCP remote tools must have format: mcp::remote::<server_id>::<tool_name>
+    if id.startswith(MCP_REMOTE_TOOL_ID_PREFIX):
+        server_id, tool_name = mcp_server_and_tool_name_from_id(id)
+        if not server_id or not tool_name:
+            raise ValueError(
+                f"Invalid remote MCP tool ID: {id}. Expected format: 'mcp::remote::<server_id>::<tool_name>'."
+            )
+        return id
+
+    # MCP local tools must have format: mcp::local::<server_id>::<tool_name>
+    if id.startswith(MCP_LOCAL_TOOL_ID_PREFIX):
+        server_id, tool_name = mcp_server_and_tool_name_from_id(id)
+        if not server_id or not tool_name:
+            raise ValueError(
+                f"Invalid local MCP tool ID: {id}. Expected format: 'mcp::local::<server_id>::<tool_name>'."
+            )
+        return id
+
+    raise ValueError(f"Invalid tool ID: {id}")
+
+
+def mcp_server_and_tool_name_from_id(id: str) -> tuple[str, str]:
+    """
+    Get the tool server ID and tool name from the ID.
+    """
+    parts = id.split("::")
+    if len(parts) != 4:
+        # Determine if it's remote or local for the error message
+        if id.startswith(MCP_REMOTE_TOOL_ID_PREFIX):
+            raise ValueError(
+                f"Invalid remote MCP tool ID: {id}. Expected format: 'mcp::remote::<server_id>::<tool_name>'."
+            )
+        elif id.startswith(MCP_LOCAL_TOOL_ID_PREFIX):
+            raise ValueError(
+                f"Invalid local MCP tool ID: {id}. Expected format: 'mcp::local::<server_id>::<tool_name>'."
+            )
+        else:
+            raise ValueError(
+                f"Invalid MCP tool ID: {id}. Expected format: 'mcp::(remote|local)::<server_id>::<tool_name>'."
+            )
+    return parts[2], parts[3]  # server_id, tool_name

kiln_ai/tools/__init__.py

@@ -0,0 +1,8 @@
+from kiln_ai.tools.base_tool import KilnTool, KilnToolInterface
+from kiln_ai.tools.tool_registry import tool_from_id
+
+__all__ = [
+    "KilnTool",
+    "KilnToolInterface",
+    "tool_from_id",
+]

kiln_ai/tools/base_tool.py

@@ -0,0 +1,82 @@
+from abc import ABC, abstractmethod
+from typing import Any, Dict
+
+from kiln_ai.datamodel.json_schema import validate_schema_dict
+from kiln_ai.datamodel.tool_id import KilnBuiltInToolId, ToolId
+
+
+class KilnToolInterface(ABC):
+    """
+    Abstract interface defining the core API that all Kiln tools must implement.
+    This ensures consistency across all tool implementations.
+    """
+
+    @abstractmethod
+    async def run(self, **kwargs) -> Any:
+        """Execute the tool with the given parameters."""
+        pass
+
+    @abstractmethod
+    async def toolcall_definition(self) -> Dict[str, Any]:
+        """Return the OpenAI-compatible tool definition for this tool."""
+        pass
+
+    @abstractmethod
+    async def id(self) -> ToolId:
+        """Return a unique identifier for this tool."""
+        pass
+
+    @abstractmethod
+    async def name(self) -> str:
+        """Return the tool name (function name) of this tool."""
+        pass
+
+    @abstractmethod
+    async def description(self) -> str:
+        """Return a description of what this tool does."""
+        pass
+
+
+class KilnTool(KilnToolInterface):
+    """
+    Base helper class that provides common functionality for tool implementations.
+    Subclasses only need to implement run() and provide tool configuration.
+    """
+
+    def __init__(
+        self,
+        tool_id: KilnBuiltInToolId,
+        name: str,
+        description: str,
+        parameters_schema: Dict[str, Any],
+    ):
+        self._id = tool_id
+        self._name = name
+        self._description = description
+        validate_schema_dict(parameters_schema)
+        self._parameters_schema = parameters_schema
+
+    async def id(self) -> KilnBuiltInToolId:
+        return self._id
+
+    async def name(self) -> str:
+        return self._name
+
+    async def description(self) -> str:
+        return self._description
+
+    async def toolcall_definition(self) -> Dict[str, Any]:
+        """Generate OpenAI-compatible tool definition."""
+        return {
+            "type": "function",
+            "function": {
+                "name": await self.name(),
+                "description": await self.description(),
+                "parameters": self._parameters_schema,
+            },
+        }
+
+    @abstractmethod
+    async def run(self, **kwargs) -> str:
+        """Subclasses must implement the actual tool logic."""
+        pass

kiln_ai/tools/built_in_tools/__init__.py

@@ -0,0 +1,13 @@
+from kiln_ai.tools.built_in_tools.math_tools import (
+    AddTool,
+    DivideTool,
+    MultiplyTool,
+    SubtractTool,
+)
+
+__all__ = [
+    "AddTool",
+    "DivideTool",
+    "MultiplyTool",
+    "SubtractTool",
+]