diff --git a/examples/vite/src/App.tsx b/examples/vite/src/App.tsx
index f617e1f1c..ced37ffe8 100644
--- a/examples/vite/src/App.tsx
+++ b/examples/vite/src/App.tsx
@@ -73,7 +73,7 @@ const App = () => {
   if (!chatClient) return <>Loading...</>;
 
   return (
-    <Chat client={chatClient}>
+    <Chat client={chatClient} isMessageAIGenerated={(message) => !!message?.ai_generated}>
       <ChatView>
         <ChatView.Selector />
         <ChatView.Channels>
diff --git a/src/components/ChannelList/ChannelList.tsx b/src/components/ChannelList/ChannelList.tsx
index 6b7ba41eb..73fd8564e 100644
--- a/src/components/ChannelList/ChannelList.tsx
+++ b/src/components/ChannelList/ChannelList.tsx
@@ -30,7 +30,7 @@ import { LoadingChannels } from '../Loading/LoadingChannels';
 import { LoadMorePaginator, LoadMorePaginatorProps } from '../LoadMore/LoadMorePaginator';
 import { NullComponent } from '../UtilityComponents';
 
-import { ChannelListContextProvider } from '../../context';
+import { ChannelListContextProvider, ChatContextValue } from '../../context';
 import { useChatContext } from '../../context/ChatContext';
 
 import type { Channel, ChannelFilters, ChannelOptions, ChannelSort, Event } from 'stream-chat';
@@ -75,6 +75,7 @@ export type ChannelListProps<
     channel: Channel<StreamChatGenerics>,
     t: TranslationContextValue['t'],
     userLanguage: TranslationContextValue['userLanguage'],
+    isMessageAIGenerated?: ChatContextValue['isMessageAIGenerated'],
   ) => string | JSX.Element;
   /** Custom UI component to display the container for the queried channels, defaults to and accepts same props as: [ChannelListMessenger](https://github.com/GetStream/stream-chat-react/blob/master/src/components/ChannelList/ChannelListMessenger.tsx) */
   List?: React.ComponentType<ChannelListMessengerProps<StreamChatGenerics>>;
diff --git a/src/components/ChannelPreview/ChannelPreview.tsx b/src/components/ChannelPreview/ChannelPreview.tsx
index 836805935..a25bd333e 100644
--- a/src/components/ChannelPreview/ChannelPreview.tsx
+++ b/src/components/ChannelPreview/ChannelPreview.tsx
@@ -82,9 +82,12 @@ export const ChannelPreview = <
     channelUpdateCount,
     getLatestMessagePreview = defaultGetLatestMessagePreview,
   } = props;
