@lobehub/chat 0.156.2 → 0.157.0

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

@@ -0,0 +1,68 @@
+import {
+  EnhancedGenerateContentResponse,
+  GenerateContentStreamResult,
+} from '@google/generative-ai';
+import { readableFromAsyncIterable } from 'ai';
+
+import { nanoid } from '@/utils/uuid';
+
+import { ChatStreamCallbacks } from '../../types';
+import {
+  StreamProtocolChunk,
+  StreamStack,
+  StreamToolCallChunkData,
+  chatStreamable,
+  createCallbacksTransformer,
+  createSSEProtocolTransformer,
+  generateToolCallId,
+} from './protocol';
+
+const transformGoogleGenerativeAIStream = (
+  chunk: EnhancedGenerateContentResponse,
+  stack: StreamStack,
+): StreamProtocolChunk => {
+  // maybe need another structure to add support for multiple choices
+  const functionCalls = chunk.functionCalls();
+
+  if (functionCalls) {
+    return {
+      data: functionCalls.map(
+        (value, index): StreamToolCallChunkData => ({
+          function: {
+            arguments: JSON.stringify(value.args),
+            name: value.name,
+          },
+          id: generateToolCallId(index, value.name),
+          index: index,
+          type: 'function',
+        }),
+      ),
+      id: stack.id,
+      type: 'tool_calls',
+    };
+  }
+  const text = chunk.text();
+
+  return {
+    data: text,
+    id: stack?.id,
+    type: 'text',
+  };
+};
+
+// only use for debug
+export const googleGenAIResultToStream = (stream: GenerateContentStreamResult) => {
+  // make the response to the streamable format
+  return readableFromAsyncIterable(chatStreamable(stream.stream));
+};
+
+export const GoogleGenerativeAIStream = (
+  rawStream: ReadableStream<EnhancedGenerateContentResponse>,
+  callbacks?: ChatStreamCallbacks,
+) => {
+  const streamStack: StreamStack = { id: 'chat_' + nanoid() };
+
+  return rawStream
+    .pipeThrough(createSSEProtocolTransformer(transformGoogleGenerativeAIStream, streamStack))
+    .pipeThrough(createCallbacksTransformer(callbacks));
+};

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

@@ -0,0 +1,7 @@
+export * from './anthropic';
+export * from './bedrock';
+export * from './google-ai';
+export * from './minimax';
+export * from './ollama';
+export * from './openai';
+export * from './protocol';

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

@@ -0,0 +1,39 @@
+import OpenAI from 'openai';
+
+import { ChatStreamCallbacks } from '../../types';
+import { transformOpenAIStream } from './openai';
+import { createCallbacksTransformer, createSSEProtocolTransformer } from './protocol';
+
+const unit8ArrayToJSONChunk = (unit8Array: Uint8Array): OpenAI.ChatCompletionChunk => {
+  const decoder = new TextDecoder();
+
+  let chunkValue = decoder.decode(unit8Array, { stream: true });
+
+  // chunkValue example:
+  // data: {"id":"028a65377137d57aaceeffddf48ae99f","choices":[{"finish_reason":"tool_calls","index":0,"delta":{"role":"assistant","tool_calls":[{"id":"call_function_7371372822","type":"function","function":{"name":"realtime-weather____fetchCurrentWeather","arguments":"{\"city\": [\"杭州\", \"北京\"]}"}}]}}],"created":155511,"model":"abab6.5s-chat","object":"chat.completion.chunk"}
+
+  // so we need to remove `data:` prefix and then parse it as JSON
+  if (chunkValue.startsWith('data:')) {
+    chunkValue = chunkValue.slice(5).trim();
+  }
+
+  try {
+    return JSON.parse(chunkValue);
+  } catch (e) {
+    console.error('minimax chunk parse error:', e);
+
+    return { raw: chunkValue } as any;
+  }
+};
+
+export const MinimaxStream = (stream: ReadableStream, callbacks?: ChatStreamCallbacks) => {
+  return stream
+    .pipeThrough(
+      createSSEProtocolTransformer((buffer) => {
+        const chunk = unit8ArrayToJSONChunk(buffer);
+
+        return transformOpenAIStream(chunk);
+      }),
+    )
+    .pipeThrough(createCallbacksTransformer(callbacks));
+};

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

