Update File Management Tools to Include Root Directory (#3112)

- Permit the specification of a `root_dir` to the read/write file tools
to specify a working directory
- Add validation for attempts to read/write outside the directory (e.g.,
through `../../` or symlinks or `/abs/path`'s that don't lie in the
correct path)
- Add some tests for all


One question is whether we should make a default root directory for
these? tradeoffs either way

langchain/tools/file_management/read.py

@@ -1,8 +1,10 @@
-from typing import Type
+from pathlib import Path
+from typing import Optional, Type
 
 from pydantic import BaseModel, Field
 
 from langchain.tools.base import BaseTool
+from langchain.tools.file_management.utils import get_validated_relative_path
 
 
 class ReadFileInput(BaseModel):
@@ -15,10 +17,19 @@ class ReadFileTool(BaseTool):
     name: str = "read_file"
     args_schema: Type[BaseModel] = ReadFileInput
     description: str = "Read file from disk"
+    root_dir: Optional[str] = None
+    """Directory to read file from.
+
+    If specified, raises an error for file_paths oustide root_dir."""
 
     def _run(self, file_path: str) -> str:
+        read_path = (
+            get_validated_relative_path(Path(self.root_dir), file_path)
+            if self.root_dir
+            else Path(file_path)
+        )
         try:
-            with open(file_path, "r", encoding="utf-8") as f:
+            with read_path.open("r", encoding="utf-8") as f:
                 content = f.read()
             return content
         except Exception as e:

langchain/tools/file_management/utils.py

@@ -0,0 +1,26 @@
+import sys
+from pathlib import Path
+
+
+def is_relative_to(path: Path, root: Path) -> bool:
+    """Check if path is relative to root."""
+    if sys.version_info >= (3, 9):
+        # No need for a try/except block in Python 3.8+.
+        return path.is_relative_to(root)
+    try:
+        path.relative_to(root)
+        return True
+    except ValueError:
+        return False
+
+
+def get_validated_relative_path(root: Path, user_path: str) -> Path:
+    """Resolve a relative path, raising an error if not within the root directory."""
+    # Note, this still permits symlinks from outside that point within the root.
+    # Further validation would be needed if those are to be disallowed.
+    root = root.resolve()
+    full_path = (root / user_path).resolve()
+
+    if not is_relative_to(full_path, root):
+        raise ValueError(f"Path {user_path} is outside of the allowed directory {root}")
+    return full_path

langchain/tools/file_management/write.py

@@ -1,9 +1,10 @@
-import os
-from typing import Type
+from pathlib import Path
+from typing import Optional, Type
 
 from pydantic import BaseModel, Field
 
 from langchain.tools.base import BaseTool
+from langchain.tools.file_management.utils import get_validated_relative_path
 
 
 class WriteFileInput(BaseModel):
@@ -17,15 +18,22 @@ class WriteFileTool(BaseTool):
     name: str = "write_file"
     args_schema: Type[BaseModel] = WriteFileInput
     description: str = "Write file to disk"
+    root_dir: Optional[str] = None
+    """Directory to write file to.
+
+    If specified, raises an error for file_paths oustide root_dir."""
 
     def _run(self, file_path: str, text: str) -> str:
+        write_path = (
+            get_validated_relative_path(Path(self.root_dir), file_path)
+            if self.root_dir
+            else Path(file_path)
+        )
         try:
-            directory = os.path.dirname(file_path)
-            if not os.path.exists(directory) and directory:
-                os.makedirs(directory)
-            with open(file_path, "w", encoding="utf-8") as f:
+            write_path.parent.mkdir(exist_ok=True, parents=False)
+            with write_path.open("w", encoding="utf-8") as f:
                 f.write(text)
-            return "File written to successfully."
+            return f"File written successfully to {file_path}."
         except Exception as e:
             return "Error: " + str(e)