From f2256423d50aca7e176bf1f2ebe1429ca4176cc5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Nuno=20G=C3=B3is?= <github@nunogois.com>
Date: Thu, 17 Oct 2024 09:50:27 +0100
Subject: [PATCH] chore: adapt UI to server-side Unleash AI chat ownership
 (#8466)

https://linear.app/unleash/issue/2-2847/adapt-unleash-ai-chat-logic-to-new-server-side-chat-ownership-logic

Adapts the Unleash AI chat logic on the UI to the new server-side chat
ownership logic.
---
 frontend/src/component/ai/AIChat.tsx          | 51 +++++++++++--------
 .../hooks/api/actions/useAIApi/useAIApi.ts    | 39 ++++++++++----
 2 files changed, 60 insertions(+), 30 deletions(-)

diff --git a/frontend/src/component/ai/AIChat.tsx b/frontend/src/component/ai/AIChat.tsx
index ac085ba4c0..9fea0bf4d2 100644
--- a/frontend/src/component/ai/AIChat.tsx
+++ b/frontend/src/component/ai/AIChat.tsx
@@ -15,6 +15,11 @@ import { AIChatMessage } from './AIChatMessage';
 import { AIChatHeader } from './AIChatHeader';
 import { Resizable } from 'component/common/Resizable/Resizable';
 
+const AI_ERROR_MESSAGE = {
+    role: 'assistant',
+    content: `I'm sorry, I'm having trouble understanding you right now. I've reported the issue to the team. Please try again later.`,
+} as const;
+
 const StyledAIIconContainer = styled('div')(({ theme }) => ({
     position: 'fixed',
     bottom: 20,
@@ -71,13 +76,6 @@ const StyledChatContent = styled('div')(({ theme }) => ({
     overflowX: 'hidden',
 }));
 
-const initialMessages: ChatMessage[] = [
-    {
-        role: 'system',
-        content: `You are an assistant that helps users interact with Unleash. You should ask the user in case you're missing any required information.`,
-    },
-];
-
 export const AIChat = () => {
     const unleashAIEnabled = useUiFlag('unleashAI');
     const {
@@ -86,9 +84,9 @@ export const AIChat = () => {
     const [open, setOpen] = useState(false);
     const [loading, setLoading] = useState(false);
     const { setToastApiError } = useToast();
-    const { chat } = useAIApi();
+    const { chat, newChat } = useAIApi();
 
-    const [messages, setMessages] = useState<ChatMessage[]>(initialMessages);
+    const [messages, setMessages] = useState<ChatMessage[]>([]);
 
     const chatEndRef = useRef<HTMLDivElement | null>(null);
 
@@ -106,26 +104,34 @@ export const AIChat = () => {
         scrollToEnd();
     }, [open]);
 
-    const onSend = async (message: string) => {
-        if (!message.trim() || loading) return;
+    const onSend = async (content: string) => {
+        if (!content.trim() || loading) return;
 
         try {
             setLoading(true);
-            const tempMessages: ChatMessage[] = [
-                ...messages,
-                { role: 'user', content: message },
-                { role: 'assistant', content: '_Unleash AI is typing..._' },
-            ];
-            setMessages(tempMessages);
-            const newMessages = await chat(tempMessages.slice(0, -1));
+            setMessages((currentMessages) => [
+                ...currentMessages,
+                { role: 'user', content },
+            ]);
+            const { messages: newMessages } = await chat(content);
             mutate(() => true);
             setMessages(newMessages);
-            setLoading(false);
         } catch (error: unknown) {
+            setMessages((currentMessages) => [
+                ...currentMessages,
+                AI_ERROR_MESSAGE,
+            ]);
             setToastApiError(formatUnknownError(error));
+        } finally {
+            setLoading(false);
         }
     };
 
+    const onNewChat = () => {
+        setMessages([]);
+        newChat();
+    };
+
     if (!unleashAIEnabled || !unleashAIAvailable) {
         return null;
     }
@@ -151,7 +157,7 @@ export const AIChat = () => {
             >
                 <StyledChat>
                     <AIChatHeader
-                        onNew={() => setMessages(initialMessages)}
+                        onNew={onNewChat}
                         onClose={() => setOpen(false)}
                     />
                     <StyledChatContent>
@@ -163,6 +169,11 @@ export const AIChat = () => {
                                 {content}
                             </AIChatMessage>
                         ))}
+                        {loading && (
+                            <AIChatMessage from='assistant'>
+                                _Unleash AI is typing..._
+                            </AIChatMessage>
+                        )}
                         <div ref={chatEndRef} />
                     </StyledChatContent>
                     <AIChatInput onSend={onSend} loading={loading} />
diff --git a/frontend/src/hooks/api/actions/useAIApi/useAIApi.ts b/frontend/src/hooks/api/actions/useAIApi/useAIApi.ts
index 0d0d5fc3f1..1c62bac3ff 100644
--- a/frontend/src/hooks/api/actions/useAIApi/useAIApi.ts
+++ b/frontend/src/hooks/api/actions/useAIApi/useAIApi.ts
@@ -1,3 +1,4 @@
+import { useState } from 'react';
 import useAPI from '../useApi/useApi';
 
 const ENDPOINT = 'api/admin/ai';
@@ -7,29 +8,47 @@ export type ChatMessage = {
     content: string;
 };
 
+type Chat = {
+    id: string;
+    userId: number;
+    createdAt: string;
+    messages: ChatMessage[];
+};
+
 export const useAIApi = () => {
     const { makeRequest, createRequest, errors, loading } = useAPI({
         propagateErrors: true,
     });
 
-    const chat = async (messages: ChatMessage[]): Promise<ChatMessage[]> => {
+    const [chatId, setChatId] = useState<string>();
+
+    const chat = async (message: string): Promise<Chat> => {
         const requestId = 'chat';
 
-        const req = createRequest(`${ENDPOINT}/chat`, {
-            method: 'POST',
-            body: JSON.stringify({
-                messages,
-            }),
-            requestId,
-        });
+        const req = createRequest(
+            `${ENDPOINT}/chat${chatId ? `/${chatId}` : ''}`,
+            {
+                method: 'POST',
+                body: JSON.stringify({
+                    message,
+                }),
+                requestId,
+            },
+        );
 
         const response = await makeRequest(req.caller, req.id);
-        const { messages: newMessages } = await response.json();
-        return newMessages;
+        const chat: Chat = await response.json();
+        setChatId(chat.id);
+        return chat;
+    };
+
+    const newChat = () => {
+        setChatId(undefined);
     };
 
     return {
         chat,
+        newChat,
         errors,
         loading,
     };