@lobehub/lobehub 2.0.0-next.37 → 2.0.0-next.39

package/src/server/services/usage/index.test.ts

@@ -0,0 +1,310 @@
+import dayjs from 'dayjs';
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { LobeChatDatabase } from '@/database/type';
+import { MessageMetadata } from '@/types/message';
+
+import { UsageRecordService } from './index';
+
+describe('UsageRecordService', () => {
+  let service: UsageRecordService;
+  let mockDb: LobeChatDatabase;
+  const userId = 'test-user-id';
+
+  // Helper function to setup query chain mock
+  const setupQueryChainMock = (mockMessages: any[]) => {
+    const mockOrderBy = vi.fn().mockResolvedValue(mockMessages);
+    const mockWhere = vi.fn().mockReturnValue({ orderBy: mockOrderBy });
+    const mockFrom = vi.fn().mockReturnValue({ where: mockWhere });
+    mockDb.select = vi.fn().mockReturnValue({ from: mockFrom });
+  };
+
+  beforeEach(() => {
+    // Create a fresh mock for each test
+    const mockOrderBy = vi.fn();
+    const mockWhere = vi.fn().mockReturnValue({ orderBy: mockOrderBy });
+    const mockFrom = vi.fn().mockReturnValue({ where: mockWhere });
+    const mockSelect = vi.fn().mockReturnValue({ from: mockFrom });
+
+    mockDb = {
+      select: mockSelect,
+    } as unknown as LobeChatDatabase;
+
+    service = new UsageRecordService(mockDb, userId);
+  });
+
+  describe('findByMonth', () => {
+    it('should return usage records for the current month when no month is provided', async () => {
+      const mockMessages = [
+        {
+          id: 'msg-1',
+          userId: userId,
+          role: 'assistant',
+          provider: 'openai',
+          model: 'gpt-4',
+          createdAt: new Date(),
+          metadata: {
+            cost: 0.05,
+            totalInputTokens: 100,
+            totalOutputTokens: 50,
+            tps: 10,
+            ttft: 500,
+          } as MessageMetadata,
+        },
+      ];
+
+      setupQueryChainMock(mockMessages);
+
+      const result = await service.findByMonth();
+
+      expect(result).toHaveLength(1);
+      expect(result[0]).toMatchObject({
+        id: 'msg-1',
+        model: 'gpt-4',
+        provider: 'openai',
+        spend: 0.05,
+        totalInputTokens: 100,
+        totalOutputTokens: 50,
+        totalTokens: 150,
+        tps: 10,
+        ttft: 500,
+        type: 'chat',
+        userId: userId,
+      });
+    });
+
+    it('should return usage records for a specific month', async () => {
+      const mockMessages = [
+        {
+          id: 'msg-1',
+          userId: userId,
+          role: 'assistant',
+          provider: 'anthropic',
+          model: 'claude-3',
+          createdAt: new Date('2024-01-15'),
+          metadata: {
+            cost: 0.03,
+            totalInputTokens: 80,
+            totalOutputTokens: 40,
+          } as MessageMetadata,
+        },
+      ];
+
+      setupQueryChainMock(mockMessages);
+
+      const result = await service.findByMonth('2024-01');
+
+      expect(result[0].model).toBe('claude-3');
+      expect(result[0].spend).toBe(0.03);
+    });
+
+    it('should handle messages with missing metadata fields', async () => {
+      const mockMessages = [
+        {
+          id: 'msg-1',
+          userId: userId,
+          role: 'assistant',
+          provider: 'openai',
+          model: 'gpt-3.5-turbo',
+          createdAt: new Date(),
+          metadata: {} as MessageMetadata,
+        },
+      ];
+
+      setupQueryChainMock(mockMessages);
+
+      const result = await service.findByMonth();
+
+      expect(result).toHaveLength(1);
+      expect(result[0]).toMatchObject({
+        spend: 0,
+        totalInputTokens: 0,
+        totalOutputTokens: 0,
+        totalTokens: 0,
+        tps: 0,
+        ttft: 0,
+      });
+    });
+
+    it('should return empty array when no messages found', async () => {
+      setupQueryChainMock([]);
+
+      const result = await service.findByMonth();
+
+      expect(result).toHaveLength(0);
+    });
+  });
+
+  describe('findAndGroupByDay', () => {
+    it('should group usage records by day for current month', async () => {
+      const date1 = dayjs().startOf('month').toDate();
+      const date2 = dayjs().startOf('month').add(1, 'day').toDate();
+
+      const mockMessages = [
+        {
+          id: 'msg-1',
+          userId: userId,
+          role: 'assistant',
+          provider: 'openai',
+          model: 'gpt-4',
+          createdAt: date1,
+          metadata: {
+            cost: 0.05,
+            totalInputTokens: 100,
+            totalOutputTokens: 50,
+          } as MessageMetadata,
+        },
+        {
+          id: 'msg-2',
+          userId: userId,
+          role: 'assistant',
+          provider: 'openai',
+          model: 'gpt-4',
+          createdAt: date1,
+          metadata: {
+            cost: 0.03,
+            totalInputTokens: 60,
+            totalOutputTokens: 30,
+          } as MessageMetadata,
+        },
+        {
+          id: 'msg-3',
+          userId: userId,
+          role: 'assistant',
+          provider: 'anthropic',
+          model: 'claude-3',
+          createdAt: date2,
+          metadata: {
+            cost: 0.02,
+            totalInputTokens: 40,
+            totalOutputTokens: 20,
+          } as MessageMetadata,
+        },
+      ];
+
+      setupQueryChainMock(mockMessages);
+
+      const result = await service.findAndGroupByDay();
+
+      expect(result.length).toBeGreaterThan(0);
+
+      // Check that days with records have correct aggregations
+      const dayWithRecords = result.find((log) => log.totalRequests > 0);
+      if (dayWithRecords) {
+        expect(dayWithRecords.totalSpend).toBeGreaterThan(0);
+        expect(dayWithRecords.totalTokens).toBeGreaterThan(0);
+        expect(dayWithRecords.records.length).toBeGreaterThan(0);
+      }
+    });
+
+    it('should pad missing days with zero values', async () => {
+      const firstDay = dayjs().startOf('month');
+
+      const mockMessages = [
+        {
+          id: 'msg-1',
+          userId: userId,
+          role: 'assistant',
+          provider: 'openai',
+          model: 'gpt-4',
+          createdAt: firstDay.toDate(),
+          metadata: {
+            cost: 0.05,
+            totalInputTokens: 100,
+            totalOutputTokens: 50,
+          } as MessageMetadata,
+        },
+      ];
+
+      setupQueryChainMock(mockMessages);
+
+      const result = await service.findAndGroupByDay();
+
+      // Should have entries for every day in the month
+      const daysInMonth = dayjs().endOf('month').date();
+      expect(result.length).toBeGreaterThanOrEqual(daysInMonth - 1);
+
+      // Check that padded days have zero values
+      const paddedDay = result.find((log) => log.totalRequests === 0);
+      if (paddedDay) {
+        expect(paddedDay.totalSpend).toBe(0);
+        expect(paddedDay.totalTokens).toBe(0);
+        expect(paddedDay.records).toHaveLength(0);
+      }
+    });
+
+    it('should calculate correct totals for days with multiple records', async () => {
+      const testDate = dayjs().startOf('month').toDate();
+
+      const mockMessages = [
+        {
+          id: 'msg-1',
+          userId: userId,
+          role: 'assistant',
+          provider: 'openai',
+          model: 'gpt-4',
+          createdAt: testDate,
+          metadata: {
+            cost: 0.05,
+            totalInputTokens: 100,
+            totalOutputTokens: 50,
+          } as MessageMetadata,
+        },
+        {
+          id: 'msg-2',
+          userId: userId,
+          role: 'assistant',
+          provider: 'openai',
+          model: 'gpt-4',
+          createdAt: testDate,
+          metadata: {
+            cost: 0.03,
+            totalInputTokens: 60,
+            totalOutputTokens: 30,
+          } as MessageMetadata,
+        },
+      ];
+
+      setupQueryChainMock(mockMessages);
+
+      const result = await service.findAndGroupByDay();
+
+      const dayLog = result.find((log) => log.totalRequests === 2);
+
+      if (dayLog) {
+        expect(dayLog.totalSpend).toBe(0.08);
+        expect(dayLog.totalTokens).toBe(240); // (100+50) + (60+30)
+        expect(dayLog.totalRequests).toBe(2);
+        expect(dayLog.records).toHaveLength(2);
+      }
+    });
+
+    it('should handle specific month parameter', async () => {
+      const mockMessages = [
+        {
+          id: 'msg-1',
+          userId: userId,
+          role: 'assistant',
+          provider: 'openai',
+          model: 'gpt-4',
+          createdAt: new Date('2024-01-15'),
+          metadata: {
+            cost: 0.05,
+            totalInputTokens: 100,
+            totalOutputTokens: 50,
+          } as MessageMetadata,
+        },
+      ];
+
+      setupQueryChainMock(mockMessages);
+
+      const result = await service.findAndGroupByDay('2024-01');
+
+      expect(result.length).toBeGreaterThan(0);
+      // All days should be from January 2024
+      result.forEach((log) => {
+        expect(log.day).toMatch(/^2024-01/);
+      });
+    });
+  });
+});

