updated implementation

app/components/chat/AssistantMessage.tsx

@@ -1,15 +1,22 @@
 import { memo } from 'react';
 import { Markdown } from './Markdown';
-import { USAGE_REGEX } from '~/utils/constants';
+import type { JSONValue } from 'ai';
 
 interface AssistantMessageProps {
   content: string;
+  annotations?: JSONValue[];
 }
 
-export const AssistantMessage = memo(({ content }: AssistantMessageProps) => {
-  const match = content.match(USAGE_REGEX);
-  const usage = match ? JSON.parse(match[1]) : null;
-  const cleanContent = content.replace(USAGE_REGEX, '').trim();
+export const AssistantMessage = memo(({ content, annotations }: AssistantMessageProps) => {
+  const filteredAnnotations = (annotations?.filter(
+    (annotation: JSONValue) => annotation && typeof annotation === 'object' && Object.keys(annotation).includes('type'),
+  ) || []) as { type: string; value: any }[];
+
+  const usage: {
+    completionTokens: number;
+    promptTokens: number;
+    totalTokens: number;
+  } = filteredAnnotations.find((annotation) => annotation.type === 'usage')?.value;
 
   return (
     <div className="overflow-hidden w-full">
@@ -18,7 +25,7 @@ export const AssistantMessage = memo(({ content }: AssistantMessageProps) => {
           Tokens: {usage.totalTokens} (prompt: {usage.promptTokens}, completion: {usage.completionTokens})
         </div>
       )}
-      <Markdown html>{cleanContent}</Markdown>
+      <Markdown html>{content}</Markdown>
     </div>
   );
 });

app/components/chat/Messages.client.tsx

@@ -65,7 +65,11 @@ export const Messages = React.forwardRef<HTMLDivElement, MessagesProps>((props:
                   </div>
                 )}
                 <div className="grid grid-col-1 w-full">
-                  {isUserMessage ? <UserMessage content={content} /> : <AssistantMessage content={content} />}
+                  {isUserMessage ? (
+                    <UserMessage content={content} />
+                  ) : (
+                    <AssistantMessage content={content} annotations={message.annotations} />
+                  )}
                 </div>
                 {!isUserMessage && (
                   <div className="flex gap-2 flex-col lg:flex-row">