From 096d27f342322849355b7754e1f5908f23912082 Mon Sep 17 00:00:00 2001
From: Erik Peterson <e@eriklp.com>
Date: Sat, 17 Jun 2023 14:53:58 -0700
Subject: [PATCH] Fix execute_command coming from plugins (#4729)

Signed-off-by: Merwane Hamadi <merwanehamadi@gmail.com>
---
 autogpt/agent/agent.py             |  5 ++---
 autogpt/app.py                     |  7 +++----
 tests/unit/test_execute_command.py | 24 ++++++++++++++++++++++++
 3 files changed, 29 insertions(+), 7 deletions(-)
 create mode 100644 tests/unit/test_execute_command.py

diff --git a/autogpt/agent/agent.py b/autogpt/agent/agent.py
index 202c124a145..8af9f49e487 100644
--- a/autogpt/agent/agent.py
+++ b/autogpt/agent/agent.py
@@ -268,9 +268,8 @@ def signal_handler(signum, frame):
                         command_name, arguments
                     )
                 command_result = execute_command(
-                    self.command_registry,
-                    command_name,
-                    arguments,
+                    command_name=command_name,
+                    arguments=arguments,
                     agent=self,
                 )
                 result = f"Command {command_name} returned: " f"{command_result}"
diff --git a/autogpt/app.py b/autogpt/app.py
index 30f4b7d6b1f..78e3a4dd206 100644
--- a/autogpt/app.py
+++ b/autogpt/app.py
@@ -3,7 +3,6 @@
 from typing import Dict
 
 from autogpt.agent.agent import Agent
-from autogpt.models.command_registry import CommandRegistry
 
 
 def is_valid_int(value: str) -> bool:
@@ -79,7 +78,6 @@ def map_command_synonyms(command_name: str):
 
 
 def execute_command(
-    command_registry: CommandRegistry,
     command_name: str,
     arguments: dict[str, str],
     agent: Agent,
@@ -89,12 +87,13 @@ def execute_command(
     Args:
         command_name (str): The name of the command to execute
         arguments (dict): The arguments for the command
+        agent (Agent): The agent that is executing the command
 
     Returns:
         str: The result of the command
     """
     try:
-        cmd = command_registry.commands.get(command_name)
+        cmd = agent.command_registry.commands.get(command_name)
 
         # If the command is found, call it with the provided arguments
         if cmd:
@@ -106,7 +105,7 @@ def execute_command(
         # TODO: Change these to take in a file rather than pasted code, if
         # non-file is given, return instructions "Input should be a python
         # filepath, write your code to file and try again
-        for command in agent.prompt.commands:
+        for command in agent.ai_config.prompt_generator.commands:
             if (
                 command_name == command["label"].lower()
                 or command_name == command["name"].lower()
diff --git a/tests/unit/test_execute_command.py b/tests/unit/test_execute_command.py
new file mode 100644
index 00000000000..fb3f043a902
--- /dev/null
+++ b/tests/unit/test_execute_command.py
@@ -0,0 +1,24 @@
+from autogpt.agent import Agent
+from autogpt.app import execute_command
+
+
+def check_plan():
+    return "hi"
+
+
+def test_execute_command_plugin(agent: Agent):
+    """Test that executing a command that came from a plugin works as expected"""
+    agent.ai_config.prompt_generator.add_command(
+        "check_plan",
+        "Read the plan.md with the next goals to achieve",
+        {},
+        check_plan,
+    )
+    command_name = "check_plan"
+    arguments = {}
+    command_result = execute_command(
+        command_name=command_name,
+        arguments=arguments,
+        agent=agent,
+    )
+    assert command_result == "hi"