From 8058f4528b89ffb10e4baf9ffe3d17209cc455a1 Mon Sep 17 00:00:00 2001 From: Saurav Panda Date: Thu, 20 Jun 2024 23:51:52 -0400 Subject: [PATCH 1/3] fix: added llm based json parsing handling --- kaizen/helpers/parser.py | 12 ++++++++++++ kaizen/reviewer/code_review.py | 12 +++++++----- 2 files changed, 19 insertions(+), 5 deletions(-) diff --git a/kaizen/helpers/parser.py b/kaizen/helpers/parser.py index e928ab3a..6c833af0 100644 --- a/kaizen/helpers/parser.py +++ b/kaizen/helpers/parser.py @@ -24,6 +24,18 @@ def extract_json(text): return parsed_data +def extract_json_with_llm_retry(provider, text, total_usage): + try: + json_data = extract_json(text) + return json_data, total_usage + except: + prompt = f"Help me fix this json data: {text}" + resp, usage = provider.chat_completion(prompt) + total_usage = provider.update_usage(total_usage, usage) + json_data = extract_json(resp) + return json_data, total_usage + + def extract_multi_json(text): start_index = text.find("[") end_index = text.rfind("]") + 1 diff --git a/kaizen/reviewer/code_review.py b/kaizen/reviewer/code_review.py index 9e17a860..f203adfc 100644 --- a/kaizen/reviewer/code_review.py +++ b/kaizen/reviewer/code_review.py @@ -71,7 +71,7 @@ def review_pull_request( if self.provider.is_inside_token_limit(PROMPT=prompt): self.logger.debug("Processing Directly from Diff") resp, usage = self.provider.chat_completion(prompt, user=user) - review_json = parser.extract_json(resp) + review_json, total_usage = parser.extract_json_with_llm_retry(self.provider, resp, total_usage) reviews = review_json["review"] total_usage = self.provider.update_usage(total_usage, usage) else: @@ -95,7 +95,7 @@ def review_pull_request( continue resp, usage = self.provider.chat_completion(prompt, user=user) total_usage = self.provider.update_usage(total_usage, usage) - review_json = parser.extract_json(resp) + review_json, total_usage = parser.extract_json_with_llm_retry(self.provider, resp, total_usage) reviews.extend(review_json["review"]) topics = self.merge_topics(reviews=reviews) @@ -137,7 +137,8 @@ def generate_pull_request_desc( self.logger.debug("Processing Directly from Diff") resp, usage = self.provider.chat_completion(prompt, user=user) total_usage = self.provider.update_usage(total_usage, usage) - desc = parser.extract_json(resp)["desc"] + json_data = parser.extract_json_with_llm_retry(self.provider, resp, total_usage) + desc = json_data["desc"] else: self.logger.debug("Processing Based on files") # We recurrsively get feedback for files and then get basic summary @@ -159,13 +160,14 @@ def generate_pull_request_desc( continue resp, usage = self.provider.chat_completion(prompt, user=user) total_usage = self.provider.update_usage(total_usage, usage) - desc_json = parser.extract_json(resp) + desc_json, total_usage = parser.extract_json_with_llm_retry(self.provider, resp, total_usage) descs.append(desc_json["desc"]) prompt = MERGE_PR_DESCRIPTION_PROMPT.format(DESCS=json.dumps(descs)) resp, usage = self.provider.chat_completion(prompt, user=user) total_usage = self.provider.update_usage(total_usage, usage) - desc = parser.extract_json(resp)["desc"] + resp, total_usage = parser.extract_json_with_llm_retry(self.provider, resp, total_usage) + desc = resp["desc"] body = output.create_pr_description(desc, pull_request_desc) prompt_cost, completion_cost = self.provider.get_usage_cost( total_usage=total_usage From 8d9b4b5ba639d50571a67d461bc0a93643af1fc6 Mon Sep 17 00:00:00 2001 From: Saurav Panda Date: Thu, 20 Jun 2024 23:54:33 -0400 Subject: [PATCH 2/3] linting fixes --- kaizen/generator/ui.py | 6 +--- kaizen/llms/prompts/work_summary_prompts.py | 2 +- kaizen/reviewer/code_review.py | 36 ++++++++++++--------- kaizen/reviewer/work_summarizer.py | 6 +--- 4 files changed, 23 insertions(+), 27 deletions(-) diff --git a/kaizen/generator/ui.py b/kaizen/generator/ui.py index e6c94477..844a351a 100644 --- a/kaizen/generator/ui.py +++ b/kaizen/generator/ui.py @@ -81,11 +81,7 @@ def generate_module_tests(self, web_content: str, test_modules: dict, web_url: s This method generates UI testing points for all modules. """ ui_tests = test_modules - total_usage = { - "prompt_tokens": 0, - "completion_tokens": 0, - "total_tokens": 0 - } + total_usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} for module in ui_tests: for test in module["tests"]: test_description = test["test_description"] diff --git a/kaizen/llms/prompts/work_summary_prompts.py b/kaizen/llms/prompts/work_summary_prompts.py index 2ea02af8..3af8af15 100644 --- a/kaizen/llms/prompts/work_summary_prompts.py +++ b/kaizen/llms/prompts/work_summary_prompts.py @@ -39,4 +39,4 @@ 7. Maintain a consistent, readable tone. PATCH DATA: {PATCH_DATA} -""" +""" diff --git a/kaizen/reviewer/code_review.py b/kaizen/reviewer/code_review.py index f203adfc..45612ff6 100644 --- a/kaizen/reviewer/code_review.py +++ b/kaizen/reviewer/code_review.py @@ -63,15 +63,13 @@ def review_pull_request( PULL_REQUEST_DESC=pull_request_desc, CODE_DIFF=diff_text, ) - total_usage = { - "prompt_tokens": 0, - "completion_tokens": 0, - "total_tokens": 0 - } + total_usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} if self.provider.is_inside_token_limit(PROMPT=prompt): self.logger.debug("Processing Directly from Diff") resp, usage = self.provider.chat_completion(prompt, user=user) - review_json, total_usage = parser.extract_json_with_llm_retry(self.provider, resp, total_usage) + review_json, total_usage = parser.extract_json_with_llm_retry( + self.provider, resp, total_usage + ) reviews = review_json["review"] total_usage = self.provider.update_usage(total_usage, usage) else: @@ -90,12 +88,16 @@ def review_pull_request( PULL_REQUEST_DESC=pull_request_desc, FILE_PATCH=patch_details, ) - if not self.provider.is_inside_token_limit(PROMPT=prompt, percentage=85): + if not self.provider.is_inside_token_limit( + PROMPT=prompt, percentage=85 + ): # TODO: Chunk this big files and process them continue resp, usage = self.provider.chat_completion(prompt, user=user) total_usage = self.provider.update_usage(total_usage, usage) - review_json, total_usage = parser.extract_json_with_llm_retry(self.provider, resp, total_usage) + review_json, total_usage = parser.extract_json_with_llm_retry( + self.provider, resp, total_usage + ) reviews.extend(review_json["review"]) topics = self.merge_topics(reviews=reviews) @@ -128,16 +130,14 @@ def generate_pull_request_desc( CODE_DIFF=diff_text, ) - total_usage = { - "prompt_tokens": 0, - "completion_tokens": 0, - "total_tokens": 0 - } + total_usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} if self.provider.is_inside_token_limit(PROMPT=prompt): self.logger.debug("Processing Directly from Diff") resp, usage = self.provider.chat_completion(prompt, user=user) total_usage = self.provider.update_usage(total_usage, usage) - json_data = parser.extract_json_with_llm_retry(self.provider, resp, total_usage) + json_data = parser.extract_json_with_llm_retry( + self.provider, resp, total_usage + ) desc = json_data["desc"] else: self.logger.debug("Processing Based on files") @@ -160,13 +160,17 @@ def generate_pull_request_desc( continue resp, usage = self.provider.chat_completion(prompt, user=user) total_usage = self.provider.update_usage(total_usage, usage) - desc_json, total_usage = parser.extract_json_with_llm_retry(self.provider, resp, total_usage) + desc_json, total_usage = parser.extract_json_with_llm_retry( + self.provider, resp, total_usage + ) descs.append(desc_json["desc"]) prompt = MERGE_PR_DESCRIPTION_PROMPT.format(DESCS=json.dumps(descs)) resp, usage = self.provider.chat_completion(prompt, user=user) total_usage = self.provider.update_usage(total_usage, usage) - resp, total_usage = parser.extract_json_with_llm_retry(self.provider, resp, total_usage) + resp, total_usage = parser.extract_json_with_llm_retry( + self.provider, resp, total_usage + ) desc = resp["desc"] body = output.create_pr_description(desc, pull_request_desc) prompt_cost, completion_cost = self.provider.get_usage_cost( diff --git a/kaizen/reviewer/work_summarizer.py b/kaizen/reviewer/work_summarizer.py index fcf85a55..889937ca 100644 --- a/kaizen/reviewer/work_summarizer.py +++ b/kaizen/reviewer/work_summarizer.py @@ -22,11 +22,7 @@ def generate_work_summaries( summaries = [] # Try to merge the files untill LLM can process the response combined_diff_data = "" - total_usage = { - "prompt_tokens": 0, - "completion_tokens": 0, - "total_tokens": 0 - } + total_usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} for file_dict in diff_file_data: temp_prompt = combined_diff_data temp_prompt += f"""\n---->\nFile Name: {file_dict["file"]}\nPatch: {file_dict["patch"]}\n Status: {file_dict["status"]}""" From 6f34531d555fcc10cd2a713f44fc8cc8f9ffcfbd Mon Sep 17 00:00:00 2001 From: Saurav Panda Date: Thu, 20 Jun 2024 23:56:36 -0400 Subject: [PATCH 3/3] Update parser.py --- kaizen/helpers/parser.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/kaizen/helpers/parser.py b/kaizen/helpers/parser.py index 6c833af0..e89fc726 100644 --- a/kaizen/helpers/parser.py +++ b/kaizen/helpers/parser.py @@ -28,7 +28,8 @@ def extract_json_with_llm_retry(provider, text, total_usage): try: json_data = extract_json(text) return json_data, total_usage - except: + except Exception as e: + print(f"Error parsing json: {e}") prompt = f"Help me fix this json data: {text}" resp, usage = provider.chat_completion(prompt) total_usage = provider.update_usage(total_usage, usage)