Skip to content

Commit

Permalink
Fixed prompt in query, typo with '[-3]' instead of [:-3]
Browse files Browse the repository at this point in the history
  • Loading branch information
antoninoLorenzo committed Jul 1, 2024
1 parent 3ef75b7 commit 257d3d0
Show file tree
Hide file tree
Showing 3 changed files with 7 additions and 44 deletions.
4 changes: 3 additions & 1 deletion ai-ops-cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -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':
Expand All @@ -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):
Expand Down
6 changes: 4 additions & 2 deletions src/agent/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
41 changes: 0 additions & 41 deletions src/agent/prompts.py
Original file line number Diff line number Diff line change
Expand Up @@ -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?"
"""

0 comments on commit 257d3d0

Please sign in to comment.