Update app.py

app.py

@@ -1,6 +1,7 @@
 import re
 import os
 import panel as pn
+from io import StringIO
 from panel.io.mime_render import exec_with_return
 from llama_index import (
     VectorStoreIndex,
@@ -13,20 +14,16 @@ from llama_index.chat_engine import ContextChatEngine
 from llama_index.embeddings import OpenAIEmbedding
 from llama_index.llms import OpenAI
 
-pn.extension("codeeditor", sizing_mode="stretch_width")
 
 SYSTEM_PROMPT = (
-    "You are a renowned data visualization expert "
-    "with a strong background in hvplot and holoviews. "
-    "Note, hvplot is built on top of holoviews; so "
-    "anything you can do with holoviews, you can do "
-    "with hvplot, but prioritize hvplot kwargs "
-    "first as its simpler. Your primary goal is "
-    "to assist the user in edit the code based on user request "
-    "using best practices. Simply provide code "
-    "in code fences (```python). You absolutely "
-    "must have `hvplot_obj` as the last line of code. FYI,"
-    "Data columns: ['sepal_length', 'sepal_width', 'petal_length', 'petal_width', 'species']"
+    "You are a data visualization pro and expert in HoloViz hvplot + holoviews. "
+    "Your primary goal is to assist the user in editing based on user requests using best practices. "
+    "Simply provide code in code fences (```python). You must have `hvplot_obj` as the last line of code. "
+    "Note, data columns are ['sepal_length', 'sepal_width', 'petal_length', 'petal_width', 'species'] and "
+    "hvplot is built on top of holoviews--anything you can do with holoviews, you can do "
+    "with hvplot. First try to use hvplot **kwargs instead of opts, e.g. `legend='top_right'` "
+    "instead of `opts(legend_position='top_right')`. If you need to use opts, you can use "
+    "concise version, e.g. `opts(xlabel='Petal Length')` vs `opts(hv.Opts(xlabel='Petal Length'))`"
 )
 
 USER_CONTENT_FORMAT = """
@@ -48,6 +45,14 @@ hvplot_obj
 """.strip()
 
 
+def exception_handler(exc):
+    if retries.value == 0:
+        chat_interface.send(f"Can't figure this out: {exc}", respond=False)
+        return
+    chat_interface.send(f"Fix this error:\n```python\n{exc}\n```")
+    retries.value = retries.value - 1
+
+
 def init_llm(event):
     api_key = event.new
     if not api_key:
@@ -110,13 +115,18 @@ def callback(content: str, user: str, instance: pn.chat.ChatInterface):
         if llm_code.splitlines()[-1].strip() != "hvplot_obj":
             llm_code += "\nhvplot_obj"
         code_editor.value = llm_code
+    retries.value = 2
 
 
 def update_plot(event):
-    try:
-        hvplot_pane.object = exec_with_return(event.new)
-    except Exception as exc:
-        chat_interface.send(f"Fix this error: {exc}")
+    with StringIO() as buf:
+        hvplot_pane.object = exec_with_return(event.new, stderr=buf)
+        buf.seek(0)
+        errors = buf.read()
+        if errors:
+            exception_handler(errors)
+
+pn.extension("codeeditor", sizing_mode="stretch_width", exception_handler=exception_handler)
 
 
 # instantiate widgets and panes
@@ -147,6 +157,8 @@ code_editor = pn.widgets.CodeEditor(
     language="python",
     sizing_mode="stretch_both",
 )
+retries = pn.widgets.IntInput(value=2, visible=False)
+error = pn.widgets.StaticText(visible=False)
 
 # watch for code changes
 api_key_input.param.watch(init_llm, "value")
@@ -170,4 +182,5 @@ template = pn.template.FastListTemplate(
     header_background="#fd7000",
     title="Chat with Plot"
 )
+
 template.servable()