@lobehub/chat 1.61.5 → 1.62.0

package/src/libs/agent-runtime/utils/streams/protocol.ts

@@ -2,8 +2,16 @@ import { ChatStreamCallbacks } from '@/libs/agent-runtime';
 
 import { AgentRuntimeErrorType } from '../../error';
 
-export interface StreamStack {
+/**
+ * context in the stream to save temporarily data
+ */
+export interface StreamContext {
   id: string;
+  /**
+   * As pplx citations is in every chunk, but we only need to return it once
+   * this flag is used to check if the pplx citation is returned,and then not return it again
+   */
+  returnedPplxCitation?: boolean;
   tool?: {
     id: string;
     index: number;
@@ -15,7 +23,20 @@ export interface StreamStack {
 export interface StreamProtocolChunk {
   data: any;
   id?: string;
-  type: 'text' | 'tool_calls' | 'data' | 'stop' | 'error' | 'reasoning';
+  type: // pure text
+  | 'text'
+    // Tools use
+    | 'tool_calls'
+    // Model Thinking
+    | 'reasoning'
+    // Search or Grounding
+    | 'citations'
+    // stop signal
+    | 'stop'
+    // Error
+    | 'error'
+    // unknown data result
+    | 'data';
 }
 
 export interface StreamToolCallChunkData {
@@ -85,16 +106,20 @@ export const convertIterableToStream = <T>(stream: AsyncIterable<T>) => {
  * Create a transformer to convert the response into an SSE format
  */
 export const createSSEProtocolTransformer = (
-  transformer: (chunk: any, stack: StreamStack) => StreamProtocolChunk,
-  streamStack?: StreamStack,
+  transformer: (chunk: any, stack: StreamContext) => StreamProtocolChunk | StreamProtocolChunk[],
+  streamStack?: StreamContext,
 ) =>
   new TransformStream({
     transform: (chunk, controller) => {
-      const { type, id, data } = transformer(chunk, streamStack || { id: '' });
+      const result = transformer(chunk, streamStack || { id: '' });
+
+      const buffers = Array.isArray(result) ? result : [result];
 
-      controller.enqueue(`id: ${id}\n`);
-      controller.enqueue(`event: ${type}\n`);
-      controller.enqueue(`data: ${JSON.stringify(data)}\n\n`);
+      buffers.forEach(({ type, id, data }) => {
+        controller.enqueue(`id: ${id}\n`);
+        controller.enqueue(`event: ${type}\n`);
+        controller.enqueue(`data: ${JSON.stringify(data)}\n\n`);
+      });
     },
   });
 

package/src/libs/agent-runtime/utils/streams/vertex-ai.ts

@@ -4,8 +4,8 @@ import { nanoid } from '@/utils/uuid';
 
 import { ChatStreamCallbacks } from '../../types';
 import {
+  StreamContext,
   StreamProtocolChunk,
-  StreamStack,
   createCallbacksTransformer,
   createSSEProtocolTransformer,
   generateToolCallId,
@@ -13,7 +13,7 @@ import {
 
 const transformVertexAIStream = (
   chunk: GenerateContentResponse,
-  stack: StreamStack,
+  stack: StreamContext,
 ): StreamProtocolChunk => {
   // maybe need another structure to add support for multiple choices
   const candidates = chunk.candidates;
@@ -67,7 +67,7 @@ export const VertexAIStream = (
   rawStream: ReadableStream<EnhancedGenerateContentResponse>,
   callbacks?: ChatStreamCallbacks,
 ) => {
-  const streamStack: StreamStack = { id: 'chat_' + nanoid() };
+  const streamStack: StreamContext = { id: 'chat_' + nanoid() };
 
   return rawStream
     .pipeThrough(createSSEProtocolTransformer(transformVertexAIStream, streamStack))

package/src/locales/default/components.ts

@@ -79,6 +79,7 @@ export default {
       file: '该模型支持上传文件读取与识别',
       functionCall: '该模型支持函数调用（Function Call）',
       reasoning: '该模型支持深度思考',
+      search: '该模型支持联网搜索',
       tokens: '该模型单个会话最多支持 {{tokens}} Tokens',
       vision: '该模型支持视觉识别',
     },

package/src/server/globalConfig/index.test.ts

@@ -0,0 +1,81 @@
+import { describe, expect, it, vi } from 'vitest';
+
+import { getAppConfig } from '@/config/app';
+import { knowledgeEnv } from '@/config/knowledge';
+import { SystemEmbeddingConfig } from '@/types/knowledgeBase';
+import { FilesConfigItem } from '@/types/user/settings/filesConfig';
+
+import { getServerDefaultAgentConfig, getServerDefaultFilesConfig } from './index';
+import { parseAgentConfig } from './parseDefaultAgent';
+import { parseFilesConfig } from './parseFilesConfig';
+
+vi.mock('@/config/app', () => ({
+  getAppConfig: vi.fn(),
+}));
+
+vi.mock('@/config/knowledge', () => ({
+  knowledgeEnv: {
+    DEFAULT_FILES_CONFIG: 'test_config',
+  },
+}));
+
+vi.mock('./parseDefaultAgent', () => ({
+  parseAgentConfig: vi.fn(),
+}));
+
+vi.mock('./parseFilesConfig', () => ({
+  parseFilesConfig: vi.fn(),
+}));
+
+describe('getServerDefaultAgentConfig', () => {
+  it('should return parsed agent config', () => {
+    const mockConfig = { key: 'value' };
+    vi.mocked(getAppConfig).mockReturnValue({
+      DEFAULT_AGENT_CONFIG: 'test_agent_config',
+    } as any);
+    vi.mocked(parseAgentConfig).mockReturnValue(mockConfig);
+
+    const result = getServerDefaultAgentConfig();
+
+    expect(parseAgentConfig).toHaveBeenCalledWith('test_agent_config');
+    expect(result).toEqual(mockConfig);
+  });
+
+  it('should return empty object if parseAgentConfig returns undefined', () => {
+    vi.mocked(getAppConfig).mockReturnValue({
+      DEFAULT_AGENT_CONFIG: 'test_agent_config',
+    } as any);
+    vi.mocked(parseAgentConfig).mockReturnValue(undefined);
+
+    const result = getServerDefaultAgentConfig();
+
+    expect(result).toEqual({});
+  });
+});
+
+describe('getServerDefaultFilesConfig', () => {
+  it('should return parsed files config', () => {
+    const mockEmbeddingModel: FilesConfigItem = {
+      model: 'test-model',
+      provider: 'test-provider',
+    };
+
+    const mockRerankerModel: FilesConfigItem = {
+      model: 'test-reranker',
+      provider: 'test-provider',
+    };
+
+    const mockConfig: SystemEmbeddingConfig = {
+      embeddingModel: mockEmbeddingModel,
+      queryMode: 'hybrid',
+      rerankerModel: mockRerankerModel,
+    };
+
+    vi.mocked(parseFilesConfig).mockReturnValue(mockConfig);
+
+    const result = getServerDefaultFilesConfig();
+
+    expect(parseFilesConfig).toHaveBeenCalledWith('test_config');
+    expect(result).toEqual(mockConfig);
+  });
+});

package/src/server/routers/lambda/user.test.ts

@@ -0,0 +1,305 @@
+// @vitest-environment node
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { enableClerk } from '@/const/auth';
+import { serverDB } from '@/database/server';
+import { MessageModel } from '@/database/server/models/message';
+import { SessionModel } from '@/database/server/models/session';
+import { UserModel, UserNotFoundError } from '@/database/server/models/user';
+import { LobeNextAuthDbAdapter } from '@/libs/next-auth/adapter';
+import { KeyVaultsGateKeeper } from '@/server/modules/KeyVaultsEncrypt';
+import { UserService } from '@/server/services/user';
+
+import { userRouter } from './user';
+
+// Mock modules
+vi.mock('@clerk/nextjs/server', () => ({
+  currentUser: vi.fn(),
+}));
+
+vi.mock('@/database/server', () => ({
+  serverDB: {},
+}));
+
+vi.mock('@/database/server/models/message');
+vi.mock('@/database/server/models/session');
+vi.mock('@/database/server/models/user');
+vi.mock('@/libs/next-auth/adapter');
+vi.mock('@/server/modules/KeyVaultsEncrypt');
+vi.mock('@/server/services/user');
+vi.mock('@/const/auth', () => ({
+  enableClerk: true,
+}));
+
+describe('userRouter', () => {
+  const mockUserId = 'test-user-id';
+  const mockCtx = {
+    userId: mockUserId,
+  };
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  describe('getUserRegistrationDuration', () => {
+    it('should return registration duration', async () => {
+      const mockDuration = { duration: 100, createdAt: '2023-01-01', updatedAt: '2023-01-02' };
+      vi.mocked(UserModel).mockImplementation(
+        () =>
+          ({
+            getUserRegistrationDuration: vi.fn().mockResolvedValue(mockDuration),
+          }) as any,
+      );
+
+      const result = await userRouter.createCaller({ ...mockCtx }).getUserRegistrationDuration();
+
+      expect(result).toEqual(mockDuration);
+      expect(UserModel).toHaveBeenCalledWith(serverDB, mockUserId);
+    });
+  });
+
+  describe('getUserSSOProviders', () => {
+    it('should return SSO providers', async () => {
+      const mockProviders = [
+        {
+          provider: 'google',
+          providerAccountId: '123',
+          userId: 'user-1',
+          type: 'oauth',
+        },
+      ];
+      vi.mocked(UserModel).mockImplementation(
+        () =>
+          ({
+            getUserSSOProviders: vi.fn().mockResolvedValue(mockProviders),
+          }) as any,
+      );
+
+      const result = await userRouter.createCaller({ ...mockCtx }).getUserSSOProviders();
+
+      expect(result).toEqual(mockProviders);
+      expect(UserModel).toHaveBeenCalledWith(serverDB, mockUserId);
+    });
+  });
+
+  describe('getUserState', () => {
+    it('should return user state', async () => {
+      const mockState = {
+        isOnboarded: true,
+        preference: { telemetry: true },
+        settings: {},
+        userId: mockUserId,
+      };
+
+      vi.mocked(UserModel).mockImplementation(
+        () =>
+          ({
+            getUserState: vi.fn().mockResolvedValue(mockState),
+          }) as any,
+      );
+
+      vi.mocked(MessageModel).mockImplementation(
+        () =>
+          ({
+            hasMoreThanN: vi.fn().mockResolvedValue(true),
+          }) as any,
+      );
+
+      vi.mocked(SessionModel).mockImplementation(
+        () =>
+          ({
+            hasMoreThanN: vi.fn().mockResolvedValue(true),
+          }) as any,
+      );
+
+      const result = await userRouter.createCaller({ ...mockCtx }).getUserState();
+
+      expect(result).toEqual({
+        isOnboard: true,
+        preference: { telemetry: true },
+        settings: {},
+        hasConversation: true,
+        canEnablePWAGuide: true,
+        canEnableTrace: true,
+        userId: mockUserId,
+      });
+    });
+
+    it('should create new user when user not found (clerk enabled)', async () => {
+      const mockClerkUser = {
+        id: mockUserId,
+        createdAt: new Date(),
+        emailAddresses: [{ id: 'email-1', emailAddress: 'test@example.com' }],
+        firstName: 'Test',
+        lastName: 'User',
+        imageUrl: 'avatar.jpg',
+        phoneNumbers: [],
+        primaryEmailAddressId: 'email-1',
+        primaryPhoneNumberId: null,
+        username: 'testuser',
+      };
+
+      const { currentUser } = await import('@clerk/nextjs/server');
+      vi.mocked(currentUser).mockResolvedValue(mockClerkUser as any);
+
+      vi.mocked(UserService).mockImplementation(
+        () =>
+          ({
+            createUser: vi.fn().mockResolvedValue({ success: true }),
+          }) as any,
+      );
+
+      vi.mocked(UserModel).mockImplementation(
+        () =>
+          ({
+            getUserState: vi
+              .fn()
+              .mockRejectedValueOnce(new UserNotFoundError())
+              .mockResolvedValueOnce({
+                isOnboarded: false,
+                preference: { telemetry: null },
+                settings: {},
+              }),
+          }) as any,
+      );
+
+      vi.mocked(MessageModel).mockImplementation(
+        () =>
+          ({
+            hasMoreThanN: vi.fn().mockResolvedValue(false),
+          }) as any,
+      );
+
+      vi.mocked(SessionModel).mockImplementation(
+        () =>
+          ({
+            hasMoreThanN: vi.fn().mockResolvedValue(false),
+          }) as any,
+      );
+
+      const result = await userRouter.createCaller({ ...mockCtx } as any).getUserState();
+
+      expect(result).toEqual({
+        isOnboard: true,
+        preference: { telemetry: null },
+        settings: {},
+        hasConversation: false,
+        canEnablePWAGuide: false,
+        canEnableTrace: false,
+        userId: mockUserId,
+      });
+    });
+  });
+
+  describe('makeUserOnboarded', () => {
+    it('should update user onboarded status', async () => {
+      vi.mocked(UserModel).mockImplementation(
+        () =>
+          ({
+            updateUser: vi.fn().mockResolvedValue({ rowCount: 1 }),
+          }) as any,
+      );
+
+      await userRouter.createCaller({ ...mockCtx }).makeUserOnboarded();
+
+      expect(UserModel).toHaveBeenCalledWith(serverDB, mockUserId);
+    });
+  });
+
+  describe('unlinkSSOProvider', () => {
+    it('should unlink SSO provider successfully', async () => {
+      const mockInput = {
+        provider: 'google',
+        providerAccountId: '123',
+      };
+
+      const mockAccount = {
+        userId: mockUserId,
+        provider: 'google',
+        providerAccountId: '123',
+        type: 'oauth',
+      };
+
+      vi.mocked(LobeNextAuthDbAdapter).mockReturnValue({
+        getAccount: vi.fn().mockResolvedValue(mockAccount),
+        unlinkAccount: vi.fn().mockResolvedValue(undefined),
+      } as any);
+
+      await expect(
+        userRouter.createCaller({ ...mockCtx }).unlinkSSOProvider(mockInput),
+      ).resolves.not.toThrow();
+    });
+
+    it('should throw error if account does not exist', async () => {
+      const mockInput = {
+        provider: 'google',
+        providerAccountId: '123',
+      };
+
+      vi.mocked(LobeNextAuthDbAdapter).mockReturnValue({
+        getAccount: vi.fn().mockResolvedValue(null),
+        unlinkAccount: vi.fn(),
+      } as any);
+
+      await expect(
+        userRouter.createCaller({ ...mockCtx }).unlinkSSOProvider(mockInput),
+      ).rejects.toThrow('The account does not exist');
+    });
+
+    it('should throw error if adapter methods are not implemented', async () => {
+      const mockInput = {
+        provider: 'google',
+        providerAccountId: '123',
+      };
+
+      vi.mocked(LobeNextAuthDbAdapter).mockReturnValue({} as any);
+
+      await expect(
+        userRouter.createCaller({ ...mockCtx }).unlinkSSOProvider(mockInput),
+      ).rejects.toThrow('The method in LobeNextAuthDbAdapter `unlinkAccount` is not implemented');
+    });
+  });
+
+  describe('updateSettings', () => {
+    it('should update settings with encrypted key vaults', async () => {
+      const mockSettings = {
+        keyVaults: { openai: { key: 'test-key' } },
+        general: { language: 'en-US' },
+      };
+
+      const mockEncryptedVaults = 'encrypted-data';
+      const mockGateKeeper = {
+        encrypt: vi.fn().mockResolvedValue(mockEncryptedVaults),
+      };
+
+      vi.mocked(KeyVaultsGateKeeper.initWithEnvKey).mockResolvedValue(mockGateKeeper as any);
+      vi.mocked(UserModel).mockImplementation(
+        () =>
+          ({
+            updateSetting: vi.fn().mockResolvedValue({ rowCount: 1 }),
+          }) as any,
+      );
+
+      await userRouter.createCaller({ ...mockCtx }).updateSettings(mockSettings);
+
+      expect(mockGateKeeper.encrypt).toHaveBeenCalledWith(JSON.stringify(mockSettings.keyVaults));
+    });
+
+    it('should update settings without key vaults', async () => {
+      const mockSettings = {
+        general: { language: 'en-US' },
+      };
+
+      vi.mocked(UserModel).mockImplementation(
+        () =>
+          ({
+            updateSetting: vi.fn().mockResolvedValue({ rowCount: 1 }),
+          }) as any,
+      );
+
+      await userRouter.createCaller({ ...mockCtx }).updateSettings(mockSettings);
+
+      expect(UserModel).toHaveBeenCalledWith(serverDB, mockUserId);
+    });
+  });
+});

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

@@ -17,7 +17,7 @@ export class NextAuthUserService {
     { providerAccountId, provider }: { provider: string; providerAccountId: string },
     data: Partial<UserItem>,
   ) => {
-    pino.info('updating user due to webhook');
+    pino.info(`updating user "${JSON.stringify({ provider, providerAccountId })}" due to webhook`);
     // 1. Find User by account
     // @ts-expect-error: Already impl in `LobeNextauthDbAdapter`
     const user = await this.adapter.getUserByAccount({
@@ -37,7 +37,7 @@ export class NextAuthUserService {
       });
     } else {
       pino.warn(
-        `[${provider}]: Webhooks handler user update for "${JSON.stringify(data)}", but no user was found by the providerAccountId.`,
+        `[${provider}]: Webhooks handler user "${JSON.stringify({ provider, providerAccountId })}" update for "${JSON.stringify(data)}", but no user was found by the providerAccountId.`,
       );
     }
     return NextResponse.json({ message: 'user updated', success: true }, { status: 200 });

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

@@ -455,7 +455,7 @@ export const generateAIChat: StateCreator<
         await messageService.updateMessageError(messageId, error);
         await refreshMessages();
       },
-      onFinish: async (content, { traceId, observationId, toolCalls, reasoning }) => {
+      onFinish: async (content, { traceId, observationId, toolCalls, reasoning, citations }) => {
         // if there is traceId, update it
         if (traceId) {
           msgTraceId = traceId;
@@ -470,15 +470,26 @@ export const generateAIChat: StateCreator<
         }
 
         // update the content after fetch result
-        await internal_updateMessageContent(
-          messageId,
-          content,
+        await internal_updateMessageContent(messageId, content, {
           toolCalls,
-          !!reasoning ? { content: reasoning, duration } : undefined,
-        );
+          reasoning: !!reasoning ? { content: reasoning, duration } : undefined,
+          search: !!citations ? { citations } : undefined,
+        });
       },
       onMessageHandle: async (chunk) => {
         switch (chunk.type) {
+          case 'citations': {
+            // if there is no citations, then stop
+            if (!chunk.citations || chunk.citations.length <= 0) return;
+
+            internal_dispatchMessage({
+              id: messageId,
+              type: 'updateMessage',
+              value: { search: { citations: chunk.citations } },
+            });
+            break;
+          }
+
           case 'text': {
             output += chunk.text;
 

package/src/store/chat/slices/message/action.ts

@@ -16,6 +16,7 @@ import {
   ChatMessage,
   ChatMessageError,
   CreateMessageParams,
+  GroundingSearch,
   MessageToolCall,
   ModelReasoning,
 } from '@/types/message';
@@ -73,8 +74,11 @@ export interface ChatMessageAction {
   internal_updateMessageContent: (
     id: string,
     content: string,
-    toolCalls?: MessageToolCall[],
-    reasoning?: ModelReasoning,
+    extra?: {
+      toolCalls?: MessageToolCall[];
+      reasoning?: ModelReasoning;
+      search?: GroundingSearch;
+    },
   ) => Promise<void>;
   /**
    * update the message error with optimistic update
@@ -272,17 +276,17 @@ export const chatMessage: StateCreator<
     await messageService.updateMessage(id, { error });
     await get().refreshMessages();
   },
-  internal_updateMessageContent: async (id, content, toolCalls, reasoning) => {
+  internal_updateMessageContent: async (id, content, extra) => {
     const { internal_dispatchMessage, refreshMessages, internal_transformToolCalls } = get();
 
     // Due to the async update method and refresh need about 100ms
     // we need to update the message content at the frontend to avoid the update flick
     // refs: https://medium.com/@kyledeguzmanx/what-are-optimistic-updates-483662c3e171
-    if (toolCalls) {
+    if (extra?.toolCalls) {
       internal_dispatchMessage({
         id,
         type: 'updateMessage',
-        value: { tools: internal_transformToolCalls(toolCalls) },
+        value: { tools: internal_transformToolCalls(extra?.toolCalls) },
       });
     } else {
       internal_dispatchMessage({ id, type: 'updateMessage', value: { content } });
@@ -290,8 +294,9 @@ export const chatMessage: StateCreator<
 
     await messageService.updateMessage(id, {
       content,
-      tools: toolCalls ? internal_transformToolCalls(toolCalls) : undefined,
-      reasoning,
+      tools: extra?.toolCalls ? internal_transformToolCalls(extra?.toolCalls) : undefined,
+      reasoning: extra?.reasoning,
+      search: extra?.search,
     });
     await refreshMessages();
   },

package/src/types/aiModel.ts

@@ -34,6 +34,11 @@ export interface ModelAbilities {
    * whether model supports reasoning
    */
   reasoning?: boolean;
+  /**
+   * whether model supports search web
+   */
+  search?: boolean;
+
   /**
    *  whether model supports vision
    */

package/src/types/message/base.ts

@@ -1,3 +1,15 @@
+export interface CitationItem {
+  id?: string;
+  onlyUrl?: boolean;
+  title?: string;
+  url: string;
+}
+
+export interface GroundingSearch {
+  citations?: CitationItem[];
+  searchQueries?: string[];
+}
+
 export interface ModelReasoning {
   content?: string;
   duration?: number;
@@ -20,6 +32,7 @@ export interface MessageItem {
   quotaId: string | null;
   reasoning: ModelReasoning | null;
   role: string;
+  search: GroundingSearch | null;
   sessionId: string | null;
   threadId: string | null;
   // jsonb type

package/src/types/message/chat.ts

@@ -2,7 +2,7 @@ import { IPluginErrorType } from '@lobehub/chat-plugin-sdk';
 
 import { ILobeAgentRuntimeErrorType } from '@/libs/agent-runtime';
 import { ErrorType } from '@/types/fetch';
-import { MessageRoleType, ModelReasoning } from '@/types/message/base';
+import { GroundingSearch, MessageRoleType, ModelReasoning } from '@/types/message/base';
 import { ChatPluginPayload, ChatToolPayload } from '@/types/message/tools';
 import { Translate } from '@/types/message/translate';
 import { MetaData } from '@/types/meta';
@@ -100,11 +100,12 @@ export interface ChatMessage {
   ragRawQuery?: string | null;
 
   reasoning?: ModelReasoning | null;
-
   /**
    * message role type
    */
   role: MessageRoleType;
+
+  search?: GroundingSearch | null;
   sessionId?: string;
   threadId?: string | null;
   tool_call_id?: string;