Update app.py

app.py

@@ -1,68 +1,182 @@
-from context import query
-from openai import AsyncOpenAI
+import asyncio
+import random
+import sqlite3
+
+import os
 import panel as pn
+import pandas as pd
+from litellm import acompletion
+
+pn.extension("perspective")
+
+
+MODELS = [
+    "mistral/mistral-tiny",
+    "mistral/mistral-small",
+    "mistral/mistral-medium",
+    "mistral/mistral-large-latest",
+]
+
+VOTING_LABELS = [
+    "👈 A is better",
+    "🤗 About the same",
+    "😓 Both not good",
+    "👉 B is better",
+]
+
+
+def set_api_key(api_key):
+    os.environ["MISTRAL_API_KEY"] = api_key
+
+
+async def respond(content, user, instance):
+    """
+    Respond to the user in the chat interface.
+    """
+    try:
+        instance.disabled = True
+        chat_label = instance.name
+        if chat_model := chat_models.get(chat_label):
+            model = chat_model
+        else:
+            # remove past history up to new message
+            instance.objects = instance.objects[-1:]
+            header_a.object = f"## Model: A"
+            header_b.object = f"## Model: B"
+            model = chat_models[chat_label] = random.choice(MODELS)
+
+        messages = instance.serialize()
+        messages.append({"role": "user", "content": content})
+
+        response = await acompletion(
+            model=model, messages=messages, stream=True, max_tokens=128
+        )
+
+        message = None
+        async for chunk in response:
+            if not chunk.choices[0].delta["content"]:
+                continue
+            message = instance.stream(
+                chunk.choices[0].delta["content"], user="Assistant", message=message
+            )
+    finally:
+        instance.disabled = False
+
 
-# 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:]
+async def forward_message(content, user, instance):
+    """
+    Send the message to the other chat interface and respond to the user in both.
+    """
+    if instance is chat_interface_a:
+        other_instance = chat_interface_b
+    else:
+        other_instance = chat_interface_a
+    other_instance.append(pn.chat.ChatMessage(content, user=user))
+
+    coroutines = [
+        respond(content, user, chat_interface)
+        for chat_interface in (chat_interface_a, chat_interface_b)
+    ]
+    await asyncio.gather(*coroutines)
+
+
+def click_vote(event):
+    """
+    Count the votes and update the voting results.
+    """
+    if len(chat_models) == 0:
+        return
+
+    voting_label = event.obj.name
+    if voting_label == VOTING_LABELS[0]:
+        chat_model = chat_models[chat_interface_a.name]
+        voting_counts[chat_model] = voting_counts.get(chat_model, 0) + 1
+    elif voting_label == VOTING_LABELS[3]:
+        chat_model = chat_models[chat_interface_b.name]
+        voting_counts[chat_model] = voting_counts.get(chat_model, 0) + 1
+    elif voting_label == VOTING_LABELS[1]:
+        chat_model_a = chat_models[chat_interface_a.name]
+        chat_model_b = chat_models[chat_interface_b.name]
+        if chat_model_a == chat_model_b:
+            voting_counts[chat_model_a] = voting_counts.get(chat_model_a, 0) + 1
+        else:
+            voting_counts[chat_model_a] = voting_counts.get(chat_model_a, 0) + 1
+            voting_counts[chat_model_b] = voting_counts.get(chat_model_b, 0) + 1
+
+    header_a.object = f"## Model: {chat_models[chat_interface_a.name]}"
+    header_b.object = f"## Model: {chat_models[chat_interface_b.name]}"
+    for chat_label in set(chat_models.keys()):
+        chat_models.pop(chat_label)
+
+    perspective.object = (
+        pd.DataFrame(voting_counts, index=["Votes"])
+        .melt(var_name="Model", value_name="Votes")
+        .set_index("Model")
     )
+    with sqlite3.connect("voting_counts.db") as conn:
+        pd.DataFrame(voting_counts.items(), columns=["Model", "Votes"]).to_sql(
+            "voting_counts", conn, if_exists="replace", index=False
+        )
 
-    openai_response = await client.chat.completions.create(
-        model=MODEL, messages=messages, temperature=0.2, stream=True
+
+# initialize
+chat_models = {}
+with sqlite3.connect("voting_counts.db") as conn:
+    conn.execute(
+        "CREATE TABLE IF NOT EXISTS voting_counts (Model TEXT PRIMARY KEY, Votes INTEGER)"
+    )
+    voting_counts = (
+        pd.read_sql("SELECT * FROM voting_counts", conn)
+        .set_index("Model")["Votes"]
+        .to_dict()
     )
 
-    message = ""
-    async for chunk in openai_response:
-        token = chunk.choices[0].delta.content
-        if token:
-            message += token
-            yield message
+# header
+api_key_input = pn.widgets.PasswordInput(placeholder="Mistral API Key")
+pn.bind(set_api_key, api_key_input)
+
+# main
+tabs = pn.Tabs()
+
+# tab 1
+chat_interface_kwargs = dict(
+    callback=forward_message,
+    show_undo=False,
+    show_rerun=False,
+    show_clear=False,
+    show_stop=False,
+    show_button_name=False,
+    callback_exception="verbose",
+)
+header_a = pn.pane.Markdown("## Model: A")
+chat_interface_a = pn.chat.ChatInterface(
+    name="A", header=header_a, **chat_interface_kwargs
+)
+header_b = pn.pane.Markdown("## Model: B")
+chat_interface_b = pn.chat.ChatInterface(
+    name="B", header=header_b, **chat_interface_kwargs
+)
+
+button_kwargs = dict(sizing_mode="stretch_width")
+button_row = pn.Row()
+for voting_label in VOTING_LABELS:
+    button = pn.widgets.Button(name=voting_label, **button_kwargs)
+    button.on_click(click_vote)
+    button_row.append(button)
+
+tabs.append(("Chat", pn.Column(pn.Row(chat_interface_a, chat_interface_b), button_row)))
 
+# tab 2
+perspective = pn.pane.Perspective(
+    pd.DataFrame(voting_counts, index=["Votes"])
+    .melt(var_name="Model", value_name="Votes")
+    .set_index("Model"),
+    sizing_mode="stretch_both",
+    editable=False,
+)
+tabs.append(("Voting Results", perspective))
 
-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()
+# layout
+pn.template.FastListTemplate(
+    title="Mistral Chat Arena", header=[api_key_input], main=[tabs]
+).servable()