@lobehub/lobehub 2.0.0-next.32 → 2.0.0-next.33

package/packages/conversation-flow/src/transformation/FlatListBuilder.ts

@@ -0,0 +1,421 @@
+import type { AssistantContentBlock, ChatToolPayloadWithResult } from '@lobechat/types';
+
+import type { Message, MessageGroupMetadata } from '../types';
+import type { BranchResolver } from './BranchResolver';
+import type { MessageCollector } from './MessageCollector';
+import type { MessageTransformer } from './MessageTransformer';
+
+/**
+ * FlatListBuilder - Builds flat message list following the active path
+ *
+ * Handles:
+ * 1. Recursive traversal following active branches
+ * 2. Creating virtual messages for Compare and AssistantGroup
+ * 3. Processing different message types with priority
+ */
+export class FlatListBuilder {
+  constructor(
+    private messageMap: Map<string, Message>,
+    private messageGroupMap: Map<string, MessageGroupMetadata>,
+    private childrenMap: Map<string | null, string[]>,
+    private branchResolver: BranchResolver,
+    private messageCollector: MessageCollector,
+    private messageTransformer: MessageTransformer,
+  ) {}
+
+  /**
+   * Generate flatList from messages array
+   * Only includes messages in the active path
+   */
+  flatten(messages: Message[]): Message[] {
+    const flatList: Message[] = [];
+    const processedIds = new Set<string>();
+
+    // Build the active path by traversing from root
+    this.buildFlatListRecursive(null, flatList, processedIds, messages);
+
+    return flatList;
+  }
+
+  /**
+   * Recursively build flatList following the active path
+   */
+  private buildFlatListRecursive(
+    parentId: string | null,
+    flatList: Message[],
+    processedIds: Set<string>,
+    allMessages: Message[],
+  ): void {
+    const children = this.childrenMap.get(parentId) ?? [];
+
+    for (const childId of children) {
+      if (processedIds.has(childId)) continue;
+
+      const message = this.messageMap.get(childId);
+      if (!message) continue;
+
+      // Priority 1: Compare message group
+      const messageGroup = message.groupId ? this.messageGroupMap.get(message.groupId) : undefined;
+
+      if (messageGroup && messageGroup.mode === 'compare' && !processedIds.has(messageGroup.id)) {
+        const groupMembers = this.messageCollector.collectGroupMembers(
+          message.groupId!,
+          allMessages,
+        );
+        const compareMessage = this.createCompareMessage(messageGroup, groupMembers);
+        flatList.push(compareMessage);
+        groupMembers.forEach((m) => processedIds.add(m.id));
+        processedIds.add(messageGroup.id);
+
+        // Continue with active column's children (if any)
+        if ((compareMessage as any).activeColumnId) {
+          this.buildFlatListRecursive(
+            (compareMessage as any).activeColumnId,
+            flatList,
+            processedIds,
+            allMessages,
+          );
+        }
+        continue;
+      }
+
+      // Priority 2: AssistantGroup (assistant + tools)
+      if (message.role === 'assistant' && message.tools && message.tools.length > 0) {
+        // Collect the entire assistant group chain
+        const assistantChain: Message[] = [];
+        const allToolMessages: Message[] = [];
+        this.messageCollector.collectAssistantChain(
+          message,
+          allMessages,
+          assistantChain,
+          allToolMessages,
+          processedIds,
+        );
+
+        // Create assistantGroup virtual message
+        const groupMessage = this.createAssistantGroupMessage(
+          assistantChain[0],
+          assistantChain,
+          allToolMessages,
+        );
+        flatList.push(groupMessage);
+
+        // Mark all as processed
+        assistantChain.forEach((m) => processedIds.add(m.id));
+        allToolMessages.forEach((m) => processedIds.add(m.id));
+
+        // Continue after the assistant chain
+        // Priority 1: If last assistant has non-tool children, continue from it
+        // Priority 2: Otherwise continue from tools (for cases where user replies to tool)
+        const lastAssistant = assistantChain.at(-1);
+        const toolIds = new Set(allToolMessages.map((t) => t.id));
+
+        const lastAssistantNonToolChildren = lastAssistant
+          ? this.childrenMap.get(lastAssistant.id)?.filter((childId) => !toolIds.has(childId))
+          : undefined;
+
+        if (
+          lastAssistantNonToolChildren &&
+          lastAssistantNonToolChildren.length > 0 &&
+          lastAssistant
+        ) {
+          // Follow-up messages exist after the last assistant (not tools)
+          this.buildFlatListRecursive(lastAssistant.id, flatList, processedIds, allMessages);
+        } else {
+          // No non-tool children of last assistant, check tools for children
+          for (const toolMsg of allToolMessages) {
+            this.buildFlatListRecursive(toolMsg.id, flatList, processedIds, allMessages);
+          }
+        }
+        continue;
+      }
+
+      // Priority 3a: Compare mode from user message metadata
+      const childMessages = this.childrenMap.get(message.id) ?? [];
+      if (this.isCompareMode(message) && childMessages.length > 1) {
+        // Add user message
+        flatList.push(message);
+        processedIds.add(message.id);
+
+        // Create compare virtual message with proper handling of AssistantGroups
+        const compareMessage = this.createCompareMessageFromChildIds(
+          message,
+          childMessages,
+          allMessages,
+          processedIds,
+        );
+        flatList.push(compareMessage);
+
+        // Continue with active column's children (if any)
+        if ((compareMessage as any).activeColumnId) {
+          this.buildFlatListRecursive(
+            (compareMessage as any).activeColumnId,
+            flatList,
+            processedIds,
+            allMessages,
+          );
+        }
+        continue;
+      }
+
+      // Priority 3b: User message with branches
+      if (message.role === 'user' && childMessages.length > 1) {
+        const activeBranchId = this.branchResolver.getActiveBranchIdFromMetadata(
+          message,
+          childMessages,
+          this.childrenMap,
+        );
+        const userWithBranches = this.createUserMessageWithBranches(message);
+        flatList.push(userWithBranches);
+        processedIds.add(message.id);
+
+        // Continue with active branch and process its message
+        const activeBranchMsg = this.messageMap.get(activeBranchId);
+        if (activeBranchMsg) {
+          flatList.push(activeBranchMsg);
+          processedIds.add(activeBranchId);
+
+          // Continue with active branch's children
+          this.buildFlatListRecursive(activeBranchId, flatList, processedIds, allMessages);
+        }
+        continue;
+      }
+
+      // Priority 3c: Assistant message with branches
+      if (message.role === 'assistant' && childMessages.length > 1) {
+        const activeBranchId = this.branchResolver.getActiveBranchIdFromMetadata(
+          message,
+          childMessages,
+          this.childrenMap,
+        );
+        // Add the assistant message itself
+        flatList.push(message);
+        processedIds.add(message.id);
+
+        // Continue with active branch and process its message
+        const activeBranchMsg = this.messageMap.get(activeBranchId);
+        if (activeBranchMsg) {
+          flatList.push(activeBranchMsg);
+          processedIds.add(activeBranchId);
+
+          // Continue with active branch's children
+          this.buildFlatListRecursive(activeBranchId, flatList, processedIds, allMessages);
+        }
+        continue;
+      }
+
+      // Priority 4: Regular message
+      flatList.push(message);
+      processedIds.add(message.id);
+
+      // Continue with children
+      this.buildFlatListRecursive(message.id, flatList, processedIds, allMessages);
+    }
+  }
+
+  /**
+   * Check if message has compare mode in metadata
+   */
+  private isCompareMode(message: Message): boolean {
+    return (message.metadata as any)?.compare === true;
+  }
+
+  /**
+   * Create compare virtual message from child IDs with AssistantGroup support
+   */
+  private createCompareMessageFromChildIds(
+    parentMessage: Message,
+    childIds: string[],
+    allMessages: Message[],
+    processedIds: Set<string>,
+  ): Message {
+    const columns: Message[][] = [];
+    const columnFirstIds: string[] = [];
+    let activeColumnId: string | undefined;
+
+    // Process each child (column)
+    for (const childId of childIds) {
+      const childMessage = this.messageMap.get(childId);
+      if (!childMessage) continue;
+
+      columnFirstIds.push(childId);
+
+      // Check if this message is marked as active column
+      if ((childMessage.metadata as any)?.activeColumn === true) {
+        activeColumnId = childId;
+      }
+
+      // Check if this child is an AssistantGroup
+      if (
+        childMessage.role === 'assistant' &&
+        childMessage.tools &&
+        childMessage.tools.length > 0
+      ) {
+        // Collect the entire assistant group chain for this column
+        const assistantChain: Message[] = [];
+        const allToolMessages: Message[] = [];
+        const columnProcessedIds = new Set<string>();
+
+        this.messageCollector.collectAssistantChain(
+          childMessage,
+          allMessages,
+          assistantChain,
+          allToolMessages,
+          columnProcessedIds,
+        );
+
+        // Create assistantGroup virtual message for this column
+        const groupMessage = this.createAssistantGroupMessage(
+          assistantChain[0],
+          assistantChain,
+          allToolMessages,
+        );
+
+        columns.push([groupMessage]);
+
+        // Mark all as processed
+        assistantChain.forEach((m) => processedIds.add(m.id));
+        allToolMessages.forEach((m) => processedIds.add(m.id));
+      } else {
+        // Regular message (not an AssistantGroup)
+        columns.push([childMessage]);
+        processedIds.add(childId);
+      }
+    }
+
+    // Generate ID with all column first message IDs
+    const columnIdsStr = columnFirstIds.join('-');
+    const compareId = `compare-${parentMessage.id}-${columnIdsStr}`;
+
+    // Calculate timestamps from first column's messages
+    const firstColumnMessages = childIds.map((id) => this.messageMap.get(id)).filter(Boolean);
+    const createdAt =
+      firstColumnMessages.length > 0
+        ? Math.min(...firstColumnMessages.map((m) => m!.createdAt))
+        : parentMessage.createdAt;
+    const updatedAt =
+      firstColumnMessages.length > 0
+        ? Math.max(...firstColumnMessages.map((m) => m!.updatedAt))
+        : parentMessage.updatedAt;
+
+    return {
+      activeColumnId,
+      columns: columns as any,
+      content: '',
+      createdAt,
+      extra: {
+        parentMessageId: parentMessage.id,
+      },
+      id: compareId,
+      meta: parentMessage.meta || {},
+      role: 'compare' as any,
+      updatedAt,
+    } as Message;
+  }
+
+  /**
+   * Create compare virtual message (for group-based compare)
+   */
+  private createCompareMessage(group: MessageGroupMetadata, members: Message[]): Message {
+    // Find active column ID from members metadata
+    const activeColumnId = members.find((msg) => (msg.metadata as any)?.activeColumn === true)?.id;
+
+    // columns contain full Message objects
+    const columns: Message[][] = members.map((msg) => [msg]);
+
+    return {
+      activeColumnId,
+      columns: columns as any,
+      content: '',
+      createdAt: Math.min(...members.map((m) => m.createdAt)),
+      extra: {
+        groupMode: group.mode,
+        parentMessageId: group.parentMessageId,
+      },
+      id: group.id,
+      meta: members[0]?.meta || {},
+      role: 'compare' as any,
+      updatedAt: Math.max(...members.map((m) => m.updatedAt)),
+    } as Message;
+  }
+
+  /**
+   * Create assistant group virtual message from entire chain
+   */
+  private createAssistantGroupMessage(
+    firstAssistant: Message,
+    assistantChain: Message[],
+    allToolMessages: Message[],
+  ): Message {
+    const children: AssistantContentBlock[] = [];
+
+    // Create tool map for lookup
+    const toolMap = new Map<string, Message>();
+    allToolMessages.forEach((tm) => {
+      if (tm.tool_call_id) {
+        toolMap.set(tm.tool_call_id, tm);
+      }
+    });
+
+    // Process each assistant in the chain
+    for (const assistant of assistantChain) {
+      // Build toolsWithResults for this assistant
+      const toolsWithResults: ChatToolPayloadWithResult[] =
+        assistant.tools?.map((tool) => {
+          const toolMsg = toolMap.get(tool.id);
+          if (toolMsg) {
+            return {
+              ...tool,
+              result: {
+                content: toolMsg.content || '',
+                error: toolMsg.error,
+                id: toolMsg.id,
+                state: toolMsg.pluginState,
+              },
+              result_msg_id: toolMsg.id,
+            };
+          }
+          return tool;
+        }) || [];
+
+      const { usage: msgUsage, performance: msgPerformance } =
+        this.messageTransformer.splitMetadata(assistant.metadata);
+
+      children.push({
+        content: assistant.content || '',
+        error: assistant.error,
+        id: assistant.id,
+        imageList:
+          assistant.imageList && assistant.imageList.length > 0 ? assistant.imageList : undefined,
+        performance: msgPerformance,
+        reasoning: assistant.reasoning || undefined,
+        tools: toolsWithResults.length > 0 ? toolsWithResults : undefined,
+        usage: msgUsage,
+      });
+    }
+
+    const aggregated = this.messageTransformer.aggregateMetadata(children);
+
+    return {
+      ...firstAssistant,
+      children,
+      content: '',
+      imageList: undefined,
+      metadata: undefined,
+      performance: aggregated.performance,
+      reasoning: undefined,
+      role: 'assistantGroup' as any,
+      tools: undefined,
+      usage: aggregated.usage,
+    };
+  }
+
+  /**
+   * Create user message with branch metadata
+   */
+  private createUserMessageWithBranches(user: Message): Message {
+    // Just return the original user message with its metadata.activeBranchId
+    // No need to add extra.branches
+    return { ...user };
+  }
+}

