fix: handle the case in which closing code tag starts with </

Sinaptik-AI · May 7, 2023 · dba400f · dba400f
1 parent f2348d0
commit dba400f
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/pandasai/llm/base.py b/pandasai/llm/base.py
@@ -83,7 +83,11 @@ def _extract_code(self, response: str, separator: str = "```") -> str:
             str: Extracted code from the response
         """
         code = response
-        match = re.search(rf"{START_CODE_TAG}(.*){END_CODE_TAG}", code, re.DOTALL)
+        match = re.search(
+            rf"{START_CODE_TAG}(.*)({END_CODE_TAG}|{END_CODE_TAG.replace('<', '</')})",
+            code,
+            re.DOTALL,
+        )
         if match:
             code = match.group(1).strip()
         if len(code.split(separator)) > 1: