Add step magics (#493)

* Add step magics

* fix

* fix test cases

* fix test cases

* Bump version

* Update pyproject.toml

* update libraries

* add partial svelte support

* unlimit call llm

* remove depscan test for debugging

Author: CTY-git

patchwork/common/context_strategy/javascript.py

@@ -26,7 +26,7 @@
         (statement_block) @node
     """.strip()
 
-_javascript_exts = [".js", ".ts"]
+_javascript_exts = [".js", ".ts", ".svelte"]
 _jsx_exts = [".jsx", ".tsx"]
 
 

patchwork/step.py

@@ -1,14 +1,66 @@
-from typing_extensions import Protocol
+import abc
+from enum import Flag, auto
 
+from patchwork.logger import logger
 
-class Step(Protocol):
-    """
-    Protocol for a Step.
-    Steps do not have to inherit from this class, but they must implement the run method.
-    The __init__ method should have a single argument, inputs, which is a dictionary of inputs.
-    This is the only opportunity to set the inputs as a class property to be used in `run`.
-    """
 
+class StepStatus(Flag):
+    COMPLETED = auto()
+    FAILED = auto()
+    SKIPPED = auto()
+
+    def __str__(self):
+        return self.name.lower()
+
+
+class Step(abc.ABC):
+    def __init__(self, inputs: dict):
+        """
+        Initializes the step.
+        :param inputs: a dictionary of inputs
+        """
+        self.__status = StepStatus.COMPLETED
+        self.__status_msg = None
+        self.__step_name = self.__class__.__name__
+        # abit of a hack to wrap the implemented run method
+        self.original_run = self.run
+        self.run = self.__managed_run
+
+    def __managed_run(self, *args, **kwargs):
+        logger.info(f"Run started {self.__step_name}")
+
+        exc = None
+        try:
+            output = self.original_run(*args, **kwargs)
+        except Exception as e:
+            exc = e
+
+        is_fail = self.__status == StepStatus.FAILED or exc is not None
+        if self.__status_msg is not None:
+            message_logger = logger.error if is_fail else logger.info
+            message_logger(f"Step {self.__step_name} message: {self.__status_msg}")
+
+        if exc is not None:
+            logger.error(f"Step {self.__step_name} failed")
+            raise exc
+
+        if is_fail:
+            raise ValueError(f"Step {self.__step_name} failed")
+
+        logger.info(f"Run {self.__status} {self.__step_name}")
+        return output
+
+    def set_status(self, status: StepStatus, msg: str = None):
+        if status not in StepStatus:
+            raise ValueError(f"Invalid status: {status}")
+        self.__status = status
+        self.__status_msg = msg
+
+    @property
+    def status(self) -> StepStatus:
+        return self.__status
+
+    @abc.abstractmethod
     def run(self) -> dict:
         """
         Runs the step.

patchwork/steps/AnalyzeImpact/AnalyzeImpact.py

@@ -2,7 +2,7 @@
 from pathlib import Path
 
 from patchwork.logger import logger
-from patchwork.step import Step
+from patchwork.step import Step, StepStatus
 
 _PURL_TO_LANGUAGE_ = {
     "pypi": "python",
@@ -76,8 +76,7 @@ def find_dependency_usage(directory, dependency, language, methods):
 
 class AnalyzeImpact(Step):
     def __init__(self, inputs: dict):
-        logger.info(f"Run started {self.__class__.__name__}")
-
+        super().__init__(inputs)
         required_keys = {"extracted_responses", "library_name", "platform_type"}
         if not all(key in inputs.keys() for key in required_keys):
             raise ValueError(f'Missing required data: "{required_keys}"')
@@ -86,6 +85,10 @@ def __init__(self, inputs: dict):
 
     def run(self) -> dict:
         extracted_responses = self.inputs["extracted_responses"]
+        if len(extracted_responses) == 0:
+            self.set_status(StepStatus.SKIPPED, "No extracted responses found")
+            return dict(files_to_patch=[])
+
         name = self.inputs["library_name"]
         platform_type = self.inputs["platform_type"]
 
@@ -142,5 +145,4 @@ def run(self) -> dict:
             )
             extracted_data.append(data)
 
-        logger.info(f"Run completed {self.__class__.__name__}")
         return dict(files_to_patch=extracted_data)

patchwork/steps/CallAPI/CallAPI.py

@@ -7,6 +7,7 @@
 
 class CallAPI(Step):
     def __init__(self, inputs):
+        super().__init__(inputs)
         self.url = inputs["url"]
         self.method = inputs["method"]
         possible_headers = inputs.get("headers", {})

patchwork/steps/CallCode2Prompt/CallCode2Prompt.py

@@ -2,8 +2,7 @@
 import subprocess
 from pathlib import Path
 
-from patchwork.logger import logger
-from patchwork.step import Step
+from patchwork.step import Step, StepStatus
 
 FOLDER_PATH = "folder_path"
 
@@ -12,8 +11,7 @@ class CallCode2Prompt(Step):
     required_keys = {FOLDER_PATH}
 
     def __init__(self, inputs: dict):
-        logger.info(f"Run started {self.__class__.__name__}")
-
+        super().__init__(inputs)
         if not all(key in inputs.keys() for key in self.required_keys):
             raise ValueError(f'Missing required data: "{self.required_keys}"')
 
@@ -28,9 +26,6 @@ def __init__(self, inputs: dict):
             with open(self.code_file_path, "a") as file:
                 pass  # No need to write anything, just create the file if it doesn't exist
 
-        # Prepare for data extraction
-        self.extracted_data = []
-
     def run(self) -> dict:
         cmd = [
             "code2prompt",
@@ -54,13 +49,13 @@ def run(self) -> dict:
             with open(self.code_file_path, "r") as file:
                 file_content = file.read()
         except FileNotFoundError:
-            logger.info(f"Unable to find file: {self.code_file_path}")
+            self.set_status(StepStatus.FAILED, f"Unable to find file: {self.code_file_path}")
+            return dict(files_to_patch=[])
 
         lines = file_content.splitlines(keepends=True)
 
-        self.extracted_data.append(
-            dict(uri=self.code_file_path, startLine=0, endLine=len(lines), fullContent=prompt_content_md)
+        return dict(
+            files_to_patch=[
+                dict(uri=self.code_file_path, startLine=0, endLine=len(lines), fullContent=prompt_content_md)
+            ]
         )
-
-        logger.info(f"Run completed {self.__class__.__name__}")
-        return dict(files_to_patch=self.extracted_data)

patchwork/steps/CallLLM/CallLLM.py

@@ -12,7 +12,7 @@
 from typing_extensions import Any, Protocol
 
 from patchwork.logger import logger
-from patchwork.step import Step
+from patchwork.step import Step, StepStatus
 
 TOKEN_URL = "https://app.patched.codes/signin"
 _DEFAULT_PATCH_URL = "https://patchwork.patched.codes/v1"
@@ -158,8 +158,7 @@ def parse_model_args(self, model_args: dict) -> dict:
 
 class CallLLM(Step):
     def __init__(self, inputs: dict):
-        logger.info(f"Run started {self.__class__.__name__}")
-
+        super().__init__(inputs)
         # Set 'openai_key' from inputs or environment if not already set
         inputs.setdefault("openai_api_key", os.environ.get("OPENAI_API_KEY"))
 
@@ -178,7 +177,7 @@ def __init__(self, inputs: dict):
         else:
             raise ValueError('Missing required data: "prompt_file" or "prompts"')
 
-        self.call_limit = int(inputs.get("max_llm_calls", 50))
+        self.call_limit = int(inputs.get("max_llm_calls", -1))
         self.model_args = {key[len("model_") :]: value for key, value in inputs.items() if key.startswith("model_")}
         self.client_args = {key[len("client_") :]: value for key, value in inputs.items() if key.startswith("client_")}
         self.save_responses_to_file = inputs.get("save_responses_to_file", None)
@@ -223,13 +222,22 @@ def __init__(self, inputs: dict):
 
     def run(self) -> dict:
         prompt_length = len(self.prompts)
+        if prompt_length == 0:
+            self.set_status(StepStatus.SKIPPED, "No prompts to process")
+            return dict(openai_responses=[])
+
         if prompt_length > self.call_limit:
             logger.debug(
                 f"Number of prompts ({prompt_length}) exceeds the call limit ({self.call_limit}). "
                 f"Only the first {self.call_limit} prompts will be processed."
             )
 
-        contents = self.llm.call(list(islice(self.prompts, self.call_limit)))
+        if self.call_limit > 0:
+            prompts = list(islice(self.prompts, self.call_limit))
+        else:
+            prompts = self.prompts
+
+        contents = self.llm.call(prompts)
 
         if self.save_responses_to_file:
             # Convert relative path to absolute path
@@ -247,5 +255,4 @@ def run(self) -> dict:
                     }
                     f.write(json.dumps(data) + "\n")
 
-        logger.info(f"Run completed {self.__class__.__name__}")
         return dict(openai_responses=contents)

patchwork/steps/CommitChanges/CommitChanges.py

@@ -9,7 +9,7 @@
 
 from patchwork.common.utils.utils import get_current_branch
 from patchwork.logger import logger
-from patchwork.step import Step
+from patchwork.step import Step, StepStatus
 
 
 @contextlib.contextmanager
@@ -100,8 +100,7 @@ class CommitChanges(Step):
     required_keys = {"modified_code_files"}
 
     def __init__(self, inputs: dict):
-        logger.info(f"Run started {self.__class__.__name__}")
-
+        super().__init__(inputs)
         if not all(key in inputs.keys() for key in self.required_keys):
             raise ValueError(f'Missing required data: "{self.required_keys}"')
 
@@ -127,7 +126,7 @@ def run(self) -> dict:
         modified_files = {Path(modified_code_file["path"]).resolve() for modified_code_file in self.modified_code_files}
         true_modified_files = modified_files.intersection(repo_changed_files.union(repo_untracked_files))
         if not self.enabled or len(true_modified_files) < 1:
-            logger.debug("Branch creation is disabled.")
+            self.set_status(StepStatus.SKIPPED, "Branch creation is disabled.")
             from_branch = get_current_branch(repo)
             from_branch_name = from_branch.name if not from_branch.is_remote() else from_branch.remote_head
             return dict(target_branch=from_branch_name)
@@ -142,7 +141,6 @@ def run(self) -> dict:
                 repo.git.add(modified_file)
                 commit_with_msg(repo, f"Patched {modified_file}")
 
-            logger.info(f"Run completed {self.__class__.__name__}")
             return dict(
                 base_branch=from_branch,
                 target_branch=to_branch,

patchwork/steps/CreateIssue/CreateIssue.py

@@ -8,16 +8,14 @@
     ScmPlatformClientProtocol,
     get_slug_from_remote_url,
 )
-from patchwork.logger import logger
 from patchwork.step import Step
 
 
 class CreateIssue(Step):
     required_keys = {"issue_title", "issue_text", "scm_url"}
 
     def __init__(self, inputs: dict):
-        logger.info(f"Run started {self.__class__.__name__}")
-
+        super().__init__(inputs)
         if not all(key in inputs.keys() for key in self.required_keys):
             raise ValueError(f'Missing required data: "{self.required_keys}"')
 

patchwork/steps/CreateIssueComment/CreateIssueComment.py

@@ -3,16 +3,14 @@
     GitlabClient,
     ScmPlatformClientProtocol,
 )
-from patchwork.logger import logger
-from patchwork.step import Step
+from patchwork.step import Step, StepStatus
 
 
 class CreateIssueComment(Step):
     required_keys = {"issue_url", "issue_text"}
 
     def __init__(self, inputs: dict):
-        logger.info(f"Run started {self.__class__.__name__}")
-
+        super().__init__(inputs)
         if not all(key in inputs.keys() for key in self.required_keys):
             raise ValueError(f'Missing required data: "{self.required_keys}"')
 
@@ -35,7 +33,7 @@ def run(self) -> dict:
             slug, issue_id = self.scm_client.get_slug_and_id_from_url(self.issue_url)
             url = self.scm_client.create_issue_comment(slug, self.issue_text, issue_id=issue_id)
         except Exception as e:
-            logger.error(e)
-            return {}
+            self.set_status(StepStatus.FAILED, f"Failed to create issue comment")
+            raise e
 
         return dict(issue_comment_url=url)

patchwork/steps/CreatePR/CreatePR.py

@@ -17,8 +17,7 @@ class CreatePR(Step):
     required_keys = {"target_branch"}
 
     def __init__(self, inputs: dict):
-        logger.info(f"Run started {self.__class__.__name__}")
-
+        super().__init__(inputs)
         if not all(key in inputs.keys() for key in self.required_keys):
             raise ValueError(f'Missing required data: "{self.required_keys}"')
 
@@ -87,7 +86,6 @@ def run(self) -> dict:
         )
 
         logger.info(f"[green]PR created at [link={url}]{url}[/link][/]", extra={"markup": True})
-        logger.info(f"Run completed {self.__class__.__name__}")
         return {"pr_url": url}