package/src/server/services/usage/index.ts

@@ -0,0 +1,164 @@
+import dayjs from 'dayjs';
+import debug from 'debug';
+import { desc, eq } from 'drizzle-orm';
+
+import { messages } from '@/database/schemas';
+import { LobeChatDatabase } from '@/database/type';
+import { genRangeWhere, genWhere } from '@/database/utils/genWhere';
+import { MessageMetadata } from '@/types/message';
+import { UsageLog, UsageRecordItem } from '@/types/usage/usageRecord';
+import { formatDate } from '@/utils/format';
+
+const log = debug('lobe-usage:service');
+
+export class UsageRecordService {
+  private userId: string;
+  private db: LobeChatDatabase;
+  constructor(db: LobeChatDatabase, userId: string) {
+    this.userId = userId;
+    this.db = db;
+  }
+
+  /**
+   * @description Find usage records by month.
+   * @param mo Month
+   * @returns UsageRecordItem[]
+   */
+  findByMonth = async (mo?: string): Promise<UsageRecordItem[]> => {
+    // 设置 startAt 和 endAt
+    let startAt: string;
+    let endAt: string;
+    if (mo) {
+      // mo 格式: "YYYY-MM"
+      startAt = dayjs(mo, 'YYYY-MM').startOf('month').format('YYYY-MM-DD');
+      endAt = dayjs(mo, 'YYYY-MM').endOf('month').format('YYYY-MM-DD');
+    } else {
+      startAt = dayjs().startOf('month').format('YYYY-MM-DD');
+      endAt = dayjs().endOf('month').format('YYYY-MM-DD');
+    }
+
+    // TODO: To extend to support other features
+    // - Functionality:
+    //  - More type of usage, e.g. image generation, file processing, summary, search engine.
+    //  - More dimension for analysis, e.g. relational analysis.
+    // - Performance: Computing asynchronously for performance.
+    // For now, we only support chat messages for normal users for PoC.
+
+    const spends = await this.db
+      .select({
+        createdAt: messages.createdAt,
+        id: messages.id,
+        metadata: messages.metadata,
+        model: messages.model,
+        provider: messages.provider,
+        role: messages.role,
+        updatedAt: messages.createdAt,
+        userId: messages.userId,
+      })
+      .from(messages)
+      .where(
+        genWhere([
+          eq(messages.userId, this.userId),
+          eq(messages.role, 'assistant'),
+          genRangeWhere([startAt, endAt], messages.createdAt, (date) => date.toDate()),
+        ]),
+      )
+      .orderBy(desc(messages.createdAt));
+    return spends.map((spend) => {
+      const metadata = spend.metadata as MessageMetadata;
+      return {
+        createdAt: spend.createdAt,
+        id: spend.id,
+        metadata: spend.metadata,
+        model: spend.model,
+        provider: spend.provider,
+        spend: metadata?.cost || 0, // Messages do not have a direct cost associated
+        totalInputTokens: metadata?.totalInputTokens || 0,
+        totalOutputTokens: metadata?.totalOutputTokens || 0,
+        totalTokens: (metadata?.totalInputTokens || 0) + (metadata?.totalOutputTokens || 0),
+        tps: metadata?.tps || 0,
+        ttft: metadata?.ttft || 0,
+        type: 'chat', // Default to 'chat' for messages
+        updatedAt: spend.createdAt,
+        userId: spend.userId,
+      } as UsageRecordItem;
+    });
+  };
+
+  findAndGroupByDay = async (mo?: string): Promise<UsageLog[]> => {
+    // 设置 startAt 和 endAt
+    let startAt: string;
+    let endAt: string;
+    if (mo) {
+      // mo 格式: "YYYY-MM"
+      startAt = dayjs(mo, 'YYYY-MM').startOf('month').format('YYYY-MM-DD');
+      endAt = dayjs(mo, 'YYYY-MM').endOf('month').format('YYYY-MM-DD');
+    } else {
+      startAt = dayjs().startOf('month').format('YYYY-MM-DD');
+      endAt = dayjs().endOf('month').format('YYYY-MM-DD');
+    }
+    const spends = await this.findByMonth(mo);
+    // Clustering by time
+    let usages = new Map<string, { date: Date; logs: UsageRecordItem[] }>();
+    spends.forEach((spend) => {
+      if (!usages.has(formatDate(spend.createdAt))) {
+        usages.set(formatDate(spend.createdAt), { date: spend.createdAt, logs: [spend] });
+        return;
+      }
+      usages.get(formatDate(spend.createdAt))?.logs.push(spend);
+    });
+    // Calculate usage
+    let usageLogs: UsageLog[] = [];
+    usages.forEach((spends, date) => {
+      const totalSpend = spends.logs.reduce((acc, spend) => acc + spend.spend, 0);
+      const totalTokens = spends.logs.reduce((acc, spend) => (spend.totalTokens || 0) + acc, 0);
+      const totalRequests = spends.logs?.length ?? 0;
+      log(
+        'date',
+        date,
+        'totalSpend',
+        totalSpend,
+        'totalTokens',
+        totalTokens,
+        'totalRequests',
+        totalRequests,
+      );
+      usageLogs.push({
+        date: spends.date.getTime(),
+        day: date,
+        records: spends.logs,
+        totalRequests,
+        totalSpend,
+        totalTokens, // Store the formatted date as a string
+      });
+    });
+    // Padding to ensure the date range is complete
+    const startDate = dayjs(startAt);
+    const endDate = dayjs(endAt);
+    const paddedUsageLogs: UsageLog[] = [];
+    // For every day in the range, check if it exists in usageLogs
+    // If exists, use it; if not, create a new log with 0 values
+    log(
+      'Padding usage logs from',
+      startDate.format('YYYY-MM-DD'),
+      'to',
+      endDate.format('YYYY-MM-DD'),
+    );
+    for (let date = startDate; date.isBefore(endDate); date = date.add(1, 'day')) {
+      const log = usageLogs.find((log) => log.day === date.format('YYYY-MM-DD'));
+      if (log) {
+        paddedUsageLogs.push(log);
+      } else {
+        paddedUsageLogs.push({
+          date: date.toDate().getTime(),
+          day: date.format('YYYY-MM-DD'),
+          records: [],
+          totalRequests: 0,
+          totalSpend: 0,
+          totalTokens: 0,
+        });
+      }
+    }
+    return paddedUsageLogs;
+  };
+}

