diff options
Diffstat (limited to 'actions.py')
| -rw-r--r-- | actions.py | 117 |
1 files changed, 77 insertions, 40 deletions
@@ -1,6 +1,24 @@ from pathlib import Path -from typing import Optional +from typing import Optional, Callable, Dict +from dataclasses import dataclass import subprocess +from llm import AgentConfig, JSON + + +@dataclass +class Action: + """Class for LLM actions. + + Should only be called on JSON dicts where "action" -> `name`. + """ + + name: str + description: str + handler: Callable[[JSON, AgentConfig], JSON] + + def __call__(self, data: JSON, config: AgentConfig) -> JSON: + return self.handler(data, config) + def _path_checked(path: Path, cwd=Path.cwd()) -> Optional[Path]: """Resolves `path`, returning `None` if it's not a subdirectory of `cwd`.""" @@ -8,6 +26,7 @@ def _path_checked(path: Path, cwd=Path.cwd()) -> Optional[Path]: cwd = cwd.resolve() return path if path.is_relative_to(cwd) else None + def _path_helper(path_str: str, callback, cwd=Path.cwd()): path = _path_checked(Path(path_str), cwd=cwd) if path is None: @@ -18,59 +37,77 @@ def _path_helper(path_str: str, callback, cwd=Path.cwd()): except Exception as e: return {"status": "error", "reason": f"Failure: {e}"} -def ls(args, workdir=Path.cwd()): + +def _ls(args: JSON, config: AgentConfig) -> JSON: '{"action": "ls", "path": PATH} - list files at PATH' + def inner(path: Path): return {"content": list(map(lambda s: str(s), path.iterdir()))} - match args: - case {"action": "ls", "path": path_str}: - return _path_helper(path_str, inner, cwd=workdir) -def read(args, workdir=Path.cwd()): + return _path_helper(args["path"], inner, cwd=config.workdir) + + +def _read(args: JSON, config: AgentConfig) -> JSON: '{"action": "read", "path": PATH} - read file at PATH' + def inner(path: Path): contents = path.read_text(encoding="utf-8") return {"content": contents} - match args: - case {"action": "read", "path": path_str}: - return _path_helper(path_str, inner, cwd=workdir) -def write(args, workdir=Path.cwd()): + return _path_helper(args["path"], inner, cwd=config.workdir) + + +def _write(args: JSON, config: AgentConfig) -> JSON: '{"action": "write", "path": PATH, "content": CONTENT} - write CONTENT into file at PATH' + def inner(path: Path, text: str): n_bytes = path.write_text(text) return {"bytes": n_bytes} - match args: - case {"action": "write", "path": path_str, "content": content}: - return _path_helper(path_str, lambda path: inner(path, content), cwd=workdir) -def tldr(args, workdir=Path.cwd()): + return _path_helper( + args["path"], lambda path: inner(path, args["content"]), cwd=config.workdir + ) + + +def _tldr(args: JSON, config: AgentConfig) -> JSON: '{"action": "tldr", "program": PROGRAM} - read tldr docs for PROGRAM' - match args: - case {"action": "tldr", "program": program_str}: - try: - result = subprocess.run( - ["tldr", program_str], - capture_output=True, - text=True, - check=True - ) - return {"status": "success", "content": result.stdout} - except subprocess.CalledProcessError: - return {"status": "error", "reason": f"Tldr for '{program_str}' not found."} - except Exception as e: - return {"status": "error", "reason": f"An unexpected error occurred: {e}"} + try: + result = subprocess.run( + ["tldr", args["program"]], capture_output=True, text=True, check=True + ) + return {"status": "success", "content": result.stdout} + except subprocess.CalledProcessError: + return {"status": "error", "reason": f"Tldr for '{args['program']}' not found."} + except Exception as e: + return {"status": "error", "reason": f"An unexpected error occurred: {e}"} + -def shell(args, workdir=Path.cwd()): +def _shell(args: JSON, config: AgentConfig) -> JSON: '{"action": "shell", "command": CMD} - Propose executing a shell command, but interrupts to ask for permission from the user.' - match args: - case {"action": "shell", "command": cmd}: - user_input = input(f"Run cmd `{cmd}`? [y/N]: ") - if user_input.lower() == 'y': - res = subprocess.run(cmd, shell=True, cwd=workdir, capture_output=True, text=True) - if res.returncode == 0: - return {"status": "success", "stdout": res.stdout} - else: - return {"status": "error", "reason": f"Shell command exited with code {res.returncode}", "stdout": res.stdout, "stderr": res.stderr} - else: - return {"status": "error", "reason": "Permission denied by user."} + user_input = input(f"Run cmd `{args['command']}`? [y/N]: ") + if user_input.lower() == "y": + res = subprocess.run( + args["command"], + shell=True, + cwd=config.workdir, + capture_output=True, + text=True, + ) + if res.returncode == 0: + return {"status": "success", "stdout": res.stdout} + else: + return { + "status": "error", + "reason": f"Shell command exited with code {res.returncode}", + "stdout": res.stdout, + "stderr": res.stderr, + } + else: + return {"status": "error", "reason": "Permission denied by user."} + + +ls = Action("ls", _ls.__doc__ or "", _ls) +read = Action("read", _read.__doc__ or "", _read) +write = Action("write", _write.__doc__ or "", _write) +tldr = Action("tldr", _tldr.__doc__ or "", _tldr) +shell = Action("shell", _shell.__doc__ or "", _shell) |