Cleaned up code

Author: DianaStrauss

src/hackingBuddyGPT/usecases/web_api_testing/simple_openapi_documentation.py

@@ -97,8 +97,9 @@ def _handle_response(self, completion, response):
             self.console.print(Panel(result[:30], title="tool"))
             result_str = self.response_handler.parse_http_status_line(result)
             self._prompt_history.append(tool_message(result_str, tool_call_id))
-            invalid_flags = ["recorded","Not a valid HTTP method" ]
-            if not result_str in invalid_flags :
+            invalid_flags = ["recorded","Not a valid HTTP method", "404" ,"Client Error: Not Found"]
+            print(f'result_str:{result_str}')
+            if not result_str in invalid_flags  or any(item in result_str for item in invalid_flags):
                 self.documentation_handler.update_openapi_spec(response, result)
                 self.documentation_handler.write_openapi_to_yaml()
                 self.prompt_engineer.schemas = self.documentation_handler.schemas

src/hackingBuddyGPT/usecases/web_api_testing/utils/documentation_handler.py

@@ -1,15 +1,32 @@
 import os
 import yaml
 from datetime import datetime
-
 from hackingBuddyGPT.capabilities.yamlFile import YAMLFile
 
-
 class DocumentationHandler:
-    """Handles the generation and updating of an OpenAPI specification document based on dynamic API responses."""
+    """
+    Handles the generation and updating of an OpenAPI specification document based on dynamic API responses.
+
+    Attributes:
+        response_handler (object): An instance of the response handler for processing API responses.
+        schemas (dict): A dictionary to store API schemas.
+        filename (str): The filename for the OpenAPI specification file.
+        openapi_spec (dict): The OpenAPI specification document structure.
+        llm_handler (object): An instance of the LLM handler for interacting with the LLM.
+        api_key (str): The API key for accessing the LLM.
+        file_path (str): The path to the directory where the OpenAPI specification file will be stored.
+        file (str): The complete path to the OpenAPI specification file.
+        _capabilities (dict): A dictionary to store capabilities related to YAML file handling.
+    """
 
     def __init__(self, llm_handler, response_handler):
-        """Initializes the handler with a template OpenAPI specification."""
+        """
+        Initializes the handler with a template OpenAPI specification.
+
+        Args:
+            llm_handler (object): An instance of the LLM handler for interacting with the LLM.
+            response_handler (object): An instance of the response handler for processing API responses.
+        """
         self.response_handler = response_handler
         self.schemas = {}
         self.filename = f"openapi_spec_{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.yaml"
@@ -29,7 +46,6 @@ def __init__(self, llm_handler, response_handler):
         current_path = os.path.dirname(os.path.abspath(__file__))
         self.file_path = os.path.join(current_path, "openapi_spec")
         self.file = os.path.join(self.file_path, self.filename)
-        yamls = []
         self._capabilities = {
             "yaml": YAMLFile()
         }
@@ -39,47 +55,46 @@ def update_openapi_spec(self, resp, result):
         Updates the OpenAPI specification based on the API response provided.
 
         Args:
-        - response: The response object containing details like the path and method which should be documented.
+            resp (object): The response object containing details like the path and method which should be documented.
+            result (str): The result of the API call.
         """
         request = resp.action
 
-
-        if request.__class__.__name__ == 'RecordNote':  # TODO check why isinstance does not work
+        if request.__class__.__name__ == 'RecordNote':  # TODO: check why isinstance does not work
             self.check_openapi_spec(resp)
         if request.__class__.__name__ == 'HTTPRequest':
             path = request.path
             method = request.method
-            print(f'method:{method}')
+            print(f'method: {method}')
             # Ensure that path and method are not None and method has no numeric characters
             if path and method:
                 # Initialize the path if not already present
                 if path not in self.openapi_spec['endpoints']:
                     self.openapi_spec['endpoints'][path] = {}
                 # Update the method description within the path
                 example, reference, self.openapi_spec = self.response_handler.parse_http_response_to_openapi_example(self.openapi_spec, result, path, method)
-                if example != None or reference != None:
+                if example is not None or reference is not None:
                     self.openapi_spec['endpoints'][path][method.lower()] = {
-                    "summary": f"{method} operation on {path}",
-                    "responses": {
-                        "200": {
-                            "description": "Successful response",
-                            "content": {
-                                "application/json": {
-                                    "schema": {
-                                        "$ref": reference
-                                    },
-                                    "examples": example
+                        "summary": f"{method} operation on {path}",
+                        "responses": {
+                            "200": {
+                                "description": "Successful response",
+                                "content": {
+                                    "application/json": {
+                                        "schema": {
+                                            "$ref": reference
+                                        },
+                                        "examples": example
+                                    }
                                 }
                             }
                         }
-
                     }
-                }
-
-
 
     def write_openapi_to_yaml(self):
-        """Writes the updated OpenAPI specification to a YAML file with a timestamped filename."""
+        """
+        Writes the updated OpenAPI specification to a YAML file with a timestamped filename.
+        """
         try:
             # Prepare data to be written to YAML
             openapi_data = {
@@ -100,20 +115,14 @@ def write_openapi_to_yaml(self):
         except Exception as e:
             raise Exception(f"Error writing YAML file: {e}")
 
-
-
-
     def check_openapi_spec(self, note):
         """
