switch to using tags to fix issue of mixing up code and tests

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)

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

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>
 """