DuoChat

.gitignore

@@ -0,0 +1 @@
+__pycache__

claude.py

@@ -0,0 +1,25 @@
+import anthropic
+
+from data_model import ChatMessage
+
+client = anthropic.Anthropic(
+  # defaults to os.environ.get("ANTHROPIC_API_KEY")
+  #   api_key="my_api_key",
+)
+
+
+def call_claude_sonnet(input: str, history: list[ChatMessage]):
+  messages = [
+    {
+      "role": "assistant" if message.role == "model" else message.role,
+      "content": message.content,
+    }
+    for message in history
+  ] + [{"role": "user", "content": input}]
+  with client.messages.stream(
+    max_tokens=1024,
+    messages=messages,
+    model="claude-3-5-sonnet-20240620",
+  ) as stream:
+    for text in stream.text_stream:
+      yield text

data_model.py

@@ -0,0 +1,37 @@
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Literal
+
+import mesop as me
+
+
+Role = Literal["user", "model"]
+
+
+@dataclass(kw_only=True)
+class ChatMessage:
+  """Chat message metadata."""
+
+  role: Role = "user"
+  content: str = ""
+  in_progress: bool = False
+
+class Models(Enum):
+  GEMINI_1_5_FLASH = "Gemini 1.5 Flash"
+  GEMINI_1_5_PRO = "Gemini 1.5 Pro"
+  CLAUDE_3_5_SONNET = "Claude 3.5 Sonnet"
+
+@dataclass
+class Conversation:
+  model: str = ""
+  messages: list[ChatMessage] = field(default_factory=list)
+
+@me.stateclass
+class State:
+  is_model_picker_dialog_open: bool
+  input: str
+  conversations: list[Conversation]
+  models: list[str]
+  gemini_api_key: str
+  claude_api_key: str
+

dialog.py

@@ -0,0 +1,65 @@
+import mesop as me
+
+
+@me.content_component
+def dialog(is_open: bool):
+  """Renders a dialog component.
+
+  The design of the dialog borrows from the Angular component dialog. So basically
+  rounded corners and some box shadow.
+
+  One current drawback is that it's not possible to close the dialog
+  by clicking on the overlay background. This is due to
+  https://github.com/google/mesop/issues/268.
+
+  Args:
+    is_open: Whether the dialog is visible or not.
+  """
+  with me.box(
+    style=me.Style(
+      background="rgba(0,0,0,0.4)",
+      display="block" if is_open else "none",
+      height="100%",
+      overflow_x="auto",
+      overflow_y="auto",
+      position="fixed",
+      width="100%",
+      z_index=1000,
+    )
+  ):
+    with me.box(
+      style=me.Style(
+        align_items="center",
+        display="grid",
+        height="100vh",
+        justify_items="center",
+      )
+    ):
+      with me.box(
+        style=me.Style(
+          background="#fff",
+          border_radius=20,
+          box_sizing="content-box",
+          box_shadow=(
+            "0 3px 1px -2px #0003, 0 2px 2px #00000024, 0 1px 5px #0000001f"
+          ),
+          margin=me.Margin.symmetric(vertical="0", horizontal="auto"),
+          padding=me.Padding.all(20),
+        )
+      ):
+        me.slot()
+
+
+@me.content_component
+def dialog_actions():
+  """Helper component for rendering action buttons so they are right aligned.
+
+  This component is optional. If you want to position action buttons differently,
+  you can just write your own Mesop markup.
+  """
+  with me.box(
+    style=me.Style(
+      display="flex", justify_content="end", margin=me.Margin(top=20)
+    )
+  ):
+    me.slot()

duo_chat.py

