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

kiln_ai/datamodel/test_extraction_model.py

@@ -468,3 +468,34 @@ def test_document_invalid_mime_type(
 )
 def test_get_kind_from_mime_type(mime_type, expected_kind):
     assert get_kind_from_mime_type(mime_type) == expected_kind
+
+
+def test_document_friendly_name(mock_project, mock_attachment_factory):
+    name = f"Test Document {uuid.uuid4()!s}"
+    document = Document(
+        name=name,
+        description=f"Test description {uuid.uuid4()!s}",
+        kind=Kind.DOCUMENT,
+        original_file=FileInfo(
+            filename=f"test_{name}.txt",
+            size=100,
+            mime_type="text/plain",
+            attachment=mock_attachment_factory(),
+        ),
+        parent=mock_project,
+    )
+    document.save_to_file()
+
+    # backward compatibility: old documents did not have name_override
+    assert document.name_override is None
+    assert document.friendly_name == name
+
+    # new documents have name_override
+    document.name_override = "Test Document Override"
+    assert document.friendly_name == "Test Document Override"
+
+    document.save_to_file()
+
+    document = Document.from_id_and_parent_path(str(document.id), mock_project.path)
+    assert document is not None
+    assert document.friendly_name == "Test Document Override"

kiln_ai/datamodel/test_task.py

@@ -254,7 +254,7 @@ def test_run_config_upgrade_old_entries():
         },
         "prompt": {
             "name": "Dazzling Unicorn",
-            "description": "Frozen copy of prompt 'simple_prompt_builder', created for evaluations.",
+            "description": "Frozen copy of prompt 'simple_prompt_builder'.",
             "generator_id": "simple_prompt_builder",
             "prompt": "Generate a joke, given a theme. The theme will be provided as a word or phrase as the input to the model. The assistant should output a joke that is funny and relevant to the theme. If a style is provided, the joke should be in that style. The output should include a setup and punchline.\n\nYour response should respect the following requirements:\n1) Keep the joke on topic. If the user specifies a theme, the joke must be related to that theme.\n2) Avoid any jokes that are offensive or inappropriate. Keep the joke clean and appropriate for all audiences.\n3) Make the joke funny and engaging. It should be something that someone would want to tell to their friends. Something clever, not just a simple pun.\n",
             "chain_of_thought_instructions": None,
@@ -296,3 +296,37 @@ def test_run_config_upgrade_old_entries():
 def test_task_name_unicode_name():
     task = Task(name="你好", instruction="Do something")
     assert task.name == "你好"
+
+
+def test_task_default_run_config_id_property(tmp_path):
+    """Test that default_run_config_id can be set and retrieved."""
+
+    # Create a task
+    task = Task(
+        name="Test Task", instruction="Test instruction", path=tmp_path / "task.kiln"
+    )
+    task.save_to_file()
+
+    # Create a run config for the task
+    run_config = TaskRunConfig(
+        name="Test Config",
+        run_config_properties=RunConfigProperties(
+            model_name="gpt-4",
+            model_provider_name="openai",
+            prompt_id=PromptGenerators.SIMPLE,
+            structured_output_mode=StructuredOutputMode.json_schema,
+        ),
+        parent=task,
+    )
+    run_config.save_to_file()
+
+    # Test None default (should be valid)
+    assert task.default_run_config_id is None
+
+    # Test setting a valid ID
+    task.default_run_config_id = "123456789012"
+    assert task.default_run_config_id == "123456789012"
+
+    # Test setting back to None
+    task.default_run_config_id = None
+    assert task.default_run_config_id is None

kiln_ai/datamodel/test_tool_id.py

@@ -8,6 +8,7 @@ from kiln_ai.datamodel.tool_id import (
     KilnBuiltInToolId,
     ToolId,
     _check_tool_id,
+    kiln_task_server_id_from_tool_id,
     mcp_server_and_tool_name_from_id,
     rag_config_id_from_id,
 )
@@ -145,6 +146,39 @@ class TestCheckToolId:
         with pytest.raises(ValueError, match="Invalid RAG tool ID"):
             _check_tool_id("kiln_tool::rag::")
 
+    def test_valid_kiln_task_tools(self):
+        """Test validation of valid Kiln task tools."""
+        valid_ids = [
+            "kiln_task::server1",
+            "kiln_task::my_server",
+            "kiln_task::test_server_123",
+            "kiln_task::server_with_underscores",
+            "kiln_task::server-with-dashes",
+            "kiln_task::server.with.dots",
+        ]
+        for tool_id in valid_ids:
+            result = _check_tool_id(tool_id)
+            assert result == tool_id
+
+    def test_invalid_kiln_task_format(self):
+        """Test validation fails for invalid Kiln task tool formats."""
+        # These IDs start with the Kiln task prefix but have invalid formats
+        kiln_task_invalid_ids = [
+            "kiln_task::",  # Missing server ID
+            "kiln_task::server::extra",  # Too many parts
+            "kiln_task::server::tool::extra",  # Too many parts
+        ]
+
+        for invalid_id in kiln_task_invalid_ids:
+            with pytest.raises(ValueError, match="Invalid Kiln task tool ID"):
+                _check_tool_id(invalid_id)
+
+    def test_kiln_task_tool_empty_server_id(self):
+        """Test that Kiln task tool with empty server ID is handled properly."""
+        # This tests the case where kiln_task_server_id_from_tool_id returns empty string which should raise an error
+        with pytest.raises(ValueError, match="Invalid Kiln task tool ID"):
+            _check_tool_id("kiln_task::")
+
 
 class TestMcpServerAndToolNameFromId:
     """Test the mcp_server_and_tool_name_from_id function."""
@@ -220,7 +254,7 @@ class TestToolIdPydanticType:
             model = self._ModelWithToolId(tool_id=tool_id.value)
             assert model.tool_id == tool_id.value
 
-    def test_valid_mcp_tools(self):
+    def test_valid_tool_ids(self):
         """Test ToolId validates MCP remote and local tools."""
         valid_ids = [
             # Remote MCP tools
@@ -232,6 +266,9 @@ class TestToolIdPydanticType:
             # RAG tools
             "kiln_tool::rag::config1",
             "kiln_tool::rag::my_rag_config",
+            # Kiln task tools
+            "kiln_task::server1",
+            "kiln_task::my_server",
         ]
 
         for tool_id in valid_ids:
@@ -249,6 +286,8 @@ class TestToolIdPydanticType:
             "mcp::local::server",
             "kiln_tool::rag::",
             "kiln_tool::rag::config::extra",
+            "kiln_task::",
+            "kiln_task::server::extra",
         ]
 
         for invalid_id in invalid_ids:
