From c8c0b9240a3fa9a5d719fcd7b1f3f2fec8e0304f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=D0=A0=D0=B0=D1=81=D1=83=D0=BB?=
 <rasul.sahibjanli@gmail.com>
Date: Thu, 1 Aug 2024 15:43:45 +0300
Subject: [PATCH] fix: fix loading of ai question

---
 src/components/App/SideBar/AiSummary/index.tsx |  7 ++++---
 src/components/App/SideBar/index.tsx           |  7 ++++++-
 src/components/App/index.tsx                   | 14 ++++++--------
 src/stores/useDataStore/index.ts               |  8 +++++++-
 src/types/index.ts                             |  4 ++++
 5 files changed, 27 insertions(+), 13 deletions(-)

diff --git a/src/components/App/SideBar/AiSummary/index.tsx b/src/components/App/SideBar/AiSummary/index.tsx
index 79ee7daa3..44c12c0d2 100644
--- a/src/components/App/SideBar/AiSummary/index.tsx
+++ b/src/components/App/SideBar/AiSummary/index.tsx
@@ -17,6 +17,7 @@ import { AiSummarySkeleton } from './AiSummarySkeleton'
 type Props = {
   question: string
   response: AIEntity
+  refId: string
 }
 
 const Title = styled(Text)`
@@ -37,7 +38,7 @@ const TitleWrapper = styled(Flex).attrs({
   overflow: hidden;
 `
 
-export const AiSummary = ({ question, response }: Props) => {
+export const AiSummary = ({ question, response, refId }: Props) => {
   const ref = useRef<HTMLDivElement>(null)
   const [collapsed, setCollapsed] = useState(false)
   const { setAiSummaryAnswer } = useAiSummaryStore((s) => s)
@@ -53,8 +54,8 @@ export const AiSummary = ({ question, response }: Props) => {
   }
 
   const handleLoaded = () => {
-    if (question) {
-      setAiSummaryAnswer(question, { hasBeenRendered: true })
+    if (refId) {
+      setAiSummaryAnswer(refId, { hasBeenRendered: true })
     }
   }
 
diff --git a/src/components/App/SideBar/index.tsx b/src/components/App/SideBar/index.tsx
index 53a0cc19a..4508a50e9 100644
--- a/src/components/App/SideBar/index.tsx
+++ b/src/components/App/SideBar/index.tsx
@@ -208,7 +208,12 @@ const Content = forwardRef<HTMLDivElement, ContentProp>(({ subViewOpen }, ref) =
         )}
         <Flex>
           {Object.keys(aiSummaryAnswers).map((i: string) => (
-            <AiSummary key={i} question={i} response={aiSummaryAnswers[i]} />
+            <AiSummary
+              key={i}
+              question={aiSummaryAnswers[i]?.question || ''}
+              refId={i}
+              response={aiSummaryAnswers[i]}
+            />
           ))}
 
           {isLoading ? <EpisodeSkeleton /> : !hasAiChats && <LatestView isSearchResult={!!searchTerm || hasAiChats} />}
diff --git a/src/components/App/index.tsx b/src/components/App/index.tsx
index b219ab1ac..c1edf87a1 100644
--- a/src/components/App/index.tsx
+++ b/src/components/App/index.tsx
@@ -9,7 +9,6 @@ import { Flex } from '~/components/common/Flex'
 import { DataRetriever } from '~/components/DataRetriever'
 import { GlobalStyle } from '~/components/GlobalStyle'
 import { Overlay } from '~/components/Universe/Overlay' // Import Overlay directly
-import { Preloader } from '~/components/Universe/Preloader' // Import Preloader directly
 import { isDevelopment } from '~/constants'
 import { useSocket } from '~/hooks/useSockets'
 import { useAiSummaryStore } from '~/stores/useAiSummaryStore'
@@ -128,23 +127,23 @@ export const App = () => {
 
   const handleAiSummaryAnswer = useCallback(
     (data: AiSummaryAnswerResponse) => {
-      if (data.question && getKeyExist(aiRefId)) {
-        setAiSummaryAnswer(aiRefId, { answer: data.answer, answerLoading: false })
+      if (data.ref_id) {
+        setAiSummaryAnswer(data.ref_id, { answer: data.answer, answerLoading: false })
       }
     },
-    [setAiSummaryAnswer, getKeyExist, aiRefId],
+    [setAiSummaryAnswer],
   )
 
   const handleAiRelevantQuestions = useCallback(
     (data: AiSummaryQuestionsResponse) => {
-      if (data.question && getKeyExist(aiRefId)) {
-        setAiSummaryAnswer(aiRefId, {
+      if (data.ref_id) {
+        setAiSummaryAnswer(data.ref_id, {
           questions: data.relevant_questions.map((i) => i.question),
           questionsLoading: false,
         })
       }
     },
-    [setAiSummaryAnswer, getKeyExist, aiRefId],
+    [setAiSummaryAnswer],
   )
 
   const handleAiSources = useCallback(
@@ -239,7 +238,6 @@ export const App = () => {
               <LazyMainToolbar />
               <LazySideBar />
               <LazyUniverse />
-              {false && <Preloader fullSize={false} />}
               <Overlay />
               <SecondarySideBar />
               <AppBar />
diff --git a/src/stores/useDataStore/index.ts b/src/stores/useDataStore/index.ts
index d8965b83f..ee0de700c 100644
--- a/src/stores/useDataStore/index.ts
+++ b/src/stores/useDataStore/index.ts
@@ -153,7 +153,6 @@ export const useDataStore = create<DataStore>()(
       }
 
       if (AISearchQuery) {
-        setAiSummaryAnswer(AISearchQuery, { answer: '', answerLoading: true, sourcesLoading: true })
         ai = { ...ai, ai_summary: String(true) }
       }
 
@@ -189,6 +188,13 @@ export const useDataStore = create<DataStore>()(
 
         if (data?.query_data?.ref_id) {
           useAiSummaryStore.setState({ aiRefId: data?.query_data?.ref_id })
+
+          setAiSummaryAnswer(data?.query_data?.ref_id, {
+            question: AISearchQuery,
+            answer: '',
+            answerLoading: true,
+            sourcesLoading: true,
+          })
         }
 
         const currentNodes = currentPage === 0 && !aiRefId ? [] : [...(existingData?.nodes || [])]
diff --git a/src/types/index.ts b/src/types/index.ts
index 844a4fa2c..5f1e6ee57 100644
--- a/src/types/index.ts
+++ b/src/types/index.ts
@@ -288,19 +288,23 @@ export type RelayUser = {
 export type AiSummaryAnswerResponse = {
   question: string
   answer: string
+  ref_id: string
 }
 
 export type AiSummarySourcesResponse = {
   question: string
   sources: { ref_id: string }[]
+  ref_id: string
 }
 
 export type AiSummaryQuestionsResponse = {
   question: string
   relevant_questions: { question: string }[]
+  ref_id: string
 }
 
 export type AIEntity = {
+  question?: string
   answer?: string
   sources?: string[]
   questions?: string[]