From 172d0c035835b9c3d80dea2d165a9d03dc7ce48c Mon Sep 17 00:00:00 2001 From: Robby <45851384+h0rv@users.noreply.github.com> Date: Mon, 29 Jul 2024 10:26:58 -0400 Subject: [PATCH 1/2] improve response cleaning The prompt for the model starts with a code block (```). When testing watsonx models (llama and granite), they would generate the closing block in the response. --- pr_agent/algo/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pr_agent/algo/utils.py b/pr_agent/algo/utils.py index c534e701..3044f447 100644 --- a/pr_agent/algo/utils.py +++ b/pr_agent/algo/utils.py @@ -557,7 +557,7 @@ def _fix_key_value(key: str, value: str): def load_yaml(response_text: str, keys_fix_yaml: List[str] = [], first_key="", last_key="") -> dict: - response_text = response_text.strip('\n').removeprefix('```yaml').rstrip('`') + response_text = response_text.strip().removeprefix("```yaml").removesuffix("```").strip() try: data = yaml.safe_load(response_text) except Exception as e: From 039d85b836f2616bb674a6b60f9fc4174b6c4c59 Mon Sep 17 00:00:00 2001 From: Robby Date: Thu, 1 Aug 2024 15:50:00 -0400 Subject: [PATCH 2/2] fix cleaning --- pr_agent/algo/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pr_agent/algo/utils.py b/pr_agent/algo/utils.py index 3044f447..bcf6afca 100644 --- a/pr_agent/algo/utils.py +++ b/pr_agent/algo/utils.py @@ -557,7 +557,7 @@ def _fix_key_value(key: str, value: str): def load_yaml(response_text: str, keys_fix_yaml: List[str] = [], first_key="", last_key="") -> dict: - response_text = response_text.strip().removeprefix("```yaml").removesuffix("```").strip() + response_text = response_text.strip('\n').removeprefix('```yaml').rstrip().removesuffix('```') try: data = yaml.safe_load(response_text) except Exception as e: