Update app.py

app.py

@@ -1,24 +1,32 @@
 import re
 import os
 import panel as pn
-from mistralai.async_client import MistralAsyncClient
-from mistralai.models.chat_completion import ChatMessage
 from panel.io.mime_render import exec_with_return
+from llama_index import (
+    VectorStoreIndex,
+    SimpleDirectoryReader,
+    ServiceContext,
+    StorageContext,
+    load_index_from_storage,
+)
+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")
 
-LLM_MODEL = "mistral-small"
-SYSTEM_MESSAGE = ChatMessage(
-    role="system",
-    content=(
-        "You are a renowned data visualization expert "
-        "with a strong background in matplotlib. "
-        "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 must have `fig` "
-        "as the last line of code"
-    ),
+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']"
 )
 
 USER_CONTENT_FORMAT = """
@@ -31,68 +39,93 @@ Code:
 ```
 """.strip()
 
-DEFAULT_MATPLOTLIB = """
-import numpy as np
-import matplotlib.pyplot as plt
-
-fig = plt.figure()
-ax = plt.axes(title="Plot Title", xlabel="X Label", ylabel="Y Label")
-
-x = np.linspace(1, 10)
-y = np.sin(x)
-z = np.cos(x)
-c = np.log(x)
-
-ax.plot(x, y, c="blue", label="sin")
-ax.plot(x, z, c="orange", label="cos")
-
-img = ax.scatter(x, c, c=c, label="log")
-plt.colorbar(img, label="Colorbar")
-plt.legend()
+DEFAULT_HVPLOT = """
+import hvplot.pandas
+from bokeh.sampledata.iris import flowers
 
-# must have fig at the end!
-fig
+hvplot_obj = flowers.hvplot(x='petal_length', y='petal_width', by='species', kind='scatter')
+hvplot_obj
 """.strip()
 
 
-async def callback(content: str, user: str, instance: pn.chat.ChatInterface):
-    if not api_key_input.value:
-        yield "Please first enter your Mistral API key"
+def init_llm(event):
+    api_key = event.new
+    if not api_key:
+        api_key = os.environ.get("OPENAI_API_KEY")
+    if not api_key:
         return
-    client = MistralAsyncClient(api_key=api_key_input.value)
+    pn.state.cache["llm"] = OpenAI(api_key=api_key)
+
+
+def create_chat_engine(llm):
+    try:
+        storage_context = StorageContext.from_defaults(persist_dir="persisted/")
+        index = load_index_from_storage(storage_context=storage_context)
+    except Exception as exc:
+        embed_model = OpenAIEmbedding()
+        service_context = ServiceContext.from_defaults(llm=llm, embed_model=embed_model)
+        documents = SimpleDirectoryReader(
+            input_dir="hvplot_docs", required_exts=[".md"], recursive=True
+        ).load_data()
+        index = VectorStoreIndex.from_documents(
+            documents, service_context=service_context, show_progress=True
+        )
+        index.storage_context.persist("persisted/")
+
+    retriever = index.as_retriever()
+    chat_engine = ContextChatEngine.from_defaults(
+        system_prompt=SYSTEM_PROMPT,
+        retriever=retriever,
+        verbose=True,
+    )
+    return chat_engine
+
 
-    # system
-    messages = [SYSTEM_MESSAGE]
+def callback(content: str, user: str, instance: pn.chat.ChatInterface):
+    if "llm" not in pn.state.cache:
+        yield "Need to set OpenAI API key first"
+        return
 
-    # history
-    messages.extend([ChatMessage(**message) for message in instance.serialize()[1:-1]])
+    if "engine" not in pn.state.cache:
+        engine = pn.state.cache["engine"] = create_chat_engine(pn.state.cache["llm"])
+    else:
+        engine = pn.state.cache["engine"]
 
     # new user contents
     user_content = USER_CONTENT_FORMAT.format(
         content=content, code=code_editor.value
     )
-    messages.append(ChatMessage(role="user", content=user_content))
 
-    # stream LLM tokens
-    message = ""
-    async for chunk in client.chat_stream(model=LLM_MODEL, messages=messages):
-        if chunk.choices[0].delta.content is not None:
-            message += chunk.choices[0].delta.content
-            yield message
+    # send user content to chat engine
+    agent_response = engine.stream_chat(user_content)
+
+    message = None
+    for chunk in agent_response.response_gen:
+        message = instance.stream(chunk, message=message, user="OpenAI")
 
     # extract code
-    llm_code = re.findall(r"```python\n(.*)\n```", message, re.DOTALL)[0]
-    if llm_code.splitlines()[-1].strip() != "fig":
-        llm_code += "\nfig"
-    code_editor.value = llm_code
+    llm_matches = re.findall(r"```python\n(.*)\n```", message.object, re.DOTALL)
+    if llm_matches:
+        llm_code = llm_matches[0]
+        if llm_code.splitlines()[-1].strip() != "hvplot_obj":
+            llm_code += "\nhvplot_obj"
+        code_editor.value = llm_code
 
 
 def update_plot(event):
-    matplotlib_pane.object = exec_with_return(event.new)
+    try:
+        hvplot_pane.object = exec_with_return(event.new)
+    except Exception as exc:
+        chat_interface.send(f"Fix this error: {exc}")
 
 
 # instantiate widgets and panes
-api_key_input = pn.widgets.PasswordInput(placeholder="Enter your MistralAI API Key")
+api_key_input = pn.widgets.PasswordInput(
+    placeholder=(
+        "Currently subsidized by Andrew, "
+        "but you can also pass your own OpenAI API Key"
+    )
+)
 chat_interface = pn.chat.ChatInterface(
     callback=callback,
     show_clear=False,
@@ -105,23 +138,24 @@ chat_interface = pn.chat.ChatInterface(
     height=650,
     callback_exception="verbose",
 )
-matplotlib_pane = pn.pane.Matplotlib(
-    exec_with_return(DEFAULT_MATPLOTLIB),
+hvplot_pane = pn.pane.HoloViews(
+    exec_with_return(DEFAULT_HVPLOT),
     sizing_mode="stretch_both",
-    tight=True,
 )
 code_editor = pn.widgets.CodeEditor(
-    value=DEFAULT_MATPLOTLIB,
+    value=DEFAULT_HVPLOT,
     language="python",
     sizing_mode="stretch_both",
 )
 
 # watch for code changes
+api_key_input.param.watch(init_llm, "value")
 code_editor.param.watch(update_plot, "value")
+api_key_input.param.trigger("value")
 
 # lay them out
 tabs = pn.Tabs(
-    ("Plot", matplotlib_pane),
+    ("Plot", hvplot_pane),
     ("Code", code_editor),
 )