Patchflow that generates architecture diagram (#1022)

* added patchflow to generate architecture

* added readme

* added tests to CI

* Update test.yml

Author: plon-Susk7

.github/workflows/test.yml

@@ -201,6 +201,13 @@ jobs:
           --github_api_key=${{ secrets.SCM_GITHUB_KEY }} \
           --base_path=tests/cicd/generate_docstring \
           --disable_telemetry
+      
+      - name : Generate Diagram
+        run: |
+          poetry run patchwork GenerateDiagram --log debug \
+          --patched_api_key=${{ secrets.PATCHED_API_KEY }} \
+          --github_api_key=${{ secrets.SCM_GITHUB_KEY }} \
+          --disable_telemetry
 
       - name: Generate UnitTests
         run: |

patchwork/patchflows/GenerateDiagram/GenerateDiagram.py

@@ -0,0 +1,62 @@
+import json
+from pathlib import Path
+import yaml
+
+from patchwork.common.utils.step_typing import validate_steps_with_inputs
+from patchwork.step import Step
+from patchwork.steps import (
+    LLM, 
+    CallCode2Prompt,
+    ModifyCode,
+    PR
+)
+
+_DEFAULT_INPUT_FILE = Path(__file__).parent / "defaults.yml"
+_DEFAULT_PROMPT_JSON = Path(__file__).parent / "default_prompt.json"
+
+class GenerateDiagram(Step):
+    def __init__(self, inputs):
+        super().__init__(inputs)
+
+        final_inputs = yaml.safe_load(_DEFAULT_INPUT_FILE.read_text())
+        if final_inputs is None:
+            final_inputs = {}
+        
+        final_inputs.update(inputs)
+
+        final_inputs["prompt_id"] = "GenerateDiagram"
+        if "folder_path" not in final_inputs.keys():
+            final_inputs["folder_path"] = Path.cwd()
+        else:
+            final_inputs["folder_path"] = Path(final_inputs["folder_path"])
+
+        print(final_inputs["folder_path"])
+
+        if "prompt_template_file" not in final_inputs:
+            final_inputs["prompt_template_file"] = _DEFAULT_PROMPT_JSON
+
+        final_inputs["pr_title"] = f"PatchWork System Architecture Diagram generated"
+        final_inputs["branch_prefix"] = f"{self.__class__.__name__.lower()}-"
+
+        validate_steps_with_inputs(
+            set(final_inputs.keys()).union({"prompt_values","files_to_patch"}), LLM, CallCode2Prompt,ModifyCode,PR
+        )
+        self.inputs = final_inputs
+
+    def run(self):
+        outputs = CallCode2Prompt(self.inputs).run()
+        new_file_name = f"diagram.md"
+        new_file_path = Path(outputs['uri']).with_name(new_file_name)
+        Path(outputs['uri']).rename(new_file_path)
+        outputs['uri'] = str(new_file_path)
+        self.inputs["response_partitions"] = {"patch": ["```", "\n", "```"]}
+        self.inputs["files_to_patch"] = self.inputs["prompt_values"] = [outputs]
+        outputs = LLM(self.inputs).run()
+        self.inputs.update(outputs)
+        outputs = ModifyCode(self.inputs).run()
+        self.inputs.update(outputs)
+        self.inputs["pr_header"] = f"This pull request from patchwork generates system architecture diagram."
+        outputs = PR(self.inputs).run()
+        self.inputs.update(outputs)
+
+        return self.inputs

patchwork/patchflows/GenerateDiagram/README.md

@@ -0,0 +1,13 @@
+## GenerateDiagram Code Overview
+
+### Inputs
+- The code reads default inputs from a YAML file.
+- It updates the default inputs with any additional inputs provided.
+- It sets up various parameters like folder path, prompt template file, PR title, branch prefix, etc.
+- It validates the inputs with specific steps required for the process.
+
+### Outputs
+- The code runs a series of steps to generate a system architecture diagram.
+- It utilizes classes like LLM, CallCode2Prompt, ModifyCode, and PR to process the inputs.
+- The final output is a set of inputs updated with the results of each step.
+- The code is designed to generate a pull request with the system architecture diagram.

patchwork/patchflows/GenerateDiagram/__init__.py

@@ -0,0 +1,15 @@
+[
+    {
+        "id": "GenerateDiagram", 
+        "prompts": [
+            {
+                "role": "system",
+                "content": "You are an experienced software architect skilled at visualizing system designs. Users will provide repository details, and you will generate a comprehensive system architecture diagram using Mermaid markdown syntax. Output the diagram code directly with no additional text, formatting, or triple quotes. The response should be ready to be pasted into an editor supporting Mermaid syntax."
+            },
+            {
+                "role": "user",
+                "content": "Repo Details: {{fullContent}}"
+            }
+        ]
+    }
+]

patchwork/patchflows/GenerateDiagram/default_prompt.json

@@ -0,0 +1,22 @@
+# CallLLM Inputs
+# openai_api_key: required-for-chatgpt
+# google_api_key: required-for-gemini
+# model: gpt-4o
+# client_base_url: https://api.openai.com/v1
+# Example HF model
+# client_base_url: https://api-inference.huggingface.co/models/codellama/CodeLlama-70b-Instruct-hf/v1
+# model: codellama/CodeLlama-70b-Instruct-hf
+# model_temperature: 0.2
+# model_top_p: 0.95
+# model_max_tokens: 2000
+
+# folder_path : path/to/folder/with/class
+
+# CommitChanges Inputs
+disable_branch: false
+
+# CreatePR Inputs
+disable_pr: false
+force_pr_creation: true
+# github_api_key: required-for-github-scm
+# gitlab_api_key: required-for-gitlab-scm

patchwork/patchflows/GenerateDiagram/defaults.yml

@@ -5,5 +5,6 @@
 from .PRReview.PRReview import PRReview
 from .ResolveIssue.ResolveIssue import ResolveIssue
 from .GenerateUnitTests.GenerateUnitTests import GenerateUnitTests
+from .GenerateDiagram.GenerateDiagram import GenerateDiagram
 
-__all__ = ["AutoFix", "DependencyUpgrade", "GenerateREADME", "PRReview", "ResolveIssue", "GenerateDocstring", "GenerateUnitTests"]
+__all__ = ["AutoFix", "DependencyUpgrade", "GenerateREADME", "PRReview", "ResolveIssue", "GenerateDocstring", "GenerateUnitTests","GenerateDiagram"]