feat: --workdir

2 files changed, 18 insertions, 22 deletions

diff --git a/actions.py b/actions.py
index 9146e53..765fe3d 100644
--- a/actions.py
+++ b/actions.py
@@ -4,12 +4,12 @@ import subprocess
 
 def _path_checked(path: Path, cwd=Path.cwd()) -> Optional[Path]:
     """Resolves `path`, returning `None` if it's not a subdirectory of `cwd`."""
-    path = path.resolve()
+    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):
-    path = _path_checked(Path(path_str))
+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:
@@ -18,33 +18,33 @@ def _path_helper(path_str: str, callback):
         except Exception as e:
             return {"status": "error", "reason": f"Failure: {e}"}
 
-def ls(args):
+def ls(args, workdir=Path.cwd()):
     '{"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)
+            return _path_helper(path_str, inner, cwd=workdir)
 
-def read(args):
+def read(args, workdir=Path.cwd()):
     '{"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)
+            return _path_helper(path_str, inner, cwd=workdir)
 
-def write(args):
+def write(args, workdir=Path.cwd()):
     '{"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))
+            return _path_helper(path_str, lambda path: inner(path, content), cwd=workdir)
 
-def tldr(args):
+def tldr(args, workdir=Path.cwd()):
     '{"action": "tldr", "program": PROGRAM} - read tldr docs for PROGRAM'
     match args:
         case {"action": "tldr", "program": program_str}:
@@ -61,4 +61,3 @@ def tldr(args):
             except Exception as e:
                 return {"status": "error", "reason": f"An unexpected error occurred: {e}"}
 
-ALL_ACTIONS = [ls, read, write, tldr]
diff --git a/llm.py b/llm.py
index 57e8e27..51234bd 100755
--- a/llm.py
+++ b/llm.py
@@ -58,14 +58,7 @@ def chat(content, state = (), query_method = direct_query, role = "user"):
     return (response, state)
 
 
-def path_checked(path: Path, cwd = Path.cwd()) -> Optional[Path]:
-    """Resolves `path`, returning `None` if it's not a subdirectory of `cwd`."""
-    path = path.resolve()
-    cwd = cwd.resolve()
-    return path if path.is_relative_to(cwd) else None
-
-
-def query_agent(actions: list, query: str, state = ()):
+def query_agent(actions: list, query: str, state = (), workdir: Path = Path.cwd()):
 
     state = state + ({"role": "user", "content": query},)
 
@@ -85,7 +78,7 @@ def query_agent(actions: list, query: str, state = ()):
             continue
 
         for f in actions:
-            result_json = f(action)
+            result_json = f(action, workdir=workdir)
             if result_json is not None:
                 replacements.append((m.start(), m.end(), action["action"]))
                 result_jsons.append(result_json)
@@ -99,13 +92,17 @@ def query_agent(actions: list, query: str, state = ()):
     replaced_res += res[i:]
 
     if len(replacements) > 0:
-        (res, state) = query_agent(actions, json.dumps(result_jsons), state)
+        (res, state) = query_agent(actions, json.dumps(result_jsons), state, workdir=workdir)
         return (replaced_res + "\n\n" + res, state)
     else:
         return (replaced_res, state)
 
 
 def main():
+    parser = argparse.ArgumentParser(description="Autonomous Coding Agent")
+    parser.add_argument("--workdir", type=Path, default=Path.cwd(), help="The directory the agent has access to (default: current directory)")
+    args = parser.parse_args()
+
     actions = [ls, read, write, tldr]
     state = ({"role": "system",
                       "content": f"""You are an autonomous coding agent that can help with programming tasks.
@@ -123,7 +120,7 @@ Once you are finished with your task, give a response with none of the above JSO
         
         query = input("> ").strip()
 
-        (response, state) = query_agent(actions, query, state)
+        (response, state) = query_agent(actions, query, state, workdir=args.workdir)
         
         print(response)