From 9de79a6a5a4081b1deb6855580f660ed99ccf974 Mon Sep 17 00:00:00 2001
From: Andrew <15331990+ahuang11@users.noreply.github.com>
Date: Mon, 29 Jul 2024 02:44:19 -0700
Subject: [PATCH] Wrap in iframe (#625)

---
 lumen/ai/agents.py    | 18 ++++++++++++------
 lumen/ai/assistant.py | 24 +++++++++++++-----------
 lumen/views/base.py   |  9 ++++++++-
 3 files changed, 33 insertions(+), 18 deletions(-)

diff --git a/lumen/ai/agents.py b/lumen/ai/agents.py
index 770f44e3..49206a96 100644
--- a/lumen/ai/agents.py
+++ b/lumen/ai/agents.py
@@ -976,6 +976,7 @@ async def answer(self, messages: list | str):
         if not analyses:
             print("NONE found...")
             return None
+
         with self.interface.add_step(title="Choosing the most relevant analysis...") as step:
             if len(analyses) > 1:
                 type_ = Literal[tuple(analyses)]
@@ -984,20 +985,25 @@ async def answer(self, messages: list | str):
                     correct_name=(type_, FieldInfo(description="The name of the analysis that is most appropriate given the user query."))
                 )
                 system_prompt = await self._system_prompt_with_context(messages, analyses)
-                analysis = (await self.llm.invoke(
+                analysis_name = (await self.llm.invoke(
                     messages,
                     system=system_prompt,
                     response_model=analysis_model,
                     allow_partial=False,
                 )).correct_name
             else:
-                analysis = list(analyses)[0]
-            step.stream(f"Selected {analysis}")
-            step.success_title = f"Selected {analysis}"
+                analysis_name = list(analyses)[0]
+            step.stream(f"Selected {analysis_name}")
+            step.success_title = f"Selected {analysis_name}"
 
         with self.interface.add_step(title="Creating view...") as step:
-            print(f"Creating view for {analysis}")
-            view = analyses[analysis](pipeline)
+            print(f"Creating view for {analysis_name}")
+            await asyncio.sleep(0.1)  # necessary to give it time to render before calling sync function...
+            analysis_callable = analyses[analysis_name]
+            if asyncio.iscoroutinefunction(analysis_callable):
+                view = await analysis_callable(pipeline)
+            else:
+                view = await asyncio.to_thread(analysis_callable, pipeline)
             spec = view.to_spec()
             step.stream(f"Generated view\n```json\n{spec}\n```")
             step.success_title = "Generated view"
diff --git a/lumen/ai/assistant.py b/lumen/ai/assistant.py
index a6c502f1..84675e03 100644
--- a/lumen/ai/assistant.py
+++ b/lumen/ai/assistant.py
@@ -178,18 +178,20 @@ async def hide_suggestions(_=None):
                 suggestion_buttons.visible = False
 
         async def use_suggestion(event):
-            contents = event.obj.name
-            if hide_after_use:
-                await hide_suggestions()
-            if analysis:
-                for agent in self.agents:
-                    if isinstance(agent, AnalysisAgent):
-                        break
+            button = event.obj
+            with button.param.update(loading=True), self.interface.active_widget.param.update(loading=True):
+                contents = button.name
+                if hide_after_use:
+                    await hide_suggestions()
+                if analysis:
+                    for agent in self.agents:
+                        if isinstance(agent, AnalysisAgent):
+                            break
+                    else:
+                        return
+                    await agent.invoke([{'role': 'user', 'content': contents}])
                 else:
-                    return
-                await agent.invoke([{'role': 'user', 'content': contents}])
-            else:
-                self.interface.send(contents)
+                    self.interface.send(contents)
 
         async def run_demo(event):
             if hide_after_use:
diff --git a/lumen/views/base.py b/lumen/views/base.py
index 679f11da..87e4fc41 100644
--- a/lumen/views/base.py
+++ b/lumen/views/base.py
@@ -4,6 +4,7 @@
 """
 from __future__ import annotations
 
+import html
 import sys
 
 from io import BytesIO, StringIO
@@ -1144,7 +1145,13 @@ def _get_params(self) -> Dict[str, Any]:
 
     def get_panel(self) -> pn.pane.HTML:
         from ydata_profiling import ProfileReport
-        return self._panel_type(ProfileReport(**self._get_params()).html)
+        report_html = ProfileReport(**self._get_params()).html
+        escaped_html = html.escape(report_html)
+        iframe = f"""
+        <iframe srcdoc="{escaped_html}" width="100%" height="100%" frameborder="0" marginheight="0" marginwidth="0">
+        </iframe>
+        """
+        return self._panel_type(iframe, min_height=700, sizing_mode="stretch_both")
 
 
 __all__ = [name for name, obj in locals().items() if isinstance(obj, type) and issubclass(obj, View)] + ["Download"]