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

kiln_ai/tools/test_tool_registry.py

@@ -1,15 +1,22 @@
+from pathlib import Path
 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 (
@@ -18,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
 
@@ -85,7 +93,6 @@ class TestToolRegistry:
             type=ToolServerType.remote_mcp,
             properties={
                 "server_url": "https://example.com",
-                "headers": {},
             },
         )
 
@@ -143,7 +150,7 @@ class TestToolRegistry:
         tool_id = f"{MCP_LOCAL_TOOL_ID_PREFIX}test_server::test_tool"
         with pytest.raises(
             ValueError,
-            match="Unable to resolve tool from id.*Requires a parent project/task",
+            match=r"Unable to resolve tool from id.*Requires a parent project/task",
         ):
             tool_from_id(tool_id, task=None)
 
@@ -155,7 +162,6 @@ class TestToolRegistry:
             type=ToolServerType.remote_mcp,
             properties={
                 "server_url": "https://example.com",
-                "headers": {},
             },
         )
 
@@ -181,6 +187,93 @@ class TestToolRegistry:
             ):
                 tool_from_id(tool_id, task=mock_task)
 
+    def test_tool_from_id_rag_tool_success(self):
+        """Test that tool_from_id works with RAG tool IDs."""
+        # Create mock RAG config
+        from unittest.mock import patch
+
+        with (
+            patch("kiln_ai.tools.tool_registry.RagConfig") as mock_rag_config_class,
+            patch("kiln_ai.tools.rag_tools.RagTool") as mock_rag_tool_class,
+        ):
+            # Setup mock RAG config
+            mock_rag_config = Mock()
+            mock_rag_config.id = "test_rag_config"
+            mock_rag_config_class.from_id_and_parent_path.return_value = mock_rag_config
+
+            # Setup mock RAG tool
+            mock_rag_tool = Mock()
+            mock_rag_tool_class.return_value = mock_rag_tool
+
+            # Create mock project
+            mock_project = Mock(spec=Project)
+            mock_project.id = "test_project_id"
+            mock_project.path = Path("/test/path")
+
+            # Create mock task with parent project
+            mock_task = Mock(spec=Task)
+            mock_task.parent_project.return_value = mock_project
+
+            # Test with RAG tool ID
+            tool_id = f"{RAG_TOOL_ID_PREFIX}test_rag_config"
+            tool = tool_from_id(tool_id, task=mock_task)
+
+            # Verify the tool is RagTool
+            assert tool == mock_rag_tool
+            mock_rag_config_class.from_id_and_parent_path.assert_called_once_with(
+                "test_rag_config", Path("/test/path")
+            )
+            mock_rag_tool_class.assert_called_once_with(tool_id, mock_rag_config)
+
+    def test_tool_from_id_rag_tool_no_task(self):
+        """Test that RAG tool ID without task raises ValueError."""
+        tool_id = f"{RAG_TOOL_ID_PREFIX}test_rag_config"
+
+        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_rag_tool_no_project(self):
+        """Test that RAG 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"{RAG_TOOL_ID_PREFIX}test_rag_config"
+
+        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_rag_config_not_found(self):
+        """Test that RAG tool ID with missing RAG config raises ValueError."""
+        from unittest.mock import patch
+
+        with patch("kiln_ai.tools.tool_registry.RagConfig") as mock_rag_config_class:
+            # Setup mock to return None (config not found)
+            mock_rag_config_class.from_id_and_parent_path.return_value = None
+
+            # Create mock project
+            mock_project = Mock(spec=Project)
+            mock_project.id = "test_project_id"
+            mock_project.path = Path("/test/path")
+
+            # Create mock task with parent project
+            mock_task = Mock(spec=Task)
+            mock_task.parent_project.return_value = mock_project
+
+            tool_id = f"{RAG_TOOL_ID_PREFIX}missing_rag_config"
+
+            with pytest.raises(
+                ValueError,
+                match="RAG config not found: missing_rag_config in project test_project_id for tool",
+            ):
+                tool_from_id(tool_id, task=mock_task)
+
     def test_all_built_in_tools_are_registered(self):
         """Test that all KilnBuiltInToolId enum members are handled by the registry."""
         for tool_id in KilnBuiltInToolId:
@@ -272,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 = [
@@ -400,13 +535,71 @@ 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"
 
         with pytest.raises(
             ValueError,
-            match="Unable to resolve tool from id.*Requires a parent project/task",
+            match=r"Unable to resolve tool from id.*Requires a parent project/task",
         ):
             tool_from_id(mcp_tool_id, task=None)
 
@@ -419,7 +612,6 @@ class TestToolRegistry:
             description="Test MCP server",
             properties={
                 "server_url": "https://example.com",
-                "headers": {},
             },
         )
 
@@ -450,7 +642,6 @@ class TestToolRegistry:
             description="Different MCP server",
             properties={
                 "server_url": "https://example.com",
-                "headers": {},
             },
         )
 
@@ -471,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,9 +1,14 @@
+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,
 )
 from kiln_ai.tools.base_tool import KilnToolInterface
 from kiln_ai.tools.built_in_tools.math_tools import (
@@ -12,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
 
@@ -35,30 +41,77 @@ 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}"
+                )
+
+            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)
+
+    elif tool_id.startswith(RAG_TOOL_ID_PREFIX):
+        project = task.parent_project() if task is not None else None
+        if project is None:
+            raise ValueError(
+                f"Unable to resolve tool from id: {tool_id}. Requires a parent project/task."
+            )
 
-        server = next(
-            (
-                server
-                for server in project.external_tool_servers()
-                if server.id == tool_server_id
-            ),
-            None,
-        )
-        if server is None:
+        rag_config_id = rag_config_id_from_id(tool_id)
+        rag_config = RagConfig.from_id_and_parent_path(rag_config_id, project.path)
+        if rag_config is None:
             raise ValueError(
-                f"External tool server not found: {tool_server_id} in project ID {project.id}"
+                f"RAG config not found: {rag_config_id} in project {project.id} for tool {tool_id}"
             )
 
-        return MCPServerTool(server, tool_name)
+        # Lazy import to avoid circular dependency
+        from kiln_ai.tools.rag_tools import RagTool
+
+        return RagTool(tool_id, rag_config)
 
     raise ValueError(f"Tool ID {tool_id} not found in tool registry")

kiln_ai/utils/__init__.py

@@ -5,8 +5,11 @@ Misc utilities used in the kiln_ai library.
 """
 
 from . import config, formatting
