Simple openapi documentation procudes an openapi yaml specification

Author: DianaStrauss

usecases/web_api_testing/prompt_engineer.py

@@ -100,7 +100,7 @@ def in_context_learning(self):
         """
         return str("\n".join(self._prompt_history[self.round]["content"] + [self.prompt]))
 
-    def chain_of_thought(self):
+    def chain_of_thought(self,zero_shot = False):
         """
         Generates a prompt using the chain-of-thought strategy. https://www.promptingguide.ai/techniques/cot
 
@@ -109,10 +109,27 @@ def chain_of_thought(self):
         Returns:
             str: The generated prompt.
         """
+
         previous_prompt = self._prompt_history[self.round]["content"]
-        chain_of_thought_steps = [
-            "Let's think step by step." # zero shot prompt
-        ]
+        if zero_shot:
+            chain_of_thought_steps = [
+                "Let's think step by step."  # zero shot prompt
+            ]
+        else:
+            chain_of_thought_steps = [
+            "Explore the API by reviewing any available documentation to learn about the API endpoints, data models, and behaviors.",
+            "Identify all available endpoints (e.g., `/posts`, `/comments`, `/albums`, etc.).",
+            "Use a tool like Postman or curl to interact with the API by making GET, POST, PUT, DELETE requests to understand the responses.",
+            "Note down the response structures, status codes, and headers for each endpoint.",
+            "For each endpoint, document the following details: URL, HTTP method (GET, POST, etc.), query parameters and path variables, expected request body structure for POST and PUT requests, response structure for successful and error responses.",
+            "Identify common data structures returned by various endpoints and define them as reusable schemas. Determine the type of each field (e.g., integer, string, array) and define common response structures as components that can be referenced in multiple endpoint definitions.",
+            "Create an OpenAPI document including metadata such as API title, version, and description, define the base URL of the API, list all endpoints, methods, parameters, and responses, and define reusable schemas, response types, and parameters.",
+            "Ensure the correctness and completeness of the OpenAPI specification by validating the syntax and completeness of the document using tools like Swagger Editor, and ensure the specification matches the actual behavior of the API.",
+            "Refine the document based on feedback and additional testing, share the draft with others, gather feedback, and make necessary adjustments. Regularly update the specification as the API evolves.",
+            "Make the OpenAPI specification available to developers by incorporating it into your API documentation site and keep the documentation up to date with API changes."
+            ]
+
+
         #if previous_prompt == "Not a valid flag":
         #    return previous_prompt
         return "\n".join([previous_prompt] + chain_of_thought_steps)

usecases/web_api_testing/simple_openapi_documentation.py

@@ -1,3 +1,5 @@
+import datetime
+import os
 import time
 from dataclasses import dataclass, field
 from typing import List, Any, Union, Dict
@@ -16,6 +18,7 @@
 from usecases import use_case
 from usecases.usecase.roundbased import RoundBasedUseCase
 import pydantic_core
+import yaml
 
 Prompt = List[Union[ChatCompletionMessage, ChatCompletionMessageParam]]
 Context = Any
@@ -30,6 +33,7 @@ class SimpleWebAPIDocumentation(RoundBasedUseCase):
     _context: Context = field(default_factory=lambda: {"notes": list()})
     _capabilities: Dict[str, Capability] = field(default_factory=dict)
     _all_http_methods_found: bool = False
+
     # Parameter specifying the pattern description for expected HTTP methods in the API response
     http_method_description: str = parameter(
         desc="Pattern description for expected HTTP methods in the API response",
@@ -47,22 +51,33 @@ class SimpleWebAPIDocumentation(RoundBasedUseCase):
         desc="Comma-separated list of HTTP methods expected to be used in the API response.",
         default="GET,POST,PUT,PATCH,DELETE"
     )
+
     def init(self):
         super().init()
+        self.openapi_spec = self.openapi_spec = {
+        "openapi": "3.0.0",
+        "info": {
+            "title": "Generated API Documentation",
+            "version": "1.0",
+            "description": "Automatically generated description of the API."
+        },
+        "servers": [{"url": "https://jsonplaceholder.typicode.com"}],
+        "paths": {}
+    }
         self._prompt_history.append(
             {
                 "role": "system",
                 "content": f"You're tasked with documenting the REST APIs of a website hosted at {self.host}. "
                            f"Your main goal is to comprehensively explore the APIs endpoints and responses, and then document your findings in form of a OpenAPI specification."
-                           f" This thorough documentation will facilitate analysis later on.\n"
+                           f"Start with an empty OpenAPI specification. This thorough documentation will facilitate analysis later on.\n"
                            f"Maintain meticulousness in documenting your observations as you traverse the APIs. This will streamline the documentation process.\n"
                            f"Avoid resorting to brute-force methods. All essential information should be accessible through the API endpoints.\n"
 
             })
         self.prompt_engineer = PromptEngineer(
-                                              strategy=PromptStrategy.CHAIN_OF_THOUGHT,
-                                              api_key=self.llm.api_key,
-                                              history=self._prompt_history)
+            strategy=PromptStrategy.CHAIN_OF_THOUGHT,
+            api_key=self.llm.api_key,
+            history=self._prompt_history)
 
         self._context["host"] = self.host
         sett = set(self.http_method_template.format(method=method) for method in self.http_methods.split(","))
@@ -73,13 +88,13 @@ def init(self):
             "http_request": HTTPRequest(self.host),
             "record_note": RecordNote(self._context["notes"]),
         }
+        self.current_time = datetime.datetime.now()
 
     def all_http_methods_found(self):
         self.console.print(Panel("All HTTP methods found! Congratulations!", title="system"))
         self._all_http_methods_found = True
 
-    def perform_round(self, turn: int):
-
+    def perform_round(self, turn: int, FINAL_ROUND=20):
 
         with self.console.status("[bold green]Asking LLM for a new command..."):
             # generate prompt
@@ -101,20 +116,24 @@ def perform_round(self, turn: int):
             self._prompt_history.append(message)
             content = completion.choices[0].message.content
 
-            #print(f'message:{message}')
+            # print(f'message:{message}')
             answer = LLMResult(content, str(prompt),
                                content, toc - tic, completion.usage.prompt_tokens,
                                completion.usage.completion_tokens)
-            #print(f'answer: {answer}')
+            # print(f'answer: {answer}')
 
         with self.console.status("[bold green]Executing that command..."):
             result = response.execute()
 
             self.console.print(Panel(result, title="tool"))
             result_str = self.parse_http_status_line(result)
             self._prompt_history.append(tool_message(result_str, tool_call_id))
+            if result_str == '200 OK':
+                self.update_openapi_spec(response )
 
         self.log_db.add_log_query(self._run_id, turn, command, result, answer)
+        # if self._all_http_methods_found or turn == FINAL_ROUND:
+        self.write_openapi_to_yaml()
         return self._all_http_methods_found
 
     def parse_http_status_line(self, status_line):
@@ -130,10 +149,74 @@ def parse_http_status_line(self, status_line):
                 status_code = parts[1]  # e.g., "200"
                 status_message = parts[2].split("\r\n")[0]  # e.g., "OK"
                 print(f'status code:{status_code}, status msg:{status_message}')
-                return str(status_code +" " + status_message  )
+                return str(status_code + " " + status_message)
             else:
                 raise ValueError("Invalid HTTP status line")
 
-
-
-
+    def has_no_numbers(self,path):
+        for char in path:
+            if char.isdigit():
+                return False
+        return True
+    def update_openapi_spec(self, response):
+        # This function should parse the request and update the OpenAPI specification
+        # For the purpose of this example, let's assume it parses JSON requests and updates paths
+            request = response.action
+            path = request.path
+            method = request.method
+            if path and method:
+                if path not in self.openapi_spec['paths'] :#and self.has_no_numbers(path):
+                    self.openapi_spec['paths'][path] = {}
+                self.openapi_spec['paths'][path][method.lower()] = {
+                    "summary": f"{method} operation on {path}",
+                    "responses": {
+                        "200": {
+                            "description": "Successful response",
+                            "content": {
+                                "application/json": {
+                                    "schema": {"type": "object"}  # Simplified for example
+                                }
+                            }
+                        }
+                    }
+                }
+
+    def write_openapi_to_yaml(self, filename='openapi_spec.yaml'):
+        """Write the OpenAPI specification to a YAML file."""
+        try:
+            openapi_data = {
+                "openapi": self.openapi_spec["openapi"],
+                "info": self.openapi_spec["info"],
+                "servers": self.openapi_spec["servers"],
+                "paths": self.openapi_spec["paths"]
+            }
+
+            # Ensure the directory exists
+            file_path = filename.split(".yaml")[0]
+            file_name = filename.split(".yaml")[0] + "_"+ self.current_time.strftime("%Y-%m-%d %H:%M:%S")+".yaml"
+            os.makedirs(file_path, exist_ok=True)
+
+            with open(os.path.join(file_path, file_name), 'w') as yaml_file:
+                yaml.dump(openapi_data, yaml_file, allow_unicode=True, default_flow_style=False)
+            self.console.print(f"[green]OpenAPI specification written to [bold]{filename}[/bold].")
+        except Exception as e:
+            raise Exception(e)
+
+            #self.console.print(f"[red]Error writing YAML file: {e}")
+    def write_openapi_to_yaml2(self, filename='openapi_spec.yaml'):
+        """Write the OpenAPI specification to a YAML file."""
+        try:
+           # self.setup_yaml()  # Configure YAML to handle complex types
+            with open(filename, 'w') as yaml_file:
+                yaml.dump(self.openapi_spec, yaml_file, allow_unicode=True, default_flow_style=False)
+            self.console.print(f"[green]OpenAPI specification written to [bold]{filename}[/bold].")
+        except TypeError as e:
+            raise Exception(e)
+            #self.console.print(f"[red]Error writing YAML file: {e}")
+
+    def represent_dict_order(self, data):
+        return self.represent_mapping('tag:yaml.org,2002:map', data.items())
+
+    def setup_yaml(self):
+        """Configure YAML to output OrderedDicts as regular dicts (helpful for better YAML readability)."""
+        yaml.add_representer(dict, self.represent_dict_order)