package/packages/conversation-flow/src/transformation/MessageCollector.ts

@@ -0,0 +1,166 @@
+import type { ContextNode, IdNode, Message, MessageNode } from '../types';
+
+/**
+ * MessageCollector - Handles collection of related messages
+ *
+ * Provides utilities for:
+ * 1. Collecting messages in a group
+ * 2. Collecting tool messages
+ * 3. Collecting assistant chains
+ * 4. Finding next messages in sequences
+ */
+export class MessageCollector {
+  constructor(
+    private messageMap: Map<string, Message>,
+    private childrenMap: Map<string | null, string[]>,
+  ) {}
+
+  /**
+   * Collect all messages belonging to a message group
+   */
+  collectGroupMembers(groupId: string, messages: Message[]): Message[] {
+    return messages.filter((m) => m.groupId === groupId);
+  }
+
+  /**
+   * Collect tool messages related to an assistant message
+   */
+  collectToolMessages(assistant: Message, messages: Message[]): Message[] {
+    const toolCallIds = new Set(assistant.tools?.map((t) => t.id) || []);
+    return messages.filter(
+      (m) => m.role === 'tool' && m.tool_call_id && toolCallIds.has(m.tool_call_id),
+    );
+  }
+
+  /**
+   * Recursively collect the entire assistant chain
+   * (assistant -> tools -> assistant -> tools -> ...)
+   */
+  collectAssistantChain(
+    currentAssistant: Message,
+    allMessages: Message[],
+    assistantChain: Message[],
+    allToolMessages: Message[],
+    processedIds: Set<string>,
+  ): void {
+    if (processedIds.has(currentAssistant.id)) return;
+
+    // Add current assistant to chain
+    assistantChain.push(currentAssistant);
+
+    // Collect its tool messages
+    const toolMessages = this.collectToolMessages(currentAssistant, allMessages);
+    allToolMessages.push(...toolMessages);
+
+    // Find next assistant after tools
+    for (const toolMsg of toolMessages) {
+      const nextMessages = allMessages.filter((m) => m.parentId === toolMsg.id);
+
+      for (const nextMsg of nextMessages) {
+        if (nextMsg.role === 'assistant' && nextMsg.tools && nextMsg.tools.length > 0) {
+          // Continue the chain
+          this.collectAssistantChain(
+            nextMsg,
+            allMessages,
+            assistantChain,
+            allToolMessages,
+            processedIds,
+          );
+          return;
+        } else if (nextMsg.role === 'assistant') {
+          // Final assistant without tools
+          assistantChain.push(nextMsg);
+          return;
+        }
+      }
+    }
+  }
+
+  /**
+   * Recursively collect assistant messages for an AssistantGroup (contextTree version)
+   */
+  collectAssistantGroupMessages(
+    message: Message,
+    idNode: IdNode,
+    children: ContextNode[],
+  ): void {
+    // Get tool message IDs if this assistant has tools
+    const toolIds = idNode.children
+      .filter((child) => {
+        const childMsg = this.messageMap.get(child.id);
+        return childMsg?.role === 'tool';
+      })
+      .map((child) => child.id);
+
+    // Add current assistant message node
+    const messageNode: MessageNode = {
+      id: message.id,
+      type: 'message',
+    };
+    if (toolIds.length > 0) {
+      messageNode.tools = toolIds;
+    }
+    children.push(messageNode);
+
+    // Find next assistant message after tools
+    for (const toolNode of idNode.children) {
+      const toolMsg = this.messageMap.get(toolNode.id);
+      if (toolMsg?.role !== 'tool') continue;
+
+      // Check if tool has an assistant child
+      if (toolNode.children.length > 0) {
+        const nextChild = toolNode.children[0];
+        const nextMsg = this.messageMap.get(nextChild.id);
+
+        if (nextMsg?.role === 'assistant') {
+          // Recursively collect this assistant and its descendants
+          this.collectAssistantGroupMessages(nextMsg, nextChild, children);
+          return; // Only follow one path
+        }
+      }
+    }
+  }
+
+  /**
+   * Find next message after tools in an assistant group
+   */
+  findNextAfterTools(assistantMsg: Message, idNode: IdNode): IdNode | null {
+    // Recursively find the last message in the assistant group
+    const lastNode = this.findLastNodeInAssistantGroup(idNode);
+    if (lastNode && lastNode.children.length > 0) {
+      return lastNode.children[0];
+    }
+    return null;
+  }
+
+  /**
+   * Find the last node in an AssistantGroup sequence
+   */
+  findLastNodeInAssistantGroup(idNode: IdNode): IdNode | null {
+    // Check if has tool children
+    const toolChildren = idNode.children.filter((child) => {
+      const childMsg = this.messageMap.get(child.id);
+      return childMsg?.role === 'tool';
+    });
+
+    if (toolChildren.length === 0) {
+      return idNode;
+    }
+
+    // Check if any tool has an assistant child
+    for (const toolNode of toolChildren) {
+      if (toolNode.children.length > 0) {
+        const nextChild = toolNode.children[0];
+        const nextMsg = this.messageMap.get(nextChild.id);
+
+        if (nextMsg?.role === 'assistant') {
+          // Continue following the assistant chain
+          return this.findLastNodeInAssistantGroup(nextChild);
+        }
+      }
+    }
+
+    // No more assistant messages, return the last tool node
+    return toolChildren.at(-1) ?? null;
+  }
+}