Fix code gen

actions.py

@@ -41,7 +41,9 @@ def execute_task(task_id: int, active_index: int, error_value, *args):
     # We need to return outputs for all tasks in the row.
     outputs = [""] * n_avail_tasks
 
-    if active_index is None or error_value:  # Active index could be 0 == not active_index
+    if (
+        active_index is None or error_value
+    ):  # Active index could be 0 == not active_index
         return outputs + [
             gr.HighlightedText.update(
                 value=error_value, visible=error_value is not None
@@ -108,6 +110,9 @@ def execute_task(task_id: int, active_index: int, error_value, *args):
             )
         ]
     except Exception as e:
+        import traceback
+
+        print(traceback.format_tb(e.__traceback__))
         outputs[active_index] = "ERROR"
         return outputs + [
             gr.HighlightedText.update(

components.py

@@ -1,3 +1,4 @@
+import re
 from abc import ABC, abstractmethod
 from typing import Any, Dict, List, Union
 
@@ -184,26 +185,25 @@ class CodeTask(TaskComponent):
                         {code_prompt}
 
                         Do't save anything to disk. Instead, the function should return the necessary data.
-                        Include all the necessary imports.
+                        Include all the necessary imports. Make sure that the package names are correct.
                         """,
                     }
                 ],
                 temperature=0,
             )
 
-            parsed_output = json.loads(
-                ai.llm.next(
-                    [
-                        {
-                            "role": "user",
-                            "content": f"""
+            raw_parsed_output = ai.llm.next(
+                [
+                    {
+                        "role": "user",
+                        "content": f"""
                         The following text should have a python function with some imports that might need to be installed:
                         {raw_prompt_output}
 
-                        Extract all the python packages that need to be installed with pip.
-                        Also extract the function and the imports as a single python script.
+                        Extract all the python packages that need to be installed with pip, nothing else.
+                        Extract the function and the imports as a single python script, nothing else.
 
-                        Write a JSON as follows:
+                        Write a JSON:
                         ```
                             {{
                                 "packages": Python list of packages to be parsed with eval(). If no packages, the list should be empty.
@@ -211,12 +211,18 @@ class CodeTask(TaskComponent):
                             }}
                         ```
                         """,
-                        }
-                    ],
-                    temperature=0,
-                )
+                    }
+                ],
+                temperature=0,
+            )
+            parsed_output = json.loads(
+                re.search("({.*})", raw_parsed_output, re.DOTALL).group(1)
             )
         except Exception as e:
+            import traceback
+
+            print(str(e))
+            print(traceback.format_exc())
             error_message = gr.HighlightedText.update(
                 value=[(str(e), "ERROR")], visible=True
             )
@@ -224,7 +230,7 @@ class CodeTask(TaskComponent):
         return (
             raw_prompt_output,
             parsed_output["packages"],
-            parsed_output["script"],
+            parsed_output["script"].replace("```python", "").replace("```", ""),
             error_message,
             accordion,
         )