Workspace file tools and WorkspaceLLMTaskWorker

docs/source/usage.rst

@@ -209,6 +209,56 @@ Example:
     initial_input = Task1WorkItem(data="start")
     main_graph.run(initial_tasks=[(subgraph_worker, initial_input)])
 
+Letting an LLM Work with Files
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+``WorkspaceLLMTaskWorker`` (in ``planai.workspace_task``) is a ``CachedLLMTaskWorker`` that
+gives the LLM a set of file tools -- ``read_file``, ``write_file``, ``edit_file``,
+``list_files``, and ``grep_files`` -- jailed to a per-job working directory. This is useful
+when many jobs run concurrently, each operating on its own directory, and you want an LLM to
+read, write, or search files without risking access outside of the directory assigned to it.
+
+The working directory is discovered automatically from the task's provenance chain: include a
+``WorkspaceTask`` (or any ``Task`` subclass with a string ``workspace`` attribute) upstream, and
+``WorkspaceLLMTaskWorker.get_workspace()`` will find the nearest one. All file paths the LLM
+uses are relative to that directory; absolute paths, ``..`` components, and symlinks that
+escape the directory are rejected.
+
+.. code-block:: python
+
+    from planai import Graph, WorkspaceLLMTaskWorker, WorkspaceTask, llm_from_config
+
+    class CodeReviewer(WorkspaceLLMTaskWorker):
+        prompt = "Review the code in this repository and write your findings to review.md."
+        output_types = [ReviewResult]
+
+        def expected_output_files(self, task):
+            # a cache hit whose files are missing (e.g. a fresh checkout) is re-executed
+            return ["review.md"]
+
+    llm = llm_from_config(provider="openai", model_name="gpt-4o")
+    reviewer = CodeReviewer(llm=llm, max_tool_rounds=40)
+
+    graph = Graph(name="Review Workflow")
+    graph.add_workers(reviewer)
+    graph.set_entry(reviewer)
+    graph.set_exit(reviewer)
+    graph.run(
+        initial_tasks=[(reviewer, WorkspaceTask(workspace="/jobs/job-123"))]
+    )
+
+Because a cache hit replays only the published output tasks and not any files the tools wrote
+on a previous run, override ``expected_output_files()`` to list the workspace-relative paths
+the worker is expected to produce; if any are missing, the cache entry is treated as a miss and
+the worker re-executes. Set ``input_globs`` to fold the content of matching workspace files
+into the cache key so that changing an input file also invalidates the cache. Pass
+``read_only=True`` to omit the ``write_file`` and ``edit_file`` tools.
+
+The underlying pieces are reusable on their own: ``Workspace`` is the sandboxing primitive,
+``make_file_tools(workspace, ...)`` builds the ``llm_interface`` ``Tool`` objects for any
+``LLMTaskWorker`` (via the ``get_tools()`` hook), and ``hash_files(workspace, globs)`` computes
+a stable content hash for a set of glob patterns.
+
 Best Practices
 --------------
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "planai"
-version = "0.6.1"
+version = "0.7.0"
 description = "A simple framework for coordinating classical compute and LLM-based tasks."
 authors = ["Niels Provos <planai@provos.org>"]
 license = "Apache-2.0"
@@ -33,7 +33,7 @@ html2text = "^2024.2.26"
 dicttoxml = "^1.7.16"
 psutil = "^6.1.1"
 waitress = "^3.0.2"
-llm-interface = "^0.1.13"
+llm-interface = ">=0.2.0,<0.3.0"
 
 
 [tool.poetry.group.dev.dependencies]

src/planai/__init__.py

@@ -25,7 +25,9 @@
 from .provenance import ProvenanceChain
 from .pydantic_dict_wrapper import PydanticDictWrapper
 from .task import Task, TaskWorker
+from .tools import Workspace, hash_files, make_file_tools
 from .user_input import UserInputRequest
+from .workspace_task import WorkspaceLLMTaskWorker, WorkspaceTask
 
 # Limit what gets imported with "from planai import *"
 __all__ = [
@@ -50,4 +52,9 @@
     "UserInputRequest",
     "tool",
     "Tool",
+    "Workspace",
+    "make_file_tools",
+    "hash_files",
+    "WorkspaceLLMTaskWorker",
+    "WorkspaceTask",
 ]

src/planai/_version.py

@@ -14,4 +14,4 @@
 
 """Version information for PlanAI."""
 
-__version__ = "0.6.1"
+__version__ = "0.7.0"

src/planai/cached_task.py

@@ -50,8 +50,19 @@ def _pre_consume_work(self, task: Task):
                 logging.error("Error getting data from cache %s: %s", cache_key, str(e))
                 result = None
 
+            cache_hit = False
+            cached_results = None
             if result is not None:
                 cached_results, _ = result
+                cache_hit = self._cache_hit_is_valid(task, cached_results)
+                if not cache_hit:
+                    logging.info(
+                        "Cache hit for %s with key: %s is no longer valid; re-executing",
+                        self.name,
+                        cache_key,
+                    )
+
+            if cache_hit:
                 logging.info("Cache hit for %s with key: %s", self.name, cache_key)
                 self._publish_cached_results(cached_results, task)
             else:
