feat(lib): add new ai prompt action (#448)

* feat: add promptAI action

* feat: fix ts

* feat: add promptAI tests

* feat: rename variable

* feat: fix tests

* feat: change test name

* feat: adjust promptAI tests

* feat: allow all types of prompt

* feat: improve tests

* feat: add todo to missing storyId payload property

* feat: add missing tests

* feat: add support to the new action in the test lib-helper

* feat: add error message when promptAI action is used in the Sandbox

* feat(SHAPE-8085): apply review feedback

* feat(SHAPE-8085): fix tests

* feat(SHAPE-8085): improve prompt ai response message

* feat(SHAPE-8085): apply review feedback

Author: demetriusfeijoo

packages/demo/src/components/NonModalView.tsx

@@ -7,6 +7,7 @@ import { ContextRequester } from './ContextRequester'
 import { UserContextRequester } from './UserContextRequester'
 import { PluginComponent } from './FieldPluginDemo'
 import { LanguageView } from './LanguageView'
+import { PromptAI } from './PromptAI'
 
 export const NonModalView: PluginComponent = (props) => (
   <Paper>
@@ -18,6 +19,7 @@ export const NonModalView: PluginComponent = (props) => (
       <UserContextRequester {...props} />
       <HeightChangeDemo {...props} />
       <LanguageView {...props} />
+      <PromptAI {...props} />
     </Stack>
   </Paper>
 )

packages/demo/src/components/PromptAI.tsx

@@ -0,0 +1,114 @@
+import {
+  Button,
+  Checkbox,
+  FormControl,
+  FormControlLabel,
+  MenuItem,
+  Stack,
+  TextField,
+  Typography,
+} from '@mui/material'
+import { useState } from 'react'
+import {
+  isPromptAIPayloadValid,
+  promptAIActionsList,
+  PromptAIResponse,
+  type PromptAIAction,
+} from '@storyblok/field-plugin'
+import type { PluginComponent } from './FieldPluginDemo'
+
+export const PromptAI: PluginComponent = (props) => {
+  const { actions } = props
+
+  const [promptQuestion, setPromptQuestion] = useState<string>('')
+  const [promptAction, setPromptAction] = useState<PromptAIAction>('prompt')
+  const [promptLanguage, setPromptLanguage] = useState<string>()
+  const [promptTone, setPromptTone] = useState<string>()
+  const [promptAIResponse, setPromptAIResponse] = useState<PromptAIResponse>()
+  const [promptBasedOnCurrentStory, setPromptBasedOnCurrentStory] =
+    useState<boolean>(false)
+
+  const onSubmit = async () => {
+    const payload = {
+      action: promptAction,
+      text: promptQuestion,
+      language: promptLanguage,
+      tone: promptTone,
+      basedOnCurrentStory: promptBasedOnCurrentStory,
+    }
+
+    if (!isPromptAIPayloadValid(payload)) {
+      console.error('Invalid Prompt AI payload')
+      return
+    }
+
+    const promptAIResponse = await actions.promptAI(payload)
+
+    setPromptAIResponse(promptAIResponse)
+  }
+
+  return (
+    <Stack
+      gap={2}
+      direction={'column'}
+    >
+      <Typography variant="subtitle1">Prompt AI</Typography>
+      <Stack gap={2}>
+        <TextField
+          label="Ask a question"
+          onChange={(e) => setPromptQuestion(e.target.value)}
+          required
+        />
+        <TextField
+          label="Action"
+          value={promptAction}
+          select
+          required
+          onChange={(e) => setPromptAction(e.target.value as PromptAIAction)}
+        >
+          {promptAIActionsList.map((promptAIAction) => (
+            <MenuItem
+              key={promptAIAction}
+              value={promptAIAction}
+            >
+              {promptAIAction}
+            </MenuItem>
+          ))}
+        </TextField>
+        <TextField
+          label="Language (optional)"
+          onChange={(e) => setPromptLanguage(e.target.value)}
+        />
+        <TextField
+          label="Tone (optional)"
+          onChange={(e) => setPromptTone(e.target.value)}
+        />
+        <FormControl>
+          <FormControlLabel
+            label="Based on the current story"
+            control={
+              <Checkbox
+                value={promptBasedOnCurrentStory}
+                onChange={(e) => setPromptBasedOnCurrentStory(e.target.checked)}
+              />
+            }
+            sx={{ ml: '-9px' }}
+          />
+        </FormControl>
+        <Typography>
+          {promptAIResponse?.ok === true &&
+            `AI Generated Text: ${promptAIResponse.answer}`}
+          {promptAIResponse?.ok === false &&
+            `AI Error: ${promptAIResponse.error}`}
+        </Typography>
+        <Button
+          variant="outlined"
+          color="secondary"
+          onClick={onSubmit}
+        >
+          Prompt
+        </Button>
+      </Stack>
+    </Stack>
+  )
+}

packages/field-plugin/src/createFieldPlugin/FieldPluginActions.ts

@@ -1,4 +1,11 @@
-import type { Asset, StoryData, UserData, ModalSize } from '../messaging'
+import type {
+  Asset,
+  StoryData,
+  UserData,
+  ModalSize,
+  PromptAIPayload,
+  PromptAIResponse,
+} from '../messaging'
 import type { FieldPluginData } from './FieldPluginData'
 
 export type SetContent<Content> = (
@@ -9,6 +16,7 @@ export type SetModalOpen<Content> = (
   modalSize?: ModalSize,
 ) => Promise<FieldPluginData<Content>>
 export type RequestContext = () => Promise<StoryData>
+export type PromptAI = (payload: PromptAIPayload) => Promise<PromptAIResponse>
 export type RequestUserContext = () => Promise<UserData>
 export type SelectAsset = () => Promise<Asset>
 export type Initialize<Content> = () => Promise<FieldPluginData<Content>>
@@ -17,6 +25,7 @@ export type FieldPluginActions<Content> = {
   setContent: SetContent<Content>
   setModalOpen: SetModalOpen<Content>
   requestContext: RequestContext
+  promptAI: PromptAI
   requestUserContext: RequestUserContext
   selectAsset: SelectAsset
 }

packages/field-plugin/src/createFieldPlugin/createPluginActions/callbackQueue.ts

@@ -1,4 +1,5 @@
-import {
+import type {
+  PromptAIResponseMessage,
   AssetSelectedMessage,
   ContextRequestMessage,
   UserContextRequestMessage,
@@ -16,6 +17,7 @@ type CallbackMap = {
   userContext: Record<CallbackId, OnMessage<UserContextRequestMessage>>
   stateChanged: Record<CallbackId, OnMessage<StateChangedMessage>>
   loaded: Record<CallbackId, OnMessage<LoadedMessage>>
+  promptAI: Record<CallbackId, OnMessage<PromptAIResponseMessage>>
 }
 type CallbackType = keyof CallbackMap
 
@@ -26,7 +28,9 @@ export const callbackQueue = () => {
     userContext: {},
     stateChanged: {},
     loaded: {},
+    promptAI: {},
   }
+
   const pushCallback = <T extends CallbackType>(
     callbackType: T,
     callback: CallbackMap[T][CallbackId],

packages/field-plugin/src/createFieldPlugin/createPluginActions/createPluginActions.test.ts

@@ -5,6 +5,8 @@ import type {
   GetUserContextMessage,
   ModalChangeMessage,
   ValueChangeMessage,
+  PromptAIPayload,
+  PluginPromptAIMessage,
 } from '../../messaging'
 import { emptyAsset } from '../../messaging/pluginMessage/containerToPluginMessage/Asset.test'
 
@@ -236,6 +238,34 @@ describe('createPluginActions', () => {
     })
   })
 
+  describe('promptAI()', () => {
+    it('send a message to the container to prompt the AI', () => {
+      const { uid, postToContainer, onUpdateState } = mock()
+      const {
+        actions: { promptAI },
+      } = createPluginActions({
+        uid,
+        postToContainer,
+        onUpdateState,
+        validateContent,
+      })
+
+      const promptAIPayload: PromptAIPayload = {
+        action: 'prompt',
+        text: 'Some text to prompt',
+      }
+
+      promptAI(promptAIPayload)
+
+      expect(postToContainer).toHaveBeenLastCalledWith(
+        expect.objectContaining({
+          event: 'promptAI',
+          promptAIPayload,
+        } satisfies Partial<PluginPromptAIMessage>),
+      )
+    })
+  })
+
   describe('requestUserContext()', () => {
     it('send a message to the container to request the user info', () => {
       const { uid, postToContainer, onUpdateState } = mock()

packages/field-plugin/src/createFieldPlugin/createPluginActions/createPluginActions.ts

@@ -4,17 +4,21 @@ import {
   assetFromAssetSelectedMessage,
   assetModalChangeMessage,
   getContextMessage,
+  getResponseFromPromptAIMessage,
   getUserContextMessage,
   heightChangeMessage,
   modalChangeMessage,
-  OnAssetSelectMessage,
-  OnContextRequestMessage,
-  OnUserContextRequestMessage,
-  OnLoadedMessage,
-  OnStateChangeMessage,
-  OnUnknownPluginMessage,
   pluginLoadedMessage,
+  getPluginPromptAIMessage,
   valueChangeMessage,
+  type OnAssetSelectMessage,
+  type OnContextRequestMessage,
+  type OnUserContextRequestMessage,
+  type OnLoadedMessage,
+  type OnStateChangeMessage,
+  type OnUnknownPluginMessage,
+  type OnPromptAIMessage,
+  type PromptAIPayload,
 } from '../../messaging'
 import { FieldPluginActions, Initialize } from '../FieldPluginActions'
 import { pluginStateFromStateChangeMessage } from './partialPluginStateFromStateChangeMessage'
@@ -57,10 +61,12 @@ export const createPluginActions: CreatePluginActions = ({
     popCallback('stateChanged', data.callbackId)?.(data)
     onUpdateState(pluginStateFromStateChangeMessage(data, validateContent))
   }
+
   const onLoaded: OnLoadedMessage = (data) => {
     popCallback('loaded', data.callbackId)?.(data)
     onUpdateState(pluginStateFromStateChangeMessage(data, validateContent))
   }
+
   const onContextRequest: OnContextRequestMessage = (data) => {
     popCallback('context', data.callbackId)?.(data)
   }
@@ -70,6 +76,11 @@ export const createPluginActions: CreatePluginActions = ({
   const onAssetSelect: OnAssetSelectMessage = (data) => {
     popCallback('asset', data.callbackId)?.(data)
   }
+
+  const onPromptAI: OnPromptAIMessage = (data) => {
+    popCallback('promptAI', data.callbackId)?.(data)
+  }
+
   const onUnknownMessage: OnUnknownPluginMessage = (data) => {
     // TODO remove side-effect, making functions in this file pure.
     //  perhaps only show this message in development mode?
@@ -88,6 +99,7 @@ export const createPluginActions: CreatePluginActions = ({
     onContextRequest,
     onUserContextRequest,
     onAssetSelect,
+    onPromptAI,
     onUnknownMessage,
   }
 
@@ -146,6 +158,16 @@ export const createPluginActions: CreatePluginActions = ({
           postToContainer(getContextMessage({ uid, callbackId }))
         })
       },
+      promptAI: (promptAIMessage: PromptAIPayload) => {
+        return new Promise((resolve) => {
+          const callbackId = pushCallback('promptAI', (message) =>
+            resolve(getResponseFromPromptAIMessage(message)),
+          )
+          postToContainer(
+            getPluginPromptAIMessage(promptAIMessage, { uid, callbackId }),
+          )
+        })
+      },
       requestUserContext: () => {
         return new Promise((resolve) => {
           const callbackId = pushCallback('userContext', (message) =>

packages/field-plugin/src/createFieldPlugin/createPluginActions/createPluginMessageListener/createPluginMessageListener.ts

@@ -1,8 +1,9 @@
-import {
+import type {
   OnAssetSelectMessage,
   OnContextRequestMessage,
   OnUserContextRequestMessage,
   OnLoadedMessage,
+  OnPromptAIMessage,
   OnStateChangeMessage,
   OnUnknownPluginMessage,
 } from '../../../messaging'
@@ -14,6 +15,7 @@ export type PluginMessageCallbacks = {
   onContextRequest: OnContextRequestMessage
   onUserContextRequest: OnUserContextRequestMessage
   onAssetSelect: OnAssetSelectMessage
+  onPromptAI: OnPromptAIMessage
   onUnknownMessage: OnUnknownPluginMessage
 }