@@ -318,3 +357,69 @@ class TestRagConfigIdFromId:
         # The validation for empty config ID happens in _check_tool_id
         result = rag_config_id_from_id("kiln_tool::rag::")
         assert result == ""
+
+
+class TestKilnTaskServerIdFromToolId:
+    """Test the kiln_task_server_id_from_tool_id function."""
+
+    def test_valid_kiln_task_ids(self):
+        """Test parsing valid Kiln task tool IDs."""
+        test_cases = [
+            ("kiln_task::server1", "server1"),
+            ("kiln_task::my_server", "my_server"),
+            ("kiln_task::test_server_123", "test_server_123"),
+            ("kiln_task::a", "a"),  # Minimal valid case
+            ("kiln_task::server_with_underscores", "server_with_underscores"),
+            ("kiln_task::server-with-dashes", "server-with-dashes"),
+            ("kiln_task::server.with.dots", "server.with.dots"),
+        ]
+
+        for tool_id, expected in test_cases:
+            result = kiln_task_server_id_from_tool_id(tool_id)
+            assert result == expected
+
+    def test_invalid_kiln_task_ids(self):
+        """Test parsing fails for invalid Kiln task tool IDs."""
+        # Test various invalid formats
+        invalid_ids = [
+            "kiln_task::",  # Empty server ID
+            "kiln_task::server::extra",  # Too many parts (3 parts)
+            "kiln_task::server::tool::extra",  # Too many parts (4 parts)
+            "wrong::server",  # Wrong prefix
+            "kiln_wrong::server",  # Wrong prefix
+            "task::server",  # Too few parts (2 parts)
+            "",  # Empty string
+            "single_part",  # Only 1 part
+            "kiln_task",  # Missing server ID
+        ]
+
+        for invalid_id in invalid_ids:
+            with pytest.raises(ValueError, match="Invalid Kiln task tool ID format"):
+                kiln_task_server_id_from_tool_id(invalid_id)
+
+    def test_kiln_task_id_with_empty_server_id(self):
+        """Test that Kiln task tool ID with empty server ID raises error."""
+        with pytest.raises(ValueError, match="Invalid Kiln task tool ID format"):
+            kiln_task_server_id_from_tool_id("kiln_task::")
+
+    def test_kiln_task_id_with_whitespace_server_id(self):
+        """Test that Kiln task tool ID with whitespace-only server ID raises error."""
+        with pytest.raises(ValueError, match="Invalid Kiln task tool ID format"):
+            kiln_task_server_id_from_tool_id("kiln_task::")
+
+    def test_kiln_task_id_with_multiple_colons(self):
+        """Test that Kiln task tool ID with multiple colons raises error."""
+        with pytest.raises(ValueError, match="Invalid Kiln task tool ID format"):
+            kiln_task_server_id_from_tool_id("kiln_task::server::extra")
+
+    def test_kiln_task_id_case_sensitivity(self):
+        """Test that Kiln task tool IDs are case sensitive."""
+        # These should work
+        result1 = kiln_task_server_id_from_tool_id("kiln_task::Server")
+        assert result1 == "Server"
+
+        result2 = kiln_task_server_id_from_tool_id("kiln_task::SERVER")
+        assert result2 == "SERVER"
+
+        result3 = kiln_task_server_id_from_tool_id("kiln_task::server")
+        assert result3 == "server"

