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

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

@@ -0,0 +1,177 @@
+import type { AssistantContentBlock, ModelPerformance, ModelUsage } from '@lobechat/types';
+
+import type { Message } from '../types';
+
+/**
+ * MessageTransformer - Handles message transformation utilities
+ *
+ * Provides utilities for:
+ * 1. Converting Message to AssistantContentBlock
+ * 2. Splitting metadata into usage and performance
+ * 3. Aggregating metadata from multiple messages
+ */
+export class MessageTransformer {
+  /**
+   * Convert a Message to AssistantContentBlock
+   */
+  messageToContentBlock(message: Message): AssistantContentBlock {
+    const { usage, performance } = this.splitMetadata(message.metadata);
+
+    return {
+      content: message.content || '',
+      error: message.error,
+      id: message.id,
+      imageList: message.imageList,
+      performance,
+      reasoning: message.reasoning || undefined,
+      tools: message.tools as any,
+      usage,
+    };
+  }
+
+  /**
+   * Split metadata into usage and performance objects
+   */
+  splitMetadata(metadata?: any): {
+    performance?: ModelPerformance;
+    usage?: ModelUsage;
+  } {
+    if (!metadata) return {};
+
+    const usage: ModelUsage = {};
+    const performance: ModelPerformance = {};
+
+    const usageFields = [
+      'acceptedPredictionTokens',
+      'cost',
+      'inputAudioTokens',
+      'inputCacheMissTokens',
+      'inputCachedTokens',
+      'inputCitationTokens',
+      'inputImageTokens',
+      'inputTextTokens',
+      'inputWriteCacheTokens',
+      'outputAudioTokens',
+      'outputImageTokens',
+      'outputReasoningTokens',
+      'outputTextTokens',
+      'rejectedPredictionTokens',
+      'totalInputTokens',
+      'totalOutputTokens',
+      'totalTokens',
+    ] as const;
+
+    let hasUsage = false;
+    usageFields.forEach((field) => {
+      if (metadata[field] !== undefined) {
+        (usage as any)[field] = metadata[field];
+        hasUsage = true;
+      }
+    });
+
+    const performanceFields = ['duration', 'latency', 'tps', 'ttft'] as const;
+    let hasPerformance = false;
+    performanceFields.forEach((field) => {
+      if (metadata[field] !== undefined) {
+        (performance as any)[field] = metadata[field];
+        hasPerformance = true;
+      }
+    });
+
+    return {
+      performance: hasPerformance ? performance : undefined,
+      usage: hasUsage ? usage : undefined,
+    };
+  }
+
+  /**
+   * Aggregate metadata from multiple children
+   * - Sums token counts and costs
+   * - Takes first ttft
+   * - Averages tps
+   * - Sums duration and latency
+   */
+  aggregateMetadata(children: AssistantContentBlock[]): {
+    performance?: ModelPerformance;
+    usage?: ModelUsage;
+  } {
+    const usage: ModelUsage = {};
+    const performance: ModelPerformance = {};
+    let hasUsageData = false;
+    let hasPerformanceData = false;
+    let tpsSum = 0;
+    let tpsCount = 0;
+
+    children.forEach((child) => {
+      if (child.usage) {
+        const tokenFields = [
+          'acceptedPredictionTokens',
+          'inputAudioTokens',
+          'inputCacheMissTokens',
+          'inputCachedTokens',
+          'inputCitationTokens',
+          'inputImageTokens',
+          'inputTextTokens',
+          'inputWriteCacheTokens',
+          'outputAudioTokens',
+          'outputImageTokens',
+          'outputReasoningTokens',
+          'outputTextTokens',
+          'rejectedPredictionTokens',
+          'totalInputTokens',
+          'totalOutputTokens',
+          'totalTokens',
+        ] as const;
+
+        tokenFields.forEach((field) => {
+          if (typeof child.usage![field] === 'number') {
+            (usage as any)[field] = ((usage as any)[field] || 0) + child.usage![field]!;
+            hasUsageData = true;
+          }
+        });
+
+        if (typeof child.usage.cost === 'number') {
+          usage.cost = (usage.cost || 0) + child.usage.cost;
+          hasUsageData = true;
+        }
+      }
+
+      if (child.performance) {
+        // Take first ttft (time to first token)
+        if (child.performance.ttft !== undefined && performance.ttft === undefined) {
+          performance.ttft = child.performance.ttft;
+          hasPerformanceData = true;
+        }
+
+        // Average tps (tokens per second)
+        if (typeof child.performance.tps === 'number') {
+          tpsSum += child.performance.tps;
+          tpsCount += 1;
+          hasPerformanceData = true;
+        }
+
+        // Sum duration
+        if (child.performance.duration !== undefined) {
+          performance.duration = (performance.duration || 0) + child.performance.duration;
+          hasPerformanceData = true;
+        }
+
+        // Sum latency
+        if (child.performance.latency !== undefined) {
+          performance.latency = (performance.latency || 0) + child.performance.latency;
+          hasPerformanceData = true;
+        }
+      }
+    });
+
+    // Calculate average tps
+    if (tpsCount > 0) {
+      performance.tps = tpsSum / tpsCount;
+    }
+
+    return {
+      performance: hasPerformanceData ? performance : undefined,
+      usage: hasUsageData ? usage : undefined,
+    };
+  }
+}