+from .lock import AsyncLockManager, shared_async_lock_manager
 
 __all__ = [
+    "AsyncLockManager",
     "config",
     "formatting",
+    "shared_async_lock_manager",
 ]

kiln_ai/utils/async_job_runner.py

@@ -1,7 +1,7 @@
 import asyncio
 import logging
 from dataclasses import dataclass
-from typing import AsyncGenerator, Awaitable, Callable, List, TypeVar
+from typing import AsyncGenerator, Awaitable, Callable, Generic, List, TypeVar
 
 logger = logging.getLogger(__name__)
 
@@ -15,29 +15,66 @@ class Progress:
     errors: int
 
 
-class AsyncJobRunner:
-    def __init__(self, concurrency: int = 1):
+class AsyncJobRunnerObserver(Generic[T]):
+    async def on_error(self, job: T, error: Exception):
+        """
+        Called when a job raises an unhandled exception.
+        """
+        pass
+
+    async def on_success(self, job: T):
+        """
+        Called when a job completes successfully.
+        """
+        pass
+
+    async def on_job_start(self, job: T):
+        """
+        Called when a job starts.
+        """
+        pass
+
+
+class AsyncJobRunner(Generic[T]):
+    def __init__(
+        self,
+        jobs: List[T],
+        run_job_fn: Callable[[T], Awaitable[bool]],
+        concurrency: int = 1,
+        observers: List[AsyncJobRunnerObserver[T]] | None = None,
+    ):
         if concurrency < 1:
             raise ValueError("concurrency must be ≥ 1")
         self.concurrency = concurrency