-            Uses OpenAI's GPT model to generate a complete OpenAPI specification based on a natural language description.
+        Uses OpenAI's GPT model to generate a complete OpenAPI specification based on a natural language description.
 
-            Parameters:
-            - api_key: str. Your OpenAI API key.
-            - file_path: str. Path to the YAML file to update.
-            - description: str. A detailed description of the entire API.
-            """
+        Args:
+            note (object): The note object containing the description of the API.
+        """
         description = self.response_handler.extract_description(note)
         from hackingBuddyGPT.usecases.web_api_testing.utils.yaml_assistant import YamlFileAssistant
         yaml_file_assistant = YamlFileAssistant(self.file_path, self.llm_handler)
         yaml_file_assistant.run(description)
-

src/hackingBuddyGPT/usecases/web_api_testing/utils/llm_handler.py

@@ -1,22 +1,63 @@
 from hackingBuddyGPT.capabilities.capability import capabilities_to_action_model
 
-
 class LLMHandler(object):
+    """
+    LLMHandler is a class responsible for managing interactions with a large language model (LLM).
+    It handles the execution of prompts and the management of created objects based on the capabilities.
+
+    Attributes:
+        llm (object): The large language model to interact with.
+        _capabilities (dict): A dictionary of capabilities that define the actions the LLM can perform.
+        created_objects (dict): A dictionary to keep track of created objects by their type.
+    """
+
     def __init__(self, llm, capabilities):
+        """
+        Initializes the LLMHandler with the specified LLM and capabilities.
+
+        Args:
+            llm (object): The large language model to interact with.
+            capabilities (dict): A dictionary of capabilities that define the actions the LLM can perform.
+        """
         self.llm = llm
         self._capabilities = capabilities
         self.created_objects = {}
 
     def call_llm(self, prompt):
-        return self.llm.instructor.chat.completions.create_with_completion(model=self.llm.model, messages=prompt, response_model=capabilities_to_action_model(self._capabilities))
+        """
+        Calls the LLM with the specified prompt and retrieves the response.
+
+        Args:
+            prompt (list): The prompt messages to send to the LLM.
+
+        Returns:
+            response (object): The response from the LLM.
+        """
+        return self.llm.instructor.chat.completions.create_with_completion(
+            model=self.llm.model,
+            messages=prompt,
+            response_model=capabilities_to_action_model(self._capabilities)
+        )
 
     def add_created_object(self, created_object, object_type):
+        """
+        Adds a created object to the dictionary of created objects, categorized by object type.
+
+        Args:
+            created_object (object): The object that was created.
+            object_type (str): The type/category of the created object.
+        """
         if object_type not in self.created_objects:
             self.created_objects[object_type] = []
         if len(self.created_objects[object_type]) < 7:
             self.created_objects[object_type].append(created_object)
 
     def get_created_objects(self):
+        """
+        Retrieves the dictionary of created objects and prints its contents.
+
+        Returns:
+            dict: The dictionary of created objects.
+        """
         print(f'created_objects: {self.created_objects}')
         return self.created_objects
-

src/hackingBuddyGPT/usecases/web_api_testing/utils/openapi_converter.py

@@ -1,21 +1,36 @@
 import os.path
-
 import yaml
 import json
 
-
 class OpenAPISpecificationConverter:
+    """
+    OpenAPISpecificationConverter is a class for converting OpenAPI specification files between YAML and JSON formats.
+
+    Attributes:
+        base_directory (str): The base directory for the output files.
+    """
+
     def __init__(self, base_directory):
+        """
+        Initializes the OpenAPISpecificationConverter with the specified base directory.
+
+        Args:
+            base_directory (str): The base directory for the output files.
+        """
         self.base_directory = base_directory
 
     def convert_file(self, input_filepath, output_directory, input_type, output_type):
         """
-        Generic method to convert files between YAML and JSON.
+        Converts files between YAML and JSON formats.
+
+        Args:
+            input_filepath (str): The path to the input file.
+            output_directory (str): The subdirectory for the output files.
+            input_type (str): The type of the input file ('yaml' or 'json').
+            output_type (str): The type of the output file ('json' or 'yaml').
 
-        :param input_filepath: Path to the input file.
-        :param output_directory: Subdirectory for the output files.
-        :param input_type: Type of the input file ('yaml' or 'json').
-        :param output_type: Type of the output file ('json' or 'yaml').
+        Returns:
+            str: The path to the converted output file, or None if an error occurred.
         """
         try:
             filename = os.path.basename(input_filepath)