package/packages/conversation-flow/src/transformation/__tests__/BranchResolver.test.ts

@@ -0,0 +1,151 @@
+import { describe, expect, it } from 'vitest';
+
+import type { IdNode, Message } from '../../types';
+import { BranchResolver } from '../BranchResolver';
+
+describe('BranchResolver', () => {
+  const resolver = new BranchResolver();
+
+  describe('getActiveBranchId', () => {
+    it('should return branch by activeBranchIndex from metadata when present', () => {
+      const message: Message = {
+        content: 'test',
+        createdAt: 0,
+        id: 'msg-1',
+        meta: {},
+        metadata: { activeBranchIndex: 1 },
+        role: 'user',
+        updatedAt: 0,
+      };
+
+      const idNode: IdNode = {
+        children: [
+          { children: [], id: 'msg-2' },
+          { children: [], id: 'msg-3' },
+        ],
+        id: 'msg-1',
+      };
+
+      expect(resolver.getActiveBranchId(message, idNode)).toBe('msg-3');
+    });
+
+    it('should infer active branch from which branch has children', () => {
+      const message: Message = {
+        content: 'test',
+        createdAt: 0,
+        id: 'msg-1',
+        meta: {},
+        role: 'user',
+        updatedAt: 0,
+      };
+
+      const idNode: IdNode = {
+        children: [
+          { children: [], id: 'msg-2' },
+          { children: [{ children: [], id: 'msg-4' }], id: 'msg-3' },
+        ],
+        id: 'msg-1',
+      };
+
+      expect(resolver.getActiveBranchId(message, idNode)).toBe('msg-3');
+    });
+
+    it('should default to first branch when no hints available', () => {
+      const message: Message = {
+        content: 'test',
+        createdAt: 0,
+        id: 'msg-1',
+        meta: {},
+        role: 'user',
+        updatedAt: 0,
+      };
+
+      const idNode: IdNode = {
+        children: [
+          { children: [], id: 'msg-2' },
+          { children: [], id: 'msg-3' },
+        ],
+        id: 'msg-1',
+      };
+
+      expect(resolver.getActiveBranchId(message, idNode)).toBe('msg-2');
+    });
+
+    it('should ignore invalid activeBranchIndex', () => {
+      const message: Message = {
+        content: 'test',
+        createdAt: 0,
+        id: 'msg-1',
+        meta: {},
+        metadata: { activeBranchIndex: 5 }, // out of bounds
+        role: 'user',
+        updatedAt: 0,
+      };
+
+      const idNode: IdNode = {
+        children: [
+          { children: [], id: 'msg-2' },
+          { children: [], id: 'msg-3' },
+        ],
+        id: 'msg-1',
+      };
+
+      // Should default to first branch
+      expect(resolver.getActiveBranchId(message, idNode)).toBe('msg-2');
+    });
+  });
+
+  describe('getActiveBranchIdFromMetadata', () => {
+    it('should return branch by activeBranchIndex from metadata', () => {
+      const message: Message = {
+        content: 'test',
+        createdAt: 0,
+        id: 'msg-1',
+        meta: {},
+        metadata: { activeBranchIndex: 1 },
+        role: 'user',
+        updatedAt: 0,
+      };
+
+      const childIds = ['msg-2', 'msg-3'];
+      const childrenMap = new Map<string | null, string[]>();
+
+      expect(resolver.getActiveBranchIdFromMetadata(message, childIds, childrenMap)).toBe('msg-3');
+    });
+
+    it('should infer active branch from which child has descendants', () => {
+      const message: Message = {
+        content: 'test',
+        createdAt: 0,
+        id: 'msg-1',
+        meta: {},
+        role: 'user',
+        updatedAt: 0,
+      };
+
+      const childIds = ['msg-2', 'msg-3'];
+      const childrenMap = new Map<string | null, string[]>([
+        ['msg-2', []],
+        ['msg-3', ['msg-4']],
+      ]);
+
+      expect(resolver.getActiveBranchIdFromMetadata(message, childIds, childrenMap)).toBe('msg-3');
+    });
+
+    it('should default to first child when no hints available', () => {
+      const message: Message = {
+        content: 'test',
+        createdAt: 0,
+        id: 'msg-1',
+        meta: {},
+        role: 'user',
+        updatedAt: 0,
+      };
+
+      const childIds = ['msg-2', 'msg-3'];
+      const childrenMap = new Map<string | null, string[]>();
+
+      expect(resolver.getActiveBranchIdFromMetadata(message, childIds, childrenMap)).toBe('msg-2');
+    });
+  });
+});