@@ -0,0 +1,77 @@
+import { ChatResponse } from 'ollama/browser';
+import { describe, expect, it, vi } from 'vitest';
+
+import * as uuidModule from '@/utils/uuid';
+
+import { OllamaStream } from './ollama';
+
+describe('OllamaStream', () => {
+  it('should transform Ollama stream to protocol stream', async () => {
+    vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1');
+
+    const mockOllamaStream: AsyncIterable<ChatResponse> = {
+      // @ts-ignore
+      async *[Symbol.asyncIterator]() {
+        yield { message: { content: 'Hello' }, done: false };
+        yield { message: { content: ' world!' }, done: false };
+        yield { message: { content: '' }, done: true };
+      },
+    };
+
+    const onStartMock = vi.fn();
+    const onTextMock = vi.fn();
+    const onTokenMock = vi.fn();
+    const onCompletionMock = vi.fn();
+
+    const protocolStream = OllamaStream(mockOllamaStream, {
+      onStart: onStartMock,
+      onText: onTextMock,
+      onToken: onTokenMock,
+      onCompletion: onCompletionMock,
+    });
+
+    const decoder = new TextDecoder();
+    const chunks = [];
+
+    // @ts-ignore
+    for await (const chunk of protocolStream) {
+      chunks.push(decoder.decode(chunk, { stream: true }));
+    }
+
+    expect(chunks).toEqual([
+      'id: chat_1\n',
+      'event: text\n',
+      `data: "Hello"\n\n`,
+      'id: chat_1\n',
+      'event: text\n',
+      `data: " world!"\n\n`,
+      'id: chat_1\n',
+      'event: stop\n',
+      `data: "finished"\n\n`,
+    ]);
+
+    expect(onStartMock).toHaveBeenCalledTimes(1);
+    expect(onTextMock).toHaveBeenNthCalledWith(1, '"Hello"');
+    expect(onTextMock).toHaveBeenNthCalledWith(2, '" world!"');
+    expect(onTokenMock).toHaveBeenCalledTimes(2);
+    expect(onCompletionMock).toHaveBeenCalledTimes(1);
+  });
+
+  it('should handle empty stream', async () => {
+    const mockOllamaStream = {
+      async *[Symbol.asyncIterator]() {},
+    };
+
+    const protocolStream = OllamaStream(mockOllamaStream);
+
+    const decoder = new TextDecoder();
+    const chunks = [];
+
+    // @ts-ignore
+    for await (const chunk of protocolStream) {
+      chunks.push(decoder.decode(chunk, { stream: true }));
+    }
+
+    expect(chunks).toEqual([]);
+  });
+});

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

@@ -0,0 +1,38 @@
+import { readableFromAsyncIterable } from 'ai';
+import { ChatResponse } from 'ollama/browser';
+
+import { ChatStreamCallbacks } from '@/libs/agent-runtime';
+import { nanoid } from '@/utils/uuid';
+
+import {
+  StreamProtocolChunk,
+  StreamStack,
+  createCallbacksTransformer,
+  createSSEProtocolTransformer,
+} from './protocol';
+
+const transformOllamaStream = (chunk: ChatResponse, stack: StreamStack): StreamProtocolChunk => {
+  // maybe need another structure to add support for multiple choices
+  if (chunk.done) {
+    return { data: 'finished', id: stack.id, type: 'stop' };
+  }
+
+  return { data: chunk.message.content, id: stack.id, type: 'text' };
+};
+
+const chatStreamable = async function* (stream: AsyncIterable<ChatResponse>) {
+  for await (const response of stream) {
+    yield response;
+  }
+};
+
+export const OllamaStream = (
+  res: AsyncIterable<ChatResponse>,
+  cb?: ChatStreamCallbacks,
+): ReadableStream<string> => {
+  const streamStack: StreamStack = { id: 'chat_' + nanoid() };
+
+  return readableFromAsyncIterable(chatStreamable(res))
+    .pipeThrough(createSSEProtocolTransformer(transformOllamaStream, streamStack))
+    .pipeThrough(createCallbacksTransformer(cb));
+};

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

