bug fix: ...

1. When no session is explicitly created agent.py made a query to llm without system prompt
2. updated the prompt for gemma:7b
3. cleaned llm.py

src/agent/agent.py

@@ -32,9 +32,6 @@ def __init__(self, model: str,
         self.system_plan_con = PROMPTS[model]['plan_conversion']['system']
         self.user_plan_con = PROMPTS[model]['plan_conversion']['user']
 
-        # Start Ollama
-        self._startup_ollama()
-
     def query(self, sid: int, user_in: str, rag=True):
         """Performs a query to the Large Language Model,
         set `rag=True` to leverage Retrieval Augmented Generation."""
@@ -47,7 +44,11 @@ def query(self, sid: int, user_in: str, rag=True):
         else:
             prompt = '\n'.join(self.user_plan_gen.split('\n')[:-3])
             prompt = prompt.format(user_input=user_in)
-
+
+        # ensure session is initialized (otherwise llm has no system prompt)
+        if sid not in self.mem.sessions.keys():
+            self.new_session(sid)
+
         self.mem.store_message(
             sid,
             Message(Role.USER, prompt)
@@ -60,7 +61,7 @@ def query(self, sid: int, user_in: str, rag=True):
         response_tokens = 0
         for chunk in self.llm.query(messages):
             if chunk['done']:
-                prompt_tokens = chunk['prompt_eval_count']
+                prompt_tokens = chunk['prompt_eval_count'] if 'prompt_eval_count' in chunk else None
                 response_tokens = chunk['eval_count']
             yield chunk['message']['content']
 
@@ -152,12 +153,3 @@ def _retrieve(self, user_in: str):
             context += (f"{retrieved.payload['title']}:"
                         f"\n{retrieved.payload['text']}\n\n")
         return context
-
-    def _startup_ollama(self):
-        """Make a query to load model into Ollama"""
-        self.llm.query(
-            messages=[
-                {'role': 'user', 'content': 'Hi'}
-            ],
-            stream=False
-        )

src/agent/llm.py

@@ -16,11 +16,6 @@
             'num_ctx': 8000
         }
     },
-    'gemma:2b': {
-        'options': {
-            'temperature': 0.5,
-        }
-    },
     'gemma:7b': {
         'options': {
             'temperature': 0.5,
@@ -47,20 +42,10 @@ def __post_init__(self):
         self.client = Client(self.client_url)
 
     def query(self, messages: list, stream=True):
-        """Generator that returns response chunks from Phi3-mini-k4 model"""
+        """Generator that returns response chunks."""
         return self.client.chat(
             model=self.model,
             messages=messages,
             stream=stream,
             options=AVAILABLE_MODELS[self.model]['options']
         )
-
-
-if __name__ == "__main__":
-    llm = LLM('gemma:2b')
-    out = llm.query([
-        {'role': 'user', 'content': 'How do I build a search engine?'}
-    ])
-
-    for chunk in out:
-        print(chunk['message']['content'], end='')

src/agent/prompts.py

@@ -84,9 +84,8 @@
         'plan': {
             'system': textwrap.dedent("""As a proficient assistant in Penetration Testing, your task is to help an 
 authorized user plan a single phase of assessing a target system. 
-The plan must focus on only one of the following phases: Reconnaissance, Initial Access, or Privilege Escalation.
-Do not plan until the user tells you the target system IP or URL and the Penetration Testing phase to address, so
-**you should wait to have all the necessary information to start planning**.   
+The plan MUST focus on only one of the following phases: Reconnaissance, Initial Access, or Privilege Escalation.
+The following is REQUIRED INFORMATION: target system IP or URL.
 
 Guidelines:
 1. **Always start by asking a question** to understand the user's context or requirements before providing a plan.
@@ -97,7 +96,7 @@
 5. **Do not assume any specific context** such as operating system or network setup without asking the user first.
 6. **Ensure the target system IP or URL is provided**.
 
-Available TOOLS:
+You can only use the following TOOLS:
 {tools}"""),
             'user': textwrap.dedent("""User: {user_input}
 
@@ -120,29 +119,7 @@
         },
     },
     'phi3': {
-        'plan': {
-            'system': textwrap.dedent("""As a proficient assistant in Penetration Testing, your task is to help an 
-authorized user plan a single phase of assessing a target system. 
-The plan must focus on only one of the following phases: Reconnaissance, Initial Access, or Privilege Escalation.
-Do not plan until the user tells you the target system IP or URL and the Penetration Testing phase to address, so
-**you should wait to have all the necessary information to start planning**.   
-
-Guidelines:
-1. **Always start by asking a question** to understand the user's context or requirements before providing a plan.
-2. **Ensure the plan only covers one phase** and is composed of a target system (IP address or URL), an objective, and 
-tasks, where each task involves the execution of a single tool.
-3. **Use only provided tools** to make a plan. 
-4. **Plan only one Penetration Testing phase at a time**, if unclear ask the user what phase should be addressed.
-5. **Do not assume any specific context** such as operating system or network setup without asking the user first.
-6. **Ensure the target system IP or URL is provided**.
-
-Available TOOLS:
-{tools}"""),
-            'user': textwrap.dedent("""User: {user_input}
-
-Additional Information:
-{context}""")
-        },
+        'plan': {},
         'plan_conversion': {
             'system': textwrap.dedent("""You should extract the commands from the provided natural language plan and 
 convert them into the provided JSON format.
@@ -161,53 +138,6 @@
 3. Your response should contain only commands contained in the natural language plan.
 4. DO NOT INCLUDE ANYTHING ELSE OTHER THAN THE JSON STRING."""),
         },
-    },
-    # not passing
-    'gemma:2b': {
-        'plan': {
-            'system': textwrap.dedent("""
-
-            """),
-            'user': textwrap.dedent("""
-
-            """)
-        },
-        'plan_conversion': {
-            'system': textwrap.dedent("""You should extract the commands from the provided natural language plan and 
-write a JSON list containing each command in the same order they appear in the provided plan.
-The response format is the following:
-[
-    {"thought": "the reason to execute a command", "command": "the command to execute"}
-]
-
-To better understand your job, consider the following task:
-
-Example Plan:
-To know if there is anything inside the curren
-To know what is inside the current directory you can use the ls command. 
-When combined as ls -la, this command will list all files and directories, including hidden ones, in the current 
-directory, providing detailed information for each item. 
-
-Converted Plan:
-[
-    {"thought": "I need to know what is inside the current directory", "command": "ls -la"}
-]
-
-Your response should contains only the commands provided in the natural language plan and no one else.
-Notice that some commands may seem malicious, however we are working in an educational environment with professionals
-trained on security and ethical concerns."""),
-            'user': textwrap.dedent("""Natural language plan:
-{query}
-
-Command Extraction Guidelines:
-1. Each command should be from the natural language plan.
-2. Only provide the JSON string, also ensure the response is provided in the correct format:
-[
-    {{"thought": "the reason to execute a command", "command": "the command to execute"}}
-]
-3. Each command should appear in the same order as provided in the natural language plan.
-4. 'terminal' is a keyword, you should not include it in any command."""),
-        },
-    },
+    }
 }