Improve Eval page

Also heavily refactor code to be a bit more maintainable

.gitignore

@@ -4,3 +4,7 @@ __pycache__
 
 # Mesop
 .env
+
+
+# Prompt Tuner App
+saved_prompts

README.md

@@ -9,5 +9,9 @@ app_port: 8080
 
 # Mesop Prompt Tuner
 
-Prompt tuner UI built using [Mesop](https://google.github.io/mesop/). This is a
-work in progress.
+_This is a work in progress._
+
+This app is built using [Mesop](https://google.github.io/mesop/). The UI/UX is
+heavily inspired by the Anthropic Console Workbench.
+
+You can test out the WIP demo on [Hugging Face Spaces](https://huggingface.co/spaces/richard-to/mesop-prompt-tuner).

components/__init__.py

@@ -7,4 +7,3 @@ from components.header import header as header
 from components.header import header_section as header_section
 from components.sidebar import icon_sidebar as icon_sidebar
 from components.sidebar import icon_menu_item as icon_menu_item
-from components.table import prompt_eval_table as prompt_eval_table

components/sidebar.py

@@ -14,6 +14,7 @@ def icon_sidebar():
     style=me.Style(
       background="#F5F8FC",
       border=me.Border.symmetric(horizontal=me.BorderSide(width=1, style="solid", color="#DEE2E6")),
+      height="100%",
     )
   ):
     me.slot()

components/table.py

@@ -1,131 +0,0 @@
-from typing import Callable
-
-import mesop as me
-
-_NUM_REQUIRED_ROWS = 3
-
-
-@me.component
-def prompt_eval_table(prompt, on_select_rating: Callable | None = None):
-  """Creates a grid table for displaying and comparing different prompt version runs."""
-  # Add a row for each variable
-  num_vars = len(prompt.variables)
-  table_size = num_vars + _NUM_REQUIRED_ROWS
-  with me.box(
-    style=me.Style(
-      border=me.Border.all(me.BorderSide(width=1, style="solid", color="#DEE2E6")),
-      display="grid",
-      grid_template_columns=f"1fr repeat({num_vars}, 20fr) 20fr 1fr"
-      if num_vars
-      else "1fr 20fr 1fr",
-      margin=me.Margin.all(15),
-      overflow_x="scroll",
-    )
-  ):
-    # Render first row. This row only displays the Prompt version.
-    for i in range(table_size):
-      with me.box(
-        style=me.Style(
-          background="#FFF",
-          border=me.Border.all(me.BorderSide(width=1, style="solid", color="#DEE2E6")),
-          color="#000",
-          font_size=15,
-          font_weight="bold",
-          padding=me.Padding.all(10),
-        )
-      ):
-        if i == num_vars + 1:
-          me.text(f"Version {prompt.version}")
-        else:
-          me.text("")
-
-    # Render second row. This row only displays the headers of the table:
-    # variable names, model response, avg rating.
-    header_row = [""] + prompt.variables + ["Model response"] + [""]
-    for i, header_text in enumerate(header_row):
-      with me.box(
-        style=me.Style(
-          background="#FFF",
-          border=me.Border.all(me.BorderSide(width=1, style="solid", color="#DEE2E6")),
-          color="#0063FF" if header_text and header_text != "Model response" else "#444",
-          font_size=13,
-          font_weight="bold",
-          padding=me.Padding.all(10),
-        )
-      ):
-        # Handle the variable header case.
-        if header_text and header_text != "Model response":
-          me.text("{{" + header_text + "}}")
-        elif i == table_size - 1:
-          avg_rating = _calculate_avg_rating_from_prompt(prompt)
-          if avg_rating is not None:
-            with me.tooltip(message="Average rating"):
-              me.text(f"{avg_rating:.2f}", style=me.Style(text_align="center"))
-          else:
-            me.text("")
-        else:
-          me.text(header_text)
-
-    # Render the data rows by going through the prompt responses.
-    for row_index, example in enumerate(prompt.responses):
-      content_row = (
-        [row_index]
-        + [example["variables"][v] for v in prompt.variables]
-        + [example["output"], example.get("rating", "")]
-      )
-      for col_index, content in enumerate(content_row):
-        if col_index == len(content_row) - 1:
-          with me.box(
-            style=me.Style(
-              background="#FFF",
-              border=me.Border.all(me.BorderSide(width=1, style="solid", color="#DEE2E6")),
-              color="#000",
-              padding=me.Padding.all(10),
-            )
-          ):
-            me.select(
-              value=content,
-              options=[
-                me.SelectOption(label="1", value="1"),
-                me.SelectOption(label="2", value="2"),
-                me.SelectOption(label="3", value="3"),
-                me.SelectOption(label="4", value="4"),
-                me.SelectOption(label="5", value="5"),
-              ],
-              on_selection_change=on_select_rating,
-              key=f"rating_{prompt.version}_{row_index}",
-              style=me.Style(width=60),
-            )
-        elif col_index == 0 or not content:
-          with me.box(
-            style=me.Style(
-              background="#FFF",
-              border=me.Border.all(me.BorderSide(width=1, style="solid", color="#DEE2E6")),
-              color="#000",
-              font_size=14,
-              padding=me.Padding.all(10),
-              text_align="center",
-            )
-          ):
-            me.text(str(content))
-        else:
-          with me.box(
-            style=me.Style(
-              background="#FFF",
-              border=me.Border.all(me.BorderSide(width=1, style="solid", color="#DEE2E6")),
-              color="#000",
-              font_size=14,
-              padding=me.Padding.all(10),
-              max_height=300,
-              min_width=300,
-              overflow_y="scroll",
-            )
-          ):
-            me.markdown(content)
-
-
-def _calculate_avg_rating_from_prompt(prompt) -> float | None:
-  ratings = [int(response["rating"]) for response in prompt.responses if response.get("rating")]
-  if ratings:
-    return sum(ratings) / float(len(ratings))
-  return None

constants.py

@@ -0,0 +1,6 @@
+DIALOG_INPUT_WIDTH = 350
+
+MODEL_TEMPERATURE_MAX = 2
+MODEL_TEMPERATURE_MIN = 0
+
+SAVED_PROMPT_DIRECTORY = "saved_prompts"

dialogs/__init__.py

@@ -0,0 +1,7 @@
+from dialogs.update_title import update_title as update_title
+from dialogs.model_settings import model_settings as model_settings
+from dialogs.prompt_variables import prompt_variables as prompt_variables
+from dialogs.prompt_version_history import prompt_version_history as prompt_version_history
+from dialogs.add_comparisons import add_comparisons as add_comparisons
+from dialogs.generate_prompt import generate_prompt as generate_prompt
+from dialogs.load_prompt import load_prompt as load_prompt

dialogs/add_comparisons.py

@@ -0,0 +1,34 @@
+import mesop as me
+
+import components as mex
+import handlers
+from constants import DIALOG_INPUT_WIDTH
+from helpers import parse_variables
+from state import State
+
+
+@me.component
+def add_comparisons():
+  state = me.state(State)
+  with mex.dialog(state.dialog_show_add_comparison):
+    variable_names = set(parse_variables(state.prompt))
+    me.text("Add Comparisons", type="headline-6")
+    me.select(
+      label="Select Versions",
+      multiple=True,
+      options=[
+        me.SelectOption(label=f"v{prompt.version}", value=str(prompt.version))
+        for prompt in state.prompts
+        if prompt.version != state.version and set(prompt.variables) == variable_names
+      ],
+      style=me.Style(width=DIALOG_INPUT_WIDTH),
+      on_selection_change=on_select_comparison,
+    )
+    with mex.dialog_actions():
+      me.button("Close", key="dialog_show_add_comparison", on_click=handlers.on_close_dialog)
+
+
+def on_select_comparison(e: me.SelectSelectionChangeEvent):
+  """Update UI to show the selected prompt version and close the dialog."""
+  state = me.state(State)
+  state.comparisons = list(map(int, e.values))

dialogs/generate_prompt.py

@@ -0,0 +1,39 @@
+import mesop as me
+
+import components as mex
+import handlers
+import llm
+from constants import DIALOG_INPUT_WIDTH
+from helpers import parse_variables
+from state import State
+
+
+@me.component
+def generate_prompt():
+  state = me.state(State)
+  with mex.dialog(state.dialog_show_generate_prompt):
+    me.text("Generate Prompt", type="headline-6")
+    me.textarea(
+      label="Describe your task",
+      value=state.prompt_gen_task_description,
+      on_blur=handlers.on_update_input,
+      key="prompt_gen_task_description",
+      style=me.Style(width=DIALOG_INPUT_WIDTH),
+    )
+    with mex.dialog_actions():
+      me.button("Close", key="dialog_show_generate_prompt", on_click=handlers.on_close_dialog)
+      me.button("Generate", type="flat", on_click=on_click_generate_prompt)
+
+
+def on_click_generate_prompt(e: me.ClickEvent):
+  """Generates an improved prompt based on the given task description and closes dialog."""
+  state = me.state(State)
+  state.prompt = llm.generate_prompt(
+    state.prompt_gen_task_description, state.model, state.model_temperature
+  )
+  variable_names = parse_variables(state.prompt)
+  for variable_name in variable_names:
+    if variable_name not in state.prompt_variables:
+      state.prompt_variables[variable_name] = ""
+
+  state.dialog_show_generate_prompt = False

dialogs/load_prompt.py

@@ -0,0 +1,36 @@
+import json
+
+import mesop as me
+
+import components as mex
+import handlers
+from constants import DIALOG_INPUT_WIDTH
+from state import State
+from state import Prompt
+
+
+@me.component
+def load_prompt():
+  state = me.state(State)
+  with mex.dialog(state.dialog_show_load):
+    with me.box(style=me.Style(width=DIALOG_INPUT_WIDTH)):
+      me.text("Upload saved prompt", type="headline-6")
+      me.uploader(
+        label="Upload",
+        accepted_file_types=["application/json"],
+        type="flat",
+        color="primary",
+        on_upload=on_upload_prompt,
+        style=me.Style(font_weight="bold"),
+      )
+      with mex.dialog_actions():
+        me.button("Close", key="dialog_show_load", on_click=handlers.on_close_dialog)
+
+
+def on_upload_prompt(e: me.UploadEvent):
+  state = me.state(State)
+  data = json.loads(e.file.getvalue())
+  data["prompts"] = [Prompt(**raw_prompt) for raw_prompt in data["prompts"]]
+  for key, value in data.items():
+    setattr(state, key, value)
+  state.dialog_show_load = False

dialogs/model_settings.py

@@ -0,0 +1,68 @@
+import mesop as me
+
+import components as mex
+import handlers
+from constants import DIALOG_INPUT_WIDTH
+from constants import MODEL_TEMPERATURE_MAX
+from constants import MODEL_TEMPERATURE_MIN
+from state import State
+
+
+@me.component
+def model_settings():
+  state = me.state(State)
+  with mex.dialog(state.dialog_show_model_settings):
+    me.text("Model Settings", type="headline-6")
+    with me.box():
+      me.select(
+        label="Model",
+        key="model",
+        options=[
+          me.SelectOption(label="Gemini 1.5 Flash", value="gemini-1.5-flash"),
+          me.SelectOption(label="Gemini 1.5 Pro", value="gemini-1.5-pro"),
+        ],
+        value=state.model,
+        style=me.Style(width=DIALOG_INPUT_WIDTH),
+        on_selection_change=handlers.on_update_selection,
+      )
+    with me.box():
+      me.text("Temperature", style=me.Style(font_weight="bold"))
+      with me.box(style=me.Style(display="flex", gap=10, width=DIALOG_INPUT_WIDTH)):
+        me.slider(
+          min=MODEL_TEMPERATURE_MIN,
+          max=MODEL_TEMPERATURE_MAX,
+          step=0.1,
+          style=me.Style(width=260),
+          on_value_change=on_slider_temperature,
+          value=state.model_temperature,
+        )
+        me.input(
+          value=state.model_temperature_input,
+          on_input=on_input_temperature,
+          style=me.Style(width=60),
+        )
+
+    with mex.dialog_actions():
+      me.button(
+        "Close",
+        key="dialog_show_model_settings",
+        on_click=handlers.on_close_dialog,
+      )
+
+
+def on_slider_temperature(e: me.SliderValueChangeEvent):
+  """Adjust temperature slider value."""
+  state = me.state(State)
+  state.model_temperature = float(e.value)
+  state.model_temperature_input = str(state.model_temperature)
+
+
+def on_input_temperature(e: me.InputEvent):
+  """Adjust temperature slider value by input."""
+  state = me.state(State)
+  try:
+    model_temperature = float(e.value)
+    if MODEL_TEMPERATURE_MIN <= model_temperature <= MODEL_TEMPERATURE_MAX:
+      state.model_temperature = model_temperature
+  except ValueError:
+    pass

dialogs/prompt_variables.py

@@ -0,0 +1,60 @@
+import mesop as me
+
+import components as mex
+import handlers
+import llm
+from constants import DIALOG_INPUT_WIDTH
+from helpers import parse_variables
+from state import State
+
+
+@me.component
+def prompt_variables():
+  state = me.state(State)
+
+  with mex.dialog(state.dialog_show_prompt_variables):
+    me.text("Prompt Variables", type="headline-6")
+    if not state.prompt_variables:
+      me.text("No variables defined in prompt.", style=me.Style(width=DIALOG_INPUT_WIDTH))
+    else:
+      with me.box(
+        style=me.Style(display="flex", justify_content="end", margin=me.Margin(bottom=15))
+      ):
+        me.button("Generate", type="flat", on_click=on_click_generate_variables)
+      variable_names = set(parse_variables(state.prompt))
+      with me.box(style=me.Style(display="flex", flex_direction="column")):
+        for name, value in state.prompt_variables.items():
+          if name not in variable_names:
+            continue
+          me.textarea(
+            label=name,
+            value=value,
+            on_blur=on_input_variable,
+            style=me.Style(width=DIALOG_INPUT_WIDTH),
+            key=name,
+          )
+
+    with mex.dialog_actions():
+      me.button("Close", on_click=handlers.on_close_dialog, key="dialog_show_prompt_variables")
+
+
+def on_input_variable(e: me.InputBlurEvent):
+  """Generic event to save input variables.
+
+  TODO: Probably should prefix the key to avoid key collisions.
+  """
+  state = me.state(State)
+  state.prompt_variables[e.key] = e.value
+
+
+def on_click_generate_variables(e: me.ClickEvent):
+  """Generates values for the given empty variables."""
+  state = me.state(State)
+  print(state.model)
+  variable_names = set(parse_variables(state.prompt))
+  generated_variables = llm.generate_variables(
+    state.prompt, variable_names, state.model, state.model_temperature
+  )
+  for name in state.prompt_variables:
+    if name in variable_names and name in generated_variables:
+      state.prompt_variables[name] = generated_variables[name]

dialogs/prompt_version_history.py

@@ -0,0 +1,47 @@
+import mesop as me
+
+import components as mex
+import handlers
+from constants import DIALOG_INPUT_WIDTH
+from helpers import find_prompt
+from state import State
+from state import Prompt
+
+
+@me.component
+def prompt_version_history():
+  state = me.state(State)
+  with mex.dialog(state.dialog_show_version_history):
+    me.text("Version history", type="headline-6")
+    me.select(
+      label="Select Version",
+      options=[
+        me.SelectOption(label=f"v{prompt.version}", value=str(prompt.version))
+        for prompt in state.prompts
+      ],
+      style=me.Style(width=DIALOG_INPUT_WIDTH),
+      on_selection_change=on_select_version,
+    )
+    with mex.dialog_actions():
+      me.button("Close", key="dialog_show_version_history", on_click=handlers.on_close_dialog)
+
+
+def on_select_version(e: me.SelectSelectionChangeEvent):
+  """Update UI to show the selected prompt version and close the dialog."""
+  state = me.state(State)
+  selected_version = int(e.value)
+  prompt = find_prompt(state.prompts, selected_version)
+  if prompt != Prompt():
+    state.prompt = prompt.prompt
+    state.version = prompt.version
+    state.system_instructions = prompt.system_instructions
+    state.model = prompt.model
+    state.model_temperature = prompt.model_temperature
+    state.model_temperature_input = str(prompt.model_temperature)
+    # If there is an existing response, select the most recent one.
+    if prompt.responses:
+      state.prompt_variables = prompt.responses[-1]["variables"]
+      state.response = prompt.responses[-1]["output"]
+    else:
+      state.response = ""
+    state.dialog_show_version_history = False

dialogs/update_title.py

@@ -0,0 +1,32 @@
+import mesop as me
+
+import components as mex
+import handlers
+from constants import DIALOG_INPUT_WIDTH
+from state import State
+
+
+@me.component
+def update_title():
+  state = me.state(State)
+  # Update prompt title dialog
+  with mex.dialog(state.dialog_show_title):
+    me.text("Update Prompt Title", type="headline-6")
+    me.input(
+      label="Title",
+      value=state.temp_title,
+      on_blur=handlers.on_update_input,
+      key="temp_title",
+      style=me.Style(width=DIALOG_INPUT_WIDTH),
+    )
+    with mex.dialog_actions():
+      me.button("Cancel", on_click=handlers.on_close_dialog, key="dialog_show_title")
+      me.button("Save", type="flat", disabled=not state.temp_title.strip(), on_click=on_save_title)
+
+
+def on_save_title(e: me.InputBlurEvent):
+  """Saves the title and closes the dialog."""
+  state = me.state(State)
+  if state.temp_title:
+    state.title = state.temp_title
+    state.dialog_show_title = False

eval_table.py

@@ -0,0 +1,179 @@
+from typing import Callable
+import hashlib
+
+import mesop as me
+
+from state import Prompt
+
+
+@me.component
+def prompt_eval_table(prompts: list[Prompt], on_select_rating: Callable | None = None):
+  data = _make_table_meta(prompts)
+  response_map = _make_response_map(prompts)
+
+  with me.box(
+    style=me.Style(
+      display="grid",
+      grid_template_columns=" ".join([row.get("size", "1fr") for row in data]),
+      overflow_x="scroll",
+    )
+  ):
+    # Render first header row. This row only displays the Prompt version.
+    for row in data:
+      with me.box(style=_HEADER_STYLE):
+        me.text(row.get("header_1", ""))
+
+    # Render second header row.
+    for row in data:
+      with me.box(style=_HEADER_STYLE):
+        if row["type"] == "model_rating":
+          avg_rating = _calculate_avg_rating_from_prompt(row["prompt"])
+          if avg_rating is not None:
+            with me.tooltip(message="Average rating"):
+              me.text(
+                f"{avg_rating:.2f}", style=me.Style(font_size=13, color="#666", text_align="center")
+              )
+        elif row["type"] == "variable":
+          me.text(row.get("header_2", ""), style=me.Style(font_size=13, color="#0063FF"))
+        else:
+          me.text(row.get("header_2", ""), style=me.Style(font_size=13, color="#aaa"))
+
+    # Render examples
+    for row_index, example in enumerate(prompts[0].responses):
+      response_key = _make_variables_key(example["variables"])
+      for row in data:
+        if row["type"] == "index":
+          with me.box(style=_INDEX_STYLE):
+            me.text(str(row_index))
+        elif row["type"] == "variable":
+          with me.box(style=_MARKDOWN_BOX_STYLE):
+            me.markdown(example["variables"][row["variable_name"]])
+        elif row["type"] == "model_response":
+          with me.box(style=_MARKDOWN_BOX_STYLE):
+            prompt_response = response_map[row["prompt"].version].get(response_key)
+            if prompt_response:
+              me.markdown(prompt_response[0]["output"])
+            else:
+              me.text("")
+        elif row["type"] == "model_rating":
+          with me.box(style=_RATING_STYLE):
+            prompt_response = response_map[row["prompt"].version].get(response_key)
+            if prompt_response:
+              me.select(
+                value=prompt_response[0].get("rating", ""),
+                options=[
+                  me.SelectOption(label="1", value="1"),
+                  me.SelectOption(label="2", value="2"),
+                  me.SelectOption(label="3", value="3"),
+                  me.SelectOption(label="4", value="4"),
+                  me.SelectOption(label="5", value="5"),
+                ],
+                on_selection_change=on_select_rating,
+                key=f"rating_{row["prompt"].version}_{prompt_response[1]}",
+                style=me.Style(width=60),
+              )
+
+
+def _hash_string(string: str) -> str:
+  encoded_string = string.encode("utf-8")
+  result = hashlib.md5(encoded_string)
+  return result.hexdigest()
+
+
+def _make_variables_key(variables: dict[str, str]) -> str:
+  return _hash_string("--".join(variables.values()))
+
+
+def _make_response_map(prompts: list[Prompt]) -> dict:
+  prompt_map = {}
+  for prompt in prompts:
+    response_map = {}
+    for response_index, response in enumerate(prompt.responses):
+      key = _make_variables_key(response["variables"])
+      response_map[key] = (response, response_index)
+    prompt_map[prompt.version] = response_map
+  return prompt_map
+
+
+def _make_table_meta(prompts: list[Prompt]) -> dict:
+  data = [
+    {
+      "type": "index",
+    }
+  ]
+  for variable in prompts[0].variables:
+    data.append(
+      {
+        "type": "variable",
+        "header_2": "{{" + variable + "}}",
+        "size": "20fr",
+        "variable_name": variable,
+      }
+    )
+
+  for i, prompt in enumerate(prompts):
+    data.append(
+      {
+        "type": "model_response",
+        "header_1": "Version " + str(prompt.version),
+        "header_2": "Model response",
+        "index": i,
+        "size": "20fr",
+        "prompt": prompt,
+      }
+    )
+    data.append(
+      {
+        "type": "model_rating",
+        "header_1": "",
+        "header_2": "Rating",
+        "prompt": prompt,
+      }
+    )
+  return data
+
+
+def _calculate_avg_rating_from_prompt(prompt: Prompt) -> float | None:
+  ratings = [int(response["rating"]) for response in prompt.responses if response.get("rating")]
+  if ratings:
+    return sum(ratings) / float(len(ratings))
+  return None
+
+
+_BORDER_SIDE = me.BorderSide(width=1, style="solid", color="#DEE2E6")
+
+_HEADER_STYLE = me.Style(
+  background="#FFF",
+  border=me.Border.all(_BORDER_SIDE),
+  color="#000",
+  font_size=15,
+  font_weight="bold",
+  padding=me.Padding.all(10),
+)
+
+_INDEX_STYLE = me.Style(
+  background="#FFF",
+  border=me.Border.all(_BORDER_SIDE),
+  color="#000",
+  font_size=14,
+  padding=me.Padding.all(10),
+  text_align="center",
+)
+
+_MARKDOWN_BOX_STYLE = me.Style(
+  background="#FFF",
+  border=me.Border.all(_BORDER_SIDE),
+  color="#000",
+  font_size=14,
+  padding=me.Padding.all(10),
+  max_height=300,
+  min_width=300,
+  overflow_y="scroll",
+)
+
+_RATING_STYLE = me.Style(
+  background="#FFF",
+  border=me.Border.all(_BORDER_SIDE),
+  color="#000",
+  padding=me.Padding.all(10),
+)

handlers.py

@@ -0,0 +1,27 @@
+import mesop as me
+from state import State
+
+
+def on_open_dialog(e: me.ClickEvent):
+  """Generic event to open a dialog."""
+  state = me.state(State)
+  setattr(state, e.key, True)
+
+
+def on_close_dialog(e: me.ClickEvent):
+  """Generic event to close a dialog."""
+  state = me.state(State)
+  setattr(state, e.key, False)
+
+
+def on_update_input(e: me.InputBlurEvent | me.InputEvent | me.InputEnterEvent):
+  """Generic event to update input values."""
+  state = me.state(State)
+  setattr(state, e.key, e.value)
+
+
+def on_update_selection(e: me.SelectSelectionChangeEvent):
+  """Generic event to update input values."""
+  state = me.state(State)
+  print("EEE", e)
+  setattr(state, e.key, e.value)

helpers.py

@@ -0,0 +1,18 @@
+import re
+from state import Prompt
+
+
+_RE_VARIABLES = re.compile(r"\{\{(\w+)\}\}")
+
+
+def parse_variables(prompt: str) -> list[str]:
+  return _RE_VARIABLES.findall(prompt)
+
+
+def find_prompt(prompts: list[Prompt], version: int) -> Prompt:
+  # We don't expect too many versions, so we'll just loop through the list to find the
+  # right version.
+  for prompt in prompts:
+    if prompt.version == version:
+      return prompt
+  return Prompt()

main.py

@@ -1,15 +1,13 @@
-from dataclasses import dataclass, field
-import re
-
 import mesop as me
 
 import components as mex
+import dialogs
+import handlers
 import llm
-
-_DIALOG_INPUT_WIDTH = 350
-
-_MODEL_TEMPERATURE_MAX = 2
-_MODEL_TEMPERATURE_MIN = 0
+from eval_table import prompt_eval_table
+from tool_sidebar import tool_sidebar
+from helpers import find_prompt, parse_variables
+from state import State, Prompt
 
 _INSTRUCTIONS = """
 - Write your prompt.
@@ -21,55 +19,6 @@ _INSTRUCTIONS = """
   new version of your prompt.
 """.strip()
 
-_RE_VARIABLES = re.compile(r"\{\{(\w+)\}\}")
-
-
-@dataclass
-class Prompt:
-  prompt: str = ""
-  model: str = ""
-  model_temperature: float = 0.0
-  system_instructions: str = ""
-  version: int = 0
-  variables: list[str] = field(default_factory=lambda: [])
-  # Storing the responses as a dict to workaround bug with lists
-  # of nested dataclass.
-  responses: list[dict] = field(default_factory=lambda: [])
-
-
-@me.stateclass
-class State:
-  # Main UI variables
-  system_prompt_card_expanded: bool = False
-  title: str = "Untitled Prompt"
-  temp_title: str
-  system_instructions: str
-  prompt: str
-  response: str
-  version: int = 0
-
-  # Prompt variables
-  prompt_variables: dict[str, str]
-
-  # Model info
-  model: str = "gemini-1.5-flash"
-  model_temperature: float = 1.0
-  model_temperature_input: str = "1.0"
-
-  # Dialogs
-  dialog_show_title: bool = False
-  dialog_show_model_settings: bool = False
-  dialog_show_prompt_variables: bool = False
-  dialog_show_generate_prompt: bool = False
-  dialog_show_version_history: bool = False
-  prompts: list[Prompt]
-
-  # LLM Generate functionality
-  prompt_gen_task_description: str
-
-  # Valid modes: Prompt or Eval
-  mode: str = "Prompt"
-
 
 @me.page(
   security_policy=me.SecurityPolicy(allowed_iframe_parents=["https://huggingface.co"]),
@@ -77,114 +26,13 @@ class State:
 def app():
   state = me.state(State)
 
-  # Update prompt title dialog
-  with mex.dialog(state.dialog_show_title):
-    me.text("Update Prompt Title", type="headline-6")
-    me.input(
-      label="Title",
-      value=state.temp_title,
-      on_blur=on_update_input,
-      key="temp_title",
-      style=me.Style(width=_DIALOG_INPUT_WIDTH),
-    )
-    with mex.dialog_actions():
-      me.button("Cancel", on_click=on_close_dialog, key="dialog_show_title")
-      me.button("Save", type="flat", disabled=not state.temp_title.strip(), on_click=on_save_title)
-
-  # Dialog for controlling Model settings
-  with mex.dialog(state.dialog_show_model_settings):
-    me.text("Model Settings", type="headline-6")
-    with me.box():
-      me.select(
-        label="Model",
-        key="model",
-        options=[
-          me.SelectOption(label="Gemini 1.5 Flash", value="gemini-1.5-flash"),
-          me.SelectOption(label="Gemini 1.5 Pro", value="gemini-1.5-pro"),
-        ],
-        value=state.model,
-        style=me.Style(width=_DIALOG_INPUT_WIDTH),
-        on_selection_change=on_update_input,
-      )
-    with me.box():
-      me.text("Temperature", style=me.Style(font_weight="bold"))
-      with me.box(style=me.Style(display="flex", gap=10, width=_DIALOG_INPUT_WIDTH)):
-        me.slider(
-          min=_MODEL_TEMPERATURE_MIN,
-          max=_MODEL_TEMPERATURE_MAX,
-          step=0.1,
-          style=me.Style(width=260),
-          on_value_change=on_slider_temperature,
-          value=state.model_temperature,
-        )
-        me.input(
-          value=state.model_temperature_input,
-          on_input=on_input_temperature,
-          style=me.Style(width=60),
-        )
-
-    with mex.dialog_actions():
-      me.button(
-        "Close",
-        key="dialog_show_model_settings",
-        on_click=on_close_dialog,
-      )
-
-  # Dialog for setting variables
-  with mex.dialog(state.dialog_show_prompt_variables):
-    me.text("Prompt Variables", type="headline-6")
-    if not state.prompt_variables:
-      me.text("No variables defined in prompt.", style=me.Style(width=_DIALOG_INPUT_WIDTH))
-    else:
-      with me.box(
-        style=me.Style(display="flex", justify_content="end", margin=me.Margin(bottom=15))
-      ):
-        me.button("Generate", type="flat", on_click=on_click_generate_variables)
-      variable_names = set(_parse_variables(state.prompt))
-      with me.box(style=me.Style(display="flex", flex_direction="column")):
-        for name, value in state.prompt_variables.items():
-          if name not in variable_names:
-            continue
-          me.textarea(
-            label=name,
-            value=value,
-            on_blur=on_input_variable,
-            style=me.Style(width=_DIALOG_INPUT_WIDTH),
-            key=name,
-          )
-
-    with mex.dialog_actions():
-      me.button("Close", on_click=on_close_dialog, key="dialog_show_prompt_variables")
-
-  # Dialog for showing prompt version history
-  with mex.dialog(state.dialog_show_version_history):
-    me.text("Version history", type="headline-6")
-    me.select(
-      label="Select Version",
-      options=[
-        me.SelectOption(label=f"v{prompt.version}", value=str(prompt.version))
-        for prompt in state.prompts
-      ],
-      style=me.Style(width=_DIALOG_INPUT_WIDTH),
-      on_selection_change=on_select_version,
-    )
-    with mex.dialog_actions():
-      me.button("Close", key="dialog_show_version_history", on_click=on_close_dialog)
-
-  # Dialog for generating a prompt with LLM assistance
-  # TODO: Integrate with LLM
-  with mex.dialog(state.dialog_show_generate_prompt):
-    me.text("Generate Prompt", type="headline-6")
-    me.textarea(
-      label="Describe your task",
-      value=state.prompt_gen_task_description,
-      on_blur=on_update_input,
-      key="prompt_gen_task_description",
-      style=me.Style(width=_DIALOG_INPUT_WIDTH),
-    )
-    with mex.dialog_actions():
-      me.button("Close", key="dialog_show_generate_prompt", on_click=on_close_dialog)
-      me.button("Generate", type="flat", on_click=on_click_generate_prompt)
+  dialogs.update_title()
+  dialogs.model_settings()
+  dialogs.prompt_variables()
+  dialogs.prompt_version_history()
+  dialogs.add_comparisons()
+  dialogs.generate_prompt()
+  dialogs.load_prompt()
 
   with me.box(
     style=me.Style(
@@ -226,7 +74,7 @@ def app():
             min_rows=2,
             placeholder="Optional tone and style instructions for the model",
             value=state.system_instructions,
-            on_blur=on_update_input,
+            on_blur=handlers.on_update_input,
             style=_STYLE_INVISIBLE_TEXTAREA,
             key="system_instructions",
           )
@@ -257,7 +105,7 @@ def app():
             "Generate prompt",
             disabled=bool(state.prompt),
             style=me.Style(background="#EBF1FD", border_radius="10"),
-            on_click=on_open_dialog,
+            on_click=handlers.on_open_dialog,
             key="dialog_show_generate_prompt",
           )
 
@@ -270,33 +118,16 @@ def app():
             me.markdown(_INSTRUCTIONS)
     else:
       # Render eval page
-      with me.box(style=me.Style(grid_column="1 / -2")):
-        prompt = _find_prompt(state.prompts, state.version)
+      with me.box(style=me.Style(grid_column="1 / -2", overflow_y="scroll")):
+        prompt = find_prompt(state.prompts, state.version)
         if prompt:
-          mex.prompt_eval_table(prompt, on_select_rating=on_select_rating)
-
-    with mex.icon_sidebar():
-      if state.mode == "Prompt":
-        mex.icon_menu_item(
-          icon="tune",
-          tooltip="Model settings",
-          key="dialog_show_model_settings",
-          on_click=on_open_dialog,
-        )
-        mex.icon_menu_item(
-          icon="data_object",
-          tooltip="Set variables",
-          key="dialog_show_prompt_variables",
-          on_click=on_open_dialog,
-        )
-      mex.icon_menu_item(
-        icon="history",
-        tooltip="Version history",
-        key="dialog_show_version_history",
-        on_click=on_open_dialog,
-      )
-      if state.mode == "Prompt":
-        mex.icon_menu_item(icon="code", tooltip="Get code")
+          with me.box(style=me.Style(margin=me.Margin.all(15))):
+            compare_prompts = [
+              prompt for prompt in state.prompts if prompt.version in state.comparisons
+            ]
+            prompt_eval_table([prompt] + compare_prompts, on_select_rating=on_select_rating)
+
+    tool_sidebar()
 
 
 # Event handlers
@@ -323,7 +154,7 @@ def on_click_run(e: me.ClickEvent):
   else:
     current_prompt_meta = Prompt()
 
-  variable_names = set(_parse_variables(state.prompt))
+  variable_names = set(parse_variables(state.prompt))
   prompt_variables = {
     name: value for name, value in state.prompt_variables.items() if name in variable_names
   }
@@ -369,94 +200,12 @@ def on_update_prompt(e: me.InputBlurEvent):
   """
   state = me.state(State)
   state.prompt = e.value.strip()
-  variable_names = _parse_variables(state.prompt)
+  variable_names = parse_variables(state.prompt)
   for variable_name in variable_names:
     if variable_name not in state.prompt_variables:
       state.prompt_variables[variable_name] = ""
 
 
-def on_save_title(e: me.InputBlurEvent):
-  """Saves the title and closes the dialog."""
-  state = me.state(State)
-  if state.temp_title:
-    state.title = state.temp_title
-    state.dialog_show_title = False
-
-
-def on_slider_temperature(e: me.SliderValueChangeEvent):
-  """Adjust temperature slider value."""
-  state = me.state(State)
-  state.model_temperature = float(e.value)
-  state.model_temperature_input = str(state.model_temperature)
-
-
-def on_input_temperature(e: me.InputEvent):
-  """Adjust temperature slider value by input."""
-  state = me.state(State)
-  try:
-    model_temperature = float(e.value)
-    if _MODEL_TEMPERATURE_MIN <= model_temperature <= _MODEL_TEMPERATURE_MAX:
-      state.model_temperature = model_temperature
-  except ValueError:
-    pass
-
-
-def on_input_variable(e: me.InputBlurEvent):
-  """Generic event to save input variables.
-
-  TODO: Probably should prefix the key to avoid key collisions.
-  """
-  state = me.state(State)
-  state.prompt_variables[e.key] = e.value
-
-
-def on_select_version(e: me.SelectSelectionChangeEvent):
-  """Update UI to show the selected prompt version and close the dialog."""
-  state = me.state(State)
-  selected_version = int(e.value)
-  prompt = _find_prompt(state.prompts, selected_version)
-  if prompt != Prompt():
-    state.prompt = prompt.prompt
-    state.version = prompt.version
-    state.system_instructions = prompt.system_instructions
-    state.model = prompt.model
-    state.model_temperature = prompt.model_temperature
-    state.model_temperature_input = str(prompt.model_temperature)
-    # If there is an existing response, select the most recent one.
-    if prompt.responses:
-      state.prompt_variables = prompt.responses[-1]["variables"]
-      state.response = prompt.responses[-1]["output"]
-    else:
-      state.response = ""
-    state.dialog_show_version_history = False
-
-
-def on_click_generate_prompt(e: me.ClickEvent):
-  """Generates an improved prompt based on the given task description and closes dialog."""
-  state = me.state(State)
-  state.prompt = llm.generate_prompt(
-    state.prompt_gen_task_description, state.model, state.model_temperature
-  )
-  variable_names = _parse_variables(state.prompt)
-  for variable_name in variable_names:
-    if variable_name not in state.prompt_variables:
-      state.prompt_variables[variable_name] = ""
-
-  state.dialog_show_generate_prompt = False
-
-
-def on_click_generate_variables(e: me.ClickEvent):
-  """Generates values for the given empty variables."""
-  state = me.state(State)
-  variable_names = set(_parse_variables(state.prompt))
-  generated_variables = llm.generate_variables(
-    state.prompt, variable_names, state.model, state.model_temperature
-  )
-  for name in state.prompt_variables:
-    if name in variable_names and name in generated_variables:
-      state.prompt_variables[name] = generated_variables[name]
-
-
 def on_click_mode_toggle(e: me.ClickEvent):
   """Toggle between Prompt and Eval modes."""
   state = me.state(State)
@@ -466,47 +215,10 @@ def on_click_mode_toggle(e: me.ClickEvent):
 def on_select_rating(e: me.SelectSelectionChangeEvent):
   state = me.state(State)
   _, prompt_version, response_index = e.key.split("_")
-  prompt = _find_prompt(state.prompts, int(prompt_version))
+  prompt = find_prompt(state.prompts, int(prompt_version))
   prompt.responses[int(response_index)]["rating"] = e.value
 
 
-# Generic event handlers
-
-
-def on_open_dialog(e: me.ClickEvent):
-  """Generic event to open a dialog."""
-  state = me.state(State)
-  setattr(state, e.key, True)
-
-
-def on_close_dialog(e: me.ClickEvent):
-  """Generic event to close a dialog."""
-  state = me.state(State)
-  setattr(state, e.key, False)
-
-
-def on_update_input(e: me.InputBlurEvent | me.SelectSelectionChangeEvent):
-  """Generic event to update input/select values."""
-  state = me.state(State)
-  setattr(state, e.key, e.value)
-
-
-# Helper functions
-
-
-def _parse_variables(prompt: str) -> list[str]:
-  return _RE_VARIABLES.findall(prompt)
-
-
-def _find_prompt(prompts: list[Prompt], version: int) -> Prompt:
-  # We don't expect too many versions, so we'll just loop through the list to find the
-  # right version.
-  for prompt in prompts:
-    if prompt.version == version:
-      return prompt
-  return Prompt()
-
-
 # Style helpers
 
 _STYLE_INVISIBLE_TEXTAREA = me.Style(

state.py

@@ -0,0 +1,57 @@
+from dataclasses import dataclass, field
+
+import mesop as me
+
+
+@dataclass
+class Prompt:
+  prompt: str = ""
+  model: str = ""
+  model_temperature: float = 0.0
+  system_instructions: str = ""
+  version: int = 0
+  variables: list[str] = field(default_factory=list)
+  # Storing the responses as a dict to workaround bug with lists
+  # of nested dataclass.
+  #
+  # Keys: output, variables, rating
+  responses: list[dict] = field(default_factory=list)
+
+
+@me.stateclass
+class State:
+  # Main UI variables
+  system_prompt_card_expanded: bool = False
+  title: str = "Untitled Prompt"
+  temp_title: str
+  system_instructions: str
+  prompt: str
+  response: str
+  version: int = 0
+
+  # Prompt variables
+  prompt_variables: dict[str, str]
+
+  # Model info
+  model: str = "gemini-1.5-flash"
+  model_temperature: float = 1.0
+  model_temperature_input: str = "1.0"
+
+  # Dialogs
+  dialog_show_title: bool = False
+  dialog_show_model_settings: bool = False
+  dialog_show_prompt_variables: bool = False
+  dialog_show_generate_prompt: bool = False
+  dialog_show_version_history: bool = False
+  dialog_show_add_comparison: bool = False
+  dialog_show_load: bool = False
+
+  prompts: list[Prompt]
+
+  # LLM Generate functionality
+  prompt_gen_task_description: str
+
+  # Valid modes: Prompt or Eval
+  mode: str = "Prompt"
+
+  comparisons: list[int]

tool_sidebar.py

@@ -0,0 +1,95 @@
+from dataclasses import asdict
+import errno
+import json
+import os
+import re
+
+import mesop as me
+
+import components as mex
+import handlers
+from constants import SAVED_PROMPT_DIRECTORY
+from state import State
+
+
+@me.component
+def tool_sidebar():
+  state = me.state(State)
+  with mex.icon_sidebar():
+    if state.mode == "Prompt":
+      mex.icon_menu_item(
+        icon="tune",
+        tooltip="Model settings",
+        key="dialog_show_model_settings",
+        on_click=handlers.on_open_dialog,
+      )
+      mex.icon_menu_item(
+        icon="data_object",
+        tooltip="Set variables",
+        key="dialog_show_prompt_variables",
+        on_click=handlers.on_open_dialog,
+      )
+    mex.icon_menu_item(
+      icon="history",
+      tooltip="Version history",
+      key="dialog_show_version_history",
+      on_click=handlers.on_open_dialog,
+    )
+
+    if state.mode == "Eval":
+      mex.icon_menu_item(
+        icon="compare",
+        tooltip="Compare versions",
+        key="dialog_show_add_comparison",
+        on_click=handlers.on_open_dialog,
+      )
+
+    mex.icon_menu_item(
+      icon="upload",
+      tooltip="Load prompt data",
+      key="dialog_show_load",
+      on_click=handlers.on_open_dialog,
+    )
+
+    mex.icon_menu_item(
+      icon="download",
+      tooltip="Download prompt data",
+      key="dialog_show_download",
+      on_click=on_click_download,
+    )
+
+
+def on_click_download(e: me.ClickEvent):
+  state = me.state(State)
+  cleaned_title = _clean_title(state.title)
+  _create_directory(SAVED_PROMPT_DIRECTORY)
+
+  with open(f"{SAVED_PROMPT_DIRECTORY}/prompt-{cleaned_title}.json", "w") as outfile:
+    output = {
+      key: value
+      for key, value in asdict(state).items()
+      if key
+      not in (
+        "temp_title",
+        "mode",
+        "comparisons",
+        "system_prompt_card_expanded",
+        "prompt_gen_task_description",
+      )
+      and not key.startswith("dialog_")
+    }
+    json.dump(output, outfile)
+
+
+def _clean_title(title: str) -> str:
+  return re.sub(r"[^a-z0-9_]", "", title.lower().replace(" ", "_"))
+
+
+def _create_directory(directory_path):
+  """Creates a directory if it doesn't exist."""
+  try:
+    os.makedirs(directory_path)
+    print(f"Directory '{directory_path}' created successfully.")
+  except OSError as e:
+    if e.errno != errno.EEXIST:
+      raise