Plan conversion prompt

src/agent/agent.py

@@ -1,4 +1,5 @@
 """Contains the class `Agent`, the core of the system."""
+import re
 import json
 from json import JSONDecodeError
 
@@ -94,26 +95,33 @@ def extract_plan(self, plan_nl):
         prompt = self.user_plan_con.format(query=plan_nl)
         messages = [
             {'role': 'system', 'content': self.system_plan_con},
-            {'role': 'system', 'content': prompt}
+            {'role': 'user', 'content': prompt}
         ]
         response = self.llm.query(messages=messages, stream=False)
+
         try:
             plan_data = json.loads(response['message']['content'])
-
-            tasks = []
-            for task in plan_data['plan']:
-                if len(task['command']) == 0:
-                    continue
-                tasks.append(Task(
-                    command=task['command'],
-                    thought=task['thought'],
-                    tool=Terminal
-                ))
-
-            return Plan(tasks)
         except JSONDecodeError:
-            print(response)  # use regex (or retry)
-            return None
+            # try extracting json from response
+            json_regex = re.compile(r'\[.*?\]', re.DOTALL)
+            json_match = json_regex.search(response['message']['content'])
+            if json_match:
+                plan_data = json.loads(json_match.group())
+            else:
+                print(f'PlanError: \n{response["message"]["content"]}')
+                return None
+
+        tasks = []
+        for task in plan_data:
+            if len(task['command']) == 0:
+                continue
+            tasks.append(Task(
+                command=task['command'],
+                thought=task['thought'],
+                tool=Terminal
+            ))
+
+        return Plan(tasks)
 
     def execute_plan(self, sid):
         """Extracts the plan from last message,

src/agent/prompts.py

@@ -97,24 +97,18 @@
         },
         'plan_conversion': {
             'system': textwrap.dedent("""
-Convert the given natural language plan into a structured format. 
-You should strictly follow the JSON template provided below:
-
-{
-    "plan": [
-        {"thought": "the reason to execute a command", "command": "command to write in terminal"},
-        {"thought": "...", "command": "..."}
-    ]
-}
-Remember: provide only the json string
+                You should extract the commands from the provided natural language plan and convert them into the provided JSON format.
+                [
+                    {"thought": "the reason to execute a command", "command": "command to write in terminal"}
+                ]
             """),
             'user': textwrap.dedent("""
-Convert this plan in the provided JSON format:
-{query}
-
-Remember: when a command starts with terminal consider only the subsequent part of the command
-Remember: when multiple commands are provided split them, when none is provided do not write it
-Remember: provide only the json string
+                Natural language plan:
+                {query}
+                
+                Command Extraction Guidelines:
+                1. Only provide the JSON string
+                2. 'terminal' is a keyword, you should not include it in any command
             """),
         },
         'routing': {

test/tests/plan_tests.json

@@ -36,27 +36,5 @@
         "commands": [
             "nmap -p 1-1024 -sV 192.168.1.252"
         ]
-    },
-    {
-        "content": [
-            "**Plan for Reconnaissance Phase:**\n\n",
-            "**Target System:** 154.142.1.1 (IP Camera).\n",
-            "**Objective:** Gather information about the operating system, services, data transmission, and potential vulnerabilities.\n",
-            "**Tasks:**\n\n",
-            "1. **Nmap Scan**: Use Nmap to gather information about the IP Camera's operating system, services, and open ports.\n",
-            "   - Execute the following command: `nmap -O -sV 154.142.1.1`\n",
-            "2. **Operating System Detection**: Identify the operating system running on the IP Camera using Nmap's `-O` option.\n",
-            "3. **Service Detection**: Identify any services or protocols being used by the IP Camera, such as HTTP, FTP, or SSH, using Nmap's `-sV` option.\n",
-            "4. **Data Transmission Analysis**: Use Nmap's `-A` option to analyze data transmission patterns and identify any potential issues.\n",
-            "5. **SSH and FTP Vulnerability Scan**: Run a vulnerability scan for SSH and FTP protocols using Nmap's `--script` option with the following commands:\n",
-            "   - `nmap --script sshvuln 154.142.1.1`\n",
-            "   - `nmap --script ftp-vuln 154.142.1.1`\n\n",
-            "Please let me know if you'd like to proceed with this plan or if you have any specific requirements or concerns.\n"
-        ],
-        "commands": [
-            "nmap -O -sV 154.142.1.1",
-            "nmap --script sshvuln 154.142.1.1",
-            "nmap --script ftp-vuln 154.142.1.1"
-        ]
     }
 ]