fix(opentrons-ai-client): Fixed protocol regenerate button (#16807)

<!--
Thanks for taking the time to open a Pull Request (PR)! Please make sure
you've read the "Opening Pull Requests" section of our Contributing
Guide:


https://github.com/Opentrons/opentrons/blob/edge/CONTRIBUTING.md#opening-pull-requests

GitHub provides robust markdown to format your PR. Links, diagrams,
pictures, and videos along with text formatting make it possible to
create a rich and informative PR. For more information on GitHub
markdown, see:


https://docs.github.com/en/get-started/writing-on-github/getting-started-with-writing-and-formatting-on-github/basic-writing-and-formatting-syntax

To ensure your code is reviewed quickly and thoroughly, please fill out
the sections below to the best of your ability!
-->

# Overview

This functionality was broken as the app didn't know to use the new
update/createnew endpoints when regenerating vs completion endpoint

<!--
Describe your PR at a high level. State acceptance criteria and how this
PR fits into other work. Link issues, PRs, and other relevant resources.
-->

## Test Plan and Hands on Testing

<!--
Describe your testing of the PR. Emphasize testing not reflected in the
code. Attach protocols, logs, screenshots and any other assets that
support your testing.
-->

## Changelog

<!--
List changes introduced by this PR considering future developers and the
end user. Give careful thought and clear documentation to breaking
changes.
-->

## Review requests

<!--
- What do you need from reviewers to feel confident this PR is ready to
merge?
- Ask questions.
-->

## Risk assessment

<!--
- Indicate the level of attention this PR needs.
- Provide context to guide reviewers.
- Discuss trade-offs, coupling, and side effects.
- Look for the possibility, even if you think it's small, that your
change may affect some other part of the system.
- For instance, changing return tip behavior may also change the
behavior of labware calibration.
- How do your unit tests and on hands on testing mitigate this PR's
risks and the risk of future regressions?
- Especially in high risk PRs, explain how you know your testing is
enough.
-->

opentrons-ai-client/src/OpentronsAIRoutes.tsx

@@ -1,6 +1,6 @@
 import { Route, Navigate, Routes } from 'react-router-dom'
 import { Landing } from './pages/Landing'
-import { UpdateProtocol } from './organisms/UpdateProtocol'
+import { UpdateProtocol } from './pages/UpdateProtocol'
 
 import type { RouteProps } from './resources/types'
 import { Chat } from './pages/Chat'

opentrons-ai-client/src/molecules/ChatDisplay/index.tsx

@@ -26,7 +26,10 @@ import { useAtom } from 'jotai'
 import {
   chatDataAtom,
   feedbackModalAtom,
+  regenerateProtocolAtom,
   scrollToBottomAtom,
+  createProtocolChatAtom,
+  updateProtocolChatAtom,
 } from '../../resources/atoms'
 import { delay } from 'lodash'
 import { useFormContext } from 'react-hook-form'
@@ -56,6 +59,9 @@ const StyledIcon = styled(Icon)`
 export function ChatDisplay({ chat, chatId }: ChatDisplayProps): JSX.Element {
   const { t } = useTranslation('protocol_generator')
   const [isCopied, setIsCopied] = useState<boolean>(false)
+  const [, setRegenerateProtocol] = useAtom(regenerateProtocolAtom)
+  const [createProtocolChat] = useAtom(createProtocolChatAtom)
+  const [updateProtocolChat] = useAtom(updateProtocolChatAtom)
   const [, setShowFeedbackModal] = useAtom(feedbackModalAtom)
   const { setValue } = useFormContext()
   const [chatdata] = useAtom(chatDataAtom)
@@ -64,9 +70,30 @@ export function ChatDisplay({ chat, chatId }: ChatDisplayProps): JSX.Element {
   const isUser = role === 'user'
 
   const setInputFieldToCorrespondingRequest = (): void => {
-    const prompt = chatdata.find(
-      chat => chat.role === 'user' && chat.requestId === requestId
-    )?.reply
+    let prompt = ''
+    if (
+      requestId.includes('NewProtocol') ||
+      requestId.includes('UpdateProtocol')
+    ) {
+      setRegenerateProtocol({
+        isCreateOrUpdateProtocol: true,
+        regenerate: true,
+      })
+      if (createProtocolChat.prompt !== '') {
+        prompt = createProtocolChat.prompt
+      } else {
+        prompt = updateProtocolChat.prompt
+      }
+    } else {
+      setRegenerateProtocol({
+        isCreateOrUpdateProtocol: false,
+        regenerate: true,
+      })
+      prompt =
+        chatdata.find(
+          chat => chat.role === 'user' && chat.requestId === requestId
+        )?.reply ?? ''
+    }
     setScrollToBottom(!scrollToBottom)
     setValue('userPrompt', prompt)
   }

opentrons-ai-client/src/molecules/InputPrompt/index.tsx

@@ -20,6 +20,7 @@ import {
   chatDataAtom,
   chatHistoryAtom,
   createProtocolChatAtom,
+  regenerateProtocolAtom,
   tokenAtom,
   updateProtocolChatAtom,
 } from '../../resources/atoms'
@@ -38,7 +39,11 @@ import {
 } from '../../resources/constants'
 
 import type { AxiosRequestConfig } from 'axios'
-import type { ChatData } from '../../resources/types'
+import type {
+  ChatData,
+  CreatePrompt,
+  UpdatePrompt,
+} from '../../resources/types'
 
 export function InputPrompt(): JSX.Element {
   const { t } = useTranslation('protocol_generator')
@@ -50,6 +55,9 @@ export function InputPrompt(): JSX.Element {
   const [sendAutoFilledPrompt, setSendAutoFilledPrompt] = useState<boolean>(
     false
   )
+  const [regenerateProtocol, setRegenerateProtocol] = useAtom(
+    regenerateProtocolAtom
+  )
 
   const [, setChatData] = useAtom(chatDataAtom)
   const [chatHistory, setChatHistory] = useAtom(chatHistoryAtom)
@@ -78,13 +86,24 @@ export function InputPrompt(): JSX.Element {
     }
   }, [watchUserPrompt])
 
+  useEffect(() => {
+    if (regenerateProtocol.regenerate) {
+      handleClick(regenerateProtocol.isCreateOrUpdateProtocol, true)
+      setRegenerateProtocol({
+        isCreateOrUpdateProtocol: false,
+        regenerate: false,
+      })
+    }
+  }, [regenerateProtocol])
+
   const handleClick = async (
-    isUpdateOrCreateRequest: boolean = false
+    isUpdateOrCreateRequest: boolean = false,
+    isRegenerateRequest: boolean = false
   ): Promise<void> => {
-    setRequestId(uuidv4() + getPreFixText(isUpdateOrCreateRequest))
-
+    const newRequestId = uuidv4() + getPreFixText(isUpdateOrCreateRequest)
+    setRequestId(newRequestId)
     const userInput: ChatData = {
-      requestId,
+      requestId: newRequestId,
       role: 'user',
       reply: watchUserPrompt,
     }
@@ -106,7 +125,7 @@ export function InputPrompt(): JSX.Element {
         method: 'POST',
         headers,
         data: isUpdateOrCreateRequest
-          ? getUpdateOrCreatePrompt()
+          ? getUpdateOrCreatePrompt(isRegenerateRequest)
           : {
               message: watchUserPrompt,
               history: chatHistory,
@@ -126,7 +145,11 @@ export function InputPrompt(): JSX.Element {
     }
   }
 
-  const getUpdateOrCreatePrompt = (): any => {
+  const getUpdateOrCreatePrompt = (
+    isRegenerateRequest: boolean
+  ): CreatePrompt | UpdatePrompt => {
+    createProtocol.regenerate = isRegenerateRequest
+    updateProtocol.regenerate = isRegenerateRequest
     return isNewProtocol ? createProtocol : updateProtocol
   }
 

opentrons-ai-client/src/pages/CreateProtocol/index.tsx

@@ -9,6 +9,8 @@ import { useEffect, useRef, useState } from 'react'
 import { PromptPreview } from '../../molecules/PromptPreview'
 import { useForm, FormProvider } from 'react-hook-form'
 import {
+  chatDataAtom,
+  chatHistoryAtom,
   createProtocolAtom,
   createProtocolChatAtom,
   headerWithMeterAtom,
@@ -55,6 +57,8 @@ export function CreateProtocol(): JSX.Element | null {
   )
   const [, setCreateProtocolChatAtom] = useAtom(createProtocolChatAtom)
   const [, setUpdateProtocolChatAtom] = useAtom(updateProtocolChatAtom)
+  const [, setChatHistoryAtom] = useAtom(chatHistoryAtom)
+  const [, setChatData] = useAtom(chatDataAtom)
   const navigate = useNavigate()
   const trackEvent = useTrackEvent()
   const [leftWidth, setLeftWidth] = useState(50)
@@ -79,8 +83,23 @@ export function CreateProtocol(): JSX.Element | null {
     },
   })
 
-  // Reset the update protocol chat atom when navigating to the create protocol page
+  // Reset the chat data atom and protocol atoms when navigating to the update protocol page
   useEffect(() => {
+    setCreateProtocolChatAtom({
+      prompt: '',
+      regenerate: false,
+      scientific_application_type: '',
+      description: '',
+      robots: 'opentrons_flex',
+      mounts: [],
+      flexGripper: false,
+      modules: [],
+      labware: [],
+      liquids: [],
+      steps: [],
+      fake: false,
+      fake_id: 0,
+    })
     setUpdateProtocolChatAtom({
       prompt: '',
       protocol_text: '',
@@ -90,6 +109,8 @@ export function CreateProtocol(): JSX.Element | null {
       fake: false,
       fake_id: 0,
     })
+    setChatHistoryAtom([])
+    setChatData([])
   }, [])
 
   useEffect(() => {

opentrons-ai-client/src/organisms/UpdateProtocol/index.tsx → opentrons-ai-client/src/pages/UpdateProtocol/index.tsx

@@ -20,9 +20,11 @@ import { Trans, useTranslation } from 'react-i18next'
 import { FileUpload } from '../../molecules/FileUpload'
 import { useNavigate } from 'react-router-dom'
 import {
+  chatHistoryAtom,
   createProtocolChatAtom,
   headerWithMeterAtom,
   updateProtocolChatAtom,
+  chatDataAtom,
 } from '../../resources/atoms'
 import { CSSTransition } from 'react-transition-group'
 import { useAtom } from 'jotai'
@@ -105,16 +107,19 @@ export function UpdateProtocol(): JSX.Element {
   const [headerState, setHeaderWithMeterAtom] = useAtom(headerWithMeterAtom)
   const [updateType, setUpdateType] = useState<DropdownOption | null>(null)
   const [detailsValue, setDetailsValue] = useState<string>('')
-  const [, setUpdatePromptAtom] = useAtom(updateProtocolChatAtom)
+  const [, setUpdateProtocolChatAtom] = useAtom(updateProtocolChatAtom)
   const [, setCreateProtocolChatAtom] = useAtom(createProtocolChatAtom)
+  const [, setChatHistoryAtom] = useAtom(chatHistoryAtom)
+  const [, setChatData] = useAtom(chatDataAtom)
   const [fileValue, setFile] = useState<File | null>(null)
   const [pythonText, setPythonTextValue] = useState<string>('')
   const [errorText, setErrorText] = useState<string | null>(null)
 
-  // Reset the create protocol chat atom when navigating to the update protocol page
+  // Reset the chat data atom and protocol atoms when navigating to the update protocol page
   useEffect(() => {
     setCreateProtocolChatAtom({
       prompt: '',
+      regenerate: false,
       scientific_application_type: '',
       description: '',
       robots: 'opentrons_flex',
@@ -127,6 +132,17 @@ export function UpdateProtocol(): JSX.Element {
       fake: false,
       fake_id: 0,
     })
+    setUpdateProtocolChatAtom({
+      prompt: '',
+      protocol_text: '',
+      regenerate: false,
+      update_type: 'adapt_python_protocol',
+      update_details: '',
+      fake: false,
+      fake_id: 0,
+    })
+    setChatHistoryAtom([])
+    setChatData([])
   }, [])
 
   useEffect(() => {
@@ -193,7 +209,9 @@ export function UpdateProtocol(): JSX.Element {
 
     const chatPrompt = `${introText}${originalCodeText}${updateTypeText}${detailsText}`
 
-    setUpdatePromptAtom({
+    console.log(chatPrompt)
+
+    setUpdateProtocolChatAtom({
       prompt: chatPrompt,
       protocol_text: pythonText,
       regenerate: false,

opentrons-ai-client/src/resources/atoms.ts

@@ -16,6 +16,7 @@ export const chatDataAtom = atom<ChatData[]>([])
 /** CreateProtocolChatAtom is for the prefilled userprompt when navigating to the chat page from Create New protocol page */
 export const createProtocolChatAtom = atom<CreatePrompt>({
   prompt: '',
+  regenerate: false,
   scientific_application_type: '',
   description: '',
   robots: 'opentrons_flex',
@@ -40,6 +41,15 @@ export const updateProtocolChatAtom = atom<UpdatePrompt>({
   fake_id: 0,
 })
 
+/** Regenerate protocol atom */
+export const regenerateProtocolAtom = atom<{
+  isCreateOrUpdateProtocol: boolean
+  regenerate: boolean
+}>({
+  isCreateOrUpdateProtocol: false,
+  regenerate: false,
+})
+
 /** Scroll to bottom of chat atom */
 export const scrollToBottomAtom = atom<boolean>(false)
 

opentrons-ai-client/src/resources/types.ts

@@ -15,6 +15,7 @@ export interface ChatData {
 export interface CreatePrompt {
   /** the prompt that is generated by the create protocol page */
   prompt: string
+  regenerate: boolean
   scientific_application_type: string
   description: string
   robots: 'opentrons_flex' | 'opentrons_ot2' | string

opentrons-ai-client/src/resources/utils/createProtocolUtils.tsx

@@ -230,6 +230,7 @@ export function generateChatPrompt(
 
   setCreateProtocolChatAtom({
     prompt,
+    regenerate: false,
     scientific_application_type: values.application.scientificApplication,
     description,
     robots: values.instruments.robot,

opentrons-ai-server/api/models/create_protocol.py

@@ -5,6 +5,7 @@
 
 class CreateProtocol(BaseModel):
     prompt: str = Field(..., description="Prompt")
+    regenerate: bool = Field(..., description="Flag to indicate if regeneration is needed")
     scientific_application_type: str = Field(..., description="Scientific application type")
     description: str = Field(..., description="Description of the protocol")
     robots: Literal["opentrons_flex", "opentrons_ot2"] = Field(..., description="List of required robots")