@lobehub/chat 1.61.6 → 1.62.1

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/services/nextAuthUser/index.test.ts

@@ -0,0 +1,109 @@
+// @vitest-environment node
+import { NextResponse } from 'next/server';
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { UserItem } from '@/database/schemas';
+import { serverDB } from '@/database/server';
+import { UserModel } from '@/database/server/models/user';
+import { pino } from '@/libs/logger';
+import { LobeNextAuthDbAdapter } from '@/libs/next-auth/adapter';
+
+import { NextAuthUserService } from './index';
+
+vi.mock('@/libs/logger', () => ({
+  pino: {
+    info: vi.fn(),
+    warn: vi.fn(),
+  },
+}));
+
+vi.mock('@/database/server/models/user');
+vi.mock('@/database/server');
+
+describe('NextAuthUserService', () => {
+  let service: NextAuthUserService;
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+    service = new NextAuthUserService();
+  });
+
+  describe('safeUpdateUser', () => {
+    const mockUser = {
+      id: 'user-123',
+      email: 'test@example.com',
+    };
+
+    const mockAccount = {
+      provider: 'github',
+      providerAccountId: '12345',
+    };
+
+    const mockUpdateData: Partial<UserItem> = {
+      avatar: 'https://example.com/avatar.jpg',
+      email: 'new@example.com',
+      fullName: 'Test User',
+    };
+
+    it('should update user when user is found', async () => {
+      const mockUserModel = {
+        updateUser: vi.fn().mockResolvedValue({}),
+      };
+
+      vi.mocked(UserModel).mockImplementation(() => mockUserModel as any);
+
+      // Mock the adapter directly on the service instance
+      service.adapter = {
+        getUserByAccount: vi.fn().mockResolvedValue(mockUser),
+      };
+
+      const response = await service.safeUpdateUser(mockAccount, mockUpdateData);
+
+      expect(pino.info).toHaveBeenCalledWith(
+        `updating user "${JSON.stringify(mockAccount)}" due to webhook`,
+      );
+
+      expect(service.adapter.getUserByAccount).toHaveBeenCalledWith(mockAccount);
+      expect(UserModel).toHaveBeenCalledWith(serverDB, mockUser.id);
+      expect(mockUserModel.updateUser).toHaveBeenCalledWith(mockUpdateData);
+
+      expect(response).toBeInstanceOf(NextResponse);
+      expect(response.status).toBe(200);
+      const data = await response.json();
+      expect(data).toEqual({ message: 'user updated', success: true });
+    });
+
+    it('should handle case when user is not found', async () => {
+      // Mock the adapter directly on the service instance
+      service.adapter = {
+        getUserByAccount: vi.fn().mockResolvedValue(null),
+      };
+
+      const response = await service.safeUpdateUser(mockAccount, mockUpdateData);
+
+      expect(pino.warn).toHaveBeenCalledWith(
+        `[${mockAccount.provider}]: Webhooks handler user "${JSON.stringify(mockAccount)}" update for "${JSON.stringify(mockUpdateData)}", but no user was found by the providerAccountId.`,
+      );
+
+      expect(UserModel).not.toHaveBeenCalled();
+
+      expect(response).toBeInstanceOf(NextResponse);
+      expect(response.status).toBe(200);
+      const data = await response.json();
+      expect(data).toEqual({ message: 'user updated', success: true });
+    });
+
+    it('should handle errors during user update', async () => {
+      const mockError = new Error('Database error');
+
+      // Mock the adapter directly on the service instance
+      service.adapter = {
+        getUserByAccount: vi.fn().mockRejectedValue(mockError),
+      };
+
+      await expect(service.safeUpdateUser(mockAccount, mockUpdateData)).rejects.toThrow(mockError);
+
+      expect(UserModel).not.toHaveBeenCalled();
+    });
+  });
+});

package/src/services/user/client.test.ts

@@ -95,4 +95,14 @@ describe('ClientService', () => {
     expect(spyOn).toHaveBeenCalledWith(newPreference);
     expect(spyOn).toHaveBeenCalledTimes(1);
   });
+
+  it('should return empty array for getUserSSOProviders', async () => {
+    const providers = await clientService.getUserSSOProviders();
+    expect(providers).toEqual([]);
+  });
+
+  it('should do nothing when unlinkSSOProvider is called', async () => {
+    const result = await clientService.unlinkSSOProvider('google', '123');
+    expect(result).toBeUndefined();
+  });
 });

package/src/services/user/server.test.ts

