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

kiln_ai/tools/test_rag_tools.py

@@ -8,6 +8,7 @@ from kiln_ai.datamodel.embedding import EmbeddingConfig
 from kiln_ai.datamodel.project import Project
 from kiln_ai.datamodel.rag import RagConfig
 from kiln_ai.datamodel.vector_store import VectorStoreConfig, VectorStoreType
+from kiln_ai.tools.base_tool import ToolCallContext
 from kiln_ai.tools.rag_tools import ChunkContext, RagTool, format_search_results
 
 
@@ -420,7 +421,7 @@ class TestRagTool:
             tool = RagTool("tool_123", mock_rag_config)
 
             # Run the tool
-            result = await tool.run("test query")
+            result = await tool.run(context=None, query="test query")
 
             # Verify the result format
             expected_result = (
@@ -500,7 +501,7 @@ class TestRagTool:
             tool = RagTool("tool_123", mock_rag_config)
 
             # Run the tool
-            result = await tool.run("hybrid query")
+            result = await tool.run(context=None, query="hybrid query")
 
             # Verify embedding generation was called
             mock_embedding_adapter.generate_embeddings.assert_called_once_with(
@@ -566,7 +567,7 @@ class TestRagTool:
             tool = RagTool("tool_123", mock_rag_config)
 
             # Run the tool
-            result = await tool.run("fts query")
+            result = await tool.run(context=None, query="fts query")
 
             # Verify the result format
             expected_result = (
@@ -629,7 +630,7 @@ class TestRagTool:
 
             # Run the tool and expect an error
             with pytest.raises(ValueError, match="No embeddings generated"):
-                await tool.run("query with no embeddings")
+                await tool.run(context=None, query="query with no embeddings")
 
     async def test_rag_tool_run_empty_search_results(
         self, mock_rag_config, mock_project
@@ -675,11 +676,91 @@ class TestRagTool:
             tool = RagTool("tool_123", mock_rag_config)
 
             # Run the tool
-            result = await tool.run("query with no results")
+            result = await tool.run(context=None, query="query with no results")
 
             # Should return empty string for no results
             assert result == ""
 
+    async def test_rag_tool_run_with_context_is_accepted(
+        self, mock_rag_config, mock_project
+    ):
+        """Ensure RagTool.run accepts and works when a ToolCallContext is provided."""
+        mock_rag_config.parent_project.return_value = mock_project
+
+        # Mock search results
+        search_results = [
+            SearchResult(
+                document_id="doc_ctx",
+                chunk_idx=3,
+                chunk_text="Context ok",
+                similarity=0.77,
+            )
+        ]
+
+        with (
+            patch("kiln_ai.tools.rag_tools.VectorStoreConfig") as mock_vs_config_class,
+            patch("kiln_ai.tools.rag_tools.EmbeddingConfig") as mock_embed_config_class,
+            patch(
+                "kiln_ai.tools.rag_tools.embedding_adapter_from_type"
+            ) as mock_adapter_factory,
+            patch(
+                "kiln_ai.tools.rag_tools.vector_store_adapter_for_config",
+                new_callable=AsyncMock,
+            ) as mock_vs_adapter_factory,
+        ):
+            # VECTOR type → embedding path taken
+            mock_vector_store_config = Mock()
+            mock_vector_store_config.store_type = VectorStoreType.LANCE_DB_VECTOR
+            mock_vs_config_class.from_id_and_parent_path.return_value = (
+                mock_vector_store_config
+            )
+
+            mock_embedding_config = Mock()
+            mock_embed_config_class.from_id_and_parent_path.return_value = (
+                mock_embedding_config
+            )
+
+            mock_embedding_adapter = AsyncMock()
+            mock_embedding_result = Mock()
+            mock_embedding_result.embeddings = [Mock(vector=[1.0])]
+            mock_embedding_adapter.generate_embeddings.return_value = (
+                mock_embedding_result
+            )
+            mock_adapter_factory.return_value = mock_embedding_adapter
+
+            mock_vector_store_adapter = AsyncMock()
+            mock_vector_store_adapter.search.return_value = search_results
+            mock_vs_adapter_factory.return_value = mock_vector_store_adapter
+
+            tool = RagTool("tool_ctx", mock_rag_config)
+
+            ctx = ToolCallContext(allow_saving=False)
+            result = await tool.run(context=ctx, query="with context")
+
+            # Works and returns formatted text
+            assert result == "[document_id: doc_ctx, chunk_idx: 3]\nContext ok\n\n"
+
+            # Normal behavior still occurs
+            mock_embedding_adapter.generate_embeddings.assert_called_once_with(
+                ["with context"]
+            )
+            mock_vector_store_adapter.search.assert_called_once()
+
+    async def test_rag_tool_run_missing_query_raises(
+        self, mock_rag_config, mock_project
+    ):
+        """Ensure RagTool.run enforces the 'if not query' guard."""
+        mock_rag_config.parent_project.return_value = mock_project
+
+        with (
+            patch("kiln_ai.tools.rag_tools.VectorStoreConfig") as mock_vs_config_class,
+        ):
+            mock_vs_config_class.from_id_and_parent_path.return_value = Mock()
+            tool = RagTool("tool_err", mock_rag_config)
+
+            with pytest.raises(KeyError, match="query"):
+                await tool.run(context=None)
+
 
 class TestRagToolNameAndDescription:
     """Test RagTool name and description functionality with tool_name and tool_description fields."""

kiln_ai/tools/test_tool_registry.py

@@ -3,15 +3,20 @@ from unittest.mock import Mock
 
 import pytest
 
-from kiln_ai.datamodel.external_tool_server import ExternalToolServer, ToolServerType
+from kiln_ai.datamodel.external_tool_server import (
+    ExternalToolServer,
+    ToolServerType,
+)
 from kiln_ai.datamodel.project import Project
 from kiln_ai.datamodel.task import Task
 from kiln_ai.datamodel.tool_id import (
+    KILN_TASK_TOOL_ID_PREFIX,
     MCP_LOCAL_TOOL_ID_PREFIX,
     MCP_REMOTE_TOOL_ID_PREFIX,
     RAG_TOOL_ID_PREFIX,
     KilnBuiltInToolId,
     _check_tool_id,
+    kiln_task_server_id_from_tool_id,
     mcp_server_and_tool_name_from_id,
 )
 from kiln_ai.tools.built_in_tools.math_tools import (
@@ -20,6 +25,7 @@ from kiln_ai.tools.built_in_tools.math_tools import (
     MultiplyTool,
     SubtractTool,
 )
+from kiln_ai.tools.kiln_task_tool import KilnTaskTool
 from kiln_ai.tools.mcp_server_tool import MCPServerTool
 from kiln_ai.tools.tool_registry import tool_from_id
 
@@ -87,7 +93,6 @@ class TestToolRegistry:
             type=ToolServerType.remote_mcp,
             properties={
                 "server_url": "https://example.com",
-                "headers": {},
             },
         )
 
@@ -157,7 +162,6 @@ class TestToolRegistry:
             type=ToolServerType.remote_mcp,
             properties={
                 "server_url": "https://example.com",
-                "headers": {},
             },
         )
 
@@ -361,6 +365,48 @@ class TestToolRegistry:
             with pytest.raises(ValueError, match=f"Invalid tool ID: {invalid_id}"):
                 _check_tool_id(invalid_id)
 
+    def test_check_tool_id_valid_kiln_task_tool_id(self):
+        """Test that _check_tool_id accepts valid Kiln task tool IDs."""
+        valid_kiln_task_ids = [
+            f"{KILN_TASK_TOOL_ID_PREFIX}server123",
+            f"{KILN_TASK_TOOL_ID_PREFIX}my_task_server",
+            f"{KILN_TASK_TOOL_ID_PREFIX}123456789",
+            f"{KILN_TASK_TOOL_ID_PREFIX}server_with_underscores",
+            f"{KILN_TASK_TOOL_ID_PREFIX}server-with-dashes",
+        ]
+
+        for tool_id in valid_kiln_task_ids:
+            result = _check_tool_id(tool_id)
+            assert result == tool_id
+
+    def test_check_tool_id_invalid_kiln_task_tool_id(self):
+        """Test that _check_tool_id rejects invalid Kiln task tool IDs."""
+        # These start with the prefix but have wrong format
+        invalid_kiln_task_format_ids = [
+            f"{KILN_TASK_TOOL_ID_PREFIX}",  # Missing server ID
+            f"{KILN_TASK_TOOL_ID_PREFIX}::",  # Empty server ID
+            f"{KILN_TASK_TOOL_ID_PREFIX}server::tool",  # Too many parts (3 instead of 2)
+            f"{KILN_TASK_TOOL_ID_PREFIX}server::tool::extra",  # Too many parts (4 instead of 2)
+        ]
+
+        for invalid_id in invalid_kiln_task_format_ids:
+            with pytest.raises(
+                ValueError, match=f"Invalid Kiln task tool ID format: {invalid_id}"
+            ):
+                _check_tool_id(invalid_id)
+
+        # These don't match the prefix - get generic error
+        invalid_generic_ids = [
+            "kiln_task:",  # Missing last colon (doesn't match full prefix)
+            "kiln:task::server",  # Wrong prefix format
+            "kiln_task_server",  # Missing colons
+            "task::server",  # Missing kiln prefix
+        ]
+
+        for invalid_id in invalid_generic_ids:
+            with pytest.raises(ValueError, match=f"Invalid tool ID: {invalid_id}"):
+                _check_tool_id(invalid_id)
+
     def test_mcp_server_and_tool_name_from_id_valid_inputs(self):
         """Test that mcp_server_and_tool_name_from_id correctly parses valid MCP tool IDs."""
         test_cases = [
@@ -489,6 +535,64 @@ class TestToolRegistry:
         assert server_id == expected_server
         assert tool_name == expected_tool
 
+    def test_kiln_task_server_id_from_tool_id_valid_inputs(self):
+        """Test that kiln_task_server_id_from_tool_id correctly parses valid Kiln task tool IDs."""
+        test_cases = [
+            ("kiln_task::server123", "server123"),
+            ("kiln_task::my_task_server", "my_task_server"),
+            ("kiln_task::123456789", "123456789"),
+            ("kiln_task::server_with_underscores", "server_with_underscores"),
+            ("kiln_task::server-with-dashes", "server-with-dashes"),
+            ("kiln_task::a", "a"),  # Minimal valid case
+            (
+                "kiln_task::very_long_server_name_with_numbers_123",
+                "very_long_server_name_with_numbers_123",
+            ),
+        ]
+
+        for tool_id, expected_server_id in test_cases:
+            result = kiln_task_server_id_from_tool_id(tool_id)
+            assert result == expected_server_id, (
+                f"Failed for {tool_id}: expected {expected_server_id}, got {result}"
+            )
+
+    def test_kiln_task_server_id_from_tool_id_invalid_inputs(self):
+        """Test that kiln_task_server_id_from_tool_id raises ValueError for invalid Kiln task tool IDs."""
+        invalid_inputs = [
+            "kiln_task::",  # Empty server ID
+            "kiln_task::server::tool",  # Too many parts (3 instead of 2)
+            "kiln_task::server::tool::extra",  # Too many parts (4 instead of 2)
+            "invalid::format",  # Wrong prefix
+            "",  # Empty string
+            "single_part",  # No separators
+            "two::parts",  # Only 2 parts but wrong prefix
+            "kiln_task",  # Missing colons
+        ]
+
+        for invalid_id in invalid_inputs:
+            with pytest.raises(
+                ValueError,
+                match=r"Invalid Kiln task tool ID format:.*Expected format.*kiln_task::<server_id>",
+            ):
+                kiln_task_server_id_from_tool_id(invalid_id)
+
+    @pytest.mark.parametrize(
+        "tool_id,expected_server_id",
+        [
+            ("kiln_task::test_server", "test_server"),
+            ("kiln_task::s", "s"),
+            ("kiln_task::long_server_name_123", "long_server_name_123"),
+            ("kiln_task::server-with-dashes", "server-with-dashes"),
+            ("kiln_task::server_with_underscores", "server_with_underscores"),
+        ],
+    )
+    def test_kiln_task_server_id_from_tool_id_parametrized(
+        self, tool_id, expected_server_id
+    ):
+        """Parametrized test for kiln_task_server_id_from_tool_id with various valid inputs."""
+        server_id = kiln_task_server_id_from_tool_id(tool_id)
+        assert server_id == expected_server_id
+
     def test_tool_from_id_mcp_missing_task_raises_error(self):
         """Test that MCP tool ID with missing task raises ValueError."""
         mcp_tool_id = f"{MCP_REMOTE_TOOL_ID_PREFIX}test_server::test_tool"
@@ -508,7 +612,6 @@ class TestToolRegistry:
             description="Test MCP server",
             properties={
                 "server_url": "https://example.com",
-                "headers": {},
             },
         )
 
@@ -539,7 +642,6 @@ class TestToolRegistry:
             description="Different MCP server",
             properties={
                 "server_url": "https://example.com",
-                "headers": {},
             },
         )
 
@@ -560,3 +662,95 @@ class TestToolRegistry:
             match="External tool server not found: nonexistent_server in project ID test_project_id",
         ):
             tool_from_id(mcp_tool_id, task=mock_task)
+
+    def test_tool_from_id_kiln_task_tool_success(self):
+        """Test that tool_from_id works with Kiln task tool IDs."""
+        # Create mock external tool server for Kiln task
+        mock_server = ExternalToolServer(
+            name="test_kiln_task_server",
+            type=ToolServerType.kiln_task,
+            description="Test Kiln task server",
+            properties={
+                "name": "test_task_tool",
+                "description": "A test task tool",
+                "task_id": "test_task_123",
+                "run_config_id": "test_config_456",
+                "is_archived": False,
+            },
+        )
+
+        # Create mock project with the external tool server
+        mock_project = Mock(spec=Project)
+        mock_project.id = "test_project_id"
+        mock_project.external_tool_servers.return_value = [mock_server]
+
+        # Create mock task with parent project
+        mock_task = Mock(spec=Task)
+        mock_task.parent_project.return_value = mock_project
+
+        # Test with Kiln task tool ID
+        tool_id = f"{KILN_TASK_TOOL_ID_PREFIX}{mock_server.id}"
+        tool = tool_from_id(tool_id, task=mock_task)
+
+        # Verify the tool is KilnTaskTool
+        assert isinstance(tool, KilnTaskTool)
+        assert tool._project_id == "test_project_id"
+        assert tool._tool_id == tool_id
+        assert tool._tool_server_model == mock_server
+
+    def test_tool_from_id_kiln_task_tool_no_task(self):
+        """Test that Kiln task tool ID without task raises ValueError."""
+        tool_id = f"{KILN_TASK_TOOL_ID_PREFIX}test_server"
+        with pytest.raises(
+            ValueError,
+            match=r"Unable to resolve tool from id.*Requires a parent project/task",
+        ):
+            tool_from_id(tool_id, task=None)
+
+    def test_tool_from_id_kiln_task_tool_no_project(self):
+        """Test that Kiln task tool ID with task but no project raises ValueError."""
+        # Create mock task without parent project
+        mock_task = Mock(spec=Task)
+        mock_task.parent_project.return_value = None
+
+        tool_id = f"{KILN_TASK_TOOL_ID_PREFIX}test_server"
+
+        with pytest.raises(
+            ValueError,
+            match=r"Unable to resolve tool from id.*Requires a parent project/task",
+        ):
+            tool_from_id(tool_id, task=mock_task)
+
+    def test_tool_from_id_kiln_task_tool_server_not_found(self):
+        """Test that Kiln task tool ID with server not found raises ValueError."""
+        # Create mock external tool server with different ID
+        mock_server = ExternalToolServer(
+            name="different_server",
+            type=ToolServerType.kiln_task,
+            description="Different Kiln task server",
+            properties={
+                "name": "different_tool",
+                "description": "A different task tool",
+                "task_id": "different_task_123",
+                "run_config_id": "different_config_456",
+                "is_archived": False,
+            },
+        )
+
+        # Create mock project with the external tool server
+        mock_project = Mock(spec=Project)
+        mock_project.id = "test_project_id"
+        mock_project.external_tool_servers.return_value = [mock_server]
+
+        # Create mock task with parent project
+        mock_task = Mock(spec=Task)
+        mock_task.parent_project.return_value = mock_project
+
+        # Use a tool ID with a server that doesn't exist in the project
+        tool_id = f"{KILN_TASK_TOOL_ID_PREFIX}nonexistent_server"
+
+        with pytest.raises(
+            ValueError,
+            match="Kiln Task External tool server not found: nonexistent_server in project ID test_project_id",
+        ):
+            tool_from_id(tool_id, task=mock_task)

kiln_ai/tools/tool_registry.py

@@ -1,10 +1,12 @@
 from kiln_ai.datamodel.rag import RagConfig
 from kiln_ai.datamodel.task import Task
 from kiln_ai.datamodel.tool_id import (
+    KILN_TASK_TOOL_ID_PREFIX,
     MCP_LOCAL_TOOL_ID_PREFIX,
     MCP_REMOTE_TOOL_ID_PREFIX,
     RAG_TOOL_ID_PREFIX,
     KilnBuiltInToolId,
+    kiln_task_server_id_from_tool_id,
     mcp_server_and_tool_name_from_id,
     rag_config_id_from_id,
 )
@@ -15,6 +17,7 @@ from kiln_ai.tools.built_in_tools.math_tools import (
     MultiplyTool,
     SubtractTool,
 )
+from kiln_ai.tools.kiln_task_tool import KilnTaskTool
 from kiln_ai.tools.mcp_server_tool import MCPServerTool
 from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
 
@@ -38,31 +41,60 @@ def tool_from_id(tool_id: str, task: Task | None = None) -> KilnToolInterface:
             case _:
                 raise_exhaustive_enum_error(typed_tool_id)
 
-    # Check MCP Server Tools
-    if tool_id.startswith((MCP_REMOTE_TOOL_ID_PREFIX, MCP_LOCAL_TOOL_ID_PREFIX)):
+    # Check if this looks like an MCP or Kiln Task tool ID that requires a project
+    is_mcp_tool = tool_id.startswith(
+        (MCP_REMOTE_TOOL_ID_PREFIX, MCP_LOCAL_TOOL_ID_PREFIX)
+    )
+    is_kiln_task_tool = tool_id.startswith(KILN_TASK_TOOL_ID_PREFIX)
+
+    if is_mcp_tool or is_kiln_task_tool:
         project = task.parent_project() if task is not None else None
-        if project is None:
+        if project is None or project.id is None:
             raise ValueError(
                 f"Unable to resolve tool from id: {tool_id}. Requires a parent project/task."
             )
 
-        # Get the tool server ID and tool name from the ID
-        tool_server_id, tool_name = mcp_server_and_tool_name_from_id(tool_id)
+        # Check MCP Server Tools
+        if is_mcp_tool:
+            # Get the tool server ID and tool name from the ID
+            tool_server_id, tool_name = mcp_server_and_tool_name_from_id(
+                tool_id
+            )  # Fixed function name
 
-        server = next(
-            (
-                server
-                for server in project.external_tool_servers()
-                if server.id == tool_server_id
-            ),
-            None,
-        )
-        if server is None:
-            raise ValueError(
-                f"External tool server not found: {tool_server_id} in project ID {project.id}"
+            server = next(
+                (
+                    server
+                    for server in project.external_tool_servers()
+                    if server.id == tool_server_id
+                ),
+                None,
+            )
+            if server is None:
+                raise ValueError(
+                    f"External tool server not found: {tool_server_id} in project ID {project.id}"
+                )
+
+            return MCPServerTool(server, tool_name)
+
+        # Check Kiln Task Tools
+        if is_kiln_task_tool:
+            server_id = kiln_task_server_id_from_tool_id(tool_id)
+
+            server = next(
+                (
+                    server
+                    for server in project.external_tool_servers()
+                    if server.id == server_id
+                ),
+                None,
             )
+            if server is None:
+                raise ValueError(
+                    f"Kiln Task External tool server not found: {server_id} in project ID {project.id}"
+                )
+
+            return KilnTaskTool(project.id, tool_id, server)
 
-        return MCPServerTool(server, tool_name)
     elif tool_id.startswith(RAG_TOOL_ID_PREFIX):
         project = task.parent_project() if task is not None else None
         if project is None:

kiln_ai/utils/filesystem.py

@@ -5,10 +5,10 @@ from pathlib import Path
 
 
 def open_folder(path: str | Path) -> None:
-    log_dir = os.path.dirname(path)
+    dir = os.path.dirname(path)
     if sys.platform.startswith("darwin"):
-        subprocess.run(["open", log_dir], check=True)
+        subprocess.run(["open", dir], check=True)
     elif sys.platform.startswith("win"):
-        os.startfile(log_dir)  # type: ignore[attr-defined]
+        os.startfile(dir)  # type: ignore[attr-defined]
     else:
-        subprocess.run(["xdg-open", log_dir], check=True)
+        subprocess.run(["xdg-open", dir], check=True)

kiln_ai/utils/open_ai_types.py

@@ -17,11 +17,11 @@ from typing import (
 )
 
 from openai.types.chat import (
+    ChatCompletionContentPartTextParam,
     ChatCompletionDeveloperMessageParam,
     ChatCompletionFunctionMessageParam,
     ChatCompletionMessageToolCallParam,
     ChatCompletionSystemMessageParam,
-    ChatCompletionToolMessageParam,
     ChatCompletionUserMessageParam,
 )
 from openai.types.chat.chat_completion_assistant_message_param import (
@@ -84,11 +84,28 @@ class ChatCompletionAssistantMessageParamWrapper(TypedDict, total=False):
     """The tool calls generated by the model, such as function calls."""
 
 
+class ChatCompletionToolMessageParamWrapper(TypedDict, total=False):
+    content: Required[Union[str, Iterable[ChatCompletionContentPartTextParam]]]
+    """The contents of the tool message."""
+
+    role: Required[Literal["tool"]]
+    """The role of the messages author, in this case `tool`."""
+
+    tool_call_id: Required[str]
+    """Tool call that this message is responding to."""
+
+    kiln_task_tool_data: Optional[str]
+    """The data for the Kiln task tool that this message is responding to.
+    
+    Formatted as `<project_id>:::<tool_id>:::<task_id>:::<run_id>`
+    """
+
+
 ChatCompletionMessageParam: TypeAlias = Union[
     ChatCompletionDeveloperMessageParam,
     ChatCompletionSystemMessageParam,
     ChatCompletionUserMessageParam,
     ChatCompletionAssistantMessageParamWrapper,
-    ChatCompletionToolMessageParam,
+    ChatCompletionToolMessageParamWrapper,
     ChatCompletionFunctionMessageParam,
 ]

kiln_ai/utils/pdf_utils.py

@@ -8,6 +8,7 @@ from contextlib import asynccontextmanager
 from pathlib import Path
 from typing import AsyncGenerator
 
+import pypdfium2
 from pypdf import PdfReader, PdfWriter
 
 
@@ -36,3 +37,23 @@ async def split_pdf_into_pages(pdf_path: Path) -> AsyncGenerator[list[Path], Non
                 page_paths.append(page_path)
 
         yield page_paths
+
+
+async def convert_pdf_to_images(pdf_path: Path, output_dir: Path) -> list[Path]:
+    image_paths = []
+
+    # note: doing this in a thread causes a segfault - but this is slow and blocking
+    # so we should try to find a better way
+    pdf = pypdfium2.PdfDocument(pdf_path)
+    try:
+        for idx, page in enumerate(pdf):
+            await asyncio.sleep(0)
+            # scale=2 is legible for ~A4 pages (research papers, etc.) - lower than this is blurry
+            bitmap = page.render(scale=2).to_pil()
+            target_path = output_dir / f"img-{pdf_path.name}-{idx}.png"
+            bitmap.save(target_path)
+            image_paths.append(target_path)
+
+        return image_paths
+    finally:
+        pdf.close()