@@ -0,0 +1,263 @@
+import { describe, expect, it, vi } from 'vitest';
+
+import { OpenAIStream } from './openai';
+
+describe('OpenAIStream', () => {
+  it('should transform OpenAI stream to protocol stream', async () => {
+    const mockOpenAIStream = new ReadableStream({
+      start(controller) {
+        controller.enqueue({
+          choices: [
+            {
+              delta: { content: 'Hello' },
+              index: 0,
+            },
+          ],
+          id: '1',
+        });
+        controller.enqueue({
+          choices: [
+            {
+              delta: { content: ' world!' },
+              index: 1,
+            },
+          ],
+          id: '1',
+        });
+        controller.enqueue({
+          choices: [
+            {
+              delta: null,
+              finish_reason: 'stop',
+              index: 2,
+            },
+          ],
+          id: '1',
+        });
+
+        controller.close();
+      },
+    });
+
+    const onStartMock = vi.fn();
+    const onTextMock = vi.fn();
+    const onTokenMock = vi.fn();
+    const onCompletionMock = vi.fn();
+
+    const protocolStream = OpenAIStream(mockOpenAIStream, {
+      onStart: onStartMock,
+      onText: onTextMock,
+      onToken: onTokenMock,
+      onCompletion: onCompletionMock,
+    });
+
+    const decoder = new TextDecoder();
+    const chunks = [];
+
+    // @ts-ignore
+    for await (const chunk of protocolStream) {
+      chunks.push(decoder.decode(chunk, { stream: true }));
+    }
+
+    expect(chunks).toEqual([
+      'id: 1\n',
+      'event: text\n',
+      `data: "Hello"\n\n`,
+      'id: 1\n',
+      'event: text\n',
+      `data: " world!"\n\n`,
+      'id: 1\n',
+      'event: stop\n',
+      `data: "stop"\n\n`,
+    ]);
+
+    expect(onStartMock).toHaveBeenCalledTimes(1);
+    expect(onTextMock).toHaveBeenNthCalledWith(1, '"Hello"');
+    expect(onTextMock).toHaveBeenNthCalledWith(2, '" world!"');
+    expect(onTokenMock).toHaveBeenCalledTimes(2);
+    expect(onCompletionMock).toHaveBeenCalledTimes(1);
+  });
+
+  it('should handle tool calls', async () => {
+    const mockOpenAIStream = new ReadableStream({
+      start(controller) {
+        controller.enqueue({
+          choices: [
+            {
+              delta: {
+                tool_calls: [
+                  {
+                    function: { name: 'tool1', arguments: '{}' },
+                    id: 'call_1',
+                    index: 0,
+                    type: 'function',
+                  },
+                  {
+                    function: { name: 'tool2', arguments: '{}' },
+                    id: 'call_2',
+                    index: 1,
+                  },
+                ],
+              },
+              index: 0,
+            },
+          ],
+          id: '2',
+        });
+
+        controller.close();
+      },
+    });
+
+    const onToolCallMock = vi.fn();
+
+    const protocolStream = OpenAIStream(mockOpenAIStream, {
+      onToolCall: onToolCallMock,
+    });
+
+    const decoder = new TextDecoder();
+    const chunks = [];
+
+    // @ts-ignore
+    for await (const chunk of protocolStream) {
+      chunks.push(decoder.decode(chunk, { stream: true }));
+    }
+
+    expect(chunks).toEqual([
+      'id: 2\n',
+      'event: tool_calls\n',
+      `data: [{"function":{"name":"tool1","arguments":"{}"},"id":"call_1","index":0,"type":"function"},{"function":{"name":"tool2","arguments":"{}"},"id":"call_2","index":1,"type":"function"}]\n\n`,
+    ]);
+
+    expect(onToolCallMock).toHaveBeenCalledTimes(1);
+  });
+
+  it('should handle empty stream', async () => {
+    const mockStream = new ReadableStream({
+      start(controller) {
+        controller.close();
+      },
+    });
+
+    const protocolStream = OpenAIStream(mockStream);
+
+    const decoder = new TextDecoder();
+    const chunks = [];
+
+    // @ts-ignore
+    for await (const chunk of protocolStream) {
+      chunks.push(decoder.decode(chunk, { stream: true }));
+    }
+
+    expect(chunks).toEqual([]);
+  });
+
+  it('should handle delta content null', async () => {
+    const mockOpenAIStream = new ReadableStream({
+      start(controller) {
+        controller.enqueue({
+          choices: [
+            {
+              delta: { content: null },
+              index: 0,
+            },
+          ],
+          id: '3',
+        });
+
+        controller.close();
+      },
+    });
+
+    const protocolStream = OpenAIStream(mockOpenAIStream);
+
+    const decoder = new TextDecoder();
+    const chunks = [];
+
+    // @ts-ignore
+    for await (const chunk of protocolStream) {
+      chunks.push(decoder.decode(chunk, { stream: true }));
+    }
+
+    expect(chunks).toEqual(['id: 3\n', 'event: data\n', `data: {"content":null}\n\n`]);
+  });
+
+  it('should handle other delta data', async () => {
+    const mockOpenAIStream = new ReadableStream({
+      start(controller) {
+        controller.enqueue({
+          choices: [
+            {
+              delta: { custom_field: 'custom_value' },
+              index: 0,
+            },
+          ],
+          id: '4',
+        });
+
+        controller.close();
+      },
+    });
+
+    const protocolStream = OpenAIStream(mockOpenAIStream);
+
+    const decoder = new TextDecoder();
+    const chunks = [];
+
+    // @ts-ignore
+    for await (const chunk of protocolStream) {
+      chunks.push(decoder.decode(chunk, { stream: true }));
+    }
+
+    expect(chunks).toEqual([
+      'id: 4\n',
+      'event: data\n',
+      `data: {"delta":{"custom_field":"custom_value"},"id":"4","index":0}\n\n`,
+    ]);
+  });
+
+  it('should handle tool calls without index and type', async () => {
+    const mockOpenAIStream = new ReadableStream({
+      start(controller) {
+        controller.enqueue({
+          choices: [
+            {
+              delta: {
+                tool_calls: [
+                  {
+                    function: { name: 'tool1', arguments: '{}' },
+                    id: 'call_1',
+                  },
+                  {
+                    function: { name: 'tool2', arguments: '{}' },
+                    id: 'call_2',
+                  },
+                ],
+              },
+              index: 0,
+            },
+          ],
+          id: '5',
+        });
+
+        controller.close();
+      },
+    });
+
+    const protocolStream = OpenAIStream(mockOpenAIStream);
+
+    const decoder = new TextDecoder();
+    const chunks = [];
+
+    // @ts-ignore
+    for await (const chunk of protocolStream) {
+      chunks.push(decoder.decode(chunk, { stream: true }));
+    }
+
+    expect(chunks).toEqual([
+      'id: 5\n',
+      'event: tool_calls\n',
+      `data: [{"function":{"name":"tool1","arguments":"{}"},"id":"call_1","index":0,"type":"function"},{"function":{"name":"tool2","arguments":"{}"},"id":"call_2","index":1,"type":"function"}]\n\n`,
+    ]);
+  });
+});

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

