diff --git a/examples/SampleApp/App.tsx b/examples/SampleApp/App.tsx
index 8781a891a1..2b547c3337 100644
--- a/examples/SampleApp/App.tsx
+++ b/examples/SampleApp/App.tsx
@@ -6,6 +6,7 @@ import { createStackNavigator } from '@react-navigation/stack';
 import { SafeAreaProvider, useSafeAreaInsets } from 'react-native-safe-area-context';
 import {
   Chat,
+  MessageType,
   OverlayProvider,
   QuickSqliteClient,
   ThemeProvider,
@@ -182,6 +183,7 @@ const DrawerNavigatorWrapper: React.FC<{
           enableOfflineSupport
           // @ts-expect-error
           ImageComponent={FastImage}
+          isMessageAIGenerated={(message: MessageType) => message.ai_generated}
         >
           <AppOverlayProvider>
             <UserSearchProvider>
diff --git a/examples/SampleApp/yarn.lock b/examples/SampleApp/yarn.lock
index 07526b0fd8..8d0629c810 100644
--- a/examples/SampleApp/yarn.lock
+++ b/examples/SampleApp/yarn.lock
@@ -6837,10 +6837,10 @@ statuses@~1.5.0:
   resolved "https://registry.yarnpkg.com/statuses/-/statuses-1.5.0.tgz#161c7dac177659fd9811f43771fa99381478628c"
   integrity sha512-OpZ3zP+jT1PI7I8nemJX4AKmAX070ZkYPVWV/AaKTJl+tXCTGyVdC1a4SL8RUQYEwk/f34ZX8UTykN68FwrqAA==
 
-stream-chat-react-native-core@5.43.2:
-  version "5.43.2"
-  resolved "https://registry.yarnpkg.com/stream-chat-react-native-core/-/stream-chat-react-native-core-5.43.2.tgz#b16add60d231509f864d7301ae281c4b4681ff7a"
-  integrity sha512-pdSaqw1aPHtxH0md7nnw/TLWPIqsb5TKFzGoWYd/cKRAw4pVDvxolSlha8vYSCwz7MmlarfsJPB0TlS9WcGsDQ==
+stream-chat-react-native-core@5.44.0:
+  version "5.44.0"
+  resolved "https://registry.yarnpkg.com/stream-chat-react-native-core/-/stream-chat-react-native-core-5.44.0.tgz#2d3f8b8ca8a15b6f15fa03bea099e706e4c88e79"
+  integrity sha512-gN7/aaDvESPc2s7d6xyCTu4vpiok1HCwcNiHxvNC2NlvL0O5xd/7240L+eoSPhECXTPZvCgVJVEra9CWMM4pXw==
   dependencies:
     "@gorhom/bottom-sheet" "^4.6.4"
     dayjs "1.10.5"
@@ -6853,7 +6853,7 @@ stream-chat-react-native-core@5.43.2:
     path "0.12.7"
     react-native-markdown-package "1.8.2"
     react-native-url-polyfill "^1.3.0"
-    stream-chat "8.45.1"
+    stream-chat "8.46.0"
 
 "stream-chat-react-native-core@link:../../package":
   version "0.0.0"
@@ -6863,21 +6863,6 @@ stream-chat-react-native-core@5.43.2:
   version "0.0.0"
   uid ""
 
-stream-chat@8.45.1:
-  version "8.45.1"
-  resolved "https://registry.yarnpkg.com/stream-chat/-/stream-chat-8.45.1.tgz#001f452520602ebffe45f1634b5f5dd1ff036ec5"
-  integrity sha512-7OMpL2RHUd+PXSWzhTUAIjjXlI9Oqc4HhUBRfc5i6dK+Ug9S5ertb7RvyzGL5N4ITpq/6ZUAwXtTqRaN9+UUkw==
-  dependencies:
-    "@babel/runtime" "^7.16.3"
-    "@types/jsonwebtoken" "~9.0.0"
-    "@types/ws" "^7.4.0"
-    axios "^1.6.0"
-    base64-js "^1.5.1"
-    form-data "^4.0.0"
-    isomorphic-ws "^4.0.1"
-    jsonwebtoken "~9.0.0"
-    ws "^7.5.10"
-
 stream-chat@8.46.0:
   version "8.46.0"
   resolved "https://registry.yarnpkg.com/stream-chat/-/stream-chat-8.46.0.tgz#416b325e05b144d0937a3527d1e622463113d605"
diff --git a/package/src/components/Channel/Channel.tsx b/package/src/components/Channel/Channel.tsx
index dffd98eff9..a2fecbec66 100644
--- a/package/src/components/Channel/Channel.tsx
+++ b/package/src/components/Channel/Channel.tsx
@@ -36,6 +36,7 @@ import { useCreateTypingContext } from './hooks/useCreateTypingContext';
 
 import { useTargetedMessage } from './hooks/useTargetedMessage';
 
+import { MessageContextValue } from '../../contexts';
 import { ChannelContextValue, ChannelProvider } from '../../contexts/channelContext/ChannelContext';
 import type { UseChannelStateValue } from '../../contexts/channelsStateContext/useChannelState';
 import { useChannelState } from '../../contexts/channelsStateContext/useChannelState';
@@ -338,6 +339,7 @@ export type ChannelPropsWithContext<
       | 'StreamingMessageView'
     >
   > &
+  Partial<Pick<MessageContextValue<StreamChatGenerics>, 'isMessageAIGenerated'>> &
   Partial<Pick<ThreadContextValue<StreamChatGenerics>, 'allowThreadMessagesInChannel'>> & {
     shouldSyncChannel: boolean;
     thread: ThreadType<StreamChatGenerics>;
@@ -532,6 +534,7 @@ const ChannelWithContext = <
     InputGiphySearch = InputGiphyCommandInputDefault,
     InputReplyStateHeader = InputReplyStateHeaderDefault,
     isAttachmentEqual,
+    isMessageAIGenerated = () => false,
     keyboardBehavior,
     KeyboardCompatibleView = KeyboardCompatibleViewDefault,
     keyboardVerticalOffset,
@@ -2424,6 +2427,7 @@ const ChannelWithContext = <
     InlineDateSeparator,
     InlineUnreadIndicator,
     isAttachmentEqual,
+    isMessageAIGenerated,
     legacyImageViewerSwipeBehaviour,
     markdownRules,
     Message,
@@ -2559,7 +2563,8 @@ export const Channel = <
 >(
   props: PropsWithChildren<ChannelProps<StreamChatGenerics>>,
 ) => {
-  const { client, enableOfflineSupport } = useChatContext<StreamChatGenerics>();
+  const { client, enableOfflineSupport, isMessageAIGenerated } =
+    useChatContext<StreamChatGenerics>();
   const { t } = useTranslationContext();
 
   const threadFromProps = props?.thread;
@@ -2605,6 +2610,7 @@ export const Channel = <
       {...props}
       shouldSyncChannel={shouldSyncChannel}
       {...{
+        isMessageAIGenerated,
         members,
         messages: props.messages || messages,
         read,
diff --git a/package/src/components/Channel/hooks/useCreateMessagesContext.ts b/package/src/components/Channel/hooks/useCreateMessagesContext.ts
index 9587c9c1bd..9d65fedba2 100644
--- a/package/src/components/Channel/hooks/useCreateMessagesContext.ts
+++ b/package/src/components/Channel/hooks/useCreateMessagesContext.ts
@@ -50,6 +50,7 @@ export const useCreateMessagesContext = <
   InlineDateSeparator,
   InlineUnreadIndicator,
   isAttachmentEqual,
+  isMessageAIGenerated,
   legacyImageViewerSwipeBehaviour,
   markdownRules,
   Message,
@@ -152,6 +153,7 @@ export const useCreateMessagesContext = <
       InlineDateSeparator,
       InlineUnreadIndicator,
       isAttachmentEqual,
+      isMessageAIGenerated,
       legacyImageViewerSwipeBehaviour,
       markdownRules,
       Message,
diff --git a/package/src/components/Chat/Chat.tsx b/package/src/components/Chat/Chat.tsx
index 06f0f329c5..9905983ca7 100644
--- a/package/src/components/Chat/Chat.tsx
+++ b/package/src/components/Chat/Chat.tsx
@@ -36,7 +36,12 @@ init();
 export type ChatProps<
   StreamChatGenerics extends DefaultStreamChatGenerics = DefaultStreamChatGenerics,
 > = Pick<ChatContextValue<StreamChatGenerics>, 'client'> &
-  Partial<Pick<ChatContextValue<StreamChatGenerics>, 'ImageComponent' | 'resizableCDNHosts'>> & {
+  Partial<
+    Pick<
+      ChatContextValue<StreamChatGenerics>,
+      'ImageComponent' | 'resizableCDNHosts' | 'isMessageAIGenerated'
+    >
+  > & {
     /**
      * When false, ws connection won't be disconnection upon backgrounding the app.
      * To receive push notifications, its necessary that user doesn't have active
@@ -149,6 +154,7 @@ const ChatWithContext = <
     enableOfflineSupport = false,
     i18nInstance,
     ImageComponent = Image,
+    isMessageAIGenerated,
     LoadingIndicator = null,
     resizableCDNHosts = ['.stream-io-cdn.com'],
     style,
@@ -250,6 +256,7 @@ const ChatWithContext = <
     connectionRecovering,
     enableOfflineSupport,
     ImageComponent,
+    isMessageAIGenerated,
     isOnline,
     mutedUsers,
     resizableCDNHosts,
diff --git a/package/src/components/Chat/hooks/useCreateChatContext.ts b/package/src/components/Chat/hooks/useCreateChatContext.ts
index a729cc9c18..1ca1fbf50f 100644
--- a/package/src/components/Chat/hooks/useCreateChatContext.ts
+++ b/package/src/components/Chat/hooks/useCreateChatContext.ts
@@ -12,6 +12,7 @@ export const useCreateChatContext = <
   connectionRecovering,
   enableOfflineSupport,
   ImageComponent,
+  isMessageAIGenerated,
   isOnline,
   mutedUsers,
   resizableCDNHosts,
@@ -33,6 +34,7 @@ export const useCreateChatContext = <
       connectionRecovering,
       enableOfflineSupport,
       ImageComponent,
+      isMessageAIGenerated,
       isOnline,
       mutedUsers,
       resizableCDNHosts,
diff --git a/package/src/components/Message/Message.tsx b/package/src/components/Message/Message.tsx
index bc5d1cadea..891c123415 100644
--- a/package/src/components/Message/Message.tsx
+++ b/package/src/components/Message/Message.tsx
@@ -139,8 +139,16 @@ export type MessagePropsWithContext<
   StreamChatGenerics extends DefaultStreamChatGenerics = DefaultStreamChatGenerics,
 > = Pick<ChannelContextValue<StreamChatGenerics>, 'channel' | 'enforceUniqueReaction' | 'members'> &
   Pick<KeyboardContextValue, 'dismissKeyboard'> &
-  Partial<Omit<MessageContextValue<StreamChatGenerics>, 'groupStyles' | 'message'>> &
-  Pick<MessageContextValue<StreamChatGenerics>, 'groupStyles' | 'message'> &
+  Partial<
+    Omit<
+      MessageContextValue<StreamChatGenerics>,
+      'groupStyles' | 'message' | 'isMessageAIGenerated'
+    >
+  > &
+  Pick<
+    MessageContextValue<StreamChatGenerics>,
+    'groupStyles' | 'message' | 'isMessageAIGenerated'
+  > &
   Pick<
     MessagesContextValue<StreamChatGenerics>,
     | 'sendReaction'
@@ -305,6 +313,11 @@ const MessageWithContext = <
     threadList = false,
     updateMessage,
   } = props;
+  const isMessageAIGenerated = messagesContext.isMessageAIGenerated;
+  const isAIGenerated = useMemo(
+    () => isMessageAIGenerated(message),
+    [message, isMessageAIGenerated],
+  );
   const { client } = chatContext;
   const {
     theme: {
@@ -457,7 +470,7 @@ const MessageWithContext = <
       case 'poll':
         return !!message.poll_id;
       case 'ai_text':
-        return !!message.ai_generated;
+        return isAIGenerated;
       case 'text':
       default:
         return !!message.text;
@@ -704,6 +717,7 @@ const MessageWithContext = <
     hasReactions,
     images: attachments.images,
     isEditedMessageOpen,
+    isMessageAIGenerated,
     isMyMessage,
     lastGroupMessage: groupStyles?.[0] === 'single' || groupStyles?.[0] === 'bottom',
     lastReceivedId,
@@ -865,8 +879,7 @@ const areEqual = <StreamChatGenerics extends DefaultStreamChatGenerics = Default
     prevMessage.text === nextMessage.text &&
     prevMessage.pinned === nextMessage.pinned &&
     `${prevMessage?.updated_at}` === `${nextMessage?.updated_at}` &&
-    prevMessage.i18n === nextMessage.i18n &&
-    prevMessage.ai_generated === nextMessage.ai_generated;
+    prevMessage.i18n === nextMessage.i18n;
 
   if (!messageEqual) return false;
 
diff --git a/package/src/components/Message/MessageSimple/MessageContent.tsx b/package/src/components/Message/MessageSimple/MessageContent.tsx
index 5368cf72d6..ffba84e11c 100644
--- a/package/src/components/Message/MessageSimple/MessageContent.tsx
+++ b/package/src/components/Message/MessageSimple/MessageContent.tsx
@@ -1,4 +1,4 @@
-import React from 'react';
+import React, { useMemo } from 'react';
 import {
   AnimatableNumericValue,
   LayoutChangeEvent,
@@ -82,6 +82,7 @@ export type MessageContentPropsWithContext<
   | 'preventPress'
   | 'showMessageStatus'
   | 'threadList'
+  | 'isMessageAIGenerated'
 > &
   Pick<
     MessagesContextValue<StreamChatGenerics>,
@@ -122,6 +123,7 @@ const MessageContentWithContext = <
     Gallery,
     groupStyles,
     hasReactions,
+    isMessageAIGenerated,
     isMyMessage,
     lastGroupMessage,
     members,
@@ -186,6 +188,11 @@ const MessageContentWithContext = <
     setMessageContentWidth(width);
   };
 
+  const isAIGenerated = useMemo(
+    () => isMessageAIGenerated(message),
+    [message, isMessageAIGenerated],
+  );
+
   const error = message.type === 'error' || message.status === MessageStatusTypes.FAILED;
 
   const groupStyle = `${alignment}_${groupStyles?.[0]?.toLowerCase?.()}`;
@@ -396,7 +403,7 @@ const MessageContentWithContext = <
                 ) : null;
               }
               case 'ai_text':
-                return message.ai_generated ? (
+                return isAIGenerated ? (
                   <StreamingMessageView
                     key={`ai_message_text_container_${messageContentOrderIndex}`}
                   />
@@ -404,7 +411,7 @@ const MessageContentWithContext = <
               case 'text':
               default:
                 return (otherAttachments.length && otherAttachments[0].actions) ||
-                  message.ai_generated ? null : (
+                  isAIGenerated ? null : (
                   <MessageTextContainer<StreamChatGenerics>
                     key={`message_text_container_${messageContentOrderIndex}`}
                   />
@@ -493,8 +500,7 @@ const areEqual = <StreamChatGenerics extends DefaultStreamChatGenerics = Default
     prevMessage.type === nextMessage.type &&
     prevMessage.text === nextMessage.text &&
     prevMessage.pinned === nextMessage.pinned &&
-    prevMessage.i18n === nextMessage.i18n &&
-    prevMessage.ai_generated === nextMessage.ai_generated;
+    prevMessage.i18n === nextMessage.i18n;
   if (!messageEqual) return false;
 
   const isPrevQuotedMessageTypeDeleted = prevMessage.quoted_message?.type === 'deleted';
@@ -577,6 +583,7 @@ export const MessageContent = <
     groupStyles,
     hasReactions,
     isEditedMessageOpen,
+    isMessageAIGenerated,
     isMyMessage,
     lastGroupMessage,
     lastReceivedId,
@@ -624,6 +631,7 @@ export const MessageContent = <
         hasReactions,
         isAttachmentEqual,
         isEditedMessageOpen,
+        isMessageAIGenerated,
         isMyMessage,
         lastGroupMessage,
         lastReceivedId,
diff --git a/package/src/components/Message/MessageSimple/MessageFooter.tsx b/package/src/components/Message/MessageSimple/MessageFooter.tsx
index 0cf33cbb31..d037995b62 100644
--- a/package/src/components/Message/MessageSimple/MessageFooter.tsx
+++ b/package/src/components/Message/MessageSimple/MessageFooter.tsx
@@ -1,4 +1,4 @@
-import React from 'react';
+import React, { useMemo } from 'react';
 import { StyleSheet, Text, View } from 'react-native';
 
 import type { Attachment } from 'stream-chat';
@@ -41,6 +41,7 @@ type MessageFooterPropsWithContext<
   | 'otherAttachments'
   | 'showMessageStatus'
   | 'lastGroupMessage'
+  | 'isMessageAIGenerated'
 > &
   Pick<
     MessagesContextValue<StreamChatGenerics>,
@@ -94,6 +95,7 @@ const MessageFooterWithContext = <
     formattedDate,
     isDeleted,
     isEditedMessageOpen,
+    isMessageAIGenerated,
     lastGroupMessage,
     members,
     message,
@@ -114,6 +116,11 @@ const MessageFooterWithContext = <
   } = useTheme();
   const { t } = useTranslationContext();
 
+  const isAIGenerated = useMemo(
+    () => isMessageAIGenerated(message),
+    [message, isMessageAIGenerated],
+  );
+
   if (isDeleted) {
     return (
       <View style={[styles.container, metaContainer]} testID='message-deleted'>
@@ -129,7 +136,7 @@ const MessageFooterWithContext = <
     return null;
   }
 
-  const isEdited = isEditedMessage(message);
+  const isEdited = isEditedMessage(message) && !isAIGenerated;
 
   return (
     <>
@@ -267,6 +274,7 @@ export const MessageFooter = <
   const {
     alignment,
     isEditedMessageOpen,
+    isMessageAIGenerated,
     lastGroupMessage,
     members,
     message,
@@ -283,6 +291,7 @@ export const MessageFooter = <
         alignment,
         deletedMessagesVisibilityType,
         isEditedMessageOpen,
+        isMessageAIGenerated,
         lastGroupMessage,
         members,
         message,
diff --git a/package/src/components/Message/MessageSimple/MessageSimple.tsx b/package/src/components/Message/MessageSimple/MessageSimple.tsx
index ea3a66d27b..ee42dece27 100644
--- a/package/src/components/Message/MessageSimple/MessageSimple.tsx
+++ b/package/src/components/Message/MessageSimple/MessageSimple.tsx
@@ -155,8 +155,7 @@ const areEqual = <StreamChatGenerics extends DefaultStreamChatGenerics = Default
     prevMessage.type === nextMessage.type &&
     prevMessage.text === nextMessage.text &&
     prevMessage.i18n === nextMessage.i18n &&
-    prevMessage.pinned === nextMessage.pinned &&
-    prevMessage.ai_generated === nextMessage.ai_generated;
+    prevMessage.pinned === nextMessage.pinned;
   if (!messageEqual) return false;
 
   const isPrevQuotedMessageTypeDeleted = prevMessage.quoted_message?.type === 'deleted';
diff --git a/package/src/components/Message/hooks/useCreateMessageContext.ts b/package/src/components/Message/hooks/useCreateMessageContext.ts
index eedd4d561e..205533a8d1 100644
--- a/package/src/components/Message/hooks/useCreateMessageContext.ts
+++ b/package/src/components/Message/hooks/useCreateMessageContext.ts
@@ -26,6 +26,7 @@ export const useCreateMessageContext = <
   hasReactions,
   images,
   isEditedMessageOpen,
+  isMessageAIGenerated,
   isMyMessage,
   lastGroupMessage,
   lastReceivedId,
@@ -78,6 +79,7 @@ export const useCreateMessageContext = <
       hasReactions,
       images,
       isEditedMessageOpen,
+      isMessageAIGenerated,
       isMyMessage,
       lastGroupMessage,
       lastReceivedId,
diff --git a/package/src/components/MessageOverlay/MessageOverlay.tsx b/package/src/components/MessageOverlay/MessageOverlay.tsx
index a168bff155..21893f887a 100644
--- a/package/src/components/MessageOverlay/MessageOverlay.tsx
+++ b/package/src/components/MessageOverlay/MessageOverlay.tsx
@@ -158,6 +158,7 @@ const MessageOverlayWithContext = <
   const halfScreenHeight = vh(50);
 
   const myMessageTheme = messagesContext?.myMessageTheme;
+  const isMessageAIGenerated = messagesContext?.isMessageAIGenerated;
   const wrapMessageInTheme = clientId === message?.user?.id && !!myMessageTheme;
 
   const [reactionListHeight, setReactionListHeight] = useState(0);
@@ -462,7 +463,7 @@ const MessageOverlayWithContext = <
                       case 'text':
                       default:
                         return (otherAttachments?.length && otherAttachments[0].actions) ||
-                          message.ai_generated ? null : (
+                          isMessageAIGenerated?.(message) ? null : (
                           <MessageTextContainer<StreamChatGenerics>
                             key={`message_text_container_${messageContentOrderIndex}`}
                             message={message}
diff --git a/package/src/contexts/chatContext/ChatContext.tsx b/package/src/contexts/chatContext/ChatContext.tsx
index 79dc046aee..54949ab653 100644
--- a/package/src/contexts/chatContext/ChatContext.tsx
+++ b/package/src/contexts/chatContext/ChatContext.tsx
@@ -4,6 +4,7 @@ import type { ImageProps } from 'react-native';
 import type { AppSettingsAPIResponse, Channel, Mute, StreamChat } from 'stream-chat';
 
 import type { DefaultStreamChatGenerics, UnknownType } from '../../types/types';
+import { MessageContextValue } from '../messageContext/MessageContext';
 import { DEFAULT_BASE_CONTEXT_VALUE } from '../utils/defaultBaseContextValue';
 
 import { getDisplayName } from '../utils/getDisplayName';
@@ -69,7 +70,7 @@ export type ChatContextValue<
    * This option allows you to specify a list of CDNs that offer image resizing.
    */
   resizableCDNHosts?: string[];
-};
+} & Partial<Pick<MessageContextValue<StreamChatGenerics>, 'isMessageAIGenerated'>>;
 
 export const ChatContext = React.createContext(DEFAULT_BASE_CONTEXT_VALUE as ChatContextValue);
 
diff --git a/package/src/contexts/messageContext/MessageContext.tsx b/package/src/contexts/messageContext/MessageContext.tsx
index e482e19f5d..6d5e6f6d30 100644
--- a/package/src/contexts/messageContext/MessageContext.tsx
+++ b/package/src/contexts/messageContext/MessageContext.tsx
@@ -85,6 +85,10 @@ export type MessageContextValue<
   images: Attachment<StreamChatGenerics>[];
   /** Boolean that determines if the edited message is pressed. */
   isEditedMessageOpen: boolean;
+  /**
+   * A factory function that determines whether a message is AI generated or not.
+   */
+  isMessageAIGenerated: (message: MessageType<StreamChatGenerics>) => boolean;
   /** Whether or not this is the active user's message */
   isMyMessage: boolean;
   /** Whether or not this is the last message in a group of messages */
diff --git a/package/src/contexts/messageInputContext/MessageInputContext.tsx b/package/src/contexts/messageInputContext/MessageInputContext.tsx
index a84bf72413..253c306a1a 100644
--- a/package/src/contexts/messageInputContext/MessageInputContext.tsx
+++ b/package/src/contexts/messageInputContext/MessageInputContext.tsx
@@ -385,7 +385,7 @@ export type InputMessageInputContextValue<
    * Defaults to and accepts same props as: [AudioRecordingButton](https://github.com/GetStream/stream-chat-react-native/blob/main/package/src/components/MessageInput/components/AudioRecorder/AudioRecordingButton.tsx)
    */
   StartAudioRecordingButton: React.ComponentType<AudioRecordingButtonProps<StreamChatGenerics>>;
-  StopMessageStreamingButton: React.ComponentType<StopMessageStreamingButtonProps>;
+  StopMessageStreamingButton: React.ComponentType<StopMessageStreamingButtonProps> | null;
   /**
    * Custom UI component to render upload progress indicator on attachment preview.
    *
diff --git a/package/src/contexts/messagesContext/MessagesContext.tsx b/package/src/contexts/messagesContext/MessagesContext.tsx
index bde9f4d4dc..3fc752a368 100644
--- a/package/src/contexts/messagesContext/MessagesContext.tsx
+++ b/package/src/contexts/messagesContext/MessagesContext.tsx
@@ -52,7 +52,7 @@ import type { ReplyProps } from '../../components/Reply/Reply';
 import type { FlatList } from '../../native';
 import type { DefaultStreamChatGenerics, UnknownType } from '../../types/types';
 import type { ReactionData } from '../../utils/utils';
-import type { Alignment } from '../messageContext/MessageContext';
+import type { Alignment, MessageContextValue } from '../messageContext/MessageContext';
 import type { SuggestionCommand } from '../suggestionsContext/SuggestionsContext';
 import type { DeepPartial } from '../themeContext/ThemeContext';
 import type { Theme } from '../themeContext/utils/theme';
@@ -73,7 +73,7 @@ export type DeletedMessagesVisibilityType = 'always' | 'never' | 'receiver' | 's
 
 export type MessagesContextValue<
   StreamChatGenerics extends DefaultStreamChatGenerics = DefaultStreamChatGenerics,
-> = {
+> = Pick<MessageContextValue<StreamChatGenerics>, 'isMessageAIGenerated'> & {
   /**
    * UI component for Attachment.
    * Defaults to: [Attachment](https://github.com/GetStream/stream-chat-react-native/blob/main/package/src/components/Attachment/Attachment.tsx)
diff --git a/package/src/utils/utils.ts b/package/src/utils/utils.ts
index 18434107aa..009ba25852 100644
--- a/package/src/utils/utils.ts
+++ b/package/src/utils/utils.ts
@@ -108,7 +108,7 @@ export const isEditedMessage = <
   StreamChatGenerics extends DefaultStreamChatGenerics = DefaultStreamChatGenerics,
 >(
   message: MessageType<StreamChatGenerics>,
-) => !!message.message_text_updated_at && !message.ai_generated;
+) => !!message.message_text_updated_at;
 
 /**
  * Default emoji search index for auto complete text input