DuoChat Codelab Part 3: Managing state & dialogs¶
In this section, we'll expand our application's state management capabilities and implement a dialog for selecting AI models. We'll use Mesop's state management system and dialog components to create an interactive model selection experience.
Expanding the State Management¶
First, let's create a data_model.py file with a more comprehensive state structure:
data_model.py
from dataclasses import dataclass, field
from typing import Literal
from enum import Enum
import mesop as me
Role = Literal["user", "model"]
@dataclass(kw_only=True)
class ChatMessage:
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 = False
input: str = ""
conversations: list[Conversation] = field(default_factory=list)
models: list[str] = field(default_factory=list)
gemini_api_key: str = ""
claude_api_key: str = ""
@me.stateclass
class ModelDialogState:
selected_models: list[str] = field(default_factory=list)
This expanded state structure allows us to manage multiple conversations, selected models, and API keys.
Implementing the Model Picker Dialog¶
Now, let's implement the model picker dialog in our main.py file. First, we'll create a new file called dialog.py with the following content, which is based on the dialog pattern from the demo gallery:
dialog.py
import mesop as me
@me.content_component
def dialog(is_open: bool):
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():
with me.box(
style=me.Style(
display="flex", justify_content="end", margin=me.Margin(top=20)
)
):
me.slot()
Now, let's update our main.py file to include the model picker dialog. Copy the following code and replace main.py with it:
main.py
# Update the imports:
import mesop as me
from data_model import State, Models, ModelDialogState
from dialog import dialog, dialog_actions
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)
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
ROOT_BOX_STYLE = me.Style(
background="#e7f2ff",
height="100%",
font_family="Inter",
display="flex",
flex_direction="column",
)
@me.page(
path="/",
stylesheets=[
"https://fonts.googleapis.com/css2?family=Inter:wght@100..900&display=swap"
],
)
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)),
)
chat_input()
def header():
with me.box(
style=me.Style(
padding=me.Padding.all(16),
),
):
me.text(
"DuoChat",
style=me.Style(
font_weight=500,
font_size=24,
color="#3D3929",
letter_spacing="0.3px",
),
)
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 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)):
me.native_textarea(
value=state.input,
placeholder="Enter a prompt",
on_blur=on_blur,
style=me.Style(
padding=me.Padding(top=16, left=16),
outline="none",
width="100%",
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 on_blur(e: me.InputBlurEvent):
state = me.state(State)
state.input = e.value
def send_prompt(e: me.ClickEvent):
state = me.state(State)
print(f"Sending prompt: {state.input}")
print(f"Selected models: {state.models}")
state.input = ""
This updated code adds the following features:
- A model picker dialog that allows users to select AI models and enter API keys.
- State management for selected models and API keys.
- A model switcher in the chat input area that opens the model picker dialog.
- Disabling of models based on whether the corresponding API key has been entered.
Running the Updated Application¶
Run the application again with mesop main.py and navigate to http://localhost:32123. You should now see a chat input area with a model switcher. Clicking on the model switcher will open the model picker dialog, allowing you to select models and enter API keys.
Troubleshooting¶
If you're having trouble, compare your code to the solution.
Next Steps¶
In the next section, we'll integrate multiple AI models into our application. We'll set up connections to Gemini and Claude, implement model-specific chat functions, and create a way to interact with multiple models.