@lobehub/chat 0.149.3 → 0.149.5

package/src/libs/agent-runtime/anthropic/index.ts

@@ -32,15 +32,18 @@ export class LobeAnthropicAI implements LobeRuntimeAI {
     const user_messages = messages.filter((m) => m.role !== 'system');
 
     try {
-      const response = await this.client.messages.create({
-        max_tokens: max_tokens || 4096,
-        messages: buildAnthropicMessages(user_messages),
-        model: model,
-        stream: true,
-        system: system_message?.content as string,
-        temperature: temperature,
-        top_p: top_p,
-      });
+      const response = await this.client.messages.create(
+        {
+          max_tokens: max_tokens || 4096,
+          messages: buildAnthropicMessages(user_messages),
+          model: model,
+          stream: true,
+          system: system_message?.content as string,
+          temperature: temperature,
+          top_p: top_p,
+        },
+        { signal: options?.signal },
+      );
 
       const [prod, debug] = response.tee();
 

package/src/libs/agent-runtime/azureOpenai/index.ts

@@ -8,7 +8,7 @@ import { OpenAIStream, StreamingTextResponse } from 'ai';
 
 import { LobeRuntimeAI } from '../BaseAI';
 import { AgentRuntimeErrorType } from '../error';
-import { ChatStreamPayload, ModelProvider } from '../types';
+import { ChatCompetitionOptions, ChatStreamPayload, ModelProvider } from '../types';
 import { AgentRuntimeError } from '../utils/createError';
 import { debugStream } from '../utils/debugStream';
 
@@ -26,7 +26,7 @@ export class LobeAzureOpenAI implements LobeRuntimeAI {
 
   baseURL: string;
 
-  async chat(payload: ChatStreamPayload) {
+  async chat(payload: ChatStreamPayload, options?: ChatCompetitionOptions) {
     // ============  1. preprocess messages   ============ //
     const { messages, model, ...params } = payload;
 
@@ -36,10 +36,9 @@ export class LobeAzureOpenAI implements LobeRuntimeAI {
       const response = await this.client.streamChatCompletions(
         model,
         messages as ChatRequestMessage[],
-        params as GetChatCompletionsOptions,
+        { ...params, abortSignal: options?.signal } as GetChatCompletionsOptions,
       );
 
-      // TODO: we need to refactor this part in the future
       const stream = OpenAIStream(response as any);
 
       const [debug, prod] = stream.tee();

package/src/libs/agent-runtime/bedrock/index.ts

@@ -68,7 +68,7 @@ export class LobeBedrockAI implements LobeRuntimeAI {
 
     try {
       // Ask Claude for a streaming chat completion given the prompt
-      const bedrockResponse = await this.client.send(command);
+      const bedrockResponse = await this.client.send(command, { abortSignal: options?.signal });
 
       // Convert the response into a friendly text-stream
       const stream = AWSBedrockStream(

package/src/libs/agent-runtime/ollama/index.ts

@@ -31,6 +31,13 @@ export class LobeOllamaAI implements LobeRuntimeAI {
 
   async chat(payload: ChatStreamPayload, options?: ChatCompetitionOptions) {
     try {
+      const abort = () => {
+        this.client.abort();
+        options?.signal?.removeEventListener('abort', abort);
+      };
+
+      options?.signal?.addEventListener('abort', abort);
+
       const response = await this.client.chat({
         messages: this.buildOllamaMessages(payload.messages),
         model: payload.model,

package/src/libs/agent-runtime/perplexity/index.ts

@@ -33,6 +33,7 @@ export class LobePerplexityAI implements LobeRuntimeAI {
       };
       const response = await this.client.chat.completions.create(
         chatPayload as unknown as OpenAI.ChatCompletionCreateParamsStreaming,
+        { signal: options?.signal },
       );
       const [prod, debug] = response.tee();
 

package/src/libs/agent-runtime/types/chat.ts

@@ -90,8 +90,9 @@ export interface ChatStreamPayload {
 }
 
 export interface ChatCompetitionOptions {
-  callback: ChatStreamCallbacks;
+  callback?: ChatStreamCallbacks;
   headers?: Record<string, any>;
+  signal?: AbortSignal;
 }
 
 export interface ChatCompletionFunctions {

package/src/libs/agent-runtime/utils/openaiCompatibleFactory/index.ts

@@ -76,6 +76,7 @@ export const LobeOpenAICompatibleFactory = ({
         const response = await this.client.chat.completions.create(postPayload, {
           // https://github.com/lobehub/lobe-chat/pull/318
           headers: { Accept: '*/*' },
+          signal: options?.signal,
         });
 
         const [prod, useForDebug] = response.tee();

package/src/services/chat.ts

@@ -15,7 +15,7 @@ import {
   preferenceSelectors,
 } from '@/store/global/selectors';
 import { useSessionStore } from '@/store/session';
-import { agentSelectors } from '@/store/session/selectors';
+import { sessionMetaSelectors } from '@/store/session/selectors';
 import { useToolStore } from '@/store/tool';
 import { pluginSelectors, toolSelectors } from '@/store/tool/selectors';
 import { ChatErrorType } from '@/types/fetch';
@@ -168,18 +168,6 @@ export function initializeWithClientStore(provider: string, payload: any) {
   });
 }
 
-/**
- * Fetch chat completion on the client side.
- * @param provider - The provider name.
- * @param payload - The payload data for the chat stream.
- * @returns A promise that resolves to the chat response.
- */
-export async function fetchOnClient(provider: string, payload: Partial<ChatStreamPayload>) {
-  const agentRuntime = await initializeWithClientStore(provider, payload);
-  const data = payload as ChatStreamPayload;
-  return await agentRuntime.chat(data);
-}
-
 class ChatService {
   createAssistantMessage = async (
     { plugins: enabledPlugins, messages, ...params }: GetChatCompletionPayload,
@@ -279,7 +267,7 @@ class ChatService {
      */
     if (enableFetchOnClient) {
       try {
-        return await fetchOnClient(provider, payload);
+        return await this.fetchOnClient({ payload, provider, signal });
       } catch (e) {
         const {
           errorType = ChatErrorType.BadRequest,
@@ -459,7 +447,7 @@ class ChatService {
   };
 
   private mapTrace(trace?: TracePayload, tag?: TraceTagMap): TracePayload {
-    const tags = agentSelectors.currentAgentMeta(useSessionStore.getState()).tags || [];
+    const tags = sessionMetaSelectors.currentAgentMeta(useSessionStore.getState()).tags || [];
 
     const enabled = preferenceSelectors.userAllowTrace(useGlobalStore.getState());
 
@@ -472,6 +460,21 @@ class ChatService {
       userId: commonSelectors.userId(useGlobalStore.getState()),
     };
   }
+
+  /**
+   * Fetch chat completion on the client side.
+
+   */
+  private fetchOnClient = async (params: {
+    payload: Partial<ChatStreamPayload>;
+    provider: string;
+    signal?: AbortSignal;
+  }) => {
+    const agentRuntime = await initializeWithClientStore(params.provider, params.payload);
+    const data = params.payload as ChatStreamPayload;
+
+    return agentRuntime.chat(data, { signal: params.signal });
+  };
 }
 
 export const chatService = new ChatService();

package/src/services/session/client.ts

@@ -1,7 +1,10 @@
 import { DeepPartial } from 'utility-types';
 
+import { INBOX_SESSION_ID } from '@/const/session';
 import { SessionModel } from '@/database/client/models/session';
 import { SessionGroupModel } from '@/database/client/models/sessionGroup';
+import { UserModel } from '@/database/client/models/user';
+import { useGlobalStore } from '@/store/global';
 import { LobeAgentConfig } from '@/types/agent';
 import {
   ChatSessionList,
@@ -40,6 +43,18 @@ export class ClientService implements ISessionService {
     return SessionModel.queryWithGroups();
   }
 
+  async getSessionConfig(id: string): Promise<LobeAgentConfig> {
+    if (!id || id === INBOX_SESSION_ID) {
+      return UserModel.getAgentConfig();
+    }
+
+    const res = await SessionModel.findById(id);
+
+    if (!res) throw new Error('Session not found');
+
+    return res.config as LobeAgentConfig;
+  }
+
   async getSessionsByType(type: 'agent' | 'group' | 'all' = 'all'): Promise<LobeSessions> {
     switch (type) {
       // TODO: add a filter to get only agents or agents
@@ -81,6 +96,10 @@ export class ClientService implements ISessionService {
   }
 
   async updateSessionConfig(activeId: string, config: DeepPartial<LobeAgentConfig>) {
+    if (activeId === INBOX_SESSION_ID) {
+      return useGlobalStore.getState().updateDefaultAgent({ config });
+    }
+
     return SessionModel.updateConfig(activeId, config);
   }
 

package/src/services/session/type.ts

@@ -28,6 +28,8 @@ export interface ISessionService {
     id: string,
     data: Partial<{ group?: SessionGroupId; pinned?: boolean }>,
   ): Promise<any>;
+
+  getSessionConfig(id: string): Promise<LobeAgentConfig>;
   updateSessionConfig(id: string, config: DeepPartial<LobeAgentConfig>): Promise<any>;
 
   removeSession(id: string): Promise<any>;

package/src/store/agent/index.ts

@@ -0,0 +1,2 @@
+export type { AgentStore } from './store';
+export { useAgentStore } from './store';

package/src/store/agent/initialState.ts

@@ -0,0 +1,7 @@
+import { AgentState, initialSessionState } from './slices/chat/initialState';
+
+export type SessionStoreState = AgentState;
+
+export const initialState: SessionStoreState = {
+  ...initialSessionState,
+};

package/src/store/agent/selectors.ts

@@ -0,0 +1 @@
+export { agentSelectors } from './slices/chat/selectors';

package/src/store/{session/slices/agent → agent/slices/chat}/action.test.ts

@@ -3,14 +3,14 @@ import * as immer from 'immer';
 import { describe, expect, it, vi } from 'vitest';
 
 import { sessionService } from '@/services/session';
+import { useAgentStore } from '@/store/agent';
+import { agentSelectors } from '@/store/agent/selectors';
 import { useGlobalStore } from '@/store/global';
-import { useSessionStore } from '@/store/session';
-import { agentSelectors, sessionSelectors } from '@/store/session/selectors';
 
 describe('AgentSlice', () => {
   describe('removePlugin', () => {
     it('should call togglePlugin with the provided id and false', async () => {
-      const { result } = renderHook(() => useSessionStore());
+      const { result } = renderHook(() => useAgentStore());
       const pluginId = 'plugin-id';
       const togglePluginMock = vi.spyOn(result.current, 'togglePlugin');
 
@@ -25,7 +25,7 @@ describe('AgentSlice', () => {
 
   describe('togglePlugin', () => {
     it('should add plugin id to plugins array if not present and open is true or undefined', async () => {
-      const { result } = renderHook(() => useSessionStore());
+      const { result } = renderHook(() => useAgentStore());
       const pluginId = 'plugin-id';
       const updateAgentConfigMock = vi.spyOn(result.current, 'updateAgentConfig');
 
@@ -43,7 +43,7 @@ describe('AgentSlice', () => {
     });
 
     it('should remove plugin id from plugins array if present and open is false', async () => {
-      const { result } = renderHook(() => useSessionStore());
+      const { result } = renderHook(() => useAgentStore());
       const pluginId = 'plugin-id';
       const updateAgentConfigMock = vi.spyOn(result.current, 'updateAgentConfig');
 
@@ -61,7 +61,7 @@ describe('AgentSlice', () => {
     });
 
     it('should not modify plugins array if plugin id is not present and open is false', async () => {
-      const { result } = renderHook(() => useSessionStore());
+      const { result } = renderHook(() => useAgentStore());
       const pluginId = 'plugin-id';
       const updateAgentConfigMock = vi.spyOn(result.current, 'updateAgentConfig');
 
@@ -79,49 +79,53 @@ describe('AgentSlice', () => {
 
   describe('updateAgentConfig', () => {
     it('should update global config if current session is inbox session', async () => {
-      const { result } = renderHook(() => useSessionStore());
+      const { result } = renderHook(() => useAgentStore());
       const config = { model: 'gpt-3.5-turbo' };
-      const updateDefaultAgentMock = vi.spyOn(useGlobalStore.getState(), 'updateDefaultAgent');
-
-      // 模拟当前会话是收件箱会话
-      vi.spyOn(sessionSelectors, 'isInboxSession').mockReturnValue(true);
+      const updateSessionConfigMock = vi.spyOn(sessionService, 'updateSessionConfig');
+      const refreshMock = vi.spyOn(result.current, 'internal_refreshAgentConfig');
 
       await act(async () => {
         await result.current.updateAgentConfig(config);
       });
 
-      expect(updateDefaultAgentMock).toHaveBeenCalledWith({ config });
-      updateDefaultAgentMock.mockRestore();
+      expect(updateSessionConfigMock).toHaveBeenCalledWith('inbox', config);
+      expect(refreshMock).toHaveBeenCalled();
+      updateSessionConfigMock.mockRestore();
+      refreshMock.mockRestore();
     });
 
     it('should update session config if current session is not inbox session', async () => {
-      const { result } = renderHook(() => useSessionStore());
+      const { result } = renderHook(() => useAgentStore());
       const config = { model: 'gpt-3.5-turbo' };
       const updateSessionConfigMock = vi.spyOn(sessionService, 'updateSessionConfig');
-      const refreshSessionsMock = vi.spyOn(result.current, 'refreshSessions');
+      const refreshMock = vi.spyOn(result.current, 'internal_refreshAgentConfig');
 
       // 模拟当前会话不是收件箱会话
-      vi.spyOn(sessionSelectors, 'isInboxSession').mockReturnValue(false);
-      vi.spyOn(sessionSelectors, 'currentSession').mockReturnValue({ id: 'session-id' } as any);
-      vi.spyOn(result.current, 'activeId', 'get').mockReturnValue('session-id');
+      act(() => {
+        useAgentStore.setState({
+          activeId: 'session-id',
+        });
+      });
 
       await act(async () => {
         await result.current.updateAgentConfig(config);
       });
 
       expect(updateSessionConfigMock).toHaveBeenCalledWith('session-id', config);
-      expect(refreshSessionsMock).toHaveBeenCalled();
+      expect(refreshMock).toHaveBeenCalled();
       updateSessionConfigMock.mockRestore();
-      refreshSessionsMock.mockRestore();
+      refreshMock.mockRestore();
     });
 
     it('should not update config if there is no current session', async () => {
-      const { result } = renderHook(() => useSessionStore());
+      const { result } = renderHook(() => useAgentStore());
       const config = { model: 'gpt-3.5-turbo' };
       const updateSessionConfigMock = vi.spyOn(sessionService, 'updateSessionConfig');
 
       // 模拟没有当前会话
-      vi.spyOn(sessionSelectors, 'currentSession').mockReturnValue(null as any);
+      act(() => {
+        useAgentStore.setState({ activeId: null as any });
+      });
 
       await act(async () => {
         await result.current.updateAgentConfig(config);
@@ -131,45 +135,4 @@ describe('AgentSlice', () => {
       updateSessionConfigMock.mockRestore();
     });
   });
-
-  describe('updateAgentMeta', () => {
-    it('should not update meta if there is no current session', async () => {
-      const { result } = renderHook(() => useSessionStore());
-      const meta = { title: 'Test Agent' };
-      const updateSessionMock = vi.spyOn(sessionService, 'updateSession');
-      const refreshSessionsMock = vi.spyOn(result.current, 'refreshSessions');
-
-      // 模拟没有当前会话
-      vi.spyOn(sessionSelectors, 'currentSession').mockReturnValue(null as any);
-
-      await act(async () => {
-        await result.current.updateAgentMeta(meta as any);
-      });
-
-      expect(updateSessionMock).not.toHaveBeenCalled();
-      expect(refreshSessionsMock).not.toHaveBeenCalled();
-      updateSessionMock.mockRestore();
-      refreshSessionsMock.mockRestore();
-    });
-
-    it('should update session meta and refresh sessions', async () => {
-      const { result } = renderHook(() => useSessionStore());
-      const meta = { title: 'Test Agent' };
-      const updateSessionMock = vi.spyOn(sessionService, 'updateSession');
-      const refreshSessionsMock = vi.spyOn(result.current, 'refreshSessions');
-
-      // 模拟有当前会话
-      vi.spyOn(sessionSelectors, 'currentSession').mockReturnValue({ id: 'session-id' } as any);
-      vi.spyOn(result.current, 'activeId', 'get').mockReturnValue('session-id');
-
-      await act(async () => {
-        await result.current.updateAgentMeta(meta);
-      });
-
-      expect(updateSessionMock).toHaveBeenCalledWith('session-id', { meta });
-      expect(refreshSessionsMock).toHaveBeenCalled();
-      updateSessionMock.mockRestore();
-      refreshSessionsMock.mockRestore();
-    });
-  });
 });

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

@@ -0,0 +1,107 @@
+import isEqual from 'fast-deep-equal';
+import { produce } from 'immer';
+import { SWRResponse, mutate } from 'swr';
+import { DeepPartial } from 'utility-types';
+import { StateCreator } from 'zustand/vanilla';
+
+import { useClientDataSWR } from '@/libs/swr';
+import { sessionService } from '@/services/session';
+import { useSessionStore } from '@/store/session';
+import { LobeAgentConfig } from '@/types/agent';
+import { merge } from '@/utils/merge';
+
+import { AgentStore } from '../../store';
+import { agentSelectors } from './selectors';
+
+/**
+ * 助手接口
+ */
+export interface AgentChatAction {
+  removePlugin: (id: string) => void;
+  togglePlugin: (id: string, open?: boolean) => Promise<void>;
+  updateAgentConfig: (config: Partial<LobeAgentConfig>) => Promise<void>;
+
+  useFetchAgentConfig: (id: string) => SWRResponse<LobeAgentConfig>;
+
+  /* eslint-disable typescript-sort-keys/interface */
+
+  internal_updateAgentConfig: (id: string, data: DeepPartial<LobeAgentConfig>) => Promise<void>;
+  internal_refreshAgentConfig: (id: string) => Promise<void>;
+  /* eslint-enable */
+}
+
+const FETCH_AGENT_CONFIG_KEY = 'FETCH_AGENT_CONFIG';
+
+export const createChatSlice: StateCreator<
+  AgentStore,
+  [['zustand/devtools', never]],
+  [],
+  AgentChatAction
+> = (set, get) => ({
+  removePlugin: async (id) => {
+    await get().togglePlugin(id, false);
+  },
+
+  togglePlugin: async (id, open) => {
+    const originConfig = agentSelectors.currentAgentConfig(get());
+
+    const config = produce(originConfig, (draft) => {
+      draft.plugins = produce(draft.plugins || [], (plugins) => {
+        const index = plugins.indexOf(id);
+        const shouldOpen = open !== undefined ? open : index === -1;
+
+        if (shouldOpen) {
+          // 如果 open 为 true 或者 id 不存在于 plugins 中，则添加它
+          if (index === -1) {
+            plugins.push(id);
+          }
+        } else {
+          // 如果 open 为 false 或者 id 存在于 plugins 中，则移除它
+          if (index !== -1) {
+            plugins.splice(index, 1);
+          }
+        }
+      });
+    });
+
+    await get().updateAgentConfig(config);
+  },
+  updateAgentConfig: async (config) => {
+    const { activeId } = get();
+
+    if (!activeId) return;
+
+    await get().internal_updateAgentConfig(activeId, config);
+  },
+
+  useFetchAgentConfig: (sessionId) =>
+    useClientDataSWR<LobeAgentConfig>(
+      [FETCH_AGENT_CONFIG_KEY, sessionId],
+      ([, id]: string[]) => sessionService.getSessionConfig(id),
+      {
+        onSuccess: (data) => {
+          if (get().isAgentConfigInit && isEqual(get().agentConfig, data)) return;
+
+          set({ agentConfig: data, isAgentConfigInit: true }, false, 'fetchAgentConfig');
+        },
+      },
+    ),
+
+  /* eslint-disable sort-keys-fix/sort-keys-fix */
+
+  internal_updateAgentConfig: async (id, data) => {
+    const prevModel = agentSelectors.currentAgentModel(get());
+    // optimistic update at frontend
+    set({ agentConfig: merge(get().agentConfig, data) }, false, 'optimistic_updateAgentConfig');
+
+    await sessionService.updateSessionConfig(id, data);
+    await get().internal_refreshAgentConfig(id);
+
+    // refresh sessions to update the agent config if the model has changed
+    if (prevModel !== data.model) await useSessionStore.getState().refreshSessions();
+  },
+
+  internal_refreshAgentConfig: async (id) => {
+    await mutate([FETCH_AGENT_CONFIG_KEY, id]);
+  },
+});

package/src/store/agent/slices/chat/initialState.ts

@@ -0,0 +1,14 @@
+import { DEFAULT_AGENT_CONFIG } from '@/const/settings';
+import { LobeAgentConfig } from '@/types/agent';
+
+export interface AgentState {
+  activeId: string;
+  agentConfig: LobeAgentConfig;
+  isAgentConfigInit: boolean;
+}
+
+export const initialSessionState: AgentState = {
+  activeId: 'inbox',
+  agentConfig: DEFAULT_AGENT_CONFIG,
+  isAgentConfigInit: false,
+};

package/src/store/agent/slices/chat/selectors.test.ts

@@ -0,0 +1,82 @@
+import { describe, expect, it } from 'vitest';
+
+import { DEFAULT_AGENT_CONFIG, DEFAUTT_AGENT_TTS_CONFIG } from '@/const/settings';
+import { AgentStore } from '@/store/agent';
+
+import { agentSelectors } from './selectors';
+
+vi.mock('i18next', () => ({
+  t: vi.fn((key) => key), // Simplified mock return value
+}));
+
+const mockSessionStore = {
+  activeId: '1',
+  agentConfig: DEFAULT_AGENT_CONFIG,
+} as AgentStore;
+
+describe('agentSelectors', () => {
+  describe('currentAgentConfig', () => {
+    it('should return the merged default and session-specific agent config', () => {
+      const config = agentSelectors.currentAgentConfig(mockSessionStore);
+      expect(config).toEqual(expect.objectContaining(mockSessionStore.agentConfig));
+    });
+  });
+
+  describe('currentAgentModel', () => {
+    it('should return the model from the agent config', () => {
+      const model = agentSelectors.currentAgentModel(mockSessionStore);
+      expect(model).toBe(mockSessionStore.agentConfig.model);
+    });
+  });
+
+  describe('hasSystemRole', () => {
+    it('should return true if the system role is defined in the agent config', () => {
+      const hasRole = agentSelectors.hasSystemRole(mockSessionStore);
+      expect(hasRole).toBe(false);
+    });
+
+    it('should return false if the system role is not defined in the agent config', () => {
+      const modifiedSessionStore = {
+        ...mockSessionStore,
+        agentConfig: {
+          ...mockSessionStore.agentConfig,
+          systemRole: 'test',
+        },
+      };
+      const hasRole = agentSelectors.hasSystemRole(modifiedSessionStore);
+      expect(hasRole).toBe(true);
+    });
+  });
+
+  describe('currentAgentTTS', () => {
+    it('should return the TTS config from the agent config', () => {
+      const ttsConfig = agentSelectors.currentAgentTTS(mockSessionStore);
+      expect(ttsConfig).toEqual(mockSessionStore.agentConfig.tts);
+    });
+
+    it('should return the default TTS config if none is defined in the agent config', () => {
+      const modifiedSessionStore = {
+        ...mockSessionStore,
+        sessions: [
+          {
+            ...mockSessionStore.agentConfig,
+            config: {
+              ...mockSessionStore.agentConfig,
+              tts: DEFAUTT_AGENT_TTS_CONFIG,
+            },
+          },
+        ],
+      };
+      const ttsConfig = agentSelectors.currentAgentTTS(modifiedSessionStore);
+      expect(ttsConfig).toEqual(DEFAUTT_AGENT_TTS_CONFIG);
+    });
+  });
+
+  describe('currentAgentTTSVoice', () => {
+    it('should return the appropriate TTS voice based on the service and language', () => {
+      const lang = 'en';
+      const ttsVoice = agentSelectors.currentAgentTTSVoice(lang)(mockSessionStore);
+      expect(ttsVoice).toBe(mockSessionStore.agentConfig.tts.voice.openai);
+    });
+  });
+});