Added more documentation and data types

Author: Diana Strauss

src/hackingBuddyGPT/usecases/web_api_testing/prompt_generation/prompts/in_context_learning_prompt.py

@@ -1,5 +1,10 @@
-from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.information.prompt_information import PromptStrategy
+from typing import List, Dict, Optional
+
+from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.information.prompt_information import PromptStrategy, \
+    PromptContext, PromptPurpose
 from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.prompts.basic_prompt import BasicPrompt
+from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.utils import PromptGenerationHelper
+
 
 class InContextLearningPrompt(BasicPrompt):
     """
@@ -12,33 +17,34 @@ class InContextLearningPrompt(BasicPrompt):
     Attributes:
         context (PromptContext): The context in which prompts are generated.
         prompt_helper (PromptHelper): A helper object for managing and generating prompts.
-        prompt (dict): A dictionary containing the prompts for each round.
+        prompt (Dict[int, Dict[str, str]]): A dictionary containing the prompts for each round.
+        round (int): The round number for which the prompt is being generated.
+        purpose (Optional[PromptPurpose]): The purpose of the prompt generation, which can be set during the process.
     """
 
-    def __init__(self, context, prompt_helper, prompt, round):
+    def __init__(self, context: PromptContext, prompt_helper: PromptGenerationHelper, prompt: Dict[int, Dict[str, str]], round: int) -> None:
         """
         Initializes the InContextLearningPrompt with a specific context, prompt helper, and initial prompt.
 
         Args:
             context (PromptContext): The context in which prompts are generated.
             prompt_helper (PromptHelper): A helper object for managing and generating prompts.
-            prompt (dict): A dictionary containing the prompts for each round.
-            round (int): Number of round.
+            prompt (Dict[int, Dict[str, str]]): A dictionary containing the prompts for each round.
+            round (int): The round number for which the prompt is being generated.
         """
         super().__init__(context, prompt_helper, PromptStrategy.IN_CONTEXT)
-        self.round = round
-        self.prompt = prompt
-        self.purpose = None
-
+        self.round: int = round
+        self.prompt: Dict[int, Dict[str, str]] = prompt
+        self.purpose: Optional[PromptPurpose] = None
 
-    def generate_prompt(self, move_type, hint, previous_prompt): # still work in progress
+    def generate_prompt(self, move_type: str, hint: Optional[str], previous_prompt: List[Dict[str, str]]) -> str:
         """
         Generates a prompt using the in-context learning strategy.
 
         Args:
             move_type (str): The type of move to generate.
-            hint (str): An optional hint to guide the prompt generation.
-            previous_prompt (list): A list of previous prompt entries, each containing a "content" key.
+            hint (Optional[str]): An optional hint to guide the prompt generation.
+            previous_prompt (List[Dict[str, str]]): A list of previous prompt entries, each containing a "content" key.
 
         Returns:
             str: The generated prompt.

src/hackingBuddyGPT/usecases/web_api_testing/prompt_generation/prompts/tree_of_thought_prompt.py

@@ -1,5 +1,8 @@
-from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.information.prompt_information import PromptStrategy, PromptContext, \
-    PromptPurpose
+from typing import List, Optional
+
+from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.information.prompt_information import (
+    PromptStrategy, PromptContext, PromptPurpose
+)
 from hackingBuddyGPT.usecases.web_api_testing.prompt_generation.prompts.basic_prompt import BasicPrompt
 
 class TreeOfThoughtPrompt(BasicPrompt):
@@ -13,52 +16,55 @@ class TreeOfThoughtPrompt(BasicPrompt):
     Attributes:
         context (PromptContext): The context in which prompts are generated.
         prompt_helper (PromptHelper): A helper object for managing and generating prompts.
+        rest_api (str): The REST API endpoint for which prompts are generated.
+        round (int): The round number for the prompt generation process.
+        purpose (Optional[PromptPurpose]): The purpose of the prompt generation, which can be set during the process.
     """
 
-    def __init__(self, context, prompt_helper, rest_api, round):
+    def __init__(self, context: PromptContext, prompt_helper: 'PromptHelper', rest_api: str, round: int) -> None:
         """
         Initializes the TreeOfThoughtPrompt with a specific context and prompt helper.
 
         Args:
             context (PromptContext): The context in which prompts are generated.
             prompt_helper (PromptHelper): A helper object for managing and generating prompts.
-            rest_api (str): The REST API endpoint
-            round (int): The round number
+            rest_api (str): The REST API endpoint.
+            round (int): The round number for the prompt generation process.
         """
         super().__init__(context, prompt_helper, PromptStrategy.TREE_OF_THOUGHT)
-        self.rest_api = rest_api
-        self.round = round
-        self.purpose = None
+        self.rest_api: str = rest_api
+        self.round: int = round
+        self.purpose: Optional[PromptPurpose] = None
 
-    def generate_prompt(self, move_type, hint, previous_prompt): # still work in progress
+    def generate_prompt(self, move_type: str, hint: str, previous_prompt: List[dict]) -> str:
         """
         Generates a prompt using the tree-of-thought strategy.
 
         Args:
             move_type (str): The type of move to generate.
             hint (str): An optional hint to guide the prompt generation.
-            previous_prompt (list): A list of previous prompt entries, each containing a "content" key.
+            previous_prompt (List[dict]): A list of previous prompt entries, each containing a "content" key.
 
         Returns:
             str: The generated prompt.
         """
         if self.context == PromptContext.DOCUMENTATION:
             tree_of_thoughts_steps = [(
-            "Imagine three different OpenAPI specification specialist.\n"
-            "All experts will write down one step of their thinking,\n"
-            "then share it with the group.\n"
-            "After that, all remaining specialists will proceed to the next step, and so on.\n"
-            "If any specialist realizes they're wrong at any point, they will leave.\n"
-            f"The question is: Create an OpenAPI specification for this REST API {self.rest_api} "
+                "Imagine three different OpenAPI specification specialists.\n"
+                "All experts will write down one step of their thinking,\n"
+                "then share it with the group.\n"
+                "After that, all remaining specialists will proceed to the next step, and so on.\n"
+                "If any specialist realizes they're wrong at any point, they will leave.\n"
+                f"The question is: Create an OpenAPI specification for this REST API {self.rest_api} "
             )]
         else:
             tree_of_thoughts_steps = [(
-            "Imagine three different Pentest experts are answering this question.\n"
-            "All experts will write down one step of their thinking,\n"
-            "then share it with the group.\n"
-            "After that, all experts will proceed to the next step, and so on.\n"
-            "If any expert realizes they're wrong at any point, they will leave.\n"
-            f"The question is: Create pentests for this REST API {self.rest_api} "
+                "Imagine three different Pentest experts are answering this question.\n"
+                "All experts will write down one step of their thinking,\n"
+                "then share it with the group.\n"
+                "After that, all experts will proceed to the next step, and so on.\n"
+                "If any expert realizes they're wrong at any point, they will leave.\n"
+                f"The question is: Create pentests for this REST API {self.rest_api} "
             )]
         self.purpose = PromptPurpose.AUTHENTICATION_AUTHORIZATION