aboutsummaryrefslogtreecommitdiff
path: root/actions.py
blob: dfeebc3aa0d5ffdf7ae92b9565f3af251ffbbc9b (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
from pathlib import Path
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]
    formatter: Callable[[JSON], str] = lambda args: (
        f"{args['action']} with "
        + ",".join([f"{arg} as {args[arg]}" for arg in args if arg != "action"])
    )

    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`."""
    path = (path if path.is_absolute() else (cwd / path)).resolve()
    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:
        return {"status": "error", "reason": f"Permission denied: {path_str}"}
    else:
        try:
            return {"status": "success", **callback(path)}
        except Exception as e:
            return {"status": "error", "reason": f"Failure: {e}"}


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()))}

    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}

    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}

    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'
    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: JSON, config: AgentConfig) -> JSON:
    '{"action": "shell", "command": CMD} - Propose executing a shell command, but interrupts to ask for permission from the user.'
    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,
        }


ls = Action("ls", _ls.__doc__ or "", _ls)
read = Action("read", _read.__doc__ or "", _read)
write = Action(
    "write",
    _write.__doc__ or "",
    _write,
    formatter=lambda args: f"write {args['path']}",
)
tldr = Action("tldr", _tldr.__doc__ or "", _tldr)
shell = Action(
    "shell",
    _shell.__doc__ or "",
    _shell,
    formatter=lambda args: f"shell command `{args['command']}`",
)