+        self.jobs = jobs
+        self.run_job_fn = run_job_fn
+        self.observers = observers or []
 
-    async def run(
-        self,
-        jobs: List[T],
-        run_job: Callable[[T], Awaitable[bool]],
-    ) -> AsyncGenerator[Progress, None]:
+    async def notify_error(self, job: T, error: Exception):
+        for observer in self.observers:
+            await observer.on_error(job, error)
+
+    async def notify_success(self, job: T):
+        for observer in self.observers:
+            await observer.on_success(job)
+
+    async def notify_job_start(self, job: T):
+        for observer in self.observers:
+            await observer.on_job_start(job)
+
+    async def run(self) -> AsyncGenerator[Progress, None]:
         """
         Runs the jobs with parallel workers and yields progress updates.
         """
         complete = 0
         errors = 0
-        total = len(jobs)
+        total = len(self.jobs)
 
         # Send initial status
         yield Progress(complete=complete, total=total, errors=errors)
 
         worker_queue: asyncio.Queue[T] = asyncio.Queue()
-        for job in jobs:
+        for job in self.jobs:
             worker_queue.put_nowait(job)
 
         # simple status queue to return progress. True=success, False=error
@@ -46,7 +83,7 @@ class AsyncJobRunner:
         workers = []
         for _ in range(self.concurrency):
             task = asyncio.create_task(
-                self._run_worker(worker_queue, status_queue, run_job),
+                self._run_worker(worker_queue, status_queue, self.run_job_fn),
             )
             workers.append(task)
 
@@ -64,7 +101,11 @@ class AsyncJobRunner:
                     else:
                         errors += 1
 
-                    yield Progress(complete=complete, total=total, errors=errors)
+                    yield Progress(
+                        complete=complete,
+                        total=total,
+                        errors=errors,
+                    )
                 except asyncio.TimeoutError:
                     # Timeout is expected, just continue to recheck worker status
                     # Don't love this but beats sentinels for reliability
@@ -82,7 +123,7 @@ class AsyncJobRunner:
         self,
         worker_queue: asyncio.Queue[T],
         status_queue: asyncio.Queue[bool],
-        run_job: Callable[[T], Awaitable[bool]],
+        run_job_fn: Callable[[T], Awaitable[bool]],
     ):
         while True:
             try:
@@ -92,13 +133,17 @@ class AsyncJobRunner:
                 break
 
             try:
-                success = await run_job(job)
-            except Exception:
+                await self.notify_job_start(job)
+                result = await run_job_fn(job)
+                if result:
+                    await self.notify_success(job)
+            except Exception as e:
                 logger.error("Job failed to complete", exc_info=True)
-                success = False
+                await self.notify_error(job, e)
+                result = False
 
             try:
-                await status_queue.put(success)
+                await status_queue.put(result)
             except Exception:
                 logger.error("Failed to enqueue status for job", exc_info=True)
             finally:

kiln_ai/utils/config.py

@@ -221,14 +221,14 @@ class Config:
             raise AttributeError(f"Config has no attribute '{name}'")
 
     @classmethod
-    def settings_dir(cls, create=True):
+    def settings_dir(cls, create=True) -> str:
         settings_dir = os.path.join(Path.home(), ".kiln_ai")
         if create and not os.path.exists(settings_dir):
             os.makedirs(settings_dir)
         return settings_dir
 
     @classmethod
-    def settings_path(cls, create=True):
+    def settings_path(cls, create=True) -> str:
         settings_dir = cls.settings_dir(create)
         return os.path.join(settings_dir, "settings.yaml")
 

kiln_ai/utils/env.py

@@ -0,0 +1,15 @@
+import os
+from contextlib import contextmanager
+
+
+@contextmanager
+def temporary_env(var_name: str, value: str):
+    old_value = os.environ.get(var_name)
+    os.environ[var_name] = value
+    try:
+        yield
+    finally:
+        if old_value is None:
+            os.environ.pop(var_name, None)  # remove if it did not exist before
+        else:
+            os.environ[var_name] = old_value