kiln_ai/datamodel/tool_id.py

@@ -3,6 +3,8 @@ from typing import Annotated
 
 from pydantic import AfterValidator
 
+from kiln_ai.datamodel.basemodel import ID_TYPE
+
 ToolId = Annotated[
     str,
     AfterValidator(lambda v: _check_tool_id(v)),
@@ -14,6 +16,7 @@ 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>
+- A Kiln task tool: kiln_task::<server_id>
 - More coming soon like kiln_project_tool::rag::RAG_CONFIG_ID
 """
 
@@ -28,6 +31,7 @@ class KilnBuiltInToolId(str, Enum):
 MCP_REMOTE_TOOL_ID_PREFIX = "mcp::remote::"
 RAG_TOOL_ID_PREFIX = "kiln_tool::rag::"
 MCP_LOCAL_TOOL_ID_PREFIX = "mcp::local::"
+KILN_TASK_TOOL_ID_PREFIX = "kiln_task::"
 
 
 def _check_tool_id(id: str) -> str:
@@ -68,6 +72,15 @@ def _check_tool_id(id: str) -> str:
             )
         return id
 
+    # Kiln task tools must have format: kiln_task::<server_id>
+    if id.startswith(KILN_TASK_TOOL_ID_PREFIX):
+        server_id = kiln_task_server_id_from_tool_id(id)
+        if not server_id:
+            raise ValueError(
+                f"Invalid Kiln task tool ID: {id}. Expected format: 'kiln_task::<server_id>'."
+            )
+        return id
+
     raise ValueError(f"Invalid tool ID: {id}")
 
 
@@ -103,3 +116,39 @@ def rag_config_id_from_id(id: str) -> str:
             f"Invalid RAG tool ID: {id}. Expected format: 'kiln_tool::rag::<rag_config_id>'."
         )
     return parts[2]
+
+
+def build_rag_tool_id(rag_config_id: ID_TYPE) -> str:
+    """Construct the tool ID for a RAG configuration."""
+
+    return f"{RAG_TOOL_ID_PREFIX}{rag_config_id}"
+
+
+def build_kiln_task_tool_id(server_id: ID_TYPE) -> str:
+    """Construct the tool ID for a Kiln task server."""
+    return f"{KILN_TASK_TOOL_ID_PREFIX}{server_id}"
+
+
+def kiln_task_server_id_from_tool_id(tool_id: str) -> str:
+    """
+    Get the server ID from the tool ID.
+    """
+    if not tool_id.startswith(KILN_TASK_TOOL_ID_PREFIX):
+        raise ValueError(
+            f"Invalid Kiln task tool ID format: {tool_id}. Expected format: 'kiln_task::<server_id>'."
+        )
+
+    # Remove prefix and split on ::
+    remaining = tool_id[len(KILN_TASK_TOOL_ID_PREFIX) :]
+    if not remaining:
+        raise ValueError(
+            f"Invalid Kiln task tool ID format: {tool_id}. Expected format: 'kiln_task::<server_id>'."
+        )
+    parts = remaining.split("::")
+
+    if len(parts) != 1 or not parts[0].strip():
+        raise ValueError(
+            f"Invalid Kiln task tool ID format: {tool_id}. Expected format: 'kiln_task::<server_id>'."
+        )
+
+    return parts[0]  # server_id

kiln_ai/tools/base_tool.py

@@ -1,10 +1,34 @@
 from abc import ABC, abstractmethod
-from typing import Any, Dict
+from dataclasses import dataclass
+from typing import Any, Dict, TypedDict
 
 from kiln_ai.datamodel.json_schema import validate_schema_dict
 from kiln_ai.datamodel.tool_id import KilnBuiltInToolId, ToolId
 
 
+class ToolFunction(TypedDict):
+    """Typed dict for the function definition within a tool call definition."""
+
+    name: str
+    description: str
+    parameters: Dict[str, Any]
+
+
+class ToolCallDefinition(TypedDict):
+    """Typed dict for OpenAI-compatible tool call definitions."""
+
+    type: str  # Must be "function"
+    function: ToolFunction
+
+
+@dataclass
+class ToolCallContext:
+    """Context passed to tools when they are called, containing information from the calling task."""
+
+    """Used for Kiln Tasks as Tools, to know if the tool call should save the task run it invoked to that task's Dataset."""
+    allow_saving: bool = True
+
+
 class KilnToolInterface(ABC):
     """
     Abstract interface defining the core API that all Kiln tools must implement.
@@ -12,12 +36,12 @@ class KilnToolInterface(ABC):
     """
 
     @abstractmethod
-    async def run(self, **kwargs) -> Any:
-        """Execute the tool with the given parameters."""
+    async def run(self, context: ToolCallContext | None = None, **kwargs) -> Any:
+        """Execute the tool with the given parameters and calling context if provided."""
         pass
 
     @abstractmethod
-    async def toolcall_definition(self) -> Dict[str, Any]:
+    async def toolcall_definition(self) -> ToolCallDefinition:
         """Return the OpenAI-compatible tool definition for this tool."""
         pass
 
@@ -65,7 +89,7 @@ class KilnTool(KilnToolInterface):
     async def description(self) -> str:
         return self._description
 
-    async def toolcall_definition(self) -> Dict[str, Any]:
+    async def toolcall_definition(self) -> ToolCallDefinition:
         """Generate OpenAI-compatible tool definition."""
         return {
             "type": "function",
@@ -77,6 +101,6 @@ class KilnTool(KilnToolInterface):
         }
 
     @abstractmethod
-    async def run(self, **kwargs) -> str:
+    async def run(self, context: ToolCallContext | None = None, **kwargs) -> Any:
         """Subclasses must implement the actual tool logic."""
         pass

kiln_ai/tools/built_in_tools/math_tools.py

@@ -27,7 +27,9 @@ class AddTool(KilnTool):
             parameters_schema=parameters_schema,
         )
 
