Update app.py

app.py

@@ -1,68 +1,125 @@
-from context import query
-from openai import AsyncOpenAI
+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
 
-# taken from fleet context
-SYSTEM_PROMPT = """
-You are an expert in Python libraries. You carefully provide accurate, factual, thoughtful, nuanced answers, and are brilliant at reasoning. If you think there might not be a correct answer, you say so.
-Each token you produce is another opportunity to use computation, therefore you always spend a few sentences explaining background context, assumptions, and step-by-step thinking BEFORE you try to answer a question.
-Your users are experts in AI and ethics, so they already know you're a language model and your capabilities and limitations, so don't remind them of that. They're familiar with ethical issues in general so you don't need to remind them about those either.
-Your users are also in a CLI environment. You are capable of writing and running code. DO NOT write hypothetical code. ALWAYS write real code that will execute and run end-to-end.
-
-Instructions:
-- Be objective, direct. Include literal information from the context, don't add any conclusion or subjective information.
-- When writing code, ALWAYS have some sort of output (like a print statement). If you're writing a function, call it at the end. Do not generate the output, because the user can run it themselves.
-- ALWAYS cite your sources. Context will be given to you after the text ### Context source_url ### with source_url being the url to the file. For example, ### Context https://example.com/docs/api.html#files ### will have a source_url of https://example.com/docs/api.html#files.
-- When you cite your source, please cite it as [num] with `num` starting at 1 and incrementing with each source cited (1, 2, 3, ...). At the bottom, have a newline-separated `num: source_url` at the end of the response. ALWAYS add a new line between sources or else the user won't be able to read it. DO NOT convert links into markdown, EVER! If you do that, the user will not be able to click on the links.
-For example:
-**Context 1**: https://example.com/docs/api.html#pdfs
-I'm a big fan of PDFs.
-**Context 2**: https://example.com/docs/api.html#csvs
-I'm a big fan of CSVs.
-### Prompt ###
-What is this person a big fan of?
-### Response ###
-This person is a big fan of PDFs[1] and CSVs[2].
-1: https://example.com/docs/api.html#pdfs
-2: https://example.com/docs/api.html#csvs
-"""
-
-pn.extension()
-
-MODEL = "gpt-3.5-turbo"
-
-
-async def answer(contents, user, instance):
-    # start with system prompt
-    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
-
-    # add context to the user input
-    context = ""
-    fleet_responses = query(contents, k=3)
-    for i, response in enumerate(fleet_responses):
-        context += f"\n\n**Context {i}**: {response['metadata']['url']}\n{response['metadata']['text']}"
-    instance.send(context, avatar="🛩️", user="Fleet Context", respond=False)
-
-    # get history of messages (skipping the intro message)
-    # and serialize fleet context messages as "user" role
-    messages.extend(
-        instance.serialize(role_names={"user": ["user", "Fleet Context"]})[1:]
-    )
+pn.extension("codeeditor", sizing_mode="stretch_width")
+
+LLM_MODEL = "mistral-medium"
+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 (```). You must have `fig` "
+        "as the last line of code"
+    ),
+)
+
+USER_CONTENT_FORMAT = """
+Request:
+{content}
+
+Code:
+```python
+{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")
 
-    openai_response = await client.chat.completions.create(
-        model=MODEL, messages=messages, temperature=0.2, stream=True
+img = ax.scatter(x, c, c=c, label="log")
+plt.colorbar(img, label="Colorbar")
+plt.legend()
+
+# must have fig at the end!
+fig
+""".strip()
+
+
+async def callback(content: str, user: str, instance: pn.chat.ChatInterface):
+    # system
+    messages = [SYSTEM_MESSAGE]
+
+    # history
+    messages.extend([ChatMessage(**message) for message in instance.serialize()[1:-1]])
+
+    # 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 openai_response:
-        token = chunk.choices[0].delta.content
-        if token:
-            message += token
+    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
 
+    # 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
+    matplotlib_pane.object = exec_with_return(llm_code)
+
+
+# instantiate widgets and panes
+client = MistralAsyncClient(api_key=os.environ["MISTRAL_API_KEY"])
+chat_interface = pn.chat.ChatInterface(
+    callback=callback,
+    show_clear=False,
+    show_undo=False,
+    show_button_name=False,
+    message_params=dict(
+        show_reaction_icons=False,
+        show_copy_icon=False,
+    ),
+    height=700,
+    callback_exception="verbose",
+)
+matplotlib_pane = pn.pane.Matplotlib(
+    exec_with_return(DEFAULT_MATPLOTLIB),
+    sizing_mode="stretch_both",
+)
+code_editor = pn.widgets.CodeEditor(
+    value=DEFAULT_MATPLOTLIB,
+    sizing_mode="stretch_both",
+)
+
+# lay them out
+tabs = pn.Tabs(
+    ("Plot", matplotlib_pane),
+    ("Code", code_editor),
+)
 
-client = AsyncOpenAI()
-intro_message = pn.chat.ChatMessage("Ask me anything about Python libraries!", user="System")
-chat_interface = pn.chat.ChatInterface(intro_message, callback=answer, callback_user="OpenAI")
-template = pn.template.FastListTemplate(main=[chat_interface], title="Panel UI of Fleet Context 🛩️")
-template.servable()
+sidebar = [chat_interface]
+main = [tabs]
+template = pn.template.FastListTemplate(
+    sidebar=sidebar,
+    main=main,
+    sidebar_width=600,
+    main_layout=None,
+    accent_base_color="#fd7000",
+    header_background="#fd7000",
+)
+template.servable()