@@ -44,11 +59,27 @@ def convert_file(self, input_filepath, output_directory, input_type, output_type
             return None
 
     def yaml_to_json(self, yaml_filepath):
-        """Convert a YAML file to a JSON file."""
+        """
+        Converts a YAML file to a JSON file.
+
+        Args:
+            yaml_filepath (str): The path to the YAML file to be converted.
+
+        Returns:
+            str: The path to the converted JSON file, or None if an error occurred.
+        """
         return self.convert_file(yaml_filepath, "json", 'yaml', 'json')
 
     def json_to_yaml(self, json_filepath):
-        """Convert a JSON file to a YAML file."""
+        """
+        Converts a JSON file to a YAML file.
+
+        Args:
+            json_filepath (str): The path to the JSON file to be converted.
+
+        Returns:
+            str: The path to the converted YAML file, or None if an error occurred.
+        """
         return self.convert_file(json_filepath, "yaml", 'json', 'yaml')
 
 
@@ -61,4 +92,5 @@ def json_to_yaml(self, json_filepath):
     json_file = converter.yaml_to_json(yaml_input)
 
     # Convert JSON to YAML
-    yaml_file = converter.json_to_yaml(json_file)
+    if json_file:
+        converter.json_to_yaml(json_file)

src/hackingBuddyGPT/usecases/web_api_testing/utils/openapi_parser.py

@@ -1,33 +1,73 @@
 import yaml
 
 class OpenAPISpecificationParser:
+    """
+    OpenAPISpecificationParser is a class for parsing and extracting information from an OpenAPI specification file.
+
+    Attributes:
+        filepath (str): The path to the OpenAPI specification YAML file.
+        api_data (dict): The parsed data from the YAML file.
+    """
+
     def __init__(self, filepath):
+        """
+        Initializes the OpenAPISpecificationParser with the specified file path.
+
+        Args:
+            filepath (str): The path to the OpenAPI specification YAML file.
+        """
         self.filepath = filepath
         self.api_data = self.load_yaml()
 
     def load_yaml(self):
-        """Load YAML data from the specified file."""
+        """
+        Loads YAML data from the specified file.
+
+        Returns:
+            dict: The parsed data from the YAML file.
+        """
         with open(self.filepath, 'r') as file:
             return yaml.safe_load(file)
 
     def get_servers(self):
-        """Retrieve the list of server URLs."""
+        """
+        Retrieves the list of server URLs from the OpenAPI specification.
+
+        Returns:
+            list: A list of server URLs.
+        """
         return [server['url'] for server in self.api_data.get('servers', [])]
 
     def get_paths(self):
-        """Retrieve all API paths and their methods."""
+        """
+        Retrieves all API paths and their methods from the OpenAPI specification.
+
+        Returns:
+            dict: A dictionary with API paths as keys and methods as values.
+        """
         paths_info = {}
         paths = self.api_data.get('paths', {})
         for path, methods in paths.items():
             paths_info[path] = {method: details for method, details in methods.items()}
         return paths_info
 
     def get_operations(self, path):
-        """Retrieve operations for a specific path."""
+        """
+        Retrieves operations for a specific path from the OpenAPI specification.
+
+        Args:
+            path (str): The API path to retrieve operations for.
+
+        Returns:
+            dict: A dictionary with methods as keys and operation details as values.
+        """
         return self.api_data['paths'].get(path, {})
 
     def print_api_details(self):
-        """Prints details of the API extracted from the OpenAPI document."""
+        """
+        Prints details of the API extracted from the OpenAPI document, including title, version, servers,
+        paths, and operations.
+        """
         print("API Title:", self.api_data['info']['title'])
         print("API Version:", self.api_data['info']['version'])
         print("Servers:", self.get_servers())
@@ -42,5 +82,6 @@ def print_api_details(self):
 # Usage example
 if __name__ == '__main__':
     openapi_parser = OpenAPISpecificationParser(
-        '/hackingBuddyGPT/usecases/web_api_testing/openapi_spec/openapi_spec_2024-06-13_17-16-25.yaml')
+        '/hackingBuddyGPT/usecases/web_api_testing/openapi_spec/openapi_spec_2024-06-13_17-16-25.yaml'
+    )
     openapi_parser.print_api_details()