-    async def run(self, a: Union[int, float], b: Union[int, float]) -> str:
+    async def run(
+        self, context=None, *, a: Union[int, float], b: Union[int, float]
+    ) -> str:
         """Add two numbers and return the result."""
         return str(a + b)
 
@@ -57,7 +59,9 @@ class SubtractTool(KilnTool):
             parameters_schema=parameters_schema,
         )
 
-    async def run(self, a: Union[int, float], b: Union[int, float]) -> str:
+    async def run(
+        self, context=None, *, a: Union[int, float], b: Union[int, float]
+    ) -> str:
         """Subtract b from a and return the result."""
         return str(a - b)
 
@@ -84,7 +88,9 @@ class MultiplyTool(KilnTool):
             parameters_schema=parameters_schema,
         )
 
-    async def run(self, a: Union[int, float], b: Union[int, float]) -> str:
+    async def run(
+        self, context=None, *, a: Union[int, float], b: Union[int, float]
+    ) -> str:
         """Multiply two numbers and return the result."""
         return str(a * b)
 
@@ -117,7 +123,9 @@ class DivideTool(KilnTool):
             parameters_schema=parameters_schema,
         )
 
-    async def run(self, a: Union[int, float], b: Union[int, float]) -> str:
+    async def run(
+        self, context=None, *, a: Union[int, float], b: Union[int, float]
+    ) -> str:
         """Divide a by b and return the result."""
         if b == 0:
             raise ZeroDivisionError("Cannot divide by zero")