@@ -0,0 +1,149 @@
+import { DeepPartial } from 'utility-types';
+import { describe, expect, it, vi } from 'vitest';
+
+import { lambdaClient } from '@/libs/trpc/client';
+import { UserInitializationState, UserPreference } from '@/types/user';
+import { UserSettings } from '@/types/user/settings';
+
+import { ServerService } from './server';
+
+vi.mock('@/libs/trpc/client', () => ({
+  lambdaClient: {
+    user: {
+      getUserRegistrationDuration: {
+        query: vi.fn(),
+      },
+      getUserState: {
+        query: vi.fn(),
+      },
+      getUserSSOProviders: {
+        query: vi.fn(),
+      },
+      unlinkSSOProvider: {
+        mutate: vi.fn(),
+      },
+      makeUserOnboarded: {
+        mutate: vi.fn(),
+      },
+      updatePreference: {
+        mutate: vi.fn(),
+      },
+      updateGuide: {
+        mutate: vi.fn(),
+      },
+      updateSettings: {
+        mutate: vi.fn(),
+      },
+      resetSettings: {
+        mutate: vi.fn(),
+      },
+    },
+  },
+}));
+
+describe('ServerService', () => {
+  const service = new ServerService();
+
+  it('should get user registration duration', async () => {
+    const mockData = {
+      createdAt: '2023-01-01',
+      duration: 100,
+      updatedAt: '2023-01-02',
+    };
+    vi.mocked(lambdaClient.user.getUserRegistrationDuration.query).mockResolvedValue(mockData);
+
+    const result = await service.getUserRegistrationDuration();
+    expect(result).toEqual(mockData);
+  });
+
+  it('should get user state', async () => {
+    const mockState: UserInitializationState = {
+      isOnboard: true,
+      preference: {
+        telemetry: true,
+      },
+      settings: {},
+    };
+    vi.mocked(lambdaClient.user.getUserState.query).mockResolvedValue(mockState);
+
+    const result = await service.getUserState();
+    expect(result).toEqual(mockState);
+  });
+
+  it('should get user SSO providers', async () => {
+    const mockProviders = [
+      {
+        provider: 'google',
+        providerAccountId: '123',
+        userId: 'user1',
+        type: 'oauth' as const,
+        access_token: 'token',
+        token_type: 'bearer' as const,
+        expires_at: 123,
+        scope: 'email profile',
+      },
+    ];
+    vi.mocked(lambdaClient.user.getUserSSOProviders.query).mockResolvedValue(mockProviders);
+
+    const result = await service.getUserSSOProviders();
+    expect(result).toEqual(mockProviders);
+  });
+
+  it('should unlink SSO provider', async () => {
+    const provider = 'google';
+    const providerAccountId = '123';
+    await service.unlinkSSOProvider(provider, providerAccountId);
+
+    expect(lambdaClient.user.unlinkSSOProvider.mutate).toHaveBeenCalledWith({
+      provider,
+      providerAccountId,
+    });
+  });
+
+  it('should make user onboarded', async () => {
+    await service.makeUserOnboarded();
+    expect(lambdaClient.user.makeUserOnboarded.mutate).toHaveBeenCalled();
+  });
+
+  it('should update user preference', async () => {
+    const preference: Partial<UserPreference> = {
+      telemetry: true,
+      useCmdEnterToSend: true,
+    };
+    await service.updatePreference(preference);
+    expect(lambdaClient.user.updatePreference.mutate).toHaveBeenCalledWith(preference);
+  });
+
+  it('should update user guide', async () => {
+    const guide = {
+      moveSettingsToAvatar: true,
+      topic: false,
+      uploadFileInKnowledgeBase: true,
+    };
+    await service.updateGuide(guide);
+    expect(lambdaClient.user.updateGuide.mutate).toHaveBeenCalledWith(guide);
+  });
+
+  it('should update user settings', async () => {
+    const settings: DeepPartial<UserSettings> = {
+      defaultAgent: {
+        config: {
+          model: 'gpt-4',
+          provider: 'openai',
+        },
+        meta: {
+          avatar: 'avatar',
+          description: 'test agent',
+        },
+      },
+    };
+    const signal = new AbortController().signal;
+    await service.updateUserSettings(settings, signal);
+    expect(lambdaClient.user.updateSettings.mutate).toHaveBeenCalledWith(settings, { signal });
+  });
+
+  it('should reset user settings', async () => {
+    await service.resetUserSettings();
+    expect(lambdaClient.user.resetSettings.mutate).toHaveBeenCalled();
+  });
+});

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;

package/src/utils/fetch/fetchSSE.ts

@@ -6,6 +6,7 @@ import { ChatErrorType } from '@/types/fetch';
 import { SmoothingParams } from '@/types/llm';
 import {
   ChatMessageError,
+  CitationItem,
   MessageToolCall,
   MessageToolCallChunk,
   MessageToolCallSchema,
@@ -20,6 +21,7 @@ type SSEFinishType = 'done' | 'error' | 'abort';
 export type OnFinishHandler = (
   text: string,
   context: {
+    citations?: CitationItem[];
     observationId?: string | null;
     reasoning?: string;
     toolCalls?: MessageToolCall[];
@@ -38,6 +40,11 @@ export interface MessageReasoningChunk {
   type: 'reasoning';
 }
 
+export interface MessageCitationsChunk {
+  citations: CitationItem[];
+  type: 'citations';
+}
+
 interface MessageToolCallsChunk {
   isAnimationActives?: boolean[];
   tool_calls: MessageToolCall[];
@@ -50,7 +57,7 @@ export interface FetchSSEOptions {
   onErrorHandle?: (error: ChatMessageError) => void;
   onFinish?: OnFinishHandler;
   onMessageHandle?: (
-    chunk: MessageTextChunk | MessageToolCallsChunk | MessageReasoningChunk,
+    chunk: MessageTextChunk | MessageToolCallsChunk | MessageReasoningChunk | MessageCitationsChunk,
   ) => void;
   smoothing?: SmoothingParams | boolean;
 }
@@ -279,6 +286,7 @@ export const fetchSSE = async (url: string, options: RequestInit & FetchSSEOptio
     startSpeed: smoothingSpeed,
   });
 
+  let citations: CitationItem[] | undefined = undefined;
   await fetchEventSource(url, {
     body: options.body,
     fetch: options?.fetcher,
@@ -350,6 +358,13 @@ export const fetchSSE = async (url: string, options: RequestInit & FetchSSEOptio
 
           break;
         }
+
+        case 'citations': {
+          citations = data;
+          options.onMessageHandle?.({ citations: data, type: 'citations' });
+          break;
+        }
+
         case 'reasoning': {
           if (textSmoothing) {
             thinkingController.pushToQueue(data);
@@ -419,6 +434,7 @@ export const fetchSSE = async (url: string, options: RequestInit & FetchSSEOptio
       }
 
       await options?.onFinish?.(output, {
+        citations,
         observationId,
         reasoning: !!thinking ? thinking : undefined,
         toolCalls,