@@ -0,0 +1,79 @@
+import { readableFromAsyncIterable } from 'ai';
+import OpenAI from 'openai';
+import type { Stream } from 'openai/streaming';
+
+import { ChatStreamCallbacks } from '../../types';
+import {
+  StreamProtocolChunk,
+  StreamProtocolToolCallChunk,
+  StreamToolCallChunkData,
+  createCallbacksTransformer,
+  createSSEProtocolTransformer,
+  generateToolCallId,
+} from './protocol';
+
+export const transformOpenAIStream = (chunk: OpenAI.ChatCompletionChunk): StreamProtocolChunk => {
+  // maybe need another structure to add support for multiple choices
+  const item = chunk.choices[0];
+
+  if (typeof item.delta?.content === 'string') {
+    return { data: item.delta.content, id: chunk.id, type: 'text' };
+  }
+
+  if (item.delta?.tool_calls) {
+    return {
+      data: item.delta.tool_calls.map(
+        (value, index): StreamToolCallChunkData => ({
+          function: value.function,
+          id: value.id || generateToolCallId(index, value.function?.name),
+
+          // mistral's tool calling don't have index and function field, it's data like:
+          // [{"id":"xbhnmTtY7","function":{"name":"lobe-image-designer____text2image____builtin","arguments":"{\"prompts\": [\"A photo of a small, fluffy dog with a playful expression and wagging tail.\", \"A watercolor painting of a small, energetic dog with a glossy coat and bright eyes.\", \"A vector illustration of a small, adorable dog with a short snout and perky ears.\", \"A drawing of a small, scruffy dog with a mischievous grin and a wagging tail.\"], \"quality\": \"standard\", \"seeds\": [123456, 654321, 111222, 333444], \"size\": \"1024x1024\", \"style\": \"vivid\"}"}}]
+
+          // minimax's tool calling don't have index field, it's data like:
+          // [{"id":"call_function_4752059746","type":"function","function":{"name":"lobe-image-designer____text2image____builtin","arguments":"{\"prompts\": [\"一个流浪的地球，背景是浩瀚"}}]
+
+          // so we need to add these default values
+          index: typeof value.index !== 'undefined' ? value.index : index,
+          type: value.type || 'function',
+        }),
+      ),
+      id: chunk.id,
+      type: 'tool_calls',
+    } as StreamProtocolToolCallChunk;
+  }
+
+  // 给定结束原因
+  if (item.finish_reason) {
+    return { data: item.finish_reason, id: chunk.id, type: 'stop' };
+  }
+
+  if (item.delta.content === null) {
+    return { data: item.delta, id: chunk.id, type: 'data' };
+  }
+
+  // 其余情况下，返回 delta 和 index
+  return {
+    data: { delta: item.delta, id: chunk.id, index: item.index },
+    id: chunk.id,
+    type: 'data',
+  };
+};
+
+const chatStreamable = async function* (stream: AsyncIterable<OpenAI.ChatCompletionChunk>) {
+  for await (const response of stream) {
+    yield response;
+  }
+};
+
+export const OpenAIStream = (
+  stream: Stream<OpenAI.ChatCompletionChunk> | ReadableStream,
+  callbacks?: ChatStreamCallbacks,
+) => {
+  const readableStream =
+    stream instanceof ReadableStream ? stream : readableFromAsyncIterable(chatStreamable(stream));
+
+  return readableStream
+    .pipeThrough(createSSEProtocolTransformer(transformOpenAIStream))
+    .pipeThrough(createCallbacksTransformer(callbacks));
+};

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

@@ -0,0 +1,100 @@
+import { ChatStreamCallbacks } from '@/libs/agent-runtime';
+
+export interface StreamStack {
+  id: string;
+}
+
+export interface StreamProtocolChunk {
+  data: any;
+  id?: string;
+  type: 'text' | 'tool_calls' | 'data' | 'stop';
+}
+
+export interface StreamToolCallChunkData {
+  function?: {
+    arguments?: string;
+    name?: string | null;
+  };
+  id: string;
+  index: number;
+  type: 'function' | string;
+}
+
+export interface StreamProtocolToolCallChunk {
+  data: StreamToolCallChunkData[];
+  id: string;
+  index: number;
+  type: 'tool_calls';
+}
+
+export const generateToolCallId = (index: number, functionName?: string) =>
+  `${functionName || 'unknown_tool_call'}_${index}`;
+
+export const chatStreamable = async function* <T>(stream: AsyncIterable<T>) {
+  for await (const response of stream) {
+    yield response;
+  }
+};
+
+export const createSSEProtocolTransformer = (
+  transformer: (chunk: any, stack: StreamStack) => StreamProtocolChunk,
+  streamStack?: StreamStack,
+) =>
+  new TransformStream({
+    transform: (chunk, controller) => {
+      const { type, id, data } = transformer(chunk, streamStack || { id: '' });
+
+      controller.enqueue(`id: ${id}\n`);
+      controller.enqueue(`event: ${type}\n`);
+      controller.enqueue(`data: ${JSON.stringify(data)}\n\n`);
+    },
+  });
+
+export function createCallbacksTransformer(cb: ChatStreamCallbacks | undefined) {
+  const textEncoder = new TextEncoder();
+  let aggregatedResponse = '';
+  let currentType = '';
+  const callbacks = cb || {};
+
+  return new TransformStream({
+    async flush(): Promise<void> {
+      if (callbacks.onCompletion) {
+        await callbacks.onCompletion(aggregatedResponse);
+      }
+
+      if (callbacks.onFinal) {
+        await callbacks.onFinal(aggregatedResponse);
+      }
+    },
+
+    async start(): Promise<void> {
+      if (callbacks.onStart) await callbacks.onStart();
+    },
+
+    async transform(chunk: string, controller): Promise<void> {
+      controller.enqueue(textEncoder.encode(chunk));
+
+      // track the type of the chunk
+      if (chunk.startsWith('event:')) {
+        currentType = chunk.split('event:')[1].trim();
+      }
+      // if the message is a data chunk, handle the callback
+      else if (chunk.startsWith('data:')) {
+        const content = chunk.split('data:')[1].trim();
+
+        switch (currentType) {
+          case 'text': {
+            await callbacks.onText?.(content);
+            await callbacks.onToken?.(JSON.parse(content));
+            break;
+          }
+
+          case 'tool_calls': {
+            // TODO: make on ToolCall callback
+            await callbacks.onToolCall?.();
+          }
+        }
+      }
+    },
+  });
+}