kiln_ai/tools/kiln_task_tool.py

@@ -0,0 +1,162 @@
+from dataclasses import dataclass
+from functools import cached_property
+from typing import Any, Dict
+
+from kiln_ai.datamodel import Task
+from kiln_ai.datamodel.external_tool_server import ExternalToolServer
+from kiln_ai.datamodel.task import TaskRunConfig
+from kiln_ai.datamodel.task_output import DataSource, DataSourceType
+from kiln_ai.datamodel.tool_id import ToolId
+from kiln_ai.tools.base_tool import (
+    KilnToolInterface,
+    ToolCallContext,
+    ToolCallDefinition,
+)
+from kiln_ai.utils.project_utils import project_from_id
+
+
+@dataclass
+class KilnTaskToolResult:
+    output: str
+    kiln_task_tool_data: str
+
+
+class KilnTaskTool(KilnToolInterface):
+    """
+    A tool that wraps a Kiln task, allowing it to be called as a function.
+
+    This tool loads a task by ID and executes it using the specified run configuration.
+    """
+
+    def __init__(
+        self,
+        project_id: str,
+        tool_id: str,
+        data_model: ExternalToolServer,
+    ):
+        self._project_id = project_id
+        self._tool_server_model = data_model
+        self._tool_id = tool_id
+
+        self._name = data_model.properties.get("name", "")
+        self._description = data_model.properties.get("description", "")
+        self._task_id = data_model.properties.get("task_id", "")
+        self._run_config_id = data_model.properties.get("run_config_id", "")
+
+    async def id(self) -> ToolId:
+        return self._tool_id
+
+    async def name(self) -> str:
+        return self._name
+
+    async def description(self) -> str:
+        return self._description
+
+    async def toolcall_definition(self) -> ToolCallDefinition:
+        """Generate OpenAI-compatible tool definition."""
+        return {
+            "type": "function",
+            "function": {
+                "name": await self.name(),
+                "description": await self.description(),
+                "parameters": self.parameters_schema,
+            },
+        }
+
+    async def run(
+        self, context: ToolCallContext | None = None, **kwargs
+    ) -> KilnTaskToolResult:
+        """Execute the wrapped Kiln task with the given parameters and calling context."""
+        if context is None:
+            raise ValueError("Context is required for running a KilnTaskTool.")
+
+        # Determine the input format
+        if self._task.input_json_schema:
+            # Structured input - pass kwargs directly
+            input = kwargs
+        else:
+            # Plaintext input - extract from 'input' parameter
+            if "input" in kwargs:
+                input = kwargs["input"]
+            else:
+                raise ValueError(f"Input not found in kwargs: {kwargs}")
+
+        # These imports are here to avoid circular chains
+        from kiln_ai.adapters.adapter_registry import adapter_for_task
+        from kiln_ai.adapters.model_adapters.base_adapter import AdapterConfig
+
+        # Create adapter and run the task using the calling task's allow_saving setting
+        adapter = adapter_for_task(
+            self._task,
+            run_config_properties=self._run_config.run_config_properties,
+            base_adapter_config=AdapterConfig(
+                allow_saving=context.allow_saving,
+                default_tags=["tool_call"],
+            ),
+        )
+        task_run = await adapter.invoke(
+            input,
+            input_source=DataSource(
+                type=DataSourceType.tool_call,
+                run_config=self._run_config.run_config_properties,
+            ),
+        )
+
+        return KilnTaskToolResult(
+            output=task_run.output.output,
+            kiln_task_tool_data=f"{self._project_id}:::{self._tool_id}:::{self._task.id}:::{task_run.id}",
+        )
+
+    @cached_property
+    def _task(self) -> Task:
+        # Load the project first
+        project = project_from_id(self._project_id)
+        if project is None:
+            raise ValueError(f"Project not found: {self._project_id}")
+
+        # Load the task from the project
+        task = Task.from_id_and_parent_path(self._task_id, project.path)
+        if task is None:
+            raise ValueError(
+                f"Task not found: {self._task_id} in project {self._project_id}"
+            )
+        return task
+
+    @cached_property
+    def _run_config(self) -> TaskRunConfig:
+        run_config = next(
+            (
+                run_config
+                for run_config in self._task.run_configs(readonly=True)
+                if run_config.id == self._run_config_id
+            ),
+            None,
+        )
+        if run_config is None:
+            raise ValueError(
+                f"Task run config not found: {self._run_config_id} for task {self._task_id} in project {self._project_id}"
+            )
+        return run_config
+
+    @cached_property
+    def parameters_schema(self) -> Dict[str, Any]:
+        if self._task.input_json_schema:
+            # Use the task's input schema directly if it exists
+            parameters_schema = self._task.input_schema()
+        else:
+            # For plaintext tasks, create a simple string input parameter
+            parameters_schema = {
+                "type": "object",
+                "properties": {
+                    "input": {
+                        "type": "string",
+                        "description": "Plaintext input for the tool.",
+                    }
+                },
+                "required": ["input"],
+            }
+        if parameters_schema is None:
+            raise ValueError(
+                f"Failed to create parameters schema for tool_id {self._tool_id}"
+            )
+        return parameters_schema