package/packages/conversation-flow/src/transformation/__tests__/ContextTreeBuilder.test.ts

@@ -0,0 +1,385 @@
+import { describe, expect, it } from 'vitest';
+
+import type { IdNode, Message, MessageGroupMetadata } from '../../types';
+import { BranchResolver } from '../BranchResolver';
+import { ContextTreeBuilder } from '../ContextTreeBuilder';
+import { MessageCollector } from '../MessageCollector';
+
+describe('ContextTreeBuilder', () => {
+  const createBuilder = (
+    messageMap: Map<string, Message>,
+    messageGroupMap: Map<string, MessageGroupMetadata> = new Map(),
+  ) => {
+    const childrenMap = new Map<string | null, string[]>();
+    const branchResolver = new BranchResolver();
+    const messageCollector = new MessageCollector(messageMap, childrenMap);
+    let nodeIdCounter = 0;
+    const generateNodeId = (prefix: string, messageId: string) =>
+      `${prefix}-${messageId}-${nodeIdCounter++}`;
+
+    return new ContextTreeBuilder(
+      messageMap,
+      messageGroupMap,
+      branchResolver,
+      messageCollector,
+      generateNodeId,
+    );
+  };
+
+  describe('transformAll', () => {
+    it('should transform regular message nodes', () => {
+      const messageMap = new Map<string, Message>([
+        [
+          'msg-1',
+          {
+            content: 'Hello',
+            createdAt: 0,
+            id: 'msg-1',
+            meta: {},
+            role: 'user',
+            updatedAt: 0,
+          },
+        ],
+        [
+          'msg-2',
+          {
+            content: 'Hi',
+            createdAt: 0,
+            id: 'msg-2',
+            meta: {},
+            role: 'assistant',
+            updatedAt: 0,
+          },
+        ],
+      ]);
+
+      const builder = createBuilder(messageMap);
+      const idNodes: IdNode[] = [
+        {
+          children: [{ children: [], id: 'msg-2' }],
+          id: 'msg-1',
+        },
+      ];
+
+      const result = builder.transformAll(idNodes);
+
+      expect(result).toHaveLength(2);
+      expect(result[0]).toEqual({ id: 'msg-1', type: 'message' });
+      expect(result[1]).toEqual({ id: 'msg-2', type: 'message' });
+    });
+
+    it('should create branch node for multiple children', () => {
+      const messageMap = new Map<string, Message>([
+        [
+          'msg-1',
+          {
+            content: 'Hello',
+            createdAt: 0,
+            id: 'msg-1',
+            meta: {},
+            metadata: { activeBranchIndex: 0 },
+            role: 'user',
+            updatedAt: 0,
+          },
+        ],
+        [
+          'msg-2',
+          {
+            content: 'Response 1',
+            createdAt: 0,
+            id: 'msg-2',
+            meta: {},
+            role: 'assistant',
+            updatedAt: 0,
+          },
+        ],
+        [
+          'msg-3',
+          {
+            content: 'Response 2',
+            createdAt: 0,
+            id: 'msg-3',
+            meta: {},
+            role: 'assistant',
+            updatedAt: 0,
+          },
+        ],
+      ]);
+
+      const builder = createBuilder(messageMap);
+      const idNodes: IdNode[] = [
+        {
+          children: [
+            { children: [], id: 'msg-2' },
+            { children: [], id: 'msg-3' },
+          ],
+          id: 'msg-1',
+        },
+      ];
+
+      const result = builder.transformAll(idNodes);
+
+      expect(result).toHaveLength(2);
+      expect(result[0]).toEqual({ id: 'msg-1', type: 'message' });
+      expect(result[1]).toMatchObject({
+        activeBranchIndex: 0,
+        branches: [[{ id: 'msg-2', type: 'message' }], [{ id: 'msg-3', type: 'message' }]],
+        parentMessageId: 'msg-1',
+        type: 'branch',
+      });
+    });
+
+    it('should create assistant group node for assistant with tools', () => {
+      const messageMap = new Map<string, Message>([
+        [
+          'msg-1',
+          {
+            content: 'Assistant with tools',
+            createdAt: 0,
+            id: 'msg-1',
+            meta: {},
+            role: 'assistant',
+            tools: [
+              {
+                apiName: 'test',
+                arguments: '{}',
+                id: 'tool-1',
+                identifier: 'test',
+                type: 'default',
+              },
+            ],
+            updatedAt: 0,
+          },
+        ],
+        [
+          'tool-1',
+          {
+            content: 'Tool result',
+            createdAt: 0,
+            id: 'tool-1',
+            meta: {},
+            role: 'tool',
+            updatedAt: 0,
+          },
+        ],
+      ]);
+
+      const builder = createBuilder(messageMap);
+      const idNodes: IdNode[] = [
+        {
+          children: [{ children: [], id: 'tool-1' }],
+          id: 'msg-1',
+        },
+      ];
+
+      const result = builder.transformAll(idNodes);
+
+      expect(result).toHaveLength(1);
+      expect(result[0]).toMatchObject({
+        children: [{ id: 'msg-1', tools: ['tool-1'], type: 'message' }],
+        id: 'msg-1',
+        type: 'assistantGroup',
+      });
+    });
+
+    it('should create compare node from message group', () => {
+      const messageMap = new Map<string, Message>([
+        [
+          'msg-1',
+          {
+            content: 'Compare 1',
+            createdAt: 0,
+            groupId: 'group-1',
+            id: 'msg-1',
+            meta: {},
+            metadata: { activeColumn: true },
+            role: 'assistant',
+            updatedAt: 0,
+          },
+        ],
+        [
+          'msg-2',
+          {
+            content: 'Compare 2',
+            createdAt: 0,
+            groupId: 'group-1',
+            id: 'msg-2',
+            meta: {},
+            role: 'assistant',
+            updatedAt: 0,
+          },
+        ],
+      ]);
+
+      const messageGroupMap = new Map<string, MessageGroupMetadata>([
+        ['group-1', { id: 'group-1', mode: 'compare' }],
+      ]);
+
+      const builder = createBuilder(messageMap, messageGroupMap);
+      const idNodes: IdNode[] = [{ children: [], id: 'msg-1' }];
+
+      const result = builder.transformAll(idNodes);
+
+      expect(result).toHaveLength(1);
+      expect(result[0]).toMatchObject({
+        activeColumnId: 'msg-1',
+        columns: [[{ id: 'msg-1', type: 'message' }], [{ id: 'msg-2', type: 'message' }]],
+        messageId: 'msg-1',
+        type: 'compare',
+      });
+    });
+
+    it('should create compare node from user message metadata', () => {
+      const messageMap = new Map<string, Message>([
+        [
+          'msg-1',
+          {
+            content: 'User message',
+            createdAt: 0,
+            id: 'msg-1',
+            meta: {},
+            metadata: { compare: true },
+            role: 'user',
+            updatedAt: 0,
+          },
+        ],
+        [
+          'msg-2',
+          {
+            content: 'Assistant 1',
+            createdAt: 0,
+            id: 'msg-2',
+            meta: {},
+            metadata: { activeColumn: true },
+            role: 'assistant',
+            updatedAt: 0,
+          },
+        ],
+        [
+          'msg-3',
+          {
+            content: 'Assistant 2',
+            createdAt: 0,
+            id: 'msg-3',
+            meta: {},
+            role: 'assistant',
+            updatedAt: 0,
+          },
+        ],
+      ]);
+
+      const builder = createBuilder(messageMap);
+      const idNodes: IdNode[] = [
+        {
+          children: [
+            { children: [], id: 'msg-2' },
+            { children: [], id: 'msg-3' },
+          ],
+          id: 'msg-1',
+        },
+      ];
+
+      const result = builder.transformAll(idNodes);
+
+      expect(result).toHaveLength(2);
+      expect(result[0]).toEqual({ id: 'msg-1', type: 'message' });
+      expect(result[1]).toMatchObject({
+        activeColumnId: 'msg-2',
+        columns: [[{ id: 'msg-2', type: 'message' }], [{ id: 'msg-3', type: 'message' }]],
+        id: 'compare-msg-1-msg-2-msg-3',
+        messageId: 'msg-1',
+        type: 'compare',
+      });
+    });
+
+    it('should handle empty node list', () => {
+      const messageMap = new Map<string, Message>();
+      const builder = createBuilder(messageMap);
+
+      const result = builder.transformAll([]);
+
+      expect(result).toHaveLength(0);
+    });
+
+    it('should handle missing message in map', () => {
+      const messageMap = new Map<string, Message>();
+      const builder = createBuilder(messageMap);
+      const idNodes: IdNode[] = [{ children: [], id: 'missing' }];
+
+      const result = builder.transformAll(idNodes);
+
+      expect(result).toHaveLength(0);
+    });
+
+    it('should continue with active column children in compare mode', () => {
+      const messageMap = new Map<string, Message>([
+        [
+          'msg-1',
+          {
+            content: 'User',
+            createdAt: 0,
+            id: 'msg-1',
+            meta: {},
+            metadata: { compare: true },
+            role: 'user',
+            updatedAt: 0,
+          },
+        ],
+        [
+          'msg-2',
+          {
+            content: 'Assistant 1',
+            createdAt: 0,
+            id: 'msg-2',
+            meta: {},
+            metadata: { activeColumn: true },
+            role: 'assistant',
+            updatedAt: 0,
+          },
+        ],
+        [
+          'msg-3',
+          {
+            content: 'Assistant 2',
+            createdAt: 0,
+            id: 'msg-3',
+            meta: {},
+            role: 'assistant',
+            updatedAt: 0,
+          },
+        ],
+        [
+          'msg-4',
+          {
+            content: 'Follow-up',
+            createdAt: 0,
+            id: 'msg-4',
+            meta: {},
+            role: 'user',
+            updatedAt: 0,
+          },
+        ],
+      ]);
+
+      const builder = createBuilder(messageMap);
+      const idNodes: IdNode[] = [
+        {
+          children: [
+            { children: [{ children: [], id: 'msg-4' }], id: 'msg-2' },
+            { children: [], id: 'msg-3' },
+          ],
+          id: 'msg-1',
+        },
+      ];
+
+      const result = builder.transformAll(idNodes);
+
+      expect(result).toHaveLength(3);
+      expect(result[0]).toEqual({ id: 'msg-1', type: 'message' });
+      expect(result[1]).toMatchObject({
+        type: 'compare',
+      });
+      expect(result[2]).toEqual({ id: 'msg-4', type: 'message' });
+    });
+  });
+});