package/src/services/chat/contextEngineering.test.ts

@@ -243,6 +243,10 @@ describe('contextEngineering', () => {
             type: 'text',
           },
         ],
+        reasoning: {
+          content: 'I need to calculate the answer to life, universe, and everything.',
+          signature: 'thinking_process',
+        },
         role: 'assistant',
       },
     ]);

package/src/services/mcp.test.ts

@@ -305,7 +305,13 @@ describe('MCPService', () => {
       mockPluginSelectors.getInstalledPluginById.mockReturnValue(() => mockPlugin);
       mockPluginSelectors.getCustomPluginById.mockReturnValue(() => null);
 
-      const mockResult = 'response data';
+      const mockResult = {
+        content: 'response data',
+        state: {
+          content: [{ text: 'response data', type: 'text' }],
+        },
+        success: true,
+      };
       vi.mocked(toolsClient.mcp.callTool.mutate).mockResolvedValue(mockResult);
 
       const payload: ChatToolPayload = {

package/src/services/mcp.ts

@@ -4,6 +4,7 @@ import { isLocalOrPrivateUrl, safeParseJSON } from '@lobechat/utils';
 import { PluginManifest } from '@lobehub/market-sdk';
 import { CallReportRequest } from '@lobehub/market-types';
 
+import { MCPToolCallResult } from '@/libs/mcp';
 import { desktopClient, toolsClient } from '@/libs/trpc/client';
 
 import { discoverService } from './discover';
@@ -44,16 +45,16 @@ class MCPService {
     const connection = plugin.customParams?.mcp;
     const settingsEntries = plugin.settings
       ? Object.entries(plugin.settings as Record<string, any>).filter(
-        ([, value]) => value !== undefined && value !== null,
-      )
+          ([, value]) => value !== undefined && value !== null,
+        )
       : [];
     const pluginSettings =
       settingsEntries.length > 0
         ? settingsEntries.reduce<Record<string, unknown>>((acc, [key, value]) => {
-          acc[key] = value;
+            acc[key] = value;
 
-          return acc;
-        }, {})
+            return acc;
+          }, {})
         : undefined;
 
     const params = {
@@ -77,7 +78,7 @@ class MCPService {
 
     const data = {
       args,
-      env: connection?.type === 'stdio' ? params.env : pluginSettings ?? connection?.env,
+      env: connection?.type === 'stdio' ? params.env : (pluginSettings ?? connection?.env),
       params,
       toolName: apiName,
     };
@@ -89,7 +90,7 @@ class MCPService {
     let success = false;
     let errorCode: string | undefined;
     let errorMessage: string | undefined;
-    let result: any;
+    let result: MCPToolCallResult | undefined;
 
     try {
       // For desktop and stdio, use the desktopClient
@@ -119,7 +120,7 @@ class MCPService {
 
       const requestSizeBytes = calculateObjectSizeBytes(inputParams);
       // 计算响应大小
-      const responseSizeBytes = success ? calculateObjectSizeBytes(result) : 0;
+      const responseSizeBytes = success && result ? calculateObjectSizeBytes(result.state) : 0;
 
       const isCustomPlugin = !!customPlugin;
       // 构造上报数据
@@ -127,10 +128,10 @@ class MCPService {
         callDurationMs,
         customPluginInfo: isCustomPlugin
           ? {
-            avatar: plugin.manifest?.meta.avatar,
-            description: plugin.manifest?.meta.description,
-            name: plugin.manifest?.meta.title,
-          }
+              avatar: plugin.manifest?.meta.avatar,
+              description: plugin.manifest?.meta.description,
+              name: plugin.manifest?.meta.title,
+            }
           : undefined,
         errorCode,
         errorMessage,

package/src/services/usage.ts

@@ -0,0 +1,13 @@
+import { lambdaClient } from "@/libs/trpc/client";
+
+class UsageService {
+    findByMonth = async (mo?: string) => {
+        return lambdaClient.usage.findByMonth.query({ mo });
+    };
+
+    findAndGroupByDay = async (mo?: string) => {
+        return lambdaClient.usage.findAndGroupByDay.query({ mo });
+    }
+}
+
+export const usageService = new UsageService();

package/src/store/chat/agents/createAgentExecutors.ts

@@ -39,10 +39,9 @@ export const createAgentExecutors = (context: {
     traceId?: string;
   };
   parentId: string;
-  parentMessageType: 'user' | 'assistant';
+  skipCreateFirstMessage?: boolean;
 }) => {
-  // 当通过 sendMessageInServer 的时候，已经有一条消息了，那么就不需要触发创建
-  let shouldSkipCreateMessage = context.parentMessageType === 'assistant';
+  let shouldSkipCreateMessage = context.skipCreateFirstMessage;
 
   const executors: Partial<Record<AgentInstruction['type'], InstructionExecutor>> = {
     /**

package/src/store/chat/slices/aiChat/actions/conversationLifecycle.ts

@@ -44,6 +44,10 @@ export interface ConversationLifecycleAction {
     id: string,
     params?: { skipTrace?: boolean; traceId?: string },
   ) => Promise<void>;
+  /**
+   * Continue generating from current assistant message
+   */
+  continueGenerationMessage: (lastBlockId: string, messageId: string) => Promise<void>;
   /**
    * Deletes an existing message and generates a new one in its place
    */
@@ -75,7 +79,12 @@ export const conversationLifecycle: StateCreator<
     }
 
     const messages = displayMessageSelectors.activeDisplayMessages(get());
-    const parentId = displayMessageSelectors.lastDisplayMessageId(get());
+    const lastDisplayMessageId = displayMessageSelectors.lastDisplayMessageId(get());
+
+    let parentId: string | undefined;
+    if (lastDisplayMessageId) {
+      parentId = displayMessageSelectors.findLastMessageId(lastDisplayMessageId)(get());
+    }
 
     const chatConfig = agentChatConfigSelectors.currentChatConfig(getAgentStoreState());
     const autoCreateThreshold =
@@ -230,6 +239,7 @@ export const conversationLifecycle: StateCreator<
         parentMessageType: 'assistant',
         ragQuery: get().internal_shouldUseRAG() ? message : undefined,
         threadId: activeThreadId,
+        skipCreateFirstMessage: true,
       });
 
       //
@@ -318,6 +328,35 @@ export const conversationLifecycle: StateCreator<
     await get().regenerateUserMessage(userId, params);
   },
 
+  continueGenerationMessage: async (id, messageId) => {
+    const message = dbMessageSelectors.getDbMessageById(id)(get());
+    if (!message) return;
+
+    try {
+      // Mark message as continuing
+      set(
+        { continuingIds: [...get().continuingIds, messageId] },
+        false,
+        'continueGenerationMessage/start',
+      );
+
+      const chats = displayMessageSelectors.mainAIChatsWithHistoryConfig(get());
+
+      await get().internal_execAgentRuntime({
+        messages: chats,
+        parentMessageId: id,
+        parentMessageType: message.role as 'assistant' | 'tool' | 'user',
+      });
+    } finally {
+      // Remove message from continuing state
+      set(
+        { continuingIds: get().continuingIds.filter((msgId) => msgId !== messageId) },
+        false,
+        'continueGenerationMessage/end',
+      );
+    }
+  },
+
   delAndRegenerateMessage: async (id) => {
     const traceId = dbMessageSelectors.getTraceIdByDbMessageId(id)(get());
     get().regenerateAssistantMessage(id, { skipTrace: true, traceId });