From c8def1989a04cb5c008fcb494e24d1468a81018a Mon Sep 17 00:00:00 2001
From: Willie Zutz <boarder2@bit-shift.com>
Date: Thu, 1 May 2025 12:17:08 -0600
Subject: [PATCH] fix(Library): Returns metadata back to original format so
 sources continue to work.

---
 src/app/api/chat/route.ts                   | 17 +++++++++++++----
 src/app/settings/page.tsx                   |  3 +--
 src/components/ChatWindow.tsx               | 20 ++++++--------------
 src/components/MessageActions/ModelInfo.tsx |  4 +++-
 src/components/MessageBox.tsx               | 17 +++++++++++------
 5 files changed, 34 insertions(+), 27 deletions(-)

diff --git a/src/app/api/chat/route.ts b/src/app/api/chat/route.ts
index 9705dc9..f21e278 100644
--- a/src/app/api/chat/route.ts
+++ b/src/app/api/chat/route.ts
@@ -112,7 +112,7 @@ const handleEmitterEvents = async (
   stream.on('end', () => {
     const endTime = Date.now();
     const duration = endTime - startTime;
-    
+
     modelStats = {
       ...modelStats,
       responseTime: duration,
@@ -135,9 +135,11 @@ const handleEmitterEvents = async (
         chatId: chatId,
         messageId: aiMessageId,
         role: 'assistant',
-        metadata: {
+        metadata: JSON.stringify({
+          createdAt: new Date(),
+          ...(sources && sources.length > 0 && { sources }),
           modelStats: modelStats,
-        },
+        }),
       })
       .execute();
   });
@@ -319,7 +321,14 @@ export const POST = async (req: Request) => {
     const writer = responseStream.writable.getWriter();
     const encoder = new TextEncoder();
 
-    handleEmitterEvents(stream, writer, encoder, aiMessageId, message.chatId, startTime);
+    handleEmitterEvents(
+      stream,
+      writer,
+      encoder,
+      aiMessageId,
+      message.chatId,
+      startTime,
+    );
     handleHistorySave(message, humanMessageId, body.focusMode, body.files);
 
     return new Response(responseStream.readable, {
diff --git a/src/app/settings/page.tsx b/src/app/settings/page.tsx
index 42fff1e..83a8831 100644
--- a/src/app/settings/page.tsx
+++ b/src/app/settings/page.tsx
@@ -546,8 +546,7 @@ const Page = () => {
                         Automatic Suggestions
                       </p>
                       <p className="text-xs text-black/60 dark:text-white/60 mt-0.5">
-                        Automatically show related suggestions after
-                        responses
+                        Automatically show related suggestions after responses
                       </p>
                     </div>
                   </div>
diff --git a/src/components/ChatWindow.tsx b/src/components/ChatWindow.tsx
index 35e9164..c142c13 100644
--- a/src/components/ChatWindow.tsx
+++ b/src/components/ChatWindow.tsx
@@ -18,10 +18,6 @@ export type ModelStats = {
   responseTime?: number;
 };
 
-export type MessageMetadata = {
-  modelStats?: ModelStats;
-};
-
 export type Message = {
   messageId: string;
   chatId: string;
@@ -30,7 +26,7 @@ export type Message = {
   role: 'user' | 'assistant';
   suggestions?: string[];
   sources?: Document[];
-  metadata?: MessageMetadata;
+  modelStats?: ModelStats;
 };
 
 export interface File {
@@ -217,6 +213,7 @@ const loadMessages = async (
   const messages = data.messages.map((msg: any) => {
     return {
       ...msg,
+      ...JSON.parse(msg.metadata),
     };
   }) as Message[];
 
@@ -445,11 +442,8 @@ const ChatWindow = ({ id }: { id?: string }) => {
               role: 'assistant',
               sources: sources,
               createdAt: new Date(),
-              metadata: {
-                // modelStats will be added when we receive messageEnd event
-                modelStats: {
-                  modelName: data.modelName,
-                },
+              modelStats: {
+                modelName: data.modelName,
               },
             },
           ]);
@@ -483,10 +477,8 @@ const ChatWindow = ({ id }: { id?: string }) => {
             if (message.messageId === data.messageId) {
               return {
                 ...message,
-                metadata: {
-                  // Include model stats if available, otherwise null
-                  modelStats: data.modelStats || null,
-                },
+                // Include model stats if available, otherwise null
+                modelStats: data.modelStats || null,
               };
             }
             return message;
diff --git a/src/components/MessageActions/ModelInfo.tsx b/src/components/MessageActions/ModelInfo.tsx
index fec80f2..fc06b87 100644
--- a/src/components/MessageActions/ModelInfo.tsx
+++ b/src/components/MessageActions/ModelInfo.tsx
@@ -63,7 +63,9 @@ const ModelInfoButton: React.FC<ModelInfoButtonProps> = ({ modelStats }) => {
               </div>
               {modelStats?.responseTime && (
                 <div className="flex justify-between">
-                  <span className="text-black/70 dark:text-white/70">Response time:</span>
+                  <span className="text-black/70 dark:text-white/70">
+                    Response time:
+                  </span>
                   <span className="text-black dark:text-white font-medium">
                     {(modelStats.responseTime / 1000).toFixed(2)}s
                   </span>
diff --git a/src/components/MessageBox.tsx b/src/components/MessageBox.tsx
index 193bff8..d8158ea 100644
--- a/src/components/MessageBox.tsx
+++ b/src/components/MessageBox.tsx
@@ -58,11 +58,15 @@ const MessageBox = ({
   const [speechMessage, setSpeechMessage] = useState(message.content);
   const [loadingSuggestions, setLoadingSuggestions] = useState(false);
   const [autoSuggestions, setAutoSuggestions] = useState(
-    localStorage.getItem('autoSuggestions')
+    localStorage.getItem('autoSuggestions'),
   );
 
   const handleLoadSuggestions = async () => {
-    if (loadingSuggestions || (message?.suggestions && message.suggestions.length > 0)) return;
+    if (
+      loadingSuggestions ||
+      (message?.suggestions && message.suggestions.length > 0)
+    )
+      return;
 
     setLoadingSuggestions(true);
     try {
@@ -202,8 +206,8 @@ const MessageBox = ({
                 <h3 className="text-black dark:text-white font-medium text-xl">
                   Answer
                 </h3>
-                {message.metadata?.modelStats && (
-                  <ModelInfoButton modelStats={message.metadata.modelStats} />
+                {message.modelStats && (
+                  <ModelInfoButton modelStats={message.modelStats} />
                 )}
               </div>
               <Markdown
@@ -251,8 +255,9 @@ const MessageBox = ({
                     <div className="flex flex-row items-center space-x-2 mt-4">
                       <Layers3 />
                       <h3 className="text-xl font-medium">Related</h3>{' '}
-                      {(!autoSuggestions || autoSuggestions === 'false') && (!message.suggestions ||
-                      message.suggestions.length === 0) ? (
+                      {(!autoSuggestions || autoSuggestions === 'false') &&
+                      (!message.suggestions ||
+                        message.suggestions.length === 0) ? (
                         <div className="bg-light-secondary dark:bg-dark-secondary">
                           <button
                             onClick={handleLoadSuggestions}