@lobehub/lobehub 2.0.0-next.8 → 2.0.0-next.9

package/packages/database/src/models/message.ts

@@ -1,3 +1,4 @@
+import { INBOX_SESSION_ID } from '@lobechat/const';
 import {
   ChatFileItem,
   ChatImageItem,
@@ -14,6 +15,7 @@ import {
   UIChatMessage,
   UpdateMessageParams,
   UpdateMessageRAGParams,
+  UpdateMessageResult,
 } from '@lobechat/types';
 import type { HeatmapsProps } from '@lobehub/charts';
 import dayjs from 'dayjs';
@@ -39,6 +41,7 @@ import {
 } from '../schemas';
 import { LobeChatDatabase } from '../type';
 import { genEndDateWhere, genRangeWhere, genStartDateWhere, genWhere } from '../utils/genWhere';
+import { groupAssistantMessages } from '../utils/groupMessages';
 import { idGenerator } from '../utils/idGenerator';
 
 export class MessageModel {
@@ -54,6 +57,7 @@ export class MessageModel {
   query = async (
     { current = 0, pageSize = 1000, sessionId, topicId, groupId }: QueryMessageParams = {},
     options: {
+      groupAssistantMessages?: boolean;
       postProcessUrl?: (path: string | null, file: { fileType: string }) => Promise<string>;
     } = {},
   ) => {
@@ -211,7 +215,7 @@ export class MessageModel {
       .from(messageQueries)
       .where(inArray(messageQueries.messageId, messageIds));
 
-    return result.map(
+    const mappedMessages = result.map(
       ({ model, provider, translate, ttsId, ttsFile, ttsContentMd5, ttsVoice, ...item }) => {
         const messageQuery = messageQueriesList.find((relation) => relation.messageId === item.id);
         return {
@@ -246,13 +250,15 @@ export class MessageModel {
               size: size!,
               url,
             })),
-
           imageList: imageList
             .filter((relation) => relation.messageId === item.id)
             // eslint-disable-next-line @typescript-eslint/no-unused-vars
             .map<ChatImageItem>(({ id, url, name }) => ({ alt: name!, id, url })),
-
           meta: {},
+
+          model,
+
+          provider,
           ragQuery: messageQuery?.rewriteQuery,
           ragQueryId: messageQuery?.id,
           ragRawQuery: messageQuery?.userQuery,
@@ -263,6 +269,10 @@ export class MessageModel {
         } as unknown as UIChatMessage;
       },
     );
+
+    // Group assistant messages with their tool results
+    const { groupAssistantMessages: useGroup = false } = options;
+    return useGroup ? groupAssistantMessages(mappedMessages) : mappedMessages;
   };
 
   findById = async (id: string) => {
@@ -560,7 +570,7 @@ export class MessageModel {
         sessionId: params.sessionId,
         topicId: params.topicId, // Get all messages
       },
-      options,
+      { ...options, groupAssistantMessages: true },
     );
 
     // 3. Return the result
@@ -589,27 +599,52 @@ export class MessageModel {
   };
   // **************** Update *************** //
 
-  update = async (id: string, { imageList, ...message }: Partial<UpdateMessageParams>) => {
-    return this.db.transaction(async (trx) => {
-      // 1. insert message files
-      if (imageList && imageList.length > 0) {
+  update = async (
+    id: string,
+    { imageList, ...message }: Partial<UpdateMessageParams>,
+    options?: {
+      postProcessUrl?: (path: string | null, file: { fileType: string }) => Promise<string>;
+      sessionId?: string | null;
+      topicId?: string | null;
+    },
+  ): Promise<UpdateMessageResult> => {
+    try {
+      await this.db.transaction(async (trx) => {
+        // 1. insert message files
+        if (imageList && imageList.length > 0) {
+          await trx
+            .insert(messagesFiles)
+            .values(
+              imageList.map((file) => ({ fileId: file.id, messageId: id, userId: this.userId })),
+            );
+        }
+
         await trx
-          .insert(messagesFiles)
-          .values(
-            imageList.map((file) => ({ fileId: file.id, messageId: id, userId: this.userId })),
-          );
+          .update(messages)
+          .set({ ...message })
+          .where(and(eq(messages.id, id), eq(messages.userId, this.userId)));
+      });
+
+      // if sessionId or topicId provided, return the updated message list
+      if (options?.sessionId !== undefined || options?.topicId !== undefined) {
+        const messageList = await this.query(
+          {
+            sessionId: options.sessionId,
+            topicId: options.topicId,
+          },
+          {
+            postProcessUrl: options.postProcessUrl,
+          },
+        );
+
+        return { messages: messageList, success: true };
       }
 
-      return trx
-        .update(messages)
-        .set({
-          ...message,
-          // TODO: need a better way to handle this
-          // TODO: but I forget why 🤡
-          role: message.role as any,
-        })
-        .where(and(eq(messages.id, id), eq(messages.userId, this.userId)));
-    });
+      return { success: true };
+    } catch (error) {
+      console.error('Update message error:', error);
+      return { success: false };
+    }
   };
 
   updateMetadata = async (id: string, metadata: Record<string, any>) => {
@@ -778,8 +813,11 @@ export class MessageModel {
 
   private genId = () => idGenerator('messages', 14);
 
-  private matchSession = (sessionId?: string | null) =>
-    sessionId ? eq(messages.sessionId, sessionId) : isNull(messages.sessionId);
+  private matchSession = (sessionId?: string | null) => {
+    if (sessionId === INBOX_SESSION_ID) return isNull(messages.sessionId);
+
+    return sessionId ? eq(messages.sessionId, sessionId) : isNull(messages.sessionId);
+  };
 
   private matchTopic = (topicId?: string | null) =>
     topicId ? eq(messages.topicId, topicId) : isNull(messages.topicId);

package/packages/database/src/utils/__tests__/groupMessages.test.ts

@@ -54,6 +54,7 @@ describe('groupAssistantMessages', () => {
           content: 'Beijing: Sunny, 25°C',
           state: { cached: true },
         },
+        result_msg_id: 'msg-2',
       });
     });
 
@@ -112,7 +113,9 @@ describe('groupAssistantMessages', () => {
       const block = result[0].children![0];
       expect(block.tools).toHaveLength(2);
       expect(block.tools![0].result?.content).toBe('Beijing: Sunny, 25°C');
+      expect(block.tools![0].result_msg_id).toBe('msg-2');
       expect(block.tools![1].result?.content).toBe('Latest tech news: AI breakthrough');
+      expect(block.tools![1].result_msg_id).toBe('msg-3');
     });
 
     it('should handle assistant message without tools', () => {
@@ -165,6 +168,7 @@ describe('groupAssistantMessages', () => {
       const block = result[0].children![0];
       expect(block.tools).toHaveLength(1);
       expect(block.tools![0].result).toBeUndefined();
+      expect(block.tools![0].result_msg_id).toBeUndefined();
     });
   });
 
@@ -366,11 +370,13 @@ describe('groupAssistantMessages', () => {
       // First child: original assistant with tool result
       expect(result[0].children![0].id).toBe('msg-1');
       expect(result[0].children![0].tools![0].result?.content).toBe('Sunny, 25°C');
+      expect(result[0].children![0].tools![0].result_msg_id).toBe('msg-2');
 
       // Second child: follow-up assistant with its own tool result
       expect(result[0].children![1].id).toBe('msg-3');
       expect(result[0].children![1].tools).toHaveLength(1);
       expect(result[0].children![1].tools![0].result?.content).toBe('Breaking news');
+      expect(result[0].children![1].tools![0].result_msg_id).toBe('msg-4');
     });
 
     it('should group multiple follow-up assistants in chain (3+ assistants)', () => {
@@ -448,9 +454,11 @@ describe('groupAssistantMessages', () => {
 
       expect(result[0].children![0].id).toBe('msg-1');
       expect(result[0].children![0].tools![0].result?.content).toBe('Result 1');
+      expect(result[0].children![0].tools![0].result_msg_id).toBe('msg-2');
 
       expect(result[0].children![1].id).toBe('msg-3');
       expect(result[0].children![1].tools![0].result?.content).toBe('Result 2');
+      expect(result[0].children![1].tools![0].result_msg_id).toBe('msg-4');
 
       expect(result[0].children![2].id).toBe('msg-5');
       expect(result[0].children![2].content).toBe('Step 3 final');
@@ -677,7 +685,6 @@ describe('groupAssistantMessages', () => {
       expect(block.content).toBe('Test');
       expect(block.tools).toHaveLength(1);
       expect(block.imageList).toHaveLength(1);
-      expect(block.fileList).toHaveLength(1);
     });
 
     it('should preserve all tool result fields', () => {
@@ -720,6 +727,7 @@ describe('groupAssistantMessages', () => {
         state: { step: 1 },
         error: null,
       });
+      expect(block.tools![0].result_msg_id).toBe('msg-2');
     });
   });
 
@@ -928,6 +936,142 @@ describe('groupAssistantMessages', () => {
       expect(result[0].performance).toBeUndefined();
       expect(result[0].metadata).toBeUndefined();
     });
+
+    it('should map reasoning field to reasoning in children blocks', () => {
+      const input: UIChatMessage[] = [
+        {
+          id: 'msg-1',
+          role: 'assistant',
+          content: 'Test response',
+          reasoning: {
+            content: 'This is my reasoning process',
+            duration: 1500,
+          },
+          tools: [
+            {
+              id: 'tool-1',
+              identifier: 'test',
+              apiName: 'test',
+              arguments: '{}',
+              type: 'default',
+            },
+          ],
+          createdAt: Date.now(),
+          updatedAt: Date.now(),
+          meta: {},
+        } as UIChatMessage,
+        {
+          id: 'msg-2',
+          role: 'tool',
+          tool_call_id: 'tool-1',
+          content: 'Tool result',
+          createdAt: Date.now(),
+          updatedAt: Date.now(),
+          meta: {},
+        } as UIChatMessage,
+        {
+          id: 'msg-3',
+          role: 'assistant',
+          content: 'Follow-up response',
+          parentId: 'msg-2',
+          reasoning: {
+            content: 'Follow-up reasoning',
+            duration: 2000,
+          },
+          createdAt: Date.now(),
+          updatedAt: Date.now(),
+          meta: {},
+        } as UIChatMessage,
+      ];
+
+      const result = groupAssistantMessages(input);
+
+      // First block should have reasoning
+      expect(result[0].children![0].reasoning).toEqual({
+        content: 'This is my reasoning process',
+        duration: 1500,
+      });
+
+      // Second block (follow-up) should also have reasoning
+      expect(result[0].children![1].reasoning).toEqual({
+        content: 'Follow-up reasoning',
+        duration: 2000,
+      });
+
+      // Group message should not have reasoning (moved to children)
+      expect(result[0].reasoning).toBeUndefined();
+    });
+
+    it('should preserve error field in children blocks', () => {
+      const input = [
+        {
+          id: 'msg-1',
+          role: 'assistant',
+          content: 'Failed to process',
+          error: {
+            type: 'InvalidAPIKey',
+            message: 'API key is invalid',
+          },
+          tools: [
+            {
+              id: 'tool-1',
+              identifier: 'test',
+              apiName: 'test',
+              arguments: '{}',
+              type: 'default',
+            },
+          ],
+          createdAt: Date.now(),
+          updatedAt: Date.now(),
+          meta: {},
+        } as unknown as UIChatMessage,
+      ] as UIChatMessage[];
+
+      const result = groupAssistantMessages(input);
+
+      // Child block should have error
+      expect(result[0].children![0].error).toEqual({
+        type: 'InvalidAPIKey',
+        message: 'API key is invalid',
+      });
+    });
+
+    it('should preserve imageList in children blocks', () => {
+      const input: UIChatMessage[] = [
+        {
+          id: 'msg-1',
+          role: 'assistant',
+          content: 'Here are the images',
+          imageList: [
+            { id: 'img-1', url: 'https://example.com/img1.jpg', alt: 'Image 1' },
+            { id: 'img-2', url: 'https://example.com/img2.jpg', alt: 'Image 2' },
+          ],
+          tools: [
+            {
+              id: 'tool-1',
+              identifier: 'test',
+              apiName: 'test',
+              arguments: '{}',
+              type: 'default',
+            },
+          ],
+          createdAt: Date.now(),
+          updatedAt: Date.now(),
+          meta: {},
+        } as UIChatMessage,
+      ];
+
+      const result = groupAssistantMessages(input);
+
+      // Child block should have imageList
+      expect(result[0].children![0].imageList).toEqual([
+        { id: 'img-1', url: 'https://example.com/img1.jpg', alt: 'Image 1' },
+        { id: 'img-2', url: 'https://example.com/img2.jpg', alt: 'Image 2' },
+      ]);
+
+      // Parent should not have imageList (moved to children)
+      expect(result[0].imageList).toBeUndefined();
+    });
   });
 
   describe('Empty and Null Cases', () => {
@@ -958,7 +1102,6 @@ describe('groupAssistantMessages', () => {
 
       // Empty arrays should become undefined
       expect(result[0].children![0].imageList).toBeUndefined();
-      expect(result[0].children![0].fileList).toBeUndefined();
     });
 
     it('should handle empty message list', () => {

package/packages/database/src/utils/groupMessages.ts

@@ -243,6 +243,7 @@ export function groupAssistantMessages(messages: UIChatMessage[]): UIChatMessage
             id: toolMsg.id,
             state: toolMsg.pluginState,
           },
+          result_msg_id: toolMsg.id,
         };
       }
 
@@ -253,10 +254,11 @@ export function groupAssistantMessages(messages: UIChatMessage[]): UIChatMessage
     const { usage: msgUsage, performance: msgPerformance } = splitMetadata(msg.metadata);
     children.push({
       content: msg.content || '',
-      fileList: msg.fileList && msg.fileList.length > 0 ? msg.fileList : undefined,
+      error: msg.error,
       id: msg.id,
       imageList: msg.imageList && msg.imageList.length > 0 ? msg.imageList : undefined,
       performance: msgPerformance,
+      reasoning: msg.reasoning || undefined,
       tools: toolsWithResults,
       usage: msgUsage,
     });
@@ -299,6 +301,7 @@ export function groupAssistantMessages(messages: UIChatMessage[]): UIChatMessage
                     id: followUpToolMsg.id,
                     state: followUpToolMsg.pluginState,
                   },
+                  result_msg_id: followUpToolMsg.id,
                 };
               }
 
@@ -311,16 +314,14 @@ export function groupAssistantMessages(messages: UIChatMessage[]): UIChatMessage
           );
           children.push({
             content: followUpMsg.content || '',
-            fileList:
-              followUpMsg.fileList && followUpMsg.fileList.length > 0
-                ? followUpMsg.fileList
-                : undefined,
+            error: followUpMsg.error,
             id: followUpMsg.id,
             imageList:
               followUpMsg.imageList && followUpMsg.imageList.length > 0
                 ? followUpMsg.imageList
                 : undefined,
             performance: followUpPerformance,
+            reasoning: followUpMsg.reasoning || undefined,
             tools: followUpToolsWithResults,
             usage: followUpUsage,
           });
@@ -347,6 +348,7 @@ export function groupAssistantMessages(messages: UIChatMessage[]): UIChatMessage
       assistantMsg.performance = aggregated.performance;
     }
     delete assistantMsg.metadata; // Clear individual metadata
+    delete assistantMsg.reasoning; // Reasoning moved to children blocks
     delete assistantMsg.tools;
     delete assistantMsg.imageList;
     delete assistantMsg.fileList;

package/packages/types/src/message/common/base.ts

@@ -1,5 +1,6 @@
 import type { ILobeAgentRuntimeErrorType } from '@lobechat/model-runtime';
 import type { IPluginErrorType } from '@lobehub/chat-plugin-sdk';
+import { z } from 'zod';
 
 import { ErrorType } from '../../fetch';
 
@@ -12,6 +13,12 @@ export interface ChatMessageError {
   type: ErrorType | IPluginErrorType | ILobeAgentRuntimeErrorType;
 }
 
+export const ChatMessageErrorSchema = z.object({
+  body: z.any().optional(),
+  message: z.string(),
+  type: z.union([z.string(), z.number()]),
+});
+
 export interface ChatCitationItem {
   id?: string;
   onlyUrl?: boolean;
@@ -24,3 +31,9 @@ export interface ModelReasoning {
   duration?: number;
   signature?: string;
 }
+
+export const ModelReasoningSchema = z.object({
+  content: z.string().optional(),
+  duration: z.number().optional(),
+  signature: z.string().optional(),
+});

package/packages/types/src/message/common/image.ts

@@ -1,9 +1,17 @@
+import { z } from 'zod';
+
 export interface ChatImageItem {
   alt: string;
   id: string;
   url: string;
 }
 
+export const ChatImageItemSchema = z.object({
+  alt: z.string(),
+  id: z.string(),
+  url: z.string(),
+});
+
 export interface ChatImageChunk {
   data: string;
   id: string;

package/packages/types/src/message/common/metadata.ts

@@ -1,4 +1,5 @@
 /* eslint-disable sort-keys-fix/sort-keys-fix , typescript-sort-keys/interface */
+import { z } from 'zod';
 
 export interface ModelTokensUsage {
   // Input tokens breakdown
@@ -38,6 +39,44 @@ export interface ModelTokensUsage {
   totalTokens?: number;
 }
 
+export const ModelUsageSchema = z.object({
+  // Input tokens breakdown
+  inputCachedTokens: z.number().optional(),
+  inputCacheMissTokens: z.number().optional(),
+  inputWriteCacheTokens: z.number().optional(),
+  inputTextTokens: z.number().optional(),
+  inputImageTokens: z.number().optional(),
+  inputAudioTokens: z.number().optional(),
+  inputCitationTokens: z.number().optional(),
+
+  // Output tokens breakdown
+  outputTextTokens: z.number().optional(),
+  outputImageTokens: z.number().optional(),
+  outputAudioTokens: z.number().optional(),
+  outputReasoningTokens: z.number().optional(),
+
+  // Prediction tokens
+  acceptedPredictionTokens: z.number().optional(),
+  rejectedPredictionTokens: z.number().optional(),
+
+  // Total tokens
+  totalInputTokens: z.number().optional(),
+  totalOutputTokens: z.number().optional(),
+  totalTokens: z.number().optional(),
+
+  // Cost
+  cost: z.number().optional(),
+});
+
+export const ModelPerformanceSchema = z.object({
+  tps: z.number().optional(),
+  ttft: z.number().optional(),
+  duration: z.number().optional(),
+  latency: z.number().optional(),
+});
+
+export const MessageMetadataSchema = ModelUsageSchema.merge(ModelPerformanceSchema);
+
 export interface ModelUsage extends ModelTokensUsage {
   /**
    * dollar

package/packages/types/src/message/common/tools.ts

@@ -16,6 +16,7 @@ export interface ChatToolPayload {
   arguments: string;
   id: string;
   identifier: string;
+  result_msg_id?: string;
   type: LobeToolRenderType;
 }
 
@@ -85,6 +86,15 @@ export const MessageToolCallSchema = z.object({
   type: z.string(),
 });
 
+export const ChatToolPayloadSchema = z.object({
+  apiName: z.string(),
+  arguments: z.string(),
+  id: z.string(),
+  identifier: z.string(),
+  result_msg_id: z.string().optional(),
+  type: z.string(),
+});
+
 /**
  * 聊天消息错误对象
  */

package/packages/types/src/message/db/params.ts

@@ -1,11 +1,20 @@
-import { GroundingSearch } from '../../search';
+/* eslint-disable sort-keys-fix/sort-keys-fix */
+import { z } from 'zod';
+
+import { GroundingSearch, GroundingSearchSchema } from '../../search';
 import {
   ChatImageItem,
+  ChatImageItemSchema,
   ChatMessageError,
+  ChatMessageErrorSchema,
   ChatToolPayload,
+  ChatToolPayloadSchema,
   MessageMetadata,
+  MessageMetadataSchema,
   MessageToolCall,
+  MessageToolCallSchema,
   ModelReasoning,
+  ModelReasoningSchema,
 } from '../common';
 import { UIChatMessage } from '../ui';
 
@@ -34,6 +43,21 @@ export interface CreateMessageResult {
   messages: UIChatMessage[];
 }
 
+/**
+ * Result type for updateMessage
+ * Contains success status and optional message list
+ */
+export interface UpdateMessageResult {
+  /**
+   * Updated message list (only present when success is true and sessionId/topicId provided)
+   */
+  messages?: UIChatMessage[];
+  /**
+   * Whether the update was successful
+   */
+  success: boolean;
+}
+
 export interface NewMessage {
   agentId?: string | null;
   clientId?: string | null;
@@ -67,12 +91,14 @@ export interface UpdateMessageParams {
   imageList?: ChatImageItem[];
   metadata?: MessageMetadata;
   model?: string;
+  observationId?: string;
   provider?: string;
   reasoning?: ModelReasoning;
   role?: string;
   search?: GroundingSearch;
   toolCalls?: MessageToolCall[];
   tools?: ChatToolPayload[] | null;
+  traceId?: string;
 }
 
 export interface NewMessageQueryParams {
@@ -81,3 +107,23 @@ export interface NewMessageQueryParams {
   rewriteQuery: string;
   userQuery: string;
 }
+
+// ========== Zod Schemas ========== //
+
+export const UpdateMessageParamsSchema = z
+  .object({
+    content: z.string().optional(),
+    error: ChatMessageErrorSchema.nullable().optional(),
+    imageList: z.array(ChatImageItemSchema).optional(),
+    metadata: MessageMetadataSchema.optional(),
+    model: z.string().optional(),
+    observationId: z.string().optional(),
+    provider: z.string().optional(),
+    reasoning: ModelReasoningSchema.optional(),
+    role: z.string().optional(),
+    search: GroundingSearchSchema.optional(),
+    toolCalls: z.array(MessageToolCallSchema).optional(),
+    tools: z.array(ChatToolPayloadSchema).nullable().optional(),
+    traceId: z.string().optional(),
+  })
+  .passthrough();

package/packages/types/src/message/ui/chat.ts

@@ -26,10 +26,11 @@ export interface ChatFileItem {
 
 export interface AssistantContentBlock {
   content: string;
-  fileList?: ChatFileItem[];
+  error?: ChatMessageError | null;
   id: string;
   imageList?: ChatImageItem[];
   performance?: ModelPerformance;
+  reasoning?: ModelReasoning;
   tools?: ChatToolPayloadWithResult[];
   usage?: ModelUsage;
 }
@@ -62,6 +63,7 @@ export interface UIChatMessage {
   imageList?: ChatImageItem[];
   meta: MetaData;
   metadata?: MessageMetadata | null;
+  model?: string | null;
   /**
    * observation id
    */
@@ -78,6 +80,7 @@ export interface UIChatMessage {
   plugin?: ChatPluginPayload;
   pluginError?: any;
   pluginState?: any;
+  provider?: string | null;
   /**
    * quoted other message's id
    */

package/packages/types/src/search.ts

@@ -1,3 +1,5 @@
+import { z } from 'zod';
+
 export type SearchMode = 'off' | 'auto' | 'on';
 
 export enum ModelSearchImplement {
@@ -27,3 +29,17 @@ export interface GroundingSearch {
   citations?: CitationItem[];
   searchQueries?: string[];
 }
+
+export const GroundingSearchSchema = z.object({
+  citations: z
+    .array(
+      z.object({
+        favicon: z.string().optional(),
+        id: z.string().optional(),
+        title: z.string().optional(),
+        url: z.string(),
+      }),
+    )
+    .optional(),
+  searchQueries: z.array(z.string()).optional(),
+});