@@ -66,6 +77,24 @@ def _pre_consume_work(self, task: Task):
 
             self.post_consume_work(task)
 
+    def _cache_hit_is_valid(
+        self, task: Task, cached_results: List[Tuple[str, Task]]
+    ) -> bool:
+        """
+        Hook for subclasses to reject a cache hit based on external state that isn't
+        captured by the cache key itself (e.g. output files on disk that a previous
+        run's cache entry references but that are no longer present). Defaults to
+        always accepting the cache hit.
+
+        Args:
+            task (Task): The input task that produced a cache hit.
+            cached_results (List[Tuple[str, Task]]): The cached (consumer_name, output_task) pairs.
+
+        Returns:
+            bool: True if the cached results should be used, False to treat this as a cache miss.
+        """
+        return True
+
     def pre_consume_work(self, task: Task):
         """
         This method is called before consuming the work item. It will be called even if the task has been cached.

src/planai/dispatcher.py

@@ -54,7 +54,6 @@
 across various workers effectively.
 """
 
-
 import logging
 import random
 import threading

src/planai/llm_task.py

@@ -85,6 +85,13 @@ class LLMTaskWorker(BaseLLMTaskWorker):
         default=False,
         description="Whether to use XML format for the data input to the LLM",
     )
+    max_tool_rounds: Optional[int] = Field(
+        default=None,
+        description=(
+            "Maximum number of tool-calling rounds to allow the LLM. Only forwarded to "
+            "generate_pydantic when tools are actually in use for a given call."
+        ),
+    )
 
     def __init__(self, **data):
         super().__init__(**data)
@@ -124,6 +131,38 @@ def _format_task(self, task: Task) -> str:
             else task.model_dump_xml()
         )
 
+    def get_tools(self, task: Task) -> Optional[List[Tool]]:
+        """
+        Returns the tools that should be made available to the LLM for this task.
+        Defaults to the static ``tools`` field, but subclasses can override this to
+        provide dynamic, per-task tools (e.g. file tools bound to a task-specific
+        working directory).
+
+        Args:
+            task (Task): The input task.
+
+        Returns:
+            Optional[List[Tool]]: The tools to make available, or None/empty for no tools.
+        """
+        return self.tools
+
+    def get_cache_salt(self, task: Task) -> Optional[str]:
+        """
+        Returns an optional salt to forward as ``cache_salt`` to generate_pydantic.
+        Only used by _invoke_llm when it returns a non-None value (and tools are in
+        use), so the default LLMTaskWorker behavior of not passing cache_salt is
+        unchanged. Subclasses (e.g. WorkspaceLLMTaskWorker) can override this to
+        invalidate the LLM's own response cache when external state, such as
+        workspace files, changes.
+
+        Args:
+            task (Task): The input task.
+
+        Returns:
+            Optional[str]: The cache salt, or None to omit it.
+        """
+        return None
+
     def _invoke_llm(self, task: Task):
         # allow subclasses to customize the prompt based on the input task
         task_prompt = self.format_prompt(task)
@@ -141,6 +180,19 @@ def extra_validation_with_task(response: BaseModel):
             assert isinstance(response, Task)
             return self.extra_validation(response, task)
 
+        tools = self.get_tools(task)
+
+        # only forward these new, tool-related kwargs when tools are actually in
+        # use for this call, so that callers/mocks that don't expect them (and
+        # don't use tools) keep working unchanged.
+        extra_kwargs: Dict[str, Any] = {}
+        if tools:
+            if self.max_tool_rounds is not None:
+                extra_kwargs["max_tool_rounds"] = self.max_tool_rounds
+            cache_salt = self.get_cache_salt(task)
+            if cache_salt is not None:
+                extra_kwargs["cache_salt"] = cache_salt
+
         response = self.llm.generate_pydantic(
             prompt_template=(
                 (PROMPT_TEMPLATE if processed_task is not None else "{instructions}")
@@ -152,7 +204,7 @@ def extra_validation_with_task(response: BaseModel):
             ),
             output_schema=self._output_type(),
             system=self.system_prompt,
-            tools=self.tools if self.tools else None,
+            tools=tools if tools else None,
             task=self._format_task(processed_task),
             temperature=self.temperature,
             instructions=task_prompt,
@@ -162,6 +214,7 @@ def extra_validation_with_task(response: BaseModel):
             debug_saver=save_debug_with_task if self.debug_mode else None,
             extra_validation=extra_validation_with_task,
             images=task.images if isinstance(task, MediaTask) else None,
+            **extra_kwargs,
         )
         assert isinstance(response, Task) or response is None
         self.post_process(response=response, input_task=task)

src/planai/provenance.py

@@ -40,7 +40,6 @@
 management of task dependencies and ordered task execution.
 """
 
-
 import logging
 import sys
 from collections import defaultdict

src/planai/tools/__init__.py

@@ -0,0 +1,18 @@
+# Copyright 2024 Niels Provos
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tools that can be handed to LLMTaskWorker subclasses to let an LLM act on files."""
+
+from .filesystem import Workspace, hash_files, make_file_tools
+
+__all__ = ["Workspace", "make_file_tools", "hash_files"]