Fix parsing issues with new GPT (#380)

* add backup parser for markdown

* use extract_tag with backup markdown parsing

* fix formatting

* fix python 3.10 issue

vision_agent/agent/vision_agent_coder_v2.py

@@ -78,7 +78,7 @@ def write_code(
         plan=plan,
     )
     response = cast(str, coder([{"role": "user", "content": prompt}], stream=False))
-    maybe_code = extract_tag(response, "code")
+    maybe_code = extract_tag(response, "code", extract_markdown="python")
 
     # if the response wasn't properly formatted with the code tags just retrun the response
     if maybe_code is None:
@@ -105,7 +105,7 @@ def write_test(
         media=media_list,
     )
     response = cast(str, tester([{"role": "user", "content": prompt}], stream=False))
-    maybe_code = extract_tag(response, "code")
+    maybe_code = extract_tag(response, "code", extract_markdown="python")
 
     # if the response wasn't properly formatted with the code tags just retrun the response
     if maybe_code is None:
@@ -151,8 +151,12 @@ def debug_code(
             fixed_code_and_test_str = cast(str, fixed_code_and_test_str)
             thoughts_tag = extract_tag(fixed_code_and_test_str, "thoughts")
             thoughts = thoughts_tag if thoughts_tag is not None else ""
-            fixed_code = extract_tag(fixed_code_and_test_str, "code")
-            fixed_test = extract_tag(fixed_code_and_test_str, "test")
+            fixed_code = extract_tag(
+                fixed_code_and_test_str, "code", extract_markdown="python"
+            )
+            fixed_test = extract_tag(
+                fixed_code_and_test_str, "test", extract_markdown="python"
+            )
 
             success = not (fixed_code is None and fixed_test is None)
 

vision_agent/agent/vision_agent_planner_v2.py

@@ -154,7 +154,7 @@ def run_multi_trial_planning(
         response3=responses[2],
     )
     response = cast(str, model.chat([{"role": "user", "content": prompt}]))
-    json_str = extract_tag(response, "json")
+    json_str = extract_tag(response, "json", extract_markdown="json")
     if json_str:
         json_data = extract_json(json_str)
         best = np.argmax([int(json_data[f"response{k}"]) for k in [1, 2, 3]])
@@ -230,7 +230,7 @@ def execute_code_action(
     while not execution.success and count <= 3:
         prompt = FIX_BUG.format(chat_history=get_planning(chat), code=code, error=obs)
         response = cast(str, model.chat([{"role": "user", "content": prompt}]))
-        new_code = extract_tag(response, "code")
+        new_code = extract_tag(response, "code", extract_markdown="python")
         if not new_code:
             continue
         else:
@@ -343,7 +343,7 @@ def create_finalize_plan(
     plan_str = cast(str, response)
     return_chat = [AgentMessage(role="planner", content=plan_str, media=None)]
 
-    plan_json = extract_tag(plan_str, "json")
+    plan_json = extract_tag(plan_str, "json", extract_markdown="json")
 
     # sometimes the planner model will refuse to answer a question becuase of some
     # safety concern, we then wont be able to parse the response so we have to send
@@ -357,7 +357,7 @@ def create_finalize_plan(
     except json.JSONDecodeError:
         return return_chat, ErrorContext(error=plan_str)
 
-    code_snippets = extract_tag(plan_str, "code")
+    code_snippets = extract_tag(plan_str, "code", extract_markdown="python")
     plan["code"] = code_snippets if code_snippets is not None else ""
     if verbose:
         _CONSOLE.print(
@@ -544,7 +544,9 @@ def generate_plan(
 
                 response = response_safeguards(response)
                 thinking = extract_tag(response, "thinking")
-                code = extract_tag(response, "execute_python")
+                code = extract_tag(
+                    response, "execute_python", extract_markdown="python"
+                )
                 finalize_plan = extract_tag(response, "finalize_plan")
                 finished = finalize_plan is not None
                 self.update_callback({"role": "planner_update", "content": response})

vision_agent/tools/planner_tools.py

@@ -130,7 +130,7 @@ def run_multi_judge(
 
     def run_judge() -> Tuple[Optional[Callable], str, str]:
         response = tool_chooser.generate(prompt, media=image_paths, temperature=1.0)
-        tool_choice_context = extract_tag(response, "json")  # type: ignore
+        tool_choice_context = extract_tag(response, "json", extract_markdown="json")  # type: ignore
         tool_choice_context_dict = extract_json(tool_choice_context)  # type: ignore
         tool, tool_thoughts, tool_docstring, _ = extract_tool_info(
             tool_choice_context_dict
@@ -275,7 +275,7 @@ def run_tool_testing(
     )
 
     response = lmm.generate(prompt, media=image_paths)
-    code = extract_tag(response, "code")  # type: ignore
+    code = extract_tag(response, "code", extract_markdown="python")  # type: ignore
     if code is None:
         raise ValueError(f"Could not extract code from response: {response}")
 
@@ -304,7 +304,7 @@ def run_tool_testing(
             media=str(image_paths),
         )
         response = cast(str, lmm.generate(prompt, media=image_paths))
-        code = extract_tag(response, "code")
+        code = extract_tag(response, "code", extract_markdown="python")
         if code is None:
             code = response
 

vision_agent/utils/agent.py

@@ -103,24 +103,37 @@ def extract_code(code: str) -> str:
     return code
 
 
-def extract_tag(
-    content: str,
-    tag: str,
-) -> Optional[str]:
+def _extract_arbitrary(content: str, start: str, stop: 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:
+    while start in remaning:
+        inner_content_i = remaning[remaning.find(start) + len(start) :]
+        if stop 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}>") :]
+        inner_content_i = inner_content_i[: inner_content_i.find(stop)]
+        remaning = remaning[remaning.find(stop) + len(stop) :]
         all_inner_content.append(inner_content_i)
 
     if len(all_inner_content) > 0:
         inner_content = "\n".join(all_inner_content)
+
+    return inner_content
+
+
+def extract_markdown(content: str, tag: str) -> Optional[str]:
+    return _extract_arbitrary(content, f"```{tag}", "```")
+
+
+def extract_tag(
+    content: str,
+    tag: str,
+    extract_markdown: Optional[str] = None,
+) -> Optional[str]:
+    inner_content = _extract_arbitrary(content, f"<{tag}>", f"</{tag}>")
+    if inner_content is None and extract_markdown is not None:
+        inner_content = _extract_arbitrary(content, f"```{extract_markdown}", "```")
     return inner_content