@@ -0,0 +1,381 @@
+import mesop as me
+
+
+import claude
+import gemini
+from data_model import ChatMessage, State, Models, Conversation
+from dialog import dialog, dialog_actions
+
+
+ROOT_BOX_STYLE = me.Style(
+    background="#e7f2ff",
+    height="100%",
+    font_family="Inter",
+    display="flex",
+    flex_direction="column",
+)
+
+darker_bg_color = "#b9e1ff"
+
+STYLESHEETS = [
+    "https://fonts.googleapis.com/css2?family=Inter:wght@100..900&display=swap"
+]
+
+
+@me.stateclass
+class ModelDialogState:
+    selected_models: list[str]
+
+
+def change_model_option(e: me.CheckboxChangeEvent):
+    s = me.state(ModelDialogState)
+    if e.checked:
+        s.selected_models.append(e.key)
+    else:
+        s.selected_models.remove(e.key)
+
+
+def set_gemini_api_key(e: me.InputBlurEvent):
+    me.state(State).gemini_api_key = e.value
+
+
+def set_claude_api_key(e: me.InputBlurEvent):
+    me.state(State).claude_api_key = e.value
+
+
+def model_picker_dialog():
+    state = me.state(State)
+    with dialog(state.is_model_picker_dialog_open):
+        with me.box(style=me.Style(display="flex", flex_direction="column", gap=12)):
+            me.text("API keys")
+            me.input(
+                label="Gemini API Key",
+                value=state.gemini_api_key,
+                on_blur=set_gemini_api_key,
+            )
+            me.input(
+                label="Claude API Key",
+                value=state.claude_api_key,
+                on_blur=set_claude_api_key,
+            )
+        me.text("Pick a model")
+        for model in Models:
+            if model.name.startswith("GEMINI"):
+                disabled = not state.gemini_api_key
+            elif model.name.startswith("CLAUDE"):
+                disabled = not state.claude_api_key
+            else:
+                disabled = False
+            me.checkbox(
+                key=model.value,
+                label=model.value,
+                checked=model.value in state.models,
+                disabled=disabled,
+                on_change=change_model_option,
+                style=me.Style(
+                    display="flex",
+                    flex_direction="column",
+                    gap=4,
+                    padding=me.Padding(top=12),
+                ),
+            )
+        with dialog_actions():
+            me.button("Cancel", on_click=close_model_picker_dialog)
+            me.button("Confirm", on_click=confirm_model_picker_dialog)
+
+
+@me.page(
+    path="/",
+    stylesheets=STYLESHEETS,
+    security_policy=me.SecurityPolicy(
+        allowed_iframe_parents=["https://huggingface.co"]
+    ),
+)
+def page():
+    model_picker_dialog()
+    with me.box(style=ROOT_BOX_STYLE):
+        header()
+        with me.box(
+            style=me.Style(
+                width="min(680px, 100%)",
+                margin=me.Margin.symmetric(
+                    horizontal="auto",
+                    vertical=36,
+                ),
+            )
+        ):
+            me.text(
+                "Chat with multiple models at once",
+                style=me.Style(
+                    font_size=20,
+                    margin=me.Margin(
+                        bottom=24,
+                    ),
+                ),
+            )
+            examples_row()
+            chat_input()
+
+
+EXAMPLES = [
+    "Create a file-lock in Python",
+    "Write an email to Congress to have free milk for all",
+    "Make a nice box shadow in CSS",
+]
+
+
+def examples_row():
+    with me.box(
+        style=me.Style(
+            display="flex", flex_direction="row", gap=16, margin=me.Margin(bottom=24)
+        )
+    ):
+        for i in EXAMPLES:
+            example(i)
+
+
+def example(text: str):
+    with me.box(
+        key=text,
+        on_click=click_example,
+        style=me.Style(
+            cursor="pointer",
+            background=darker_bg_color,
+            width="215px",
+            height=160,
+            font_weight=500,
+            line_height="1.5",
+            padding=me.Padding.all(16),
+            border_radius=16,
+            border=me.Border.all(me.BorderSide(width=1, color="blue", style="none")),
+        ),
+    ):
+        me.text(text)
+
+
+def click_example(e: me.ClickEvent):
+    state = me.state(State)
+    state.input = e.key
+
+
+def header():
+    def navigate_home(e: me.ClickEvent):
+        me.navigate("/")
+        state = me.state(State)
+        state.conversations = []
+
+    with me.box(
+        on_click=navigate_home,
+        style=me.Style(
+            cursor="pointer",
+            padding=me.Padding.all(16),
+        ),
+    ):
+        me.text(
+            "DuoChat",
+            style=me.Style(
+                font_weight=500,
+                font_size=24,
+                color="#3D3929",
+                letter_spacing="0.3px",
+            ),
+        )
+
+
+def chat_input():
+    state = me.state(State)
+
+    with me.box(
+        style=me.Style(
+            border_radius=16,
+            padding=me.Padding.all(8),
+            background="white",
+            display="flex",
+            width="100%",
+        )
+    ):
+        with me.box(
+            style=me.Style(
+                flex_grow=1,
+            )
+        ):
+            with me.box():
+                me.native_textarea(
+                    value=state.input,
+                    autosize=True,
+                    min_rows=4,
+                    max_rows=8,
+                    placeholder="Enter a prompt",
+                    on_blur=on_blur,
+                    style=me.Style(
+                        padding=me.Padding(top=16, left=16),
+                        outline="none",
+                        width="100%",
+                        overflow_y="auto",
+                        border=me.Border.all(
+                            me.BorderSide(style="none"),
+                        ),
+                    ),
+                )
+                with me.box(
+                    style=me.Style(
+                        display="flex",
+                        padding=me.Padding(left=12, bottom=12),
+                        cursor="pointer",
+                    ),
+                    on_click=switch_model,
+                ):
+                    me.text(
+                        "Model:",
+                        style=me.Style(font_weight=500, padding=me.Padding(right=6)),
+                    )
+                    if state.models:
+                        me.text(", ".join(state.models))
+                    else:
+                        me.text("(no model selected)")
+        with me.content_button(
+            type="icon", on_click=send_prompt, disabled=not state.models
+        ):
+            me.icon("send")
+
+
+def switch_model(e: me.ClickEvent):
+    state = me.state(State)
+    state.is_model_picker_dialog_open = True
+    dialog_state = me.state(ModelDialogState)
+    dialog_state.selected_models = state.models[:]
+
+
+def close_model_picker_dialog(e: me.ClickEvent):
+    state = me.state(State)
+    state.is_model_picker_dialog_open = False
+
+
+def confirm_model_picker_dialog(e: me.ClickEvent):
+    dialog_state = me.state(ModelDialogState)
+    state = me.state(State)
+    state.is_model_picker_dialog_open = False
+    state.models = dialog_state.selected_models
+
+
+def on_blur(e: me.InputBlurEvent):
+    state = me.state(State)
+    state.input = e.value
+
+
+def send_prompt(e: me.ClickEvent):
+    state = me.state(State)
+    if not state.conversations:
+        me.navigate("/conversation")
+        for model in state.models:
+            state.conversations.append(Conversation(model=model, messages=[]))
+    input = state.input
+    state.input = ""
+
+    for conversation in state.conversations:
+        model = conversation.model
+        messages = conversation.messages
+        history = messages[:]
+        messages.append(ChatMessage(role="user", content=input))
+        messages.append(ChatMessage(role="model", in_progress=True))
+        yield
+        me.scroll_into_view(key="end_of_messages")
+        if model == Models.GEMINI_1_5_FLASH.value:
+            llm_response = gemini.send_prompt_flash(input, history)
+        elif model == Models.GEMINI_1_5_PRO.value:
+            llm_response = gemini.send_prompt_pro(input, history)
+        elif model == Models.CLAUDE_3_5_SONNET.value:
+            llm_response = claude.call_claude_sonnet(input, history)  # type: ignore
+        else:
+            raise Exception("Unhandled model", model)
+        for chunk in llm_response:  # type: ignore
+            messages[-1].content += chunk
+            yield
+        messages[-1].in_progress = False
+        yield
+
+
+@me.page(path="/conversation", stylesheets=STYLESHEETS)
+def conversation_page():
+    state = me.state(State)
+    model_picker_dialog()
+    with me.box(style=ROOT_BOX_STYLE):
+        header()
+
+        models = len(state.conversations)
+        models_px = models * 680
+        with me.box(
+            style=me.Style(
+                width=f"min({models_px}px, calc(100% - 32px))",
+                display="grid",
+                gap=16,
+                grid_template_columns=f"repeat({models}, 1fr)",
+                flex_grow=1,
+                overflow_y="hidden",
+                margin=me.Margin.symmetric(horizontal="auto"),
+                padding=me.Padding.symmetric(horizontal=16),
+            )
+        ):
+            for conversation in state.conversations:
+                model = conversation.model
+                messages = conversation.messages
+                with me.box(
+                    style=me.Style(
+                        overflow_y="auto",
+                    )
+                ):
+                    me.text("Model: " + model, style=me.Style(font_weight=500))
+
+                    for message in messages:
+                        if message.role == "user":
+                            user_message(message.content)
+                        else:
+                            model_message(message)
+                    if messages and model == state.conversations[-1]:
+                        me.box(
+                            key="end_of_messages",
+                            style=me.Style(
+                                margin=me.Margin(
+                                    bottom="50vh" if messages[-1].in_progress else 0
+                                )
+                            ),
+                        )
+        with me.box(
+            style=me.Style(
+                display="flex",
+                justify_content="center",
+            )
+        ):
+            with me.box(
+                style=me.Style(
+                    width="min(680px, 100%)",
+                    padding=me.Padding(top=24, bottom=24),
+                )
+            ):
+                chat_input()
+
+
+def user_message(content: str):
+    with me.box(
+        style=me.Style(
+            background=darker_bg_color,
+            padding=me.Padding.all(16),
+            margin=me.Margin.symmetric(vertical=16),
+            border_radius=16,
+        )
+    ):
+        me.text(content)
+
+
+def model_message(message: ChatMessage):
+    with me.box(
+        style=me.Style(
+            background="#fff",
+            padding=me.Padding.all(16),
+            border_radius=16,
+            margin=me.Margin.symmetric(vertical=16),
+        )
+    ):
+        me.markdown(message.content)
+        if message.in_progress:
+            me.progress_spinner()

