import os import logging from typing import Type, Literal from pydantic import BaseModel, Field # Ensure openwebui.tools.base is available in your environment # Usually installed as part of open-webui backend dependencies try: from openwebui.tools.base import BaseTool except ImportError: # Fallback for environments where openwebui might not be directly in PYTHONPATH # Adjust the path based on your actual installation if needed import sys # Example: Assuming tools are siblings to a potential 'openwebui' package dir # sys.path.append(os.path.join(os.path.dirname(__file__), '..')) try: from openwebui.tools.base import BaseTool except ImportError: raise ImportError("Could not import BaseTool from openwebui.tools.base. " "Ensure open-webui backend is installed correctly.") logging.basicConfig(level=logging.INFO) log = logging.getLogger(__name__) # --- Configuration --- # Define the directory where files will be created/written. # IMPORTANT: This path is relative to where the Open WebUI *backend* is running. # For security, it's strongly recommended to keep this within a dedicated, # non-critical subdirectory. ALLOWED_WRITE_DIR_NAME = "created_files" # --- End Configuration --- # Ensure the base directory for writing exists backend_cwd = os.getcwd() # Get current working directory of the backend process ALLOWED_WRITE_DIR = os.path.abspath(os.path.join(backend_cwd, ALLOWED_WRITE_DIR_NAME)) if not os.path.exists(ALLOWED_WRITE_DIR): try: os.makedirs(ALLOWED_WRITE_DIR) log.info(f"Created allowed write directory: {ALLOWED_WRITE_DIR}") except OSError as e: log.error(f"Error creating directory {ALLOWED_WRITE_DIR}: {e}. File writing will likely fail.") # You might want to raise an exception here to prevent the tool from loading # raise RuntimeError(f"Failed to create required directory: {ALLOWED_WRITE_DIR}") from e class FileWriterArgs(BaseModel): """Arguments for the FileWriterTool.""" filename: str = Field( ..., description="The name of the file to create or write to. " "Should only be a filename, not a complex path. " "E.g., 'my_notes.txt', 'script.py'." ) content: str = Field( ..., description="The text content to write into the file." ) mode: Literal['w', 'a'] = Field( default='w', description="Write mode: 'w' to overwrite the file if it exists (or create if not), " "'a' to append to the file if it exists (or create if not)." ) class Tools: """ Container class for Open WebUI tools. This class itself might not be directly instantiated by the user, but Open WebUI may look inside it to discover tools. """ class FileWriterTool(BaseTool): """ A tool to write or append text content to files within a designated safe directory. Can also create new files if they don't exist. """ name: str = "file_writer" description: str = ( "Writes or appends text content to a specified file within a secure directory. " "Creates the file if it doesn't exist. Use 'w' mode to overwrite, 'a' to append. " "Specify only the filename (e.g., 'report.md'), not paths like '../' or '/etc'." ) args_schema: Type[BaseModel] = FileWriterArgs # Reference the standalone args schema def _is_safe_path(self, filename: str) -> tuple[bool, str | None]: """ Validates the filename and constructs the full, safe path. Returns (is_safe, full_path) or (is_safe, None). Uses the globally defined ALLOWED_WRITE_DIR. """ # Basic sanitization: remove leading/trailing whitespace and slashes sanitized_filename = filename.strip().lstrip('/').lstrip('\\') # Disallow path traversal or absolute paths attempts in the filename itself if ".." in sanitized_filename or os.path.isabs(sanitized_filename) or \ '/' in sanitized_filename or '\\' in sanitized_filename: log.warning(f"Attempted unsafe filename pattern: {filename}") return False, None # Construct the full path using the global ALLOWED_WRITE_DIR full_path = os.path.join(ALLOWED_WRITE_DIR, sanitized_filename) # Final check: Ensure the resolved absolute path is truly within the allowed directory abs_full_path = os.path.abspath(full_path) abs_allowed_dir = os.path.abspath(ALLOWED_WRITE_DIR) if os.path.commonpath([abs_full_path, abs_allowed_dir]) != abs_allowed_dir: log.warning(f"Resolved path '{abs_full_path}' is outside allowed directory '{abs_allowed_dir}'. Original filename: '{filename}'") return False, None # Check if the target is a directory - we don't want to overwrite directories if os.path.isdir(abs_full_path): log.warning(f"Attempted to write to a path that is a directory: {abs_full_path}") return False, None return True, abs_full_path def __call__(self, args: FileWriterArgs) -> str: """ Executes the file writing operation. """ log.info(f"FileWriterTool called with args: filename='{args.filename}', mode='{args.mode}'") is_safe, full_path = self._is_safe_path(args.filename) if not is_safe: return (f"Error: Invalid or unsafe filename provided: '{args.filename}'. " f"Only simple filenames are allowed within the designated directory. " f"Do not use paths like '/', '\\', or '..'.") try: # Write/Append to the file with open(full_path, mode=args.mode, encoding='utf-8') as f: f.write(args.content) action = "appended to" if args.mode == 'a' else "written to" log.info(f"Successfully {action} file: {full_path}") return f"Successfully {action} file: '{args.filename}' in the designated directory." except IOError as e: log.error(f"Error writing to file {full_path}: {e}") return f"Error: Could not write to file '{args.filename}'. Reason: {e}" except Exception as e: log.exception(f"An unexpected error occurred while writing file {full_path}: {e}") return f"Error: An unexpected error occurred while processing file '{args.filename}'. Check logs." # You could potentially add other tools here as nested classes within Tools # class AnotherTool(BaseTool): # ... # Example of how Open WebUI *might* discover and use the tool # (This is illustrative and depends on Open WebUI's internal logic) # # if __name__ == "__main__": # # Open WebUI might inspect the Tools class # available_tool_classes = { # name: tool_class for name, tool_class in Tools.__dict__.items() # if isinstance(tool_class, type) and issubclass(tool_class, BaseTool) and tool_class is not BaseTool # } # print(f"Discovered tool classes: {available_tool_classes}") # # # Then instantiate and use it # if 'FileWriterTool' in available_tool_classes: # FileWriterToolClass = available_tool_classes['FileWriterTool'] # tool_instance = FileWriterToolClass() # Instantiate the tool # print(f"Instantiated tool: {tool_instance.name}") # # # --- Test Cases (Similar to before) --- # args1 = FileWriterArgs(filename="test_file_nested.txt", content="Hello from nested tool!") # result1 = tool_instance(args1) # Call the instance # print(f"Test 1 Result: {result1}") # # args2 = FileWriterArgs(filename="../unsafe_nested.txt", content="Still unsafe.") # result2 = tool_instance(args2) # print(f"Test 2 Result: {result2}")