Auto fix json output from LLMs (patched-codes#1000)

* Autofix fix malformed json output from llm models

* bump version

* update lockfile and lint

* fix weird dependency issue

Author: CTY-git

patchwork/steps/PR/PR.py

@@ -16,6 +16,8 @@ def __handle_modified_code_files(self):
         if self.inputs.get("modified_code_files") is not None:
             return
 
+        self.inputs["modified_code_files"] = []
+
         input_modified_files = self.inputs.get("modified_files")
         if input_modified_files is None or len(input_modified_files) < 1:
             return

patchwork/steps/SimplifiedLLM/SimplifiedLLM.py

@@ -1,6 +1,8 @@
 import json
 from functools import partial
 
+from json_repair import repair_json
+
 from patchwork.common.client.llm.utils import example_json_to_schema
 from patchwork.common.utils.utils import RetryData, exclude_none_dict, retry
 from patchwork.logger import logger
@@ -13,13 +15,6 @@
 from patchwork.steps.SimplifiedLLM.typed import SimplifiedLLMInputs
 
 
-def json_loads(s: str) -> dict:
-    try:
-        return json.loads(s, strict=False)
-    except json.JSONDecodeError:
-        return dict()
-
-
 class SimplifiedLLM(Step):
     def __init__(self, inputs):
         super().__init__(inputs)
@@ -40,6 +35,20 @@ def __record_status_or_raise(self, retry_data: RetryData, step: Step):
         else:
             raise ValueError(step.status_message)
 
+    @staticmethod
+    def __json_loads(json_str: str) -> dict:
+        try:
+            return json.loads(json_str, strict=False)
+        except json.JSONDecodeError as e:
+            logger.debug(f"Json to decode: \n{json_str}\nError: \n{e}")
+
+        try:
+            json_str = repair_json(json_str, skip_json_loads=True)
+            return json.loads(json_str, strict=False)
+        except json.JSONDecodeError as e:
+            logger.debug(f"Json to decode: \n{json_str}\nError: \n{e}")
+            raise e
+
     def __retry_unit(self, prepare_prompt_outputs, call_llm_inputs, retry_data: RetryData):
         call_llm = CallLLM(call_llm_inputs)
         call_llm_outputs = call_llm.run()
@@ -49,7 +58,7 @@ def __retry_unit(self, prepare_prompt_outputs, call_llm_inputs, retry_data: Retr
             json_responses = []
             for response in call_llm_outputs.get("openai_responses"):
                 try:
-                    json_response = json.loads(response, strict=False)
+                    json_response = self.__json_loads(response)
                     json_responses.append(json_response)
                 except json.JSONDecodeError as e:
                     logger.error(f"Json to decode: \n{response}\nError: \n{e}")