FINISH

anpl/anpl.py

@@ -145,7 +145,7 @@ black, blue, red, green, yellow, grey, pink, orange, teal, maroon = range(10)
         return [f.name for f in self.funs.values() if not f.code]
     
     def user_known_funs(self) -> list[str]:
-        return [f.name for f in self.funs.values()]
+        return [f.name for f in self.funs.values() if f.code_from_user or f.prompt_from_user]
 
     def to_python(self, name: Optional[str] = None, for_user=False) -> str:
         funs_name = reversed(self.find_dependencies(name or self.entry))

robotA.py

@@ -71,7 +71,7 @@ is_correct = anpl_check(anpl, anpl.entry)
 
 while not is_correct:
     system_info("[red]ANPL WRONG[/red] Here is the anpl program.")
-    print_anpl(anpl)
+    print_anpl(anpl, for_user=True)
 
     cmd = Prompt.ask(sys_str + "Which command would you like to do? [1] Trace [2] Edit [3] Resynthesis [4] Remove IO [5] Quit", choices=["1", "2", "3", "4", "5"])
     if cmd == "5":

robotB.py

@@ -6,12 +6,22 @@ from anpl.sandbox import import_module_from_string, timeout
 import numpy as np
 import time
 from copy import deepcopy
+import re
 
 history = []
+code = ""
 def print_msg(message):
     role, text = message["role"], message["content"]    
     rich.print(f"[blue]{role}[/blue]:")
-    print(text)
+    pattern = r'```(.*?)```'
+    try:
+        global code
+        code = re.findall(pattern, text, flags=re.DOTALL)[0].strip()
+    except Exception as e:
+        print(e)
+        code = ""
+    text_without_code = re.sub(pattern, '', text, flags=re.DOTALL)
+    print(text_without_code)
 
 def print_history():
     for i, message in enumerate(history):    
@@ -24,7 +34,7 @@ logger = Logger(task_id, "B")
 
 is_correct = False
 while not is_correct:
-    cmd = Prompt.ask(sys_str + "Which command would you like to do? [1] Chat [2] Remove history [3] Check Code [4] Quit", choices=["1", "2", "3", "4"])
+    cmd = Prompt.ask(sys_str + "Which command would you like to do? [1] Chat [2] Remove history [3] Quit", choices=["1", "2", "3"])
     
     if cmd == "1":
         system_info("Please input your description or IO examples")
@@ -62,35 +72,6 @@ while not is_correct:
         else:
             logger.log("user", "remove", "no history")
             system_info("[red]No history[/red]")
-
-    elif cmd == "3":
-        system_info("Code will be executed from main function. The signature of main function should be `def main(input_grid):`")
-        system_info("Please enter your code")
-        code = multiline_input()
-        logger.log("user", "check", code)
-        try:
-            m = import_module_from_string(code)
-            inp_t = deepcopy(inp)
-            f = timeout(timeout=1)(m.main)
-            out = f(inp_t)
-        except Exception as e:
-            logger.log("system", "check", f"crash: {e}")
-            system_info("[red]Crash[/red]")
-            print_error(e)
-            continue
-        if np.array_equal(out, real_out):
-            logger.log("system", "check", f"correct")
-            system_info("[green]Code CORRECT[/green]")
-            logger.save(code)
-            is_correct = True
-        else:
-            logger.log("system", "check", f"wrong")
-            system_info("[red]Code WRONG[/red]")
-            rich.print("The output is")
-            rich.print("[green]Visual Output[/green]")
-            rich.print(rich_dumps(out))
-            rich.print("[green]Textual Output[/green]")
-            print(" ".join(out.__repr__().split()))
     else:
         quit_time = time.time()
         if quit_time - logger.start_time < 30 * 60:
@@ -100,4 +81,32 @@ while not is_correct:
         else:
             break
 
+    # print("here is code")
+    # print(code)
+    # print("code end")
+    logger.log("system", "check", code)
+    try:
+        m = import_module_from_string(code)
+        inp_t = deepcopy(inp)
+        f = timeout(timeout=1)(m.main)
+        out = f(inp_t)
+    except Exception as e:
+        logger.log("system", "check", f"crash: {e}")
+        system_info("[red]Crash[/red]")
+        print_error(e)
+        continue
+    if np.array_equal(out, real_out):
+        logger.log("system", "check", f"correct")
+        system_info("[green]Code CORRECT[/green]")
+        logger.save(code)
+        is_correct = True
+    else:
+        logger.log("system", "check", f"wrong")
+        system_info("[red]Code WRONG[/red]")
+        rich.print("The output is")
+        rich.print("[green]Visual Output[/green]")
+        rich.print(rich_dumps(out))
+        rich.print("[green]Textual Output[/green]")
+        print(" ".join(out.__repr__().split()))
+
 logger.log("system", "exit", str(is_correct))

utils.py

@@ -71,7 +71,7 @@ sys_str = "[bold red]SYSTEM: [/bold red]"
 def system_info(text):
     print(sys_str + text)
 
-def print_anpl(anpl, for_user=False):
+def print_anpl(anpl, for_user):
     print(Syntax(anpl.to_python(for_user=for_user), "python"))
 
 def print_text_IOExamples(ios: list[IOExample]):