kiln_ai/tools/mcp_server_tool.py

@@ -1,11 +1,13 @@
-from typing import Any, Dict
-
 from mcp.types import CallToolResult, TextContent
 from mcp.types import Tool as MCPTool
 
 from kiln_ai.datamodel.external_tool_server import ExternalToolServer
 from kiln_ai.datamodel.tool_id import MCP_REMOTE_TOOL_ID_PREFIX, ToolId
-from kiln_ai.tools.base_tool import KilnToolInterface
+from kiln_ai.tools.base_tool import (
+    KilnToolInterface,
+    ToolCallContext,
+    ToolCallDefinition,
+)
 from kiln_ai.tools.mcp_session_manager import MCPSessionManager
 
 
@@ -26,7 +28,7 @@ class MCPServerTool(KilnToolInterface):
         await self._load_tool_properties()
         return self._description
 
-    async def toolcall_definition(self) -> Dict[str, Any]:
+    async def toolcall_definition(self) -> ToolCallDefinition:
         """Generate OpenAI-compatible tool definition."""
         await self._load_tool_properties()
         return {
@@ -38,7 +40,7 @@ class MCPServerTool(KilnToolInterface):
             },
         }
 
-    async def run(self, **kwargs) -> Any:
+    async def run(self, context: ToolCallContext | None = None, **kwargs) -> str:
         result = await self._call_tool(**kwargs)
 
         if result.isError: