@lobehub/chat 0.156.2 → 0.157.0

package/src/utils/fetch.test.ts

@@ -1,8 +1,10 @@
+import { fetchEventSource } from '@microsoft/fetch-event-source';
+import { FetchEventSourceInit } from '@microsoft/fetch-event-source';
 import { afterEach, describe, expect, it, vi } from 'vitest';
 
 import { ErrorResponse } from '@/types/fetch';
 
-import { getMessageError } from './fetch';
+import { fetchSSE, getMessageError, parseToolCalls } from './fetch';
 
 // 模拟 i18next
 vi.mock('i18next', () => ({
@@ -39,6 +41,10 @@ const createMockResponse = (body: any, ok: boolean, status: number = 200) => ({
   },
 });
 
+vi.mock('@microsoft/fetch-event-source', () => ({
+  fetchEventSource: vi.fn(),
+}));
+
 // 在每次测试后清理所有模拟
 afterEach(() => {
   vi.restoreAllMocks();
@@ -77,3 +83,241 @@ describe('getMessageError', () => {
     expect(mockResponse.json).toHaveBeenCalled();
   });
 });
+
+describe('parseToolCalls', () => {
+  it('should create add new item', () => {
+    const chunk = [
+      { index: 0, id: '1', type: 'function', function: { name: 'func', arguments: '' } },
+    ];
+
+    const result = parseToolCalls([], chunk);
+    expect(result).toEqual([
+      { id: '1', type: 'function', function: { name: 'func', arguments: '' } },
+    ]);
+  });
+
+  it('should update arguments if there is a toolCall', () => {
+    const origin = [{ id: '1', type: 'function', function: { name: 'func', arguments: '' } }];
+
+    const chunk1 = [{ index: 0, function: { arguments: '{"lo' } }];
+
+    const result1 = parseToolCalls(origin, chunk1);
+    expect(result1).toEqual([
+      { id: '1', type: 'function', function: { name: 'func', arguments: '{"lo' } },
+    ]);
+
+    const chunk2 = [{ index: 0, function: { arguments: 'cation\\": \\"Hangzhou\\"}' } }];
+    const result2 = parseToolCalls(result1, chunk2);
+
+    expect(result2).toEqual([
+      {
+        id: '1',
+        type: 'function',
+        function: { name: 'func', arguments: '{"location\\": \\"Hangzhou\\"}' },
+      },
+    ]);
+  });
+
+  it('should add a new tool call if the index is different', () => {
+    const origin = [
+      {
+        id: '1',
+        type: 'function',
+        function: { name: 'func', arguments: '{"location\\": \\"Hangzhou\\"}' },
+      },
+    ];
+
+    const chunk = [
+      {
+        index: 1,
+        id: '2',
+        type: 'function',
+        function: { name: 'func', arguments: '' },
+      },
+    ];
+
+    const result1 = parseToolCalls(origin, chunk);
+    expect(result1).toEqual([
+      {
+        id: '1',
+        type: 'function',
+        function: { name: 'func', arguments: '{"location\\": \\"Hangzhou\\"}' },
+      },
+      { id: '2', type: 'function', function: { name: 'func', arguments: '' } },
+    ]);
+  });
+
+  it('should update correct arguments if there are multi tool calls', () => {
+    const origin = [
+      {
+        id: '1',
+        type: 'function',
+        function: { name: 'func', arguments: '{"location\\": \\"Hangzhou\\"}' },
+      },
+      { id: '2', type: 'function', function: { name: 'func', arguments: '' } },
+    ];
+
+    const chunk = [{ index: 1, function: { arguments: '{"location\\": \\"Beijing\\"}' } }];
+
+    const result1 = parseToolCalls(origin, chunk);
+    expect(result1).toEqual([
+      {
+        id: '1',
+        type: 'function',
+        function: { name: 'func', arguments: '{"location\\": \\"Hangzhou\\"}' },
+      },
+      {
+        id: '2',
+        type: 'function',
+        function: { name: 'func', arguments: '{"location\\": \\"Beijing\\"}' },
+      },
+    ]);
+  });
+});
+
+describe('fetchSSE', () => {
+  it('should handle text event correctly', async () => {
+    const mockOnMessageHandle = vi.fn();
+    const mockOnFinish = vi.fn();
+
+    (fetchEventSource as any).mockImplementationOnce(
+      (url: string, options: FetchEventSourceInit) => {
+        options.onopen!({ clone: () => ({ ok: true, headers: new Headers() }) } as any);
+        options.onmessage!({ event: 'text', data: JSON.stringify('Hello') } as any);
+        options.onmessage!({ event: 'text', data: JSON.stringify(' World') } as any);
+      },
+    );
+
+    await fetchSSE('/', { onMessageHandle: mockOnMessageHandle, onFinish: mockOnFinish });
+
+    expect(mockOnMessageHandle).toHaveBeenNthCalledWith(1, { text: 'Hello', type: 'text' });
+    expect(mockOnMessageHandle).toHaveBeenNthCalledWith(2, { text: ' World', type: 'text' });
+    expect(mockOnFinish).toHaveBeenCalledWith('Hello World', {
+      observationId: null,
+      toolCalls: undefined,
+      traceId: null,
+      type: 'done',
+    });
+  });
+
+  it('should handle tool_calls event correctly', async () => {
+    const mockOnMessageHandle = vi.fn();
+    const mockOnFinish = vi.fn();
+
+    (fetchEventSource as any).mockImplementationOnce(
+      (url: string, options: FetchEventSourceInit) => {
+        options.onopen!({ clone: () => ({ ok: true, headers: new Headers() }) } as any);
+        options.onmessage!({
+          event: 'tool_calls',
+          data: JSON.stringify([
+            { index: 0, id: '1', type: 'function', function: { name: 'func1', arguments: 'arg1' } },
+          ]),
+        } as any);
+        options.onmessage!({
+          event: 'tool_calls',
+          data: JSON.stringify([
+            { index: 1, id: '2', type: 'function', function: { name: 'func2', arguments: 'arg2' } },
+          ]),
+        } as any);
+      },
+    );
+
+    await fetchSSE('/', { onMessageHandle: mockOnMessageHandle, onFinish: mockOnFinish });
+
+    expect(mockOnMessageHandle).toHaveBeenNthCalledWith(1, {
+      tool_calls: [{ id: '1', type: 'function', function: { name: 'func1', arguments: 'arg1' } }],
+      type: 'tool_calls',
+    });
+    expect(mockOnMessageHandle).toHaveBeenNthCalledWith(2, {
+      tool_calls: [
+        { id: '1', type: 'function', function: { name: 'func1', arguments: 'arg1' } },
+        { id: '2', type: 'function', function: { name: 'func2', arguments: 'arg2' } },
+      ],
+      type: 'tool_calls',
+    });
+    expect(mockOnFinish).toHaveBeenCalledWith('', {
+      observationId: null,
+      toolCalls: [
+        { id: '1', type: 'function', function: { name: 'func1', arguments: 'arg1' } },
+        { id: '2', type: 'function', function: { name: 'func2', arguments: 'arg2' } },
+      ],
+      traceId: null,
+      type: 'done',
+    });
+  });
+
+  it('should call onAbort when AbortError is thrown', async () => {
+    const mockOnAbort = vi.fn();
+
+    (fetchEventSource as any).mockImplementationOnce(
+      (url: string, options: FetchEventSourceInit) => {
+        options.onmessage!({ event: 'text', data: JSON.stringify('Hello') } as any);
+        options.onerror!({ name: 'AbortError' });
+      },
+    );
+
+    await fetchSSE('/', { onAbort: mockOnAbort });
+
+    expect(mockOnAbort).toHaveBeenCalledWith('Hello');
+  });
+
+  it('should call onErrorHandle when other error is thrown', async () => {
+    const mockOnErrorHandle = vi.fn();
+    const mockError = new Error('Unknown error');
+
+    (fetchEventSource as any).mockImplementationOnce(
+      (url: string, options: FetchEventSourceInit) => {
+        options.onerror!(mockError);
+      },
+    );
+
+    await fetchSSE('/', { onErrorHandle: mockOnErrorHandle });
+
+    expect(mockOnErrorHandle).not.toHaveBeenCalled();
+  });
+
+  it('should call onErrorHandle when response is not ok', async () => {
+    const mockOnErrorHandle = vi.fn();
+
+    (fetchEventSource as any).mockImplementationOnce(
+      (url: string, options: FetchEventSourceInit) => {
+        const res = new Response(JSON.stringify({ errorType: 'SomeError' }), {
+          status: 400,
+          statusText: 'Error',
+        });
+
+        options.onopen!(res as any);
+      },
+    );
+
+    await fetchSSE('/', { onErrorHandle: mockOnErrorHandle });
+
+    expect(mockOnErrorHandle).toHaveBeenCalledWith({
+      body: undefined,
+      message: 'translated_response.SomeError',
+      type: 'SomeError',
+    });
+  });
+
+  it('should call onMessageHandle with full text if no message event', async () => {
+    const mockOnMessageHandle = vi.fn();
+    const mockOnFinish = vi.fn();
+
+    (fetchEventSource as any).mockImplementationOnce(
+      (url: string, options: FetchEventSourceInit) => {
+        const res = new Response('Hello World', { status: 200, statusText: 'OK' });
+        options.onopen!(res as any);
+      },
+    );
+
+    await fetchSSE('/', { onMessageHandle: mockOnMessageHandle, onFinish: mockOnFinish });
+
+    expect(mockOnMessageHandle).toHaveBeenCalledWith({ text: 'Hello World', type: 'text' });
+    expect(mockOnFinish).toHaveBeenCalledWith('Hello World', {
+      observationId: null,
+      toolCalls: undefined,
+      traceId: null,
+      type: 'done',
+    });
+  });
+});

package/src/utils/fetch.ts

@@ -1,8 +1,15 @@
+import { fetchEventSource } from '@microsoft/fetch-event-source';
 import { t } from 'i18next';
+import { produce } from 'immer';
 
 import { LOBE_CHAT_OBSERVATION_ID, LOBE_CHAT_TRACE_ID } from '@/const/trace';
 import { ErrorResponse, ErrorType } from '@/types/fetch';
-import { ChatMessageError } from '@/types/message';
+import {
+  ChatMessageError,
+  MessageToolCall,
+  MessageToolCallChunk,
+  MessageToolCallSchema,
+} from '@/types/message';
 
 export const getMessageError = async (response: Response) => {
   let chatMessageError: ChatMessageError;
@@ -32,68 +39,137 @@ export type OnFinishHandler = (
   text: string,
   context: {
     observationId?: string | null;
+    toolCalls?: MessageToolCall[];
     traceId?: string | null;
     type?: SSEFinishType;
   },
 ) => Promise<void>;
 
+export interface MessageTextChunk {
+  text: string;
+  type: 'text';
+}
+
+interface MessageToolCallsChunk {
+  tool_calls: MessageToolCall[];
+  type: 'tool_calls';
+}
+
 export interface FetchSSEOptions {
+  fetcher?: typeof fetch;
   onAbort?: (text: string) => Promise<void>;
   onErrorHandle?: (error: ChatMessageError) => void;
   onFinish?: OnFinishHandler;
-  onMessageHandle?: (text: string) => void;
+  onMessageHandle?: (chunk: MessageTextChunk | MessageToolCallsChunk) => void;
 }
 
+export const parseToolCalls = (origin: MessageToolCall[], value: MessageToolCallChunk[]) =>
+  produce(origin, (draft) => {
+    if (draft.length === 0) {
+      draft.push(...value.map((item) => MessageToolCallSchema.parse(item)));
+    } else {
+      value.forEach(({ index, ...item }) => {
+        if (!draft?.[index]) {
+          draft?.splice(index, 0, MessageToolCallSchema.parse(item));
+        } else {
+          if (item.function?.arguments) {
+            draft[index].function.arguments += item.function.arguments;
+          }
+        }
+      });
+    }
+  });
+
 /**
  * Fetch data using stream method
  */
-export const fetchSSE = async (fetchFn: () => Promise<Response>, options: FetchSSEOptions = {}) => {
-  const response = await fetchFn();
-
-  // 如果不 ok 说明有请求错误
-  if (!response.ok) {
-    const chatMessageError = await getMessageError(response);
-
-    options.onErrorHandle?.(chatMessageError);
-    return;
-  }
-
-  const returnRes = response.clone();
-
-  const data = response.body;
-
-  if (!data) return;
+// eslint-disable-next-line no-undef
+export const fetchSSE = async (url: string, options: RequestInit & FetchSSEOptions = {}) => {
   let output = '';
-  const reader = data.getReader();
-  const decoder = new TextDecoder();
+  let toolCalls: undefined | MessageToolCall[];
+  let triggerOnMessageHandler = false;
 
-  let done = false;
   let finishedType: SSEFinishType = 'done';
+  let response!: Response;
 
-  while (!done) {
-    try {
-      const { value, done: doneReading } = await reader.read();
-      done = doneReading;
-      const chunkValue = decoder.decode(value, { stream: true });
-
-      output += chunkValue;
-      options.onMessageHandle?.(chunkValue);
-    } catch (error) {
-      done = true;
-
-      if ((error as TypeError).name === 'AbortError') {
-        finishedType = 'abort';
-        options?.onAbort?.(output);
-      } else {
-        finishedType = 'error';
-        console.error(error);
-      }
+  try {
+    await fetchEventSource(url, {
+      body: options.body,
+      fetch: options?.fetcher,
+      headers: options.headers as Record<string, string>,
+      method: options.method,
+      onerror: (error) => {
+        if ((error as TypeError).name === 'AbortError') {
+          finishedType = 'abort';
+          options?.onAbort?.(output);
+        } else {
+          finishedType = 'error';
+          console.error(error);
+        }
+        throw new Error('Fetch error');
+        // options.onErrorHandle()
+      },
+      onmessage: (ev) => {
+        triggerOnMessageHandler = true;
+        let data;
+        try {
+          data = JSON.parse(ev.data);
+        } catch (e) {
+          console.warn('parse error, fallback to stream', e);
+          options.onMessageHandle?.({ text: data, type: 'text' });
+          return;
+        }
+
+        switch (ev.event) {
+          case 'text': {
+            output += data;
+            options.onMessageHandle?.({ text: data, type: 'text' });
+            break;
+          }
+
+          case 'tool_calls': {
+            if (!toolCalls) {
+              toolCalls = [];
+            }
+
+            toolCalls = parseToolCalls(toolCalls, data);
+
+            options.onMessageHandle?.({
+              tool_calls: toolCalls,
+              type: 'tool_calls',
+            });
+          }
+        }
+      },
+      onopen: async (res) => {
+        response = res.clone();
+
+        // 如果不 ok 说明有请求错误
+        if (!response.ok) {
+          const chatMessageError = await getMessageError(res);
+
+          options.onErrorHandle?.(chatMessageError);
+          return;
+        }
+      },
+
+      signal: options.signal,
+    });
+  } catch {}
+
+  // only call onFinish when response is available
+  // so like abort, we don't need to call onFinish
+  if (response) {
+    // if there is no onMessageHandler, we should call onHandleMessage first
+    if (!triggerOnMessageHandler) {
+      output = await response.clone().text();
+      options.onMessageHandle?.({ text: output, type: 'text' });
     }
-  }
 
-  const traceId = response.headers.get(LOBE_CHAT_TRACE_ID);
-  const observationId = response.headers.get(LOBE_CHAT_OBSERVATION_ID);
-  await options?.onFinish?.(output, { observationId, traceId, type: finishedType });
+    const traceId = response.headers.get(LOBE_CHAT_TRACE_ID);
+    const observationId = response.headers.get(LOBE_CHAT_OBSERVATION_ID);
+    await options?.onFinish?.(output, { observationId, toolCalls, traceId, type: finishedType });
+  }
 
-  return returnRes;
+  return response;
 };

package/src/utils/toolCall.ts

@@ -0,0 +1,21 @@
+import { Md5 } from 'ts-md5';
+
+import { PLUGIN_SCHEMA_API_MD5_PREFIX, PLUGIN_SCHEMA_SEPARATOR } from '@/const/plugin';
+
+export const genToolCallingName = (identifier: string, name: string, type: string = 'default') => {
+  const pluginType = type && type !== 'default' ? `${PLUGIN_SCHEMA_SEPARATOR + type}` : '';
+
+  // 将插件的 identifier 作为前缀，避免重复
+  let apiName = identifier + PLUGIN_SCHEMA_SEPARATOR + name + pluginType;
+
+  // OpenAI GPT function_call name can't be longer than 64 characters
+  // So we need to use md5 to shorten the name
+  // and then find the correct apiName in response by md5
+  if (apiName.length >= 64) {
+    const md5Content = PLUGIN_SCHEMA_API_MD5_PREFIX + Md5.hashStr(name).toString();
+
+    apiName = identifier + PLUGIN_SCHEMA_SEPARATOR + md5Content + pluginType;
+  }
+
+  return apiName;
+};

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

@@ -1,26 +0,0 @@
-import { ReactNode, memo } from 'react';
-
-import { isFunctionMessageAtStart } from '@/const/message';
-import { useChatStore } from '@/store/chat';
-import { chatSelectors } from '@/store/chat/selectors';
-import { ChatMessage } from '@/types/message';
-
-import Inspector from '../Plugins/Inspector';
-import { DefaultMessage } from './Default';
-
-export const AssistantMessage = memo<
-  ChatMessage & {
-    editableContent: ReactNode;
-  }
->(({ id, plugin, content, ...props }) => {
-  const fcProps = useChatStore(chatSelectors.getFunctionMessageProps({ content, id, plugin }));
-
-  if (!isFunctionMessageAtStart(content))
-    return <DefaultMessage content={content} id={id} {...props} />;
-
-  return (
-    <div id={id}>
-      <Inspector {...fcProps} />
-    </div>
-  );
-});

package/src/features/Conversation/Messages/Function.tsx

@@ -1,35 +0,0 @@
-import isEqual from 'fast-deep-equal';
-import { memo, useState } from 'react';
-import { Flexbox } from 'react-layout-kit';
-
-import { useChatStore } from '@/store/chat';
-import { chatSelectors } from '@/store/chat/selectors';
-import { ChatMessage } from '@/types/message';
-
-import Inspector from '../Plugins/Inspector';
-import PluginRender from '../Plugins/Render';
-
-export const FunctionMessage = memo<ChatMessage>(({ id, content, plugin }) => {
-  const fcProps = useChatStore(
-    chatSelectors.getFunctionMessageProps({ content, id, plugin }),
-    isEqual,
-  );
-
-  const [showRender, setShow] = useState(true);
-
-  return (
-    <Flexbox gap={12} id={id} width={'100%'}>
-      <Inspector showRender={showRender} {...fcProps} setShow={setShow} />
-      {showRender && (
-        <PluginRender
-          content={content}
-          id={id}
-          identifier={plugin?.identifier}
-          loading={fcProps.loading}
-          payload={fcProps.command}
-          type={fcProps.type}
-        />
-      )}
-    </Flexbox>
-  );
-});

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

@@ -1,31 +0,0 @@
-// copy from https://github.com/vercel/ai/discussions/539#discussioncomment-8193721
-// and I have remove the unnecessary code
-import {
-  type AIStreamCallbacksAndOptions,
-  createCallbacksTransformer,
-  createStreamDataTransformer,
-  readableFromAsyncIterable,
-} from 'ai';
-import { ChatResponse } from 'ollama/browser';
-
-// A modified version of the streamable function specifically for chat messages
-const chatStreamable = async function* (stream: AsyncIterable<ChatResponse>) {
-  for await (const response of stream) {
-    if (response.message) {
-      yield response.message;
-    }
-    if (response.done) {
-      // Additional final response data can be handled here if necessary
-      return;
-    }
-  }
-};
-
-export const OllamaStream = (
-  res: AsyncIterable<ChatResponse>,
-  cb?: AIStreamCallbacksAndOptions,
-): ReadableStream<string> => {
-  return readableFromAsyncIterable(chatStreamable(res))
-    .pipeThrough(createCallbacksTransformer(cb) as any)
-    .pipeThrough(createStreamDataTransformer(cb?.experimental_streamData));
-};