From 257d3d00e40b706d3a979d88d56d4768d5632e9b Mon Sep 17 00:00:00 2001
From: antoninoLorenzo <collab.beatbuster@gmail.com>
Date: Mon, 1 Jul 2024 17:45:56 +0200
Subject: [PATCH] Fixed prompt in query, typo with '[-3]' instead of [:-3]

---
 ai-ops-cli.py        |  4 +++-
 src/agent/agent.py   |  6 ++++--
 src/agent/prompts.py | 41 -----------------------------------------
 3 files changed, 7 insertions(+), 44 deletions(-)

diff --git a/ai-ops-cli.py b/ai-ops-cli.py
index 418421b..849b70c 100644
--- a/ai-ops-cli.py
+++ b/ai-ops-cli.py
@@ -127,6 +127,7 @@ def chat(self, print_name=True):
         if print_name:
             print(f'+ {self.current_session["name"]} ({self.current_session["sid"]})')
         query_url = f'{self.api_url}/session/{self.current_session["sid"]}/query'
+        
         while True:
             q = input('user: ')
             if q == '-1':
@@ -137,10 +138,11 @@ def chat(self, print_name=True):
                     json={'query': q},
                     headers=None,
                     stream=True) as resp:
+                resp.raise_for_status()
                 print('assistant: ')
                 for chunk in resp.iter_content():
                     if chunk:
-                        print(chunk.decode(), end='')
+                        print(chunk.decode(), end='', flush=True)
                 print()
 
     def execute_plan(self):
diff --git a/src/agent/agent.py b/src/agent/agent.py
index 5e2f951..7fd36da 100644
--- a/src/agent/agent.py
+++ b/src/agent/agent.py
@@ -42,8 +42,10 @@ def query(self, sid: int, user_in: str, rag=True, stream=True):
                 context=context
             )
         else:
-            prompt = '\n'.join(self.user_plan_gen.split('\n')[-3])
-            prompt.format(user_input=user_in)
+            print('Cutted: \n', self.user_plan_gen.split('\n')[:-3])
+            prompt = '\n'.join(self.user_plan_gen.split('\n')[:-3])
+            print(f'Prompt ({len(prompt)}):\n{prompt}')
+            prompt = prompt.format(user_input=user_in)
 
         self.mem.store_message(
             sid,
diff --git a/src/agent/prompts.py b/src/agent/prompts.py
index 1f3d568..c032820 100644
--- a/src/agent/prompts.py
+++ b/src/agent/prompts.py
@@ -96,44 +96,3 @@
     }
 }
 
-asd = """
-                To better understand what your job is, consider the following interaction between user and assistant:
-                **Initial Interaction**
-                Assistant: "Please specify the Penetration Testing phase you would like to plan."
-                
-                User: "Reconnaissance."
-                
-                Assistant: "Thank you. For the Reconnaissance phase, could you please provide more details about the 
-                target environment? For example, is it a web application, a network, or an individual host?"
-                
-                User: "It's a web application."
-                
-                **Planning Phase**
-                Assistant: "Great. Here is a preliminary plan for the Reconnaissance phase:
-                1. Overview:
-                    - Identify the target web application's domain and IP address.
-                    - Gather information about the web server, technologies used, and possible entry points.
-                
-                2. Commands to Execute:
-                    - whois example.com - To gather domain registration details.
-                    - nslookup example.com - To find the IP address of the domain.
-                    - whatweb example.com - To identify the technologies used by the web application.
-                
-                Does this meet your requirements, or is there anything specific you would like to add or modify?"
-                
-                **Adjusting Based on User Feedback**
-                User: "Could you also include a step for finding subdomains?"
-                
-                Assistant: "Certainly. Here is the updated plan for the Reconnaissance phase:
-                1. Overview:
-                    - Identify the target web application's domain and IP address.
-                    - Gather information about the web server, technologies used, and possible entry points.
-                    - Discover subdomains.
-                2. Commands to Execute:
-                    - whois example.com - To gather domain registration details.
-                    - nslookup example.com - To find the IP address of the domain.
-                    - whatweb example.com - To identify the technologies used by the web application.
-                    - sublist3r -d example.com - To discover subdomains.
-                
-                Does this meet your requirements?"
-"""