Merge pull request #669 from miurla/feat/enhance-message-metadata

feat: enhance message metadata with searchMode and modelId

Author: miurla

lib/streaming/create-chat-stream-response.ts

@@ -128,14 +128,6 @@ export async function createChatStreamResponse(
           searchMode
         })
 
-        // Write metadata including traceId at the start of streaming
-        if (parentTraceId) {
-          writer.write({
-            type: 'message-metadata',
-            messageMetadata: { traceId: parentTraceId }
-          })
-        }
-
         // Filter out reasoning parts from messages before converting to model messages
         // OpenAI API requires reasoning messages to be followed by assistant messages
         const filteredMessages = filterReasoningParts(messagesToModel)
@@ -215,7 +207,9 @@ export async function createChatStreamResponse(
         chatId,
         userId,
         titlePromise,
-        parentTraceId
+        parentTraceId,
+        searchMode,
+        context.modelId
       )
     }
   })

lib/streaming/helpers/persist-stream-results.ts

@@ -2,6 +2,7 @@ import { UIMessage } from 'ai'
 
 import { upsertMessage } from '@/lib/actions/chat'
 import { updateChatTitle } from '@/lib/db/actions'
+import { SearchMode } from '@/lib/types/search'
 import { perfTime } from '@/lib/utils/perf-logging'
 import { retryDatabaseOperation } from '@/lib/utils/retry'
 
@@ -12,14 +13,16 @@ export async function persistStreamResults(
   chatId: string,
   userId: string,
   titlePromise?: Promise<string>,
-  parentTraceId?: string
+  parentTraceId?: string,
+  searchMode?: SearchMode,
+  modelId?: string
 ) {
-  // Attach metadata to the response message if we have a traceId
-  if (parentTraceId) {
-    responseMessage.metadata = {
-      ...(responseMessage.metadata || {}),
-      traceId: parentTraceId
-    }
+  // Attach metadata to the response message
+  responseMessage.metadata = {
+    ...(responseMessage.metadata || {}),
+    ...(parentTraceId && { traceId: parentTraceId }),
+    ...(searchMode && { searchMode }),
+    ...(modelId && { modelId })
   }
 
   // Wait for title generation if it was started

lib/types/ai.ts

@@ -5,6 +5,7 @@ import { fetchTool } from '@/lib/tools/fetch'
 import { askQuestionTool } from '@/lib/tools/question'
 import { searchTool } from '@/lib/tools/search'
 import { createTodoTools, type TodoItem } from '@/lib/tools/todo'
+import type { SearchMode } from '@/lib/types/search'
 
 // Re-export TodoItem for external use
 export type { TodoItem }
@@ -13,6 +14,8 @@ export type { TodoItem }
 export interface UIMessageMetadata {
   traceId?: string
   feedbackScore?: number | null
+  searchMode?: SearchMode
+  modelId?: string
   [key: string]: any
 }