From 4e14cfd966f29a74523ff66a2cea92a95bdd7b3a Mon Sep 17 00:00:00 2001
From: Jawad Laraqui <waziers@users.noreply.github.com>
Date: Wed, 22 May 2024 22:09:32 -0400
Subject: [PATCH] fix(assistantState): cache full state when reloading
 assistant (#39)

- Add examples dependency to submit handler (last state loaded)
- Center data-load spinner and wrap both UIs

Co-authored-by: giancarlo tricerri <giancarlo@datadriven.works>
---
 .../src/pages/ExploreAssistantPage/index.tsx  | 12 +++----
 .../src/pages/ExploreBasePage/index.tsx       | 35 +++++++++++++++++++
 .../src/pages/ExploreChatPage/index.tsx       | 11 ++----
 3 files changed, 43 insertions(+), 15 deletions(-)
 create mode 100644 explore-assistant-extension/src/pages/ExploreBasePage/index.tsx

diff --git a/explore-assistant-extension/src/pages/ExploreAssistantPage/index.tsx b/explore-assistant-extension/src/pages/ExploreAssistantPage/index.tsx
index 9b0af37c..3afd8260 100644
--- a/explore-assistant-extension/src/pages/ExploreAssistantPage/index.tsx
+++ b/explore-assistant-extension/src/pages/ExploreAssistantPage/index.tsx
@@ -3,7 +3,6 @@ import {
   Button,
   FieldTextArea,
   Heading,
-  Layout,
   Paragraph,
   Section,
   Space,
@@ -26,6 +25,7 @@ import SamplePrompts from '../../components/SamplePrompts'
 import PromptHistory from '../../components/PromptHistory'
 import { RootState } from '../../store'
 import useFetchData from '../../hooks/useSendVertexMessage'
+import ExploreBasePage from '../ExploreBasePage/'
 
 const ExploreAssistantPage = () => {
   const dispatch = useDispatch()
@@ -33,7 +33,7 @@ const ExploreAssistantPage = () => {
   const [textAreaValue, setTextAreaValue] = React.useState<string>('')
   const { extensionSDK } = useContext(ExtensionContext)
 
-  const { exploreUrl, isQuerying, history } = useSelector(
+  const { exploreUrl, isQuerying, history, examples} = useSelector(
     (state: RootState) => state.assistant,
   )
 
@@ -80,7 +80,7 @@ const ExploreAssistantPage = () => {
         JSON.stringify(updatedHistory),
       )
     },
-    [history],
+    [history, examples],
   )
 
   const handleSubmit = useCallback(async () => {
@@ -97,8 +97,7 @@ const ExploreAssistantPage = () => {
   }
 
   return (
-    <>
-      <Layout height={'100%'} hasAside={true}>
+      <ExploreBasePage>
         <Aside
           paddingX={'u8'}
           paddingY={'u4'}
@@ -157,8 +156,7 @@ const ExploreAssistantPage = () => {
             </Space>
           )}
         </Section>
-      </Layout>
-    </>
+      </ExploreBasePage>
   )
 }
 
diff --git a/explore-assistant-extension/src/pages/ExploreBasePage/index.tsx b/explore-assistant-extension/src/pages/ExploreBasePage/index.tsx
new file mode 100644
index 00000000..506822ae
--- /dev/null
+++ b/explore-assistant-extension/src/pages/ExploreBasePage/index.tsx
@@ -0,0 +1,35 @@
+import {
+  Layout,
+  Spinner,
+  Space,
+  SpaceVertical
+} from '@looker/components'
+import React, { useEffect } from 'react'
+import { useSelector } from 'react-redux'
+import { RootState } from '../../store'
+
+const ExploreBasePage = ({ children }: { children: React.ReactNode }) => {
+  const [loading, setLoading] = React.useState<boolean>(true)
+  const { examples } =
+    useSelector((state: RootState) => state.assistant)
+
+
+  useEffect(() => {
+    if (
+      examples.exploreGenerationExamples.length > 0 &&
+      examples.exploreRefinementExamples.length > 0
+    ) {
+      setLoading(false)
+    }
+  }, [examples])
+
+  return (
+    <Layout height={'100%'} hasAside={true}>
+      {loading ? (
+        <Space><SpaceVertical align={'center'} ><Spinner color="key" /></SpaceVertical></Space>
+      ) : (children)}
+    </Layout>
+  )
+}
+
+export default ExploreBasePage
diff --git a/explore-assistant-extension/src/pages/ExploreChatPage/index.tsx b/explore-assistant-extension/src/pages/ExploreChatPage/index.tsx
index f24b9c0b..43ff6d21 100644
--- a/explore-assistant-extension/src/pages/ExploreChatPage/index.tsx
+++ b/explore-assistant-extension/src/pages/ExploreChatPage/index.tsx
@@ -5,12 +5,10 @@ import {
   FieldTextArea,
   Heading,
   Icon,
-  Layout,
   Section,
   Space,
   SpaceVertical,
   Span,
-  Spinner,
 } from '@looker/components'
 import React, { FormEvent, useCallback, useEffect } from 'react'
 import { useDispatch, useSelector } from 'react-redux'
@@ -28,6 +26,7 @@ import Chat from '../../components/Chat'
 import { ArrowBackIosSharp } from '@material-ui/icons'
 import GeminiLogo from '../../components/GeminiLogo'
 import { ExploreEmbed } from '../../components/ExploreEmbed'
+import ExploreBasePage from '../ExploreBasePage/'
 
 const ExploreChatPage = () => {
   const dispatch = useDispatch()
@@ -104,10 +103,7 @@ const ExploreChatPage = () => {
   }
 
   return (
-    <Layout height={'100%'} hasAside={true}>
-      {loading ? (
-        <Spinner color="key" />
-      ) : (
+    <ExploreBasePage>
         <>
           <Section>
             {exploreUrl != '' ? (
@@ -176,8 +172,7 @@ const ExploreChatPage = () => {
             )}
           </Aside>
         </>
-      )}
-    </Layout>
+    </ExploreBasePage>
   )
 }