@lobehub/chat 1.0.12 → 1.0.13

package/CHANGELOG.md

@@ -2,6 +2,31 @@
 
 # Changelog
 
+### [Version 1.0.13](https://github.com/lobehub/lobe-chat/compare/v1.0.12...v1.0.13)
+
+<sup>Released on **2024-06-19**</sup>
+
+#### 🐛 Bug Fixes
+
+- **misc**: Fix and improve tool calling.
+
+<br/>
+
+<details>
+<summary><kbd>Improvements and Fixes</kbd></summary>
+
+#### What's fixed
+
+- **misc**: Fix and improve tool calling, closes [#2941](https://github.com/lobehub/lobe-chat/issues/2941) ([ea06536](https://github.com/lobehub/lobe-chat/commit/ea06536))
+
+</details>
+
+<div align="right">
+
+[![](https://img.shields.io/badge/-BACK_TO_TOP-151515?style=flat-square)](#readme-top)
+
+</div>
+
 ### [Version 1.0.12](https://github.com/lobehub/lobe-chat/compare/v1.0.11...v1.0.12)
 
 <sup>Released on **2024-06-19**</sup>

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@lobehub/chat",
-  "version": "1.0.12",
+  "version": "1.0.13",
   "description": "Lobe Chat - an open-source, high-performance chatbot framework that supports speech synthesis, multimodal, and extensible Function Call plugin system. Supports one-click free deployment of your private ChatGPT/LLM web application.",
   "keywords": [
     "framework",
@@ -92,7 +92,7 @@
   },
   "dependencies": {
     "@ant-design/icons": "^5.3.7",
-    "@anthropic-ai/sdk": "^0.21.0",
+    "@anthropic-ai/sdk": "^0.23.0",
     "@auth/core": "0.28.0",
     "@aws-sdk/client-bedrock-runtime": "^3.583.0",
     "@aws-sdk/client-s3": "^3.583.0",
@@ -109,7 +109,7 @@
     "@lobehub/chat-plugins-gateway": "^1.9.0",
     "@lobehub/icons": "^1.23.0",
     "@lobehub/tts": "^1.24.1",
-    "@lobehub/ui": "^1.143.3",
+    "@lobehub/ui": "^1.143.5",
     "@microsoft/fetch-event-source": "^2.0.1",
     "@neondatabase/serverless": "^0.9.3",
     "@next/third-parties": "^14.2.3",
@@ -154,6 +154,7 @@
     "nuqs": "^1.17.4",
     "ollama": "^0.5.1",
     "openai": "^4.47.1",
+    "partial-json": "^0.1.7",
     "pg": "^8.11.5",
     "pino": "^9.1.0",
     "polished": "^4.3.1",

package/src/features/Conversation/Error/ErrorJsonViewer.tsx

@@ -5,7 +5,7 @@ import { Flexbox } from 'react-layout-kit';
 import { ChatMessageError } from '@/types/message';
 
 interface ErrorJSONViewerProps {
-  error?: ChatMessageError;
+  error?: ChatMessageError | null;
   id: string;
 }
 

package/src/features/Conversation/Messages/Assistant/ToolCalls/index.tsx

@@ -1,4 +1,4 @@
-import { Avatar, Highlighter, Icon } from '@lobehub/ui';
+import { Avatar, Icon } from '@lobehub/ui';
 import isEqual from 'fast-deep-equal';
 import { Loader2, LucideChevronDown, LucideChevronRight, LucideToyBrick } from 'lucide-react';
 import { CSSProperties, memo, useState } from 'react';
@@ -10,6 +10,7 @@ import { chatSelectors } from '@/store/chat/selectors';
 import { pluginHelpers, useToolStore } from '@/store/tool';
 import { toolSelectors } from '@/store/tool/selectors';
 
+import Arguments from '../../components/Arguments';
 import { useStyles } from './style';
 
 export interface InspectorProps {
@@ -64,7 +65,7 @@ const CallItem = memo<InspectorProps>(
           </Flexbox>
           <Icon icon={open ? LucideChevronDown : LucideChevronRight} />
         </Flexbox>
-        {(open || loading) && <Highlighter language={'json'}>{requestArgs}</Highlighter>}
+        {(open || loading) && <Arguments arguments={requestArgs} />}
       </Flexbox>
     );
   },

package/src/features/Conversation/Messages/Tool/index.tsx

@@ -1,4 +1,3 @@
-import { Snippet } from '@lobehub/ui';
 import { memo, useState } from 'react';
 import { Flexbox } from 'react-layout-kit';
 
@@ -7,10 +6,11 @@ import { chatSelectors } from '@/store/chat/selectors';
 import { ChatMessage } from '@/types/message';
 
 import PluginRender from '../../Plugins/Render';
+import Arguments from '../components/Arguments';
 import Inspector from './Inspector';
 
 export const ToolMessage = memo<ChatMessage>(({ id, content, plugin }) => {
-  const loading = useChatStore(chatSelectors.isMessageGenerating(id));
+  const loading = useChatStore(chatSelectors.isPluginApiInvoking(id));
 
   const [showRender, setShow] = useState(plugin?.type !== 'default');
 
@@ -35,9 +35,7 @@ export const ToolMessage = memo<ChatMessage>(({ id, content, plugin }) => {
           type={plugin?.type}
         />
       ) : (
-        <Flexbox>
-          <Snippet>{plugin?.arguments || ''}</Snippet>
-        </Flexbox>
+        <Arguments arguments={plugin?.arguments} />
       )}
     </Flexbox>
   );

package/src/features/Conversation/Messages/components/Arguments.tsx

@@ -0,0 +1,20 @@
+import { Highlighter } from '@lobehub/ui';
+import { memo } from 'react';
+
+import { useYamlArguments } from '../hooks/useYamlArguments';
+
+export interface ArgumentsProps {
+  arguments?: string;
+}
+
+const Arguments = memo<ArgumentsProps>(({ arguments: args = '' }) => {
+  const yaml = useYamlArguments(args);
+
+  return (
+    <Highlighter language={'yaml'} showLanguage={false}>
+      {yaml}
+    </Highlighter>
+  );
+});
+
+export default Arguments;

package/src/features/Conversation/Messages/hooks/useYamlArguments.ts

@@ -0,0 +1,11 @@
+import { Allow, parse } from 'partial-json';
+import { stringify } from 'yaml';
+
+export const useYamlArguments = (args: string) => {
+  try {
+    const obj = parse(args, Allow.OBJ);
+    return stringify(obj);
+  } catch {
+    return args;
+  }
+};

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

@@ -22,10 +22,6 @@ beforeEach(() => {
 
   // 使用 vi.spyOn 来模拟 chat.completions.create 方法
   vi.spyOn(instance['client'].messages, 'create').mockReturnValue(new ReadableStream() as any);
-
-  vi.spyOn(instance['client'].beta.tools.messages, 'create').mockReturnValue({
-    content: [],
-  } as any);
 });
 
 afterEach(() => {
@@ -260,35 +256,9 @@ describe('LobeAnthropicAI', () => {
         });
 
         // Assert
-        expect(instance['client'].beta.tools.messages.create).toHaveBeenCalled();
+        expect(instance['client'].messages.create).toHaveBeenCalled();
         expect(spyOn).toHaveBeenCalledWith(tools);
       });
-
-      it('should handle text and tool_use content correctly in transformResponseToStream', async () => {
-        // Arrange
-        const mockResponse = {
-          content: [
-            { type: 'text', text: 'Hello' },
-            { type: 'tool_use', id: 'tool1', name: 'tool1', input: 'input1' },
-          ],
-        };
-        // @ts-ignore
-        vi.spyOn(instance, 'transformResponseToStream').mockReturnValue(new ReadableStream());
-        vi.spyOn(instance['client'].beta.tools.messages, 'create').mockResolvedValue(
-          mockResponse as any,
-        );
-
-        // Act
-        await instance.chat({
-          messages: [{ content: 'Hello', role: 'user' }],
-          model: 'claude-3-haiku-20240307',
-          temperature: 0,
-          tools: [{ function: { name: 'tool1', description: 'desc1' }, type: 'function' }],
-        });
-
-        // Assert
-        expect(instance['transformResponseToStream']).toHaveBeenCalledWith(mockResponse);
-      });
     });
 
     describe('Error', () => {

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

@@ -30,41 +30,20 @@ export class LobeAnthropicAI implements LobeRuntimeAI {
   async chat(payload: ChatStreamPayload, options?: ChatCompetitionOptions) {
     try {
       const anthropicPayload = this.buildAnthropicPayload(payload);
-
-      // if there is no tool, we can use the normal chat API
-      if (!anthropicPayload.tools || anthropicPayload.tools.length === 0) {
-        const response = await this.client.messages.create(
-          { ...anthropicPayload, stream: true },
-          {
-            signal: options?.signal,
-          },
-        );
-
-        const [prod, debug] = response.tee();
-
-        if (process.env.DEBUG_ANTHROPIC_CHAT_COMPLETION === '1') {
-          debugStream(debug.toReadableStream()).catch(console.error);
-        }
-
-        return StreamingResponse(AnthropicStream(prod, options?.callback), {
-          headers: options?.headers,
-        });
-      }
-
-      // or we should call the tool API
-      const response = await this.client.beta.tools.messages.create(
-        { ...anthropicPayload, stream: false },
-        { signal: options?.signal },
+      const response = await this.client.messages.create(
+        { ...anthropicPayload, stream: true },
+        {
+          signal: options?.signal,
+        },
       );
 
+      const [prod, debug] = response.tee();
+
       if (process.env.DEBUG_ANTHROPIC_CHAT_COMPLETION === '1') {
-        console.log('\n[no stream response]\n');
-        console.log(JSON.stringify(response) + '\n');
+        debugStream(debug.toReadableStream()).catch(console.error);
       }
 
-      const stream = this.transformResponseToStream(response);
-
-      return StreamingResponse(AnthropicStream(stream, options?.callback), {
+      return StreamingResponse(AnthropicStream(prod, options?.callback), {
         headers: options?.headers,
       });
     } catch (error) {
@@ -118,43 +97,10 @@ export class LobeAnthropicAI implements LobeRuntimeAI {
       model,
       system: system_message?.content as string,
       temperature,
-      // TODO: Anthropic sdk don't have tools interface currently
-      // @ts-ignore
       tools: buildAnthropicTools(tools),
       top_p,
     } satisfies Anthropic.MessageCreateParams;
   }
-
-  private transformResponseToStream = (response: Anthropic.Beta.Tools.ToolsBetaMessage) => {
-    return new ReadableStream<Anthropic.MessageStreamEvent>({
-      start(controller) {
-        response.content.forEach((content) => {
-          switch (content.type) {
-            case 'text': {
-              controller.enqueue({
-                delta: { text: content.text, type: 'text_delta' },
-                type: 'content_block_delta',
-              } as Anthropic.ContentBlockDeltaEvent);
-              break;
-            }
-            case 'tool_use': {
-              controller.enqueue({
-                delta: {
-                  tool_use: { id: content.id, input: content.input, name: content.name },
-                  type: 'tool_use',
-                },
-                type: 'content_block_delta',
-              } as any);
-            }
-          }
-        });
-
-        controller.enqueue({ type: 'message_stop' } as Anthropic.MessageStopEvent);
-
-        controller.close();
-      },
-    });
-  };
 }
 
 export default LobeAnthropicAI;

package/src/libs/agent-runtime/openai/index.test.ts

@@ -49,7 +49,7 @@ describe('LobeOpenAI', () => {
     });
 
     describe('Error', () => {
-      it('should return OpenAIBizError with an openai error response when OpenAI.APIError is thrown', async () => {
+      it('should return ProviderBizError with an openai error response when OpenAI.APIError is thrown', async () => {
         // Arrange
         const apiError = new OpenAI.APIError(
           400,
@@ -79,7 +79,7 @@ describe('LobeOpenAI', () => {
               error: { message: 'Bad Request' },
               status: 400,
             },
-            errorType: 'OpenAIBizError',
+            errorType: 'ProviderBizError',
             provider: 'openai',
           });
         }
@@ -89,11 +89,11 @@ describe('LobeOpenAI', () => {
         try {
           new LobeOpenAI({});
         } catch (e) {
-          expect(e).toEqual({ errorType: 'NoOpenAIAPIKey' });
+          expect(e).toEqual({ errorType: 'InvalidProviderAPIKey' });
         }
       });
 
-      it('should return OpenAIBizError with the cause when OpenAI.APIError is thrown with cause', async () => {
+      it('should return ProviderBizError with the cause when OpenAI.APIError is thrown with cause', async () => {
         // Arrange
         const errorInfo = {
           stack: 'abc',
@@ -119,13 +119,13 @@ describe('LobeOpenAI', () => {
               cause: { message: 'api is undefined' },
               stack: 'abc',
             },
-            errorType: 'OpenAIBizError',
+            errorType: 'ProviderBizError',
             provider: 'openai',
           });
         }
       });
 
-      it('should return OpenAIBizError with an cause response with desensitize Url', async () => {
+      it('should return ProviderBizError with an cause response with desensitize Url', async () => {
         // Arrange
         const errorInfo = {
           stack: 'abc',
@@ -155,7 +155,7 @@ describe('LobeOpenAI', () => {
               cause: { message: 'api is undefined' },
               stack: 'abc',
             },
-            errorType: 'OpenAIBizError',
+            errorType: 'ProviderBizError',
             provider: 'openai',
           });
         }

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

@@ -1,4 +1,3 @@
-import { AgentRuntimeErrorType } from '../error';
 import { ModelProvider } from '../types';
 import { LobeOpenAICompatibleFactory } from '../utils/openaiCompatibleFactory';
 
@@ -7,10 +6,5 @@ export const LobeOpenAI = LobeOpenAICompatibleFactory({
   debug: {
     chatCompletion: () => process.env.DEBUG_OPENAI_CHAT_COMPLETION === '1',
   },
-  errorType: {
-    bizError: AgentRuntimeErrorType.OpenAIBizError,
-    invalidAPIKey: AgentRuntimeErrorType.NoOpenAIAPIKey,
-  },
-
   provider: ModelProvider.OpenAI,
 });

package/src/libs/agent-runtime/utils/anthropicHelpers.ts

@@ -112,9 +112,9 @@ export const buildAnthropicMessages = (
 
 export const buildAnthropicTools = (tools?: OpenAI.ChatCompletionTool[]) =>
   tools?.map(
-    (tool): Anthropic.Beta.Tools.Tool => ({
+    (tool): Anthropic.Tool => ({
       description: tool.function.description,
-      input_schema: tool.function.parameters as Anthropic.Beta.Tools.Tool.InputSchema,
+      input_schema: tool.function.parameters as Anthropic.Tool.InputSchema,
       name: tool.function.name,
     }),
   );

package/src/libs/agent-runtime/utils/streams/anthropic.test.ts

@@ -99,27 +99,67 @@ describe('AnthropicStream', () => {
   });
 
   it('should handle tool use event and ReadableStream input', async () => {
-    const toolUseEvent = {
-      type: 'content_block_delta',
-      delta: {
-        type: 'tool_use',
-        tool_use: {
-          id: 'tool_use_1',
-          name: 'example_tool',
-          input: { arg1: 'value1' },
+    const streams = [
+      {
+        type: 'message_start',
+        message: {
+          id: 'msg_017aTuY86wNxth5TE544yqJq',
+          type: 'message',
+          role: 'assistant',
+          model: 'claude-3-sonnet-20240229',
+          content: [],
+          stop_reason: null,
+          stop_sequence: null,
+          usage: { input_tokens: 457, output_tokens: 1 },
         },
       },
-    };
+      { type: 'content_block_start', index: 0, content_block: { type: 'text', text: '' } },
+      { type: 'content_block_delta', index: 0, delta: { type: 'text_delta', text: '好' } },
+      { type: 'content_block_delta', index: 0, delta: { type: 'text_delta', text: '的:' } },
+
+      { type: 'content_block_stop', index: 0 },
+      {
+        type: 'content_block_start',
+        index: 1,
+        content_block: {
+          type: 'tool_use',
+          id: 'toolu_01WdYWxYFQ8iu5iZq1Dy9Saf',
+          name: 'realtime-weather____fetchCurrentWeather',
+          input: {},
+        },
+      },
+      {
+        type: 'content_block_delta',
+        index: 1,
+        delta: { type: 'input_json_delta', partial_json: '' },
+      },
+      {
+        type: 'content_block_delta',
+        index: 1,
+        delta: { type: 'input_json_delta', partial_json: '{"city": "' },
+      },
+      {
+        type: 'content_block_delta',
+        index: 1,
+        delta: { type: 'input_json_delta', partial_json: '杭' },
+      },
+      {
+        type: 'content_block_delta',
+        index: 1,
+        delta: { type: 'input_json_delta', partial_json: '州"}' },
+      },
+      { type: 'content_block_stop', index: 1 },
+      {
+        type: 'message_delta',
+        delta: { stop_reason: 'tool_use', stop_sequence: null },
+        usage: { output_tokens: 83 },
+      },
+    ];
 
     const mockReadableStream = new ReadableStream({
       start(controller) {
-        controller.enqueue({
-          type: 'message_start',
-          message: { id: 'message_1', metadata: {} },
-        });
-        controller.enqueue(toolUseEvent);
-        controller.enqueue({
-          type: 'message_stop',
+        streams.forEach((chunk) => {
+          controller.enqueue(chunk);
         });
         controller.close();
       },
@@ -139,19 +179,53 @@ describe('AnthropicStream', () => {
       chunks.push(decoder.decode(chunk, { stream: true }));
     }
 
-    expect(chunks).toEqual([
-      'id: message_1\n',
-      'event: data\n',
-      `data: {"id":"message_1","metadata":{}}\n\n`,
-      'id: message_1\n',
-      'event: tool_calls\n',
-      `data: [{"function":{"arguments":"{\\"arg1\\":\\"value1\\"}","name":"example_tool"},"id":"tool_use_1","index":0,"type":"function"}]\n\n`,
-      'id: message_1\n',
-      'event: stop\n',
-      `data: "message_stop"\n\n`,
-    ]);
+    expect(chunks).toEqual(
+      [
+        'id: msg_017aTuY86wNxth5TE544yqJq',
+        'event: data',
+        'data: {"id":"msg_017aTuY86wNxth5TE544yqJq","type":"message","role":"assistant","model":"claude-3-sonnet-20240229","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":457,"output_tokens":1}}\n',
+        'id: msg_017aTuY86wNxth5TE544yqJq',
+        'event: data',
+        'data: ""\n',
+        'id: msg_017aTuY86wNxth5TE544yqJq',
+        'event: text',
+        'data: "好"\n',
+        'id: msg_017aTuY86wNxth5TE544yqJq',
+        'event: text',
+        'data: "的:"\n',
+        'id: msg_017aTuY86wNxth5TE544yqJq',
+        'event: data',
+        'data: {"type":"content_block_stop","index":0}\n',
+        // Tool calls
+        'id: msg_017aTuY86wNxth5TE544yqJq',
+        'event: tool_calls',
+        `data: [{"function":{"arguments":"","name":"realtime-weather____fetchCurrentWeather"},"id":"toolu_01WdYWxYFQ8iu5iZq1Dy9Saf","index":0,"type":"function"}]\n`,
+        'id: msg_017aTuY86wNxth5TE544yqJq',
+        'event: tool_calls',
+        `data: [{"function":{"arguments":""},"index":0,"type":"function"}]\n`,
+        'id: msg_017aTuY86wNxth5TE544yqJq',
+        'event: tool_calls',
+        `data: [{"function":{"arguments":"{\\"city\\": \\""},"index":0,"type":"function"}]\n`,
+        'id: msg_017aTuY86wNxth5TE544yqJq',
+        'event: tool_calls',
+
+        `data: [{"function":{"arguments":"杭"},"index":0,"type":"function"}]\n`,
+        'id: msg_017aTuY86wNxth5TE544yqJq',
+        'event: tool_calls',
+
+        `data: [{"function":{"arguments":"州\\"}"},"index":0,"type":"function"}]\n`,
+
+        'id: msg_017aTuY86wNxth5TE544yqJq',
+        'event: data',
+        'data: {"type":"content_block_stop","index":1}\n',
+
+        'id: msg_017aTuY86wNxth5TE544yqJq',
+        'event: stop',
+        'data: "tool_use"\n',
+      ].map((item) => `${item}\n`),
+    );
 
-    expect(onToolCallMock).toHaveBeenCalledTimes(1);
+    expect(onToolCallMock).toHaveBeenCalledTimes(5);
   });
 
   it('should handle ReadableStream input', async () => {

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

@@ -22,27 +22,39 @@ export const transformAnthropicStream = (
       stack.id = chunk.message.id;
       return { data: chunk.message, id: chunk.message.id, type: 'data' };
     }
+    case 'content_block_start': {
+      if (chunk.content_block.type === 'tool_use') {
+        const toolChunk = chunk.content_block;
 
-    // case 'content_block_start': {
-    //   return { data: chunk.content_block.text, id: stack.id, type: 'data' };
-    // }
+        const toolCall: StreamToolCallChunkData = {
+          function: {
+            arguments: '',
+            name: toolChunk.name,
+          },
+          id: toolChunk.id,
+          index: 0,
+          type: 'function',
+        };
+
+        stack.tool = { id: toolChunk.id, index: 0, name: toolChunk.name };
+
+        return { data: [toolCall], id: stack.id, type: 'tool_calls' };
+      }
+
+      return { data: chunk.content_block.text, id: stack.id, type: 'data' };
+    }
 
     case 'content_block_delta': {
-      switch (chunk.delta.type as string) {
-        default:
+      switch (chunk.delta.type) {
         case 'text_delta': {
           return { data: chunk.delta.text, id: stack.id, type: 'text' };
         }
 
-        // TODO: due to anthropic currently don't support streaming tool calling
-        // we need to add this new `tool_use` type to support streaming
-        // and maybe we need to update it when the feature is available
-        case 'tool_use': {
-          const delta = (chunk.delta as any).tool_use as Anthropic.Beta.Tools.ToolUseBlock;
+        case 'input_json_delta': {
+          const delta = chunk.delta.partial_json;
 
           const toolCall: StreamToolCallChunkData = {
-            function: { arguments: JSON.stringify(delta.input), name: delta.name },
-            id: delta.id,
+            function: { arguments: delta },
             index: 0,
             type: 'function',
           };
@@ -53,7 +65,12 @@ export const transformAnthropicStream = (
             type: 'tool_calls',
           } as StreamProtocolToolCallChunk;
         }
+
+        default: {
+          break;
+        }
       }
+      return { data: chunk, id: stack.id, type: 'data' };
     }
 
     case 'message_delta': {

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

@@ -2,6 +2,11 @@ import { ChatStreamCallbacks } from '@/libs/agent-runtime';
 
 export interface StreamStack {
   id: string;
+  tool?: {
+    id: string;
+    index: number;
+    name: string;
+  };
 }
 
 export interface StreamProtocolChunk {
@@ -15,7 +20,7 @@ export interface StreamToolCallChunkData {
     arguments?: string;
     name?: string | null;
   };
-  id: string;
+  id?: string;
   index: number;
   type: 'function' | string;
 }

package/src/services/message/server.ts

@@ -44,6 +44,10 @@ export class ServerService implements IMessageService {
     return lambdaClient.message.update.mutate({ id, value: { error } });
   }
 
+  async updateMessagePluginError(id: string, error: ChatMessageError): Promise<any> {
+    return lambdaClient.message.update.mutate({ id, value: { pluginError: error } });
+  }
+
   updateMessage(id: string, message: Partial<ChatMessage>): Promise<any> {
     return lambdaClient.message.update.mutate({ id, value: message });
   }

package/src/services/message/type.ts

@@ -17,7 +17,7 @@ export interface CreateMessageParams
   traceId?: string;
   topicId?: string;
   content: string;
-  error?: ChatMessageError;
+  error?: ChatMessageError | null;
   role: MessageRoleType;
 }
 

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

@@ -19,12 +19,14 @@ import { agentSelectors } from '@/store/agent/selectors';
 import { chatHelpers } from '@/store/chat/helpers';
 import { messageMapKey } from '@/store/chat/slices/message/utils';
 import { ChatStore } from '@/store/chat/store';
-import { ChatMessage, MessageToolCall } from '@/types/message';
+import { ChatMessage, ChatMessageError, MessageToolCall } from '@/types/message';
 import { TraceEventPayloads } from '@/types/trace';
 import { setNamespace } from '@/utils/storeDebug';
 import { nanoid } from '@/utils/uuid';
 
+import type { ChatStoreState } from '../../initialState';
 import { chatSelectors, topicSelectors } from '../../selectors';
+import { preventLeavingFn, toggleBooleanList } from '../../utils';
 import { MessageDispatch, messagesReducer } from './reducer';
 
 const n = setNamespace('m');
@@ -84,6 +86,12 @@ export interface ChatMessageAction {
     id?: string,
     action?: string,
   ) => AbortController | undefined;
+  internal_toggleLoadingArrays: (
+    key: keyof ChatStoreState,
+    loading: boolean,
+    id?: string,
+    action?: string,
+  ) => AbortController | undefined;
   internal_toggleToolCallingStreaming: (id: string, streaming: boolean[] | undefined) => void;
   internal_toggleMessageLoading: (loading: boolean, id: string) => void;
   /**
@@ -123,6 +131,7 @@ export interface ChatMessageAction {
     content: string,
     toolCalls?: MessageToolCall[],
   ) => Promise<void>;
+  internal_updateMessageError: (id: string, error: ChatMessageError | null) => Promise<void>;
   internal_createMessage: (
     params: CreateMessageParams,
     context?: { tempMessageId?: string; skipRefresh?: boolean },
@@ -136,24 +145,6 @@ export interface ChatMessageAction {
 const getAgentConfig = () => agentSelectors.currentAgentConfig(useAgentStore.getState());
 const getAgentChatConfig = () => agentSelectors.currentAgentChatConfig(useAgentStore.getState());
 
-const preventLeavingFn = (e: BeforeUnloadEvent) => {
-  // set returnValue to trigger alert modal
-  // Note: No matter what value is set, the browser will display the standard text
-  e.returnValue = '你有正在生成中的请求，确定要离开吗？';
-};
-
-const toggleBooleanList = (ids: string[], id: string, loading: boolean) => {
-  return produce(ids, (draft) => {
-    if (loading) {
-      draft.push(id);
-    } else {
-      const index = draft.indexOf(id);
-
-      if (index >= 0) draft.splice(index, 1);
-    }
-  });
-};
-
 export const chatMessage: StateCreator<
   ChatStore,
   [['zustand/devtools', never]],
@@ -591,35 +582,7 @@ export const chatMessage: StateCreator<
     };
   },
   internal_toggleChatLoading: (loading, id, action) => {
-    if (loading) {
-      window.addEventListener('beforeunload', preventLeavingFn);
-
-      const abortController = new AbortController();
-      set(
-        {
-          abortController,
-          chatLoadingIds: toggleBooleanList(get().messageLoadingIds, id!, loading),
-        },
-        false,
-        action,
-      );
-
-      return abortController;
-    } else {
-      if (!id) {
-        set({ abortController: undefined, chatLoadingIds: [] }, false, action);
-      } else
-        set(
-          {
-            abortController: undefined,
-            chatLoadingIds: toggleBooleanList(get().messageLoadingIds, id, loading),
-          },
-          false,
-          action,
-        );
-
-      window.removeEventListener('beforeunload', preventLeavingFn);
-    }
+    return get().internal_toggleLoadingArrays('chatLoadingIds', loading, id, action);
   },
   internal_toggleMessageLoading: (loading, id) => {
     set(
@@ -684,6 +647,11 @@ export const chatMessage: StateCreator<
     await internal_coreProcessMessage(contextMessages, latestMsg.id, { traceId });
   },
 
+  internal_updateMessageError: async (id, error) => {
+    get().internal_dispatchMessage({ id, type: 'updateMessages', value: { error } });
+    await messageService.updateMessage(id, { error });
+    await get().refreshMessages();
+  },
   internal_updateMessageContent: async (id, content, toolCalls) => {
     const { internal_dispatchMessage, refreshMessages, internal_transformToolCalls } = get();
 
@@ -762,4 +730,36 @@ export const chatMessage: StateCreator<
         .catch();
     }
   },
+
+  internal_toggleLoadingArrays: (key, loading, id, action) => {
+    if (loading) {
+      window.addEventListener('beforeunload', preventLeavingFn);
+
+      const abortController = new AbortController();
+      set(
+        {
+          abortController,
+          [key]: toggleBooleanList(get()[key] as string[], id!, loading),
+        },
+        false,
+        action,
+      );
+
+      return abortController;
+    } else {
+      if (!id) {
+        set({ abortController: undefined, [key]: [] }, false, action);
+      } else
+        set(
+          {
+            abortController: undefined,
+            [key]: toggleBooleanList(get()[key] as string[], id, loading),
+          },
+          false,
+          action,
+        );
+
+      window.removeEventListener('beforeunload', preventLeavingFn);
+    }
+  },
 });

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

@@ -26,6 +26,7 @@ export interface ChatMessageState {
    */
   messagesInit: boolean;
   messagesMap: Record<string, ChatMessage[]>;
+  pluginApiLoadingIds: string[];
   /**
    * the tool calling stream ids
    */
@@ -41,5 +42,6 @@ export const initialMessageState: ChatMessageState = {
   messageLoadingIds: [],
   messagesInit: false,
   messagesMap: {},
+  pluginApiLoadingIds: [],
   toolCallingStreamIds: {},
 };

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

@@ -129,6 +129,8 @@ const isHasMessageLoading = (s: ChatStore) => s.messageLoadingIds.length > 0;
 const isCreatingMessage = (s: ChatStore) => s.isCreatingMessage;
 
 const isMessageGenerating = (id: string) => (s: ChatStore) => s.chatLoadingIds.includes(id);
+const isPluginApiInvoking = (id: string) => (s: ChatStore) => s.pluginApiLoadingIds.includes(id);
+
 const isToolCallStreaming = (id: string, index: number) => (s: ChatStore) => {
   const isLoading = s.toolCallingStreamIds[id];
 
@@ -155,6 +157,7 @@ export const chatSelectors = {
   isMessageEditing,
   isMessageGenerating,
   isMessageLoading,
+  isPluginApiInvoking,
   isToolCallStreaming,
   latestMessage,
   showInboxWelcome,

package/src/store/chat/slices/plugin/action.test.ts

@@ -112,7 +112,7 @@ describe('ChatPluginAction', () => {
 
       vi.spyOn(storeState, 'refreshMessages');
       vi.spyOn(storeState, 'triggerAIMessage').mockResolvedValue(undefined);
-      vi.spyOn(storeState, 'internal_toggleChatLoading').mockReturnValue(undefined);
+      vi.spyOn(storeState, 'internal_togglePluginApiCalling').mockReturnValue(undefined);
 
       const runSpy = vi.spyOn(chatService, 'runPluginApi').mockResolvedValue({
         text: pluginApiResponse,
@@ -125,7 +125,7 @@ describe('ChatPluginAction', () => {
         await result.current.invokeDefaultTypePlugin(messageId, pluginPayload);
       });
 
-      expect(storeState.internal_toggleChatLoading).toHaveBeenCalledWith(
+      expect(storeState.internal_togglePluginApiCalling).toHaveBeenCalledWith(
         true,
         messageId,
         expect.any(String),
@@ -135,7 +135,7 @@ describe('ChatPluginAction', () => {
         content: pluginApiResponse,
       });
       expect(storeState.refreshMessages).toHaveBeenCalled();
-      expect(storeState.internal_toggleChatLoading).toHaveBeenCalledWith(
+      expect(storeState.internal_togglePluginApiCalling).toHaveBeenCalledWith(
         false,
         'message-id',
         'plugin/fetchPlugin/end',
@@ -150,7 +150,7 @@ describe('ChatPluginAction', () => {
       const storeState = useChatStore.getState();
       vi.spyOn(storeState, 'refreshMessages');
       vi.spyOn(storeState, 'triggerAIMessage').mockResolvedValue(undefined);
-      vi.spyOn(storeState, 'internal_toggleChatLoading').mockReturnValue(undefined);
+      vi.spyOn(storeState, 'internal_togglePluginApiCalling').mockReturnValue(undefined);
 
       vi.spyOn(chatService, 'runPluginApi').mockRejectedValue(error);
 
@@ -159,7 +159,7 @@ describe('ChatPluginAction', () => {
         await result.current.invokeDefaultTypePlugin(messageId, pluginPayload);
       });
 
-      expect(storeState.internal_toggleChatLoading).toHaveBeenCalledWith(
+      expect(storeState.internal_togglePluginApiCalling).toHaveBeenCalledWith(
         true,
         messageId,
         expect.any(String),
@@ -167,7 +167,7 @@ describe('ChatPluginAction', () => {
       expect(chatService.runPluginApi).toHaveBeenCalledWith(pluginPayload, { trace: {} });
       expect(messageService.updateMessageError).toHaveBeenCalledWith(messageId, error);
       expect(storeState.refreshMessages).toHaveBeenCalled();
-      expect(storeState.internal_toggleChatLoading).toHaveBeenCalledWith(
+      expect(storeState.internal_togglePluginApiCalling).toHaveBeenCalledWith(
         false,
         'message-id',
         'plugin/fetchPlugin/end',
@@ -467,7 +467,7 @@ describe('ChatPluginAction', () => {
       });
 
       useChatStore.setState({
-        internal_toggleChatLoading: vi.fn(),
+        internal_togglePluginApiCalling: vi.fn(),
         internal_updateMessageContent: vi.fn(),
         text2image: vi.fn(),
       });
@@ -491,12 +491,12 @@ describe('ChatPluginAction', () => {
       );
 
       // Verify that loading was toggled correctly
-      expect(result.current.internal_toggleChatLoading).toHaveBeenCalledWith(
+      expect(result.current.internal_togglePluginApiCalling).toHaveBeenCalledWith(
         true,
         messageId,
         expect.any(String),
       );
-      expect(result.current.internal_toggleChatLoading).toHaveBeenCalledWith(false);
+      expect(result.current.internal_togglePluginApiCalling).toHaveBeenCalledWith(false);
       expect(useChatStore.getState().text2image).toHaveBeenCalled();
     });
 
@@ -516,7 +516,7 @@ describe('ChatPluginAction', () => {
         });
 
         useChatStore.setState({
-          internal_toggleChatLoading: vi.fn(),
+          internal_togglePluginApiCalling: vi.fn(),
           text2image: vi.fn(),
           internal_updateMessageContent: vi.fn(),
         });
@@ -540,12 +540,12 @@ describe('ChatPluginAction', () => {
       );
 
       // Verify that loading was toggled correctly
-      expect(result.current.internal_toggleChatLoading).toHaveBeenCalledWith(
+      expect(result.current.internal_togglePluginApiCalling).toHaveBeenCalledWith(
         true,
         messageId,
         expect.any(String),
       );
-      expect(result.current.internal_toggleChatLoading).toHaveBeenCalledWith(false);
+      expect(result.current.internal_togglePluginApiCalling).toHaveBeenCalledWith(false);
       expect(useChatStore.getState().text2image).not.toHaveBeenCalled();
     });
 
@@ -563,7 +563,7 @@ describe('ChatPluginAction', () => {
       });
 
       useChatStore.setState({
-        internal_toggleChatLoading: vi.fn(),
+        internal_togglePluginApiCalling: vi.fn(),
         internal_updateMessageContent: vi.fn(),
         text2image: vi.fn(),
         refreshMessages: vi.fn(),
@@ -576,12 +576,12 @@ describe('ChatPluginAction', () => {
       });
 
       // Verify that loading was toggled correctly
-      expect(result.current.internal_toggleChatLoading).toHaveBeenCalledWith(
+      expect(result.current.internal_togglePluginApiCalling).toHaveBeenCalledWith(
         true,
         messageId,
         expect.any(String),
       );
-      expect(result.current.internal_toggleChatLoading).toHaveBeenCalledWith(false);
+      expect(result.current.internal_togglePluginApiCalling).toHaveBeenCalledWith(false);
 
       // Verify that the message content was not updated
       expect(result.current.internal_updateMessageContent).not.toHaveBeenCalled();

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

@@ -1,3 +1,4 @@
+/* eslint-disable sort-keys-fix/sort-keys-fix, typescript-sort-keys/interface */
 import { PluginErrorType } from '@lobehub/chat-plugin-sdk';
 import { t } from 'i18next';
 import { Md5 } from 'ts-md5';
@@ -25,11 +26,6 @@ export interface ChatPluginAction {
     triggerAiMessage?: boolean,
   ) => Promise<void>;
 
-  internal_callPluginApi: (id: string, payload: ChatToolPayload) => Promise<string | undefined>;
-  internal_invokeDifferentTypePlugin: (id: string, payload: ChatToolPayload) => Promise<any>;
-  internal_transformToolCalls: (toolCalls: MessageToolCall[]) => ChatToolPayload[];
-  internal_updatePluginError: (id: string, error: any) => Promise<void>;
-
   invokeBuiltinTool: (id: string, payload: ChatToolPayload) => Promise<void>;
   invokeDefaultTypePlugin: (id: string, payload: any) => Promise<string | undefined>;
   invokeMarkdownTypePlugin: (id: string, payload: ChatToolPayload) => Promise<void>;
@@ -41,6 +37,16 @@ export interface ChatPluginAction {
 
   triggerToolCalls: (id: string) => Promise<void>;
   updatePluginState: (id: string, key: string, value: any) => Promise<void>;
+
+  internal_callPluginApi: (id: string, payload: ChatToolPayload) => Promise<string | undefined>;
+  internal_invokeDifferentTypePlugin: (id: string, payload: ChatToolPayload) => Promise<any>;
+  internal_togglePluginApiCalling: (
+    loading: boolean,
+    id?: string,
+    action?: string,
+  ) => AbortController | undefined;
+  internal_transformToolCalls: (toolCalls: MessageToolCall[]) => ChatToolPayload[];
+  internal_updatePluginError: (id: string, error: any) => Promise<void>;
 }
 
 export const chatPlugin: StateCreator<
@@ -69,124 +75,17 @@ export const chatPlugin: StateCreator<
 
     if (triggerAiMessage) await triggerAIMessage({ parentId: id });
   },
-  internal_callPluginApi: async (id, payload) => {
-    const { internal_updateMessageContent, refreshMessages, internal_toggleChatLoading } = get();
-    let data: string;
-
-    try {
-      const abortController = internal_toggleChatLoading(
-        true,
-        id,
-        n('fetchPlugin/start') as string,
-      );
-
-      const message = chatSelectors.getMessageById(id)(get());
-
-      const res = await chatService.runPluginApi(payload, {
-        signal: abortController?.signal,
-        trace: { observationId: message?.observationId, traceId: message?.traceId },
-      });
-      data = res.text;
-
-      // save traceId
-      if (res.traceId) {
-        await messageService.updateMessage(id, { traceId: res.traceId });
-      }
-    } catch (error) {
-      console.log(error);
-      const err = error as Error;
-
-      // ignore the aborted request error
-      if (!err.message.includes('The user aborted a request.')) {
-        await messageService.updateMessageError(id, error as any);
-        await refreshMessages();
-      }
-
-      data = '';
-    }
-
-    internal_toggleChatLoading(false, id, n('fetchPlugin/end') as string);
-    // 如果报错则结束了
-    if (!data) return;
-
-    await internal_updateMessageContent(id, data);
-
-    return data;
-  },
-
-  internal_invokeDifferentTypePlugin: async (id, payload) => {
-    switch (payload.type) {
-      case 'standalone': {
-        return await get().invokeStandaloneTypePlugin(id, payload);
-      }
-
-      case 'markdown': {
-        return await get().invokeMarkdownTypePlugin(id, payload);
-      }
-
-      case 'builtin': {
-        return await get().invokeBuiltinTool(id, payload);
-      }
-
-      default: {
-        return await get().invokeDefaultTypePlugin(id, payload);
-      }
-    }
-  },
-
-  internal_transformToolCalls: (toolCalls) => {
-    return toolCalls
-      .map((toolCall): ChatToolPayload | null => {
-        let payload: ChatToolPayload;
-
-        const [identifier, apiName, type] = toolCall.function.name.split(PLUGIN_SCHEMA_SEPARATOR);
-
-        if (!apiName) return null;
-
-        payload = {
-          apiName,
-          arguments: toolCall.function.arguments,
-          id: toolCall.id,
-          identifier,
-          type: (type ?? 'default') as any,
-        };
-
-        // if the apiName is md5, try to find the correct apiName in the plugins
-        if (apiName.startsWith(PLUGIN_SCHEMA_API_MD5_PREFIX)) {
-          const md5 = apiName.replace(PLUGIN_SCHEMA_API_MD5_PREFIX, '');
-          const manifest = pluginSelectors.getPluginManifestById(identifier)(
-            useToolStore.getState(),
-          );
-
-          const api = manifest?.api.find((api) => Md5.hashStr(api.name).toString() === md5);
-          if (api) {
-            payload.apiName = api.name;
-          }
-        }
-
-        return payload;
-      })
-      .filter(Boolean) as ChatToolPayload[];
-  },
-
-  internal_updatePluginError: async (id, error) => {
-    const { refreshMessages } = get();
-
-    await messageService.updateMessage(id, { pluginError: error });
-    await refreshMessages();
-  },
-
   invokeBuiltinTool: async (id, payload) => {
-    const { internal_toggleChatLoading, internal_updateMessageContent } = get();
+    const { internal_togglePluginApiCalling, internal_updateMessageContent } = get();
     const params = JSON.parse(payload.arguments);
-    internal_toggleChatLoading(true, id, n('invokeBuiltinTool') as string);
+    internal_togglePluginApiCalling(true, id, n('invokeBuiltinTool') as string);
     let data;
     try {
       data = await useToolStore.getState().invokeBuiltinTool(payload.apiName, params);
     } catch (error) {
       console.log(error);
     }
-    internal_toggleChatLoading(false);
+    internal_togglePluginApiCalling(false);
 
     if (!data) return;
 
@@ -248,6 +147,8 @@ export const chatPlugin: StateCreator<
     const message = chatSelectors.getMessageById(id)(get());
     if (!message || message.role !== 'tool' || !message.plugin) return;
 
+    get().internal_updateMessageError(id, null);
+
     const payload: ChatToolPayload = { ...message.plugin, id: message.tool_call_id! };
 
     await get().internal_invokeDifferentTypePlugin(id, payload);
@@ -301,4 +202,116 @@ export const chatPlugin: StateCreator<
     await messageService.updateMessagePluginState(id, { [key]: value });
     await refreshMessages();
   },
+
+  internal_callPluginApi: async (id, payload) => {
+    const { internal_updateMessageContent, refreshMessages, internal_togglePluginApiCalling } =
+      get();
+    let data: string;
+
+    try {
+      const abortController = internal_togglePluginApiCalling(
+        true,
+        id,
+        n('fetchPlugin/start') as string,
+      );
+
+      const message = chatSelectors.getMessageById(id)(get());
+
+      const res = await chatService.runPluginApi(payload, {
+        signal: abortController?.signal,
+        trace: { observationId: message?.observationId, traceId: message?.traceId },
+      });
+      data = res.text;
+
+      // save traceId
+      if (res.traceId) {
+        await messageService.updateMessage(id, { traceId: res.traceId });
+      }
+    } catch (error) {
+      console.log(error);
+      const err = error as Error;
+
+      // ignore the aborted request error
+      if (!err.message.includes('The user aborted a request.')) {
+        await messageService.updateMessageError(id, error as any);
+        await refreshMessages();
+      }
+
+      data = '';
+    }
+
+    internal_togglePluginApiCalling(false, id, n('fetchPlugin/end') as string);
+    // 如果报错则结束了
+    if (!data) return;
+
+    await internal_updateMessageContent(id, data);
+
+    return data;
+  },
+
+  internal_invokeDifferentTypePlugin: async (id, payload) => {
+    switch (payload.type) {
+      case 'standalone': {
+        return await get().invokeStandaloneTypePlugin(id, payload);
+      }
+
+      case 'markdown': {
+        return await get().invokeMarkdownTypePlugin(id, payload);
+      }
+
+      case 'builtin': {
+        return await get().invokeBuiltinTool(id, payload);
+      }
+
+      default: {
+        return await get().invokeDefaultTypePlugin(id, payload);
+      }
+    }
+  },
+
+  internal_togglePluginApiCalling: (loading, id, action) => {
+    return get().internal_toggleLoadingArrays('pluginApiLoadingIds', loading, id, action);
+  },
+
+  internal_transformToolCalls: (toolCalls) => {
+    return toolCalls
+      .map((toolCall): ChatToolPayload | null => {
+        let payload: ChatToolPayload;
+
+        const [identifier, apiName, type] = toolCall.function.name.split(PLUGIN_SCHEMA_SEPARATOR);
+
+        if (!apiName) return null;
+
+        payload = {
+          apiName,
+          arguments: toolCall.function.arguments,
+          id: toolCall.id,
+          identifier,
+          type: (type ?? 'default') as any,
+        };
+
+        // if the apiName is md5, try to find the correct apiName in the plugins
+        if (apiName.startsWith(PLUGIN_SCHEMA_API_MD5_PREFIX)) {
+          const md5 = apiName.replace(PLUGIN_SCHEMA_API_MD5_PREFIX, '');
+          const manifest = pluginSelectors.getPluginManifestById(identifier)(
+            useToolStore.getState(),
+          );
+
+          const api = manifest?.api.find((api) => Md5.hashStr(api.name).toString() === md5);
+          if (api) {
+            payload.apiName = api.name;
+          }
+        }
+
+        return payload;
+      })
+      .filter(Boolean) as ChatToolPayload[];
+  },
+  internal_updatePluginError: async (id, error) => {
+    const { refreshMessages } = get();
+
+    get().internal_dispatchMessage({ id, type: 'updateMessages', value: { error } });
+    await messageService.updateMessage(id, { pluginError: error });
+    await refreshMessages();
+  },
 });

package/src/store/chat/utils/index.ts

@@ -0,0 +1,19 @@
+import { produce } from 'immer';
+
+export const preventLeavingFn = (e: BeforeUnloadEvent) => {
+  // set returnValue to trigger alert modal
+  // Note: No matter what value is set, the browser will display the standard text
+  e.returnValue = '你有正在生成中的请求，确定要离开吗？';
+};
+
+export const toggleBooleanList = (ids: string[], id: string, loading: boolean) => {
+  return produce(ids, (draft) => {
+    if (loading) {
+      draft.push(id);
+    } else {
+      const index = draft.indexOf(id);
+
+      if (index >= 0) draft.splice(index, 1);
+    }
+  });
+};

package/src/types/message/index.ts

@@ -32,7 +32,7 @@ export * from './tools';
 
 export interface ChatMessage extends BaseDataModel {
   content: string;
-  error?: ChatMessageError;
+  error?: ChatMessageError | null;
   // 扩展字段
   extra?: {
     fromModel?: string;