From d0e8573ba794459ad482987aba5d34496e0c653b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Demetrius=20Feij=C3=B3o?= <demetrius.feijoo@storyblok.com>
Date: Mon, 13 Jan 2025 21:09:13 -0300
Subject: [PATCH] feat(SHAPE-8085): apply review feedback

---
 packages/demo/src/components/PromptAI.tsx     | 32 ++++++++++++++++---
 .../PromptAIResponseMessage.ts                | 10 +++---
 .../MessageToContainer.ts                     |  2 +-
 .../PromptAIMessage.ts                        |  3 +-
 4 files changed, 34 insertions(+), 13 deletions(-)

diff --git a/packages/demo/src/components/PromptAI.tsx b/packages/demo/src/components/PromptAI.tsx
index 8a8a4fdf..8b95c2b9 100644
--- a/packages/demo/src/components/PromptAI.tsx
+++ b/packages/demo/src/components/PromptAI.tsx
@@ -1,4 +1,13 @@
-import { Button, MenuItem, Stack, TextField, Typography } from '@mui/material'
+import {
+  Button,
+  Checkbox,
+  FormControl,
+  FormLabel,
+  MenuItem,
+  Stack,
+  TextField,
+  Typography,
+} from '@mui/material'
 import { useState } from 'react'
 import {
   isPromptAIPayloadValid,
@@ -14,7 +23,9 @@ export const PromptAI: PluginComponent = (props) => {
   const [promptAction, setPromptAction] = useState<PromptAIAction>('prompt')
   const [promptLanguage, setPromptLanguage] = useState<string>()
   const [promptTone, setPromptTone] = useState<string>()
-  const [promptOutput, setPromptOutput] = useState<string>()
+  const [promptAIGeneratedText, setPromptAIGeneratedText] = useState<string>()
+  const [promptBasedOnCurrentStory, setPromptBasedOnCurrentStory] =
+    useState<boolean>(false)
 
   const onSubmit = async () => {
     const payload = {
@@ -22,6 +33,7 @@ export const PromptAI: PluginComponent = (props) => {
       text: promptQuestion,
       language: promptLanguage,
       tone: promptTone,
+      basedOnCurrentStory: promptBasedOnCurrentStory,
     }
 
     if (!isPromptAIPayloadValid(payload)) {
@@ -29,9 +41,9 @@ export const PromptAI: PluginComponent = (props) => {
       return
     }
 
-    const output = await actions.promptAI(payload)
+    const promptAIGeneratedText = await actions.promptAI(payload)
 
-    setPromptOutput(output)
+    setPromptAIGeneratedText(promptAIGeneratedText)
   }
 
   return (
@@ -70,7 +82,17 @@ export const PromptAI: PluginComponent = (props) => {
           label="Tone (optional)"
           onChange={(e) => setPromptTone(e.target.value)}
         />
-        <Typography>Output: {promptOutput}</Typography>
+        <FormControl>
+          <FormLabel htmlFor="based-on-current-story-checkbox">
+            Based on the current story:
+          </FormLabel>
+          <Checkbox
+            id="based-on-current-story-checkbox"
+            value={promptBasedOnCurrentStory}
+            onChange={(e) => setPromptBasedOnCurrentStory(e.target.checked)}
+          />
+        </FormControl>
+        <Typography>AI Generated Text: {promptAIGeneratedText}</Typography>
         <Button
           variant="outlined"
           color="secondary"
diff --git a/packages/field-plugin/src/messaging/pluginMessage/containerToPluginMessage/PromptAIResponseMessage.ts b/packages/field-plugin/src/messaging/pluginMessage/containerToPluginMessage/PromptAIResponseMessage.ts
index 73f89df2..4e00ab15 100644
--- a/packages/field-plugin/src/messaging/pluginMessage/containerToPluginMessage/PromptAIResponseMessage.ts
+++ b/packages/field-plugin/src/messaging/pluginMessage/containerToPluginMessage/PromptAIResponseMessage.ts
@@ -5,7 +5,7 @@ import { isMessageToPlugin, type MessageToPlugin } from './MessageToPlugin'
  * The object returned when calling the "prompt-ai" action.
  */
 export type PromptAIResponseMessage = MessageToPlugin<'prompt-ai'> & {
-  output: string
+  aiGeneratedText: string
 }
 
 export const isPromptAIMessage = (
@@ -13,12 +13,12 @@ export const isPromptAIMessage = (
 ): data is PromptAIResponseMessage =>
   isMessageToPlugin(data) &&
   data.action === 'prompt-ai' &&
-  hasKey(data, 'output') &&
-  typeof data.output === 'string'
+  hasKey(data, 'aiGeneratedText') &&
+  typeof data.aiGeneratedText === 'string'
 
 export const getResponseFromPromptAIMessage = (
   message: PromptAIResponseMessage,
 ): string => {
-  const { output } = message
-  return output
+  const { aiGeneratedText } = message
+  return aiGeneratedText
 }
diff --git a/packages/field-plugin/src/messaging/pluginMessage/pluginToContainerMessage/MessageToContainer.ts b/packages/field-plugin/src/messaging/pluginMessage/pluginToContainerMessage/MessageToContainer.ts
index b213ab98..50d0e1a7 100644
--- a/packages/field-plugin/src/messaging/pluginMessage/pluginToContainerMessage/MessageToContainer.ts
+++ b/packages/field-plugin/src/messaging/pluginMessage/pluginToContainerMessage/MessageToContainer.ts
@@ -1,7 +1,7 @@
 import { hasKey } from '../../../utils'
 
 export type MessageToContainer<Event extends string> = {
-  action: 'plugin-changed'
+  action: 'plugin-changed' | 'prompt-ai'
   uid: string
   event: Event
   callbackId?: string
diff --git a/packages/field-plugin/src/messaging/pluginMessage/pluginToContainerMessage/PromptAIMessage.ts b/packages/field-plugin/src/messaging/pluginMessage/pluginToContainerMessage/PromptAIMessage.ts
index 6997c66f..5957b176 100644
--- a/packages/field-plugin/src/messaging/pluginMessage/pluginToContainerMessage/PromptAIMessage.ts
+++ b/packages/field-plugin/src/messaging/pluginMessage/pluginToContainerMessage/PromptAIMessage.ts
@@ -36,8 +36,7 @@ export const promptAIActionsList: PromptAIAction[] = [
 export type PromptAIPayload = {
   action: PromptAIAction
   text: string
-  // todo: pass the storyId to the container when accepted by the API
-  // storyId?: number
+  basedOnCurrentStory?: boolean
   language?: string
   textLength?: string
   tone?: string