From 2dd68132a21fd389ec0f4fc53efabc1f4ced19bc Mon Sep 17 00:00:00 2001
From: Dillon Laird <dillonalaird@gmail.com>
Date: Thu, 20 Jun 2024 09:43:03 -0700
Subject: [PATCH] retry failed json decode (#143)

---
 vision_agent/agent/vision_agent.py | 30 +++++++++++++++++++++---------
 1 file changed, 21 insertions(+), 9 deletions(-)

diff --git a/vision_agent/agent/vision_agent.py b/vision_agent/agent/vision_agent.py
index d3bbcecb..55112436 100644
--- a/vision_agent/agent/vision_agent.py
+++ b/vision_agent/agent/vision_agent.py
@@ -324,16 +324,28 @@ def debug_code(
             "status": "started",
         }
     )
-    fixed_code_and_test = extract_json(
-        debugger(
-            FIX_BUG.format(
-                code=code,
-                tests=test,
-                result="\n".join(result.text().splitlines()[-50:]),
-                feedback=format_memory(working_memory + new_working_memory),
+
+    fixed_code_and_test = {"code": "", "test": "", "reflections": ""}
+    success = False
+    count = 0
+    while not success and count < 3:
+        try:
+            fixed_code_and_test = extract_json(
+                debugger(
+                    FIX_BUG.format(
+                        code=code,
+                        tests=test,
+                        result="\n".join(result.text().splitlines()[-50:]),
+                        feedback=format_memory(working_memory + new_working_memory),
+                    )
+                )
             )
-        )
-    )
+            success = True
+        except Exception as e:
+            _LOGGER.exception(f"Error while extracting JSON: {e}")
+
+        count += 1
+
     old_code = code
     old_test = test