diff --git a/vision_agent/agent/agent_utils.py b/vision_agent/agent/agent_utils.py
index e07ec619..cb7e1b44 100644
--- a/vision_agent/agent/agent_utils.py
+++ b/vision_agent/agent/agent_utils.py
@@ -92,6 +92,27 @@ def extract_code(code: str) -> str:
     return code
 
 
+def extract_tag(
+    content: str,
+    tag: str,
+) -> Optional[str]:
+    inner_content = None
+    remaning = content
+    all_inner_content = []
+
+    while f"<{tag}>" in remaning:
+        inner_content_i = remaning[remaning.find(f"<{tag}>") + len(f"<{tag}>") :]
+        if f"</{tag}>" not in inner_content_i:
+            break
+        inner_content_i = inner_content_i[: inner_content_i.find(f"</{tag}>")]
+        remaning = remaning[remaning.find(f"</{tag}>") + len(f"</{tag}>") :]
+        all_inner_content.append(inner_content_i)
+
+    if len(all_inner_content) > 0:
+        inner_content = "\n".join(all_inner_content)
+    return inner_content
+
+
 def remove_installs_from_code(code: str) -> str:
     pattern = r"\n!pip install.*?(\n|\Z)\n"
     code = re.sub(pattern, "", code, flags=re.DOTALL)
diff --git a/vision_agent/agent/vision_agent_coder.py b/vision_agent/agent/vision_agent_coder.py
index bc2de2a3..a265ec22 100644
--- a/vision_agent/agent/vision_agent_coder.py
+++ b/vision_agent/agent/vision_agent_coder.py
@@ -15,6 +15,7 @@
     DefaultImports,
     extract_code,
     extract_json,
+    extract_tag,
     format_memory,
     print_code,
     remove_installs_from_code,
@@ -260,7 +261,9 @@ def debug_code(
         }
     )
 
-    fixed_code_and_test = {"code": "", "test": "", "reflections": ""}
+    fixed_code = None
+    fixed_test = None
+    thoughts = ""
     success = False
     count = 0
     while not success and count < 3:
@@ -283,21 +286,16 @@ def debug_code(
                 stream=False,
             )
             fixed_code_and_test_str = cast(str, fixed_code_and_test_str)
-            fixed_code_and_test = extract_json(fixed_code_and_test_str)
-            code = extract_code(fixed_code_and_test_str)
-            if (
-                "which_code" in fixed_code_and_test
-                and fixed_code_and_test["which_code"] == "test"
-            ):
-                fixed_code_and_test["code"] = ""
-                fixed_code_and_test["test"] = code
-            else:  # for everything else always assume it's updating code
-                fixed_code_and_test["code"] = strip_function_calls(code)
-                fixed_code_and_test["test"] = ""
-            if "which_code" in fixed_code_and_test:
-                del fixed_code_and_test["which_code"]
-
-            success = True
+            thoughts = extract_tag(fixed_code_and_test_str, "thoughts")
+            thoughts = thoughts if thoughts is not None else ""
+            fixed_code = extract_tag(fixed_code_and_test_str, "code")
+            fixed_test = extract_tag(fixed_code_and_test_str, "test")
+
+            if fixed_code is None and fixed_test is None:
+                success = False
+            else:
+                success = True
+
         except Exception as e:
             _LOGGER.exception(f"Error while extracting JSON: {e}")
 
@@ -306,15 +304,15 @@ def debug_code(
     old_code = code
     old_test = test
 
-    if fixed_code_and_test["code"].strip() != "":
-        code = fixed_code_and_test["code"]
-    if fixed_code_and_test["test"].strip() != "":
-        test = fixed_code_and_test["test"]
+    if fixed_code is not None and fixed_code.strip() != "":
+        code = fixed_code
+    if fixed_test is not None and fixed_test.strip() != "":
+        test = fixed_test
 
     new_working_memory.append(
         {
             "code": f"{code}\n{test}",
-            "feedback": fixed_code_and_test["reflections"],
+            "feedback": cast(str, thoughts),
             "edits": get_diff(f"{old_code}\n{old_test}", f"{code}\n{test}"),
         }
     )
@@ -350,7 +348,7 @@ def debug_code(
     if verbosity == 2:
         print_code("Code and test after attempted fix:", code, test)
         _LOGGER.info(
-            f"Reflection: {fixed_code_and_test['reflections']}\nCode execution result after attempted fix: {result.text(include_logs=True)}"
+            f"Reflection: {thoughts}\nCode execution result after attempted fix: {result.text(include_logs=True)}"
         )
 
     return code, test, result
diff --git a/vision_agent/agent/vision_agent_coder_prompts.py b/vision_agent/agent/vision_agent_coder_prompts.py
index 66eb4c29..ffb83cc2 100644
--- a/vision_agent/agent/vision_agent_coder_prompts.py
+++ b/vision_agent/agent/vision_agent_coder_prompts.py
@@ -238,35 +238,29 @@ def find_text(image_path: str, text: str) -> str:
 {docstring}
 
 **Instructions**:
-Please re-complete the code to fix the error message. Here is the previous version:
-```python
+Please re-complete the code to fix the error message. Here is the current version of the CODE:
+<code>
 {code}
-```
+</code>
 
-When we run this test code:
-```python
+When we run the TEST code:
+<test>
 {tests}
-```
+</test>
 
 It raises this error:
-```
+<error>
 {result}
-```
+</error>
 
 This is previous feedback provided on the code:
 {feedback}
 
-Please fix the bug by correcting the error. Return the following JSON object followed by the fixed code in the below format:
-```json
-{{
-    "reflections": str # any thoughts you have about the bug and how you fixed it
-    "which_code": str # the code that was fixed, can only be 'code' or 'test'
-}}
-```
+Please fix the bug by correcting the error. Return thoughts you have about the bug and how you fixed in <thoughts> tags followed by the fixed CODE in <code> tags and the fixed TEST in <test> tags. For example:
 
-```python
-# Your fixed code here
-```
+<thoughts>Your thoughts here...</thoughts>
+<code># your fixed code here</code>
+<test># your fixed test here</test>
 """