gemini.py

@@ -0,0 +1,71 @@
+import os
+
+import google.generativeai as genai
+import mesop as me
+from typing_extensions import TypedDict
+
+from data_model import ChatMessage, State
+
+
+
+
+class Pad(TypedDict):
+  title: str
+  content: str
+
+
+class Output(TypedDict):
+  intro: str
+  pad: Pad
+  conclusion: str
+
+
+# Create the model
+# See https://ai.google.dev/api/python/google/generativeai/GenerativeModel
+generation_config = {
+  "temperature": 1,
+  "top_p": 0.95,
+  "top_k": 64,
+  "max_output_tokens": 8192,
+  "response_mime_type": "text/plain",
+  # "response_schema": Output,
+}
+
+def configure_gemini():
+  state = me.state(State)
+  genai.configure(api_key=state.gemini_api_key)
+
+def send_prompt_pro(prompt: str, history: list[ChatMessage]):
+  configure_gemini()
+
+  model = genai.GenerativeModel(
+    model_name="gemini-1.5-pro-latest",
+    generation_config=generation_config,
+  )
+
+  chat_session = model.start_chat(
+    history=[
+      {"role": message.role, "parts": [message.content]} for message in history
+    ]
+  )
+
+  for chunk in chat_session.send_message(prompt, stream=True):
+    yield chunk.text
+
+
+def send_prompt_flash(prompt: str, history: list[ChatMessage]):
+  configure_gemini()
+
+  model = genai.GenerativeModel(
+    model_name="gemini-1.5-flash-latest",
+    generation_config=generation_config,
+  )
+
+  chat_session = model.start_chat(
+    history=[
+      {"role": message.role, "parts": [message.content]} for message in history
+    ]
+  )
+
+  for chunk in chat_session.send_message(prompt, stream=True):
+    yield chunk.text

requirements.txt

@@ -0,0 +1,4 @@
+mesop
+gunicorn
+anthropic
+google-generativeai