From 055c4b74bfe137d3a75b4e316aa9cccd58201815 Mon Sep 17 00:00:00 2001
From: levanion <levanion1998@gmail.com>
Date: Thu, 29 Feb 2024 18:31:01 +0400
Subject: [PATCH] fix: voice on chat

---
 apps/server/agents/conversational/conversational.py    |  4 +++-
 apps/ui/src/modals/AIChatModal/hooks/useChatSocket.ts  |  1 +
 .../src/modals/AIChatModal/hooks/useUpdateChatCache.ts | 10 +++++++++-
 3 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/apps/server/agents/conversational/conversational.py b/apps/server/agents/conversational/conversational.py
index 5c0638928..bfffd32a1 100644
--- a/apps/server/agents/conversational/conversational.py
+++ b/apps/server/agents/conversational/conversational.py
@@ -113,9 +113,11 @@ async def run(
 
             yield res
 
-        history.create_ai_message(
+        ai_message = history.create_ai_message(
             res,
             human_message_id,
             agent_with_configs.agent.id,
             voice_url,
         )
+
+        chat_pubsub_service.send_chat_message(chat_message=ai_message)
diff --git a/apps/ui/src/modals/AIChatModal/hooks/useChatSocket.ts b/apps/ui/src/modals/AIChatModal/hooks/useChatSocket.ts
index 79ccc5a3e..32770cf74 100644
--- a/apps/ui/src/modals/AIChatModal/hooks/useChatSocket.ts
+++ b/apps/ui/src/modals/AIChatModal/hooks/useChatSocket.ts
@@ -70,6 +70,7 @@ const useChatSocket = ({ userId, createdChatId }: ChatSocketProps) => {
       }
       if (data.type === 'CHAT_MESSAGE_ADDED') {
         upsertChatMessageInCache(data.chat_message, {
+          replaceCache: data.chat_message.message.data.type === 'ai' ? true : false,
           agentId: data.agent_id,
           teamId: data.team_id,
           chatId: data.chat_id,
diff --git a/apps/ui/src/modals/AIChatModal/hooks/useUpdateChatCache.ts b/apps/ui/src/modals/AIChatModal/hooks/useUpdateChatCache.ts
index 393296f39..d1e2a8c12 100644
--- a/apps/ui/src/modals/AIChatModal/hooks/useUpdateChatCache.ts
+++ b/apps/ui/src/modals/AIChatModal/hooks/useUpdateChatCache.ts
@@ -13,16 +13,19 @@ const useUpdateChatCache = () => {
 
   const upsertChatMessageInCache = (
     newChatMessage: Record<string, unknown>,
+
     {
       localChatMessageRefId,
       agentId,
       teamId,
       chatId,
+      replaceCache,
     }: {
       localChatMessageRefId?: Nullable<string>
       agentId?: Nullable<string>
       teamId?: Nullable<string>
       chatId?: Nullable<string>
+      replaceCache?: boolean
     } = {},
   ) => {
     let queryVariables = omitBy(
@@ -45,7 +48,7 @@ const useUpdateChatCache = () => {
       { query: CHAT_MESSAGES_GQL, variables: queryVariables },
       data => {
         const chatMessages = data?.chatMessages || []
-        const newChatMessages = [...chatMessages]
+        let newChatMessages = [...chatMessages]
         newChatMessage = {
           __typename: 'ChatMessage',
           parent: null,
@@ -54,6 +57,11 @@ const useUpdateChatCache = () => {
           ...newChatMessage,
         }
 
+        if (replaceCache) {
+          const lastItemId = newChatMessages[newChatMessages.length - 1].id
+          newChatMessages = newChatMessages.filter((message: any) => message.id !== lastItemId)
+        }
+
         if (localChatMessageRefId && user.id === newChatMessage.sender_user_id) {
           // If the message is from the current user, we need to update the local message