feat: --workdir

1 files changed, 8 insertions, 11 deletions

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)