-  const { channel: activeChannel, client, setActiveChannel } = useChatContext<StreamChatGenerics>(
-    'ChannelPreview',
-  );
+  const {
+    channel: activeChannel,
+    client,
+    isMessageAIGenerated,
+    setActiveChannel,
+  } = useChatContext<StreamChatGenerics>('ChannelPreview');
   const { t, userLanguage } = useTranslationContext('ChannelPreview');
   const { displayImage, displayTitle, groupChannelDisplayInfo } = useChannelPreviewInfo({
     channel,
@@ -162,7 +165,12 @@ export const ChannelPreview = <
 
   if (!Preview) return null;
 
-  const latestMessagePreview = getLatestMessagePreview(channel, t, userLanguage);
+  const latestMessagePreview = getLatestMessagePreview(
+    channel,
+    t,
+    userLanguage,
+    isMessageAIGenerated,
+  );
 
   return (
     <Preview
diff --git a/src/components/ChannelPreview/utils.tsx b/src/components/ChannelPreview/utils.tsx
index 85184f21a..4b79ddd0b 100644
--- a/src/components/ChannelPreview/utils.tsx
+++ b/src/components/ChannelPreview/utils.tsx
@@ -7,6 +7,7 @@ import type { Channel, PollVote, TranslationLanguages, UserResponse } from 'stre
 import type { TranslationContextValue } from '../../context/TranslationContext';
 
 import type { DefaultStreamChatGenerics } from '../../types/types';
+import { ChatContextValue } from '../../context';
 
 export const renderPreviewText = (text: string) => <ReactMarkdown skipHtml>{text}</ReactMarkdown>;
 
@@ -32,6 +33,7 @@ export const getLatestMessagePreview = <
   channel: Channel<StreamChatGenerics>,
   t: TranslationContextValue['t'],
   userLanguage: TranslationContextValue['userLanguage'] = 'en',
+  isMessageAIGenerated?: ChatContextValue<StreamChatGenerics>['isMessageAIGenerated'],
 ): string | JSX.Element => {
   const latestMessage = channel.state.latestMessages[channel.state.latestMessages.length - 1];
 
@@ -77,7 +79,7 @@ export const getLatestMessagePreview = <
   }
 
   if (previewTextToRender) {
-    return latestMessage.ai_generated
+    return isMessageAIGenerated?.(latestMessage)
       ? previewTextToRender
       : renderPreviewText(previewTextToRender);
   }
diff --git a/src/components/Chat/Chat.tsx b/src/components/Chat/Chat.tsx
index 6e997cd44..b908ee47b 100644
--- a/src/components/Chat/Chat.tsx
+++ b/src/components/Chat/Chat.tsx
@@ -12,6 +12,7 @@ import type { StreamChat } from 'stream-chat';
 import type { SupportedTranslations } from '../../i18n/types';
 import type { Streami18n } from '../../i18n/Streami18n';
 import type { DefaultStreamChatGenerics } from '../../types/types';
+import type { MessageContextValue } from '../../context';
 
 export type ChatProps<
   StreamChatGenerics extends DefaultStreamChatGenerics = DefaultStreamChatGenerics
@@ -36,7 +37,7 @@ export type ChatProps<
    * Note: requires importing `stream-chat-react/css/v2/emoji-replacement.css` style sheet
    */
   useImageFlagEmojisOnWindows?: boolean;
-};
+} & Partial<Pick<MessageContextValue<StreamChatGenerics>, 'isMessageAIGenerated'>>;
 
 /**
  * Wrapper component for a StreamChat application. Chat needs to be placed around any other chat components
@@ -54,6 +55,7 @@ export const Chat = <
     defaultLanguage,
     i18nInstance,
     initialNavOpen = true,
+    isMessageAIGenerated,
     theme = 'messaging light',
     useImageFlagEmojisOnWindows = false,
   } = props;
@@ -79,6 +81,7 @@ export const Chat = <
     closeMobileNav,
     customClasses,
     getAppSettings,
+    isMessageAIGenerated,
     latestMessageDatesByChannels,
     mutes,
     navOpen,
diff --git a/src/components/Chat/hooks/useCreateChatContext.ts b/src/components/Chat/hooks/useCreateChatContext.ts
index 31924d9c7..d9cec2b0b 100644
--- a/src/components/Chat/hooks/useCreateChatContext.ts
+++ b/src/components/Chat/hooks/useCreateChatContext.ts
@@ -15,6 +15,7 @@ export const useCreateChatContext = <
     closeMobileNav,
     customClasses,
     getAppSettings,
+    isMessageAIGenerated,
     latestMessageDatesByChannels,
     mutes,
     navOpen,
@@ -41,6 +42,7 @@ export const useCreateChatContext = <
       closeMobileNav,
       customClasses,
       getAppSettings,
+      isMessageAIGenerated,
       latestMessageDatesByChannels,
       mutes,
       navOpen,
@@ -58,6 +60,7 @@ export const useCreateChatContext = <
       getAppSettings,
       mutedUsersLength,
       navOpen,
+      isMessageAIGenerated,
     ],
   );
 
diff --git a/src/components/Message/Message.tsx b/src/components/Message/Message.tsx
index 1ed4a60ce..786c9046a 100644
--- a/src/components/Message/Message.tsx
+++ b/src/components/Message/Message.tsx
@@ -76,7 +76,7 @@ const MessageWithContext = <
     userRoles,
   } = props;
 
-  const { client } = useChatContext('Message');
+  const { client, isMessageAIGenerated } = useChatContext('Message');
   const { read } = useChannelStateContext('Message');
   const { Message: contextMessage } = useComponentContext<StreamChatGenerics>('Message');
 
@@ -159,6 +159,7 @@ const MessageWithContext = <
     editing,
     getMessageActions: messageActionsHandler,
     handleEdit: setEdit,
+    isMessageAIGenerated,
     isMyMessage: () => isMyMessage,
     messageIsUnread,
     onUserClick,
diff --git a/src/components/Message/MessageSimple.tsx b/src/components/Message/MessageSimple.tsx
index fdd3cfff0..e82200fed 100644
--- a/src/components/Message/MessageSimple.tsx
+++ b/src/components/Message/MessageSimple.tsx
@@ -1,4 +1,4 @@
-import React, { useState } from 'react';
+import React, { useMemo, useState } from 'react';
 import clsx from 'clsx';
 
 import { MessageErrorIcon } from './icons';
@@ -56,6 +56,7 @@ const MessageSimpleWithContext = <
     handleOpenThread,
     handleRetry,
     highlighted,
+    isMessageAIGenerated,
     isMyMessage,
     message,
     onUserClick,
@@ -88,6 +89,10 @@ const MessageSimpleWithContext = <
 
   const hasAttachment = messageHasAttachments(message);
   const hasReactions = messageHasReactions(message);
+  const isAIGenerated = useMemo(() => isMessageAIGenerated?.(message), [
+    isMessageAIGenerated,
+    message,
+  ]);
 
   if (message.customType === CUSTOM_MESSAGE_TYPE.date) {
     return null;
@@ -101,7 +106,7 @@ const MessageSimpleWithContext = <
   const showReplyCountButton = !threadList && !!message.reply_count;
   const allowRetry = message.status === 'failed' && message.errorStatusCode !== 403;
   const isBounced = isMessageBounced(message);
-  const isEdited = isMessageEdited(message);
+  const isEdited = isMessageEdited(message) && !isAIGenerated;
 
   let handleClick: (() => void) | undefined = undefined;
 
@@ -187,7 +192,7 @@ const MessageSimpleWithContext = <
               {message.attachments?.length && !message.quoted_message ? (
                 <Attachment actionHandler={handleAction} attachments={message.attachments} />
               ) : null}
-              {message.ai_generated ? (
+              {isAIGenerated ? (
                 <StreamedMessageText message={message} renderText={renderText} />
               ) : (
                 <MessageText message={message} renderText={renderText} />
diff --git a/src/components/Message/utils.tsx b/src/components/Message/utils.tsx
index d8132d8b7..2a01bb31f 100644
--- a/src/components/Message/utils.tsx
+++ b/src/components/Message/utils.tsx
@@ -494,6 +494,5 @@ export const isMessageBounced = <
 export const isMessageEdited = <
   StreamChatGenerics extends DefaultStreamChatGenerics = DefaultStreamChatGenerics
 >(
-  message: Pick<StreamMessage<StreamChatGenerics>, 'message_text_updated_at'> &
-    Partial<Pick<StreamMessage<StreamChatGenerics>, 'ai_generated'>>,
-) => !!message.message_text_updated_at && !message.ai_generated;
+  message: Pick<StreamMessage<StreamChatGenerics>, 'message_text_updated_at'>,
+) => !!message.message_text_updated_at;
diff --git a/src/context/ChatContext.tsx b/src/context/ChatContext.tsx
index fecc05e07..514942a5d 100644
--- a/src/context/ChatContext.tsx
+++ b/src/context/ChatContext.tsx
@@ -56,7 +56,8 @@ export type ChatContextValue<
    */
   customClasses?: CustomClasses;
   navOpen?: boolean;
-} & Required<Pick<ChatProps<StreamChatGenerics>, 'theme' | 'client'>>;
+} & Partial<Pick<ChatProps<StreamChatGenerics>, 'isMessageAIGenerated'>> &
+  Required<Pick<ChatProps<StreamChatGenerics>, 'theme' | 'client'>>;
 
 export const ChatContext = React.createContext<ChatContextValue | undefined>(undefined);
 
diff --git a/src/context/MessageContext.tsx b/src/context/MessageContext.tsx
index 3cd6bbfdf..74e6d7a77 100644
--- a/src/context/MessageContext.tsx
+++ b/src/context/MessageContext.tsx
@@ -106,6 +106,10 @@ export type MessageContextValue<
   highlighted?: boolean;
   /** Whether the threaded message is the first in the thread list */
   initialMessage?: boolean;
+  /**
+   * A factory function that determines whether a message is AI generated or not.
+   */
+  isMessageAIGenerated?: (message: StreamMessage<StreamChatGenerics>) => boolean;
   /** Latest message id on current channel */
   lastReceivedId?: string | null;
   /** DOMRect object for parent MessageList component */