Enhance mediator workflow before platform launch

frontend/src/components/experimenter/experimenter_panel.scss

@@ -92,5 +92,17 @@
 
 .action-bar {
   @include common.flex-row;
-  justify-content: end;
+  @include typescale.label-small;
+  gap: common.$spacing-small;
+  justify-content: space-between;
+}
+
+.checkbox-wrapper {
+  @include common.flex-row-align-center;
+  gap: common.$spacing-small;
+  overflow-wrap: break-word;
+
+  md-checkbox {
+    flex-shrink: 0;
+  }
 }

frontend/src/components/experimenter/experimenter_panel.ts

@@ -5,6 +5,8 @@ import '../../pair-components/tooltip';
 import "./experimenter_data_editor";
 import './experimenter_manual_chat';
 
+import '@material/web/checkbox/checkbox.js';
+
 import {MobxLitElement} from '@adobe/lit-mobx';
 import {CSSResultGroup, html, nothing} from 'lit';
 import {customElement, state} from 'lit/decorators.js';
@@ -182,6 +184,15 @@ export class Panel extends MobxLitElement {
         stageId, {...mediator, prompt}, index
       );
     };
+    const updateJSON = () => {
+      const responseConfig = {
+        ...mediator.responseConfig,
+        isJSON: !mediator.responseConfig.isJSON,
+      };
+      this.mediatorEditor.updateMediator(
+        stageId, {...mediator, responseConfig}, index
+      );
+    };
 
     return html`
       <div class="mediator">
@@ -199,6 +210,17 @@ export class Panel extends MobxLitElement {
           >
           </pr-textarea>
           <div class="action-bar">
+            <div class="checkbox-wrapper">
+              <md-checkbox
+                touch-target="wrapper"
+                ?checked=${mediator.responseConfig.isJSON}
+                @click=${updateJSON}
+              >
+              </md-checkbox>
+              <div>
+                Parse as JSON
+              </div>
+            </div>
             <pr-button
               color="secondary"
               padding="small"

frontend/src/components/stages/chat_editor.scss

@@ -25,6 +25,7 @@
 }
 
 .description {
+  @include typescale.label-small;
   color: var(--md-sys-color-outline);
   font-style: italic;
 }

frontend/src/components/stages/chat_editor.ts

@@ -155,6 +155,13 @@ export class ChatEditor extends MobxLitElement {
         chat history (last 10 messages) and sent to the model
         (i.e., chat history + custom prompt => response)
       </div>
+      <div class="description">
+        <b>If JSON parsing enabled:</b> Make sure to
+        include appropriate instructions/examples in your prompt to
+        avoid parsing errors (if the specified message field is non-empty,
+        its contents will be turned into a chat message).
+        <b>If disabled:</b> non-empty responses will be turned into messages.
+      </div>
       <pr-textarea
         placeholder="Custom prompt for mediator"
         variant="outlined"
@@ -267,8 +274,7 @@ export class ChatEditor extends MobxLitElement {
         >
         </md-checkbox>
         <div>
-          Parse mediator response as JSON (tip: include appropriate
-          instructions/examples in prompt so that valid JSON is returned)
+          Parse mediator response as JSON
         </div>
       </div>
       ${!config.isJSON ? nothing : html`

frontend/src/services/mediator.editor.ts

@@ -29,6 +29,7 @@ export class MediatorEditor extends Service {
   // Experiment ID
   @observable experimentId: string|null = null;
   // Stage ID to chat config
+  // TODO: Map from stage ID to MediatorConfig list?
   @observable configMap: Record<string, ChatStageConfig> = {};
 
   setExperimentId(id: string) {

functions/src/stages/chat.triggers.ts

@@ -66,9 +66,22 @@ export const createMediatorMessage = onDocumentCreated(
       const response = await getGeminiAPIResponse(apiKeys.geminiKey, prompt);
 
       // Add mediator message if non-empty
-      const parsed = JSON.parse(response.text);
-      const isJSON = mediator.responseConfig.isJSON;
-      const message = isJSON ? (parsed[mediator.responseConfig.messageField] ?? '') : response.text;
+      let message = response.text;
+      let parsed = '';
+
+      if (mediator.responseConfig.isJSON) {
+        // Reset message to empty before trying to fill with JSON response
+        message = '';
+
+        try {
+          // TODO: Hack to get rid of markdown ticks surrounding {} ?
+          parsed = JSON.parse(response.text);
+        } catch {
+          // Response is already logged in console during Gemini API call
+          console.log('Could not parse JSON!');
+        }
+        message = parsed[mediator.responseConfig.messageField] ?? '';
+      }
 
       if (message.trim() === '') break;
       mediatorMessages.push({ mediator, parsed, message });