update circuit synthesizer UX

Author: lsarra-ax

src/axiomatic/magic.py

@@ -1,7 +1,10 @@
+from IPython import get_ipython  # type: ignore
+from IPython.core.magic import register_line_cell_magic, register_line_magic  # type: ignore
+from IPython.display import HTML, display  # type: ignore
+
 import platformdirs  # type: ignore
 import os
-
-
+import sys
 from . import Axiomatic
 
 
@@ -18,6 +21,7 @@ def __init__(self):
             with open(f"{self.folder}/api_key", "r") as f:
                 self.api_key = f.read()
             self.ax = Axiomatic(api_key=self.api_key)
+        self.query = ""
 
     def ax_api(self, query):
         folder = platformdirs.user_config_dir("axiomatic")
@@ -27,53 +31,63 @@ def ax_api(self, query):
             self.ax = Axiomatic(api_key=self.api)
             print("API key set.")
 
-    def ax_query(self, query, cell=None):
-        from IPython import get_ipython  # type: ignore
-        from IPython.core.magic import register_line_cell_magic, register_line_magic  # type: ignore
-        from IPython.display import HTML, display  # type: ignore
-
+    def axquery(self, query, cell=None):
         if self.api_key:
             if cell:
                 # REFINE
-                feedback = ""  # TODO: add feedback according to interface
-                result = self.ax.pic.refine(query=query, code=cell, feedback=feedback)
+                try:
+                    exec(cell)
+                    feedback = "Code executed successfully."
+                except Exception as e:
+                    feedback = f"Errors:\n{e}"
+                print(feedback)
+                current_query = f"{self.query}\n{query}"
+                result = self.ax.pic.refine(
+                    query=current_query, code=cell, feedback=feedback
+                )
             else:
                 # GENERATE FROM SCRATCH
+                self.query = query
                 result = self.ax.pic.generate(query=query)
 
             # Process output
-            output = result.thought_text.replace("\n", "<br>")
-            html_content = (
-                "<div style='font-family: Arial, sans-serif; line-height: 1.5;'>"
-            )
-            html_content += (
-                f"<div style='color: #6EB700;'><strong>AX:</strong> {output}</div>"
-            )
+            pre_thought = result.raw_content.split("<thought>")[0]
+            thought = result.thought_text.replace("\n", "<br>")
+            if not thought:
+                output = result.raw_content.replace("\n", "<br>")
+            else:
+                output = pre_thought + "<br>" + thought
+            html_content = f"""<div style='font-family: Arial, sans-serif; line-height: 1.5;'>"
+<div style='color: #6EB700;'><strong>AX:</strong> {output}</div>"""
             display(HTML(html_content))
 
             # Process code
             # remove last three lines (saving file)
             if result.code:
                 code = "\n".join(result.code.split("\n")[:-3] + ["c"])
-                try:
+                if "google.colab" in sys.modules:
                     # When running in colab
                     from google.colab import _frontend  # type: ignore
 
                     _frontend.create_scratch_cell(
-                        f"""# {query}\n{code}""", bottom_pane=True
+                        f"""# {query}\n# %%ax_fix\n{code}""", bottom_pane=True
                     )
-                except Exception as e:
+                else:
                     # When running in jupyter
-                    get_ipython().set_next_input(f"{code}", replace=False)
+                    get_ipython().set_next_input(f"# %%ax_fix\n{code}", replace=False)
 
         else:
             print(
                 "Please set your Axiomatic API key first with the command %ax_api API_KEY and restart the kernel. Request the api key at our Customer Service."
             )
-
+    def ax_query(self, query, cell=None):
+        # Updates the target circuit query
+        self.query=query
+        return self.axquery(query, cell)
+    
     def ax_fix(self, query, cell=None):
-        # Just dummy at the moment
-        return self.ax_query(query, cell)
+        # Runs again without updating the query
+        return self.axquery(query, cell)
 
 
 def ax_help(value: str):