@lobehub/chat 0.156.1 → 0.157.0

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

@@ -0,0 +1,197 @@
+import type { Stream } from '@anthropic-ai/sdk/streaming';
+import { describe, expect, it, vi } from 'vitest';
+
+import { AnthropicStream } from './anthropic';
+
+describe('AnthropicStream', () => {
+  it('should transform Anthropic stream to protocol stream', async () => {
+    // @ts-ignore
+    const mockAnthropicStream: Stream = {
+      [Symbol.asyncIterator]() {
+        let count = 0;
+        return {
+          next: async () => {
+            switch (count) {
+              case 0:
+                count++;
+                return {
+                  done: false,
+                  value: {
+                    type: 'message_start',
+                    message: { id: 'message_1', metadata: {} },
+                  },
+                };
+              case 1:
+                count++;
+                return {
+                  done: false,
+                  value: {
+                    type: 'content_block_delta',
+                    delta: { type: 'text_delta', text: 'Hello' },
+                  },
+                };
+              case 2:
+                count++;
+                return {
+                  done: false,
+                  value: {
+                    type: 'content_block_delta',
+                    delta: { type: 'text_delta', text: ' world!' },
+                  },
+                };
+              case 3:
+                count++;
+                return {
+                  done: false,
+                  value: {
+                    type: 'message_delta',
+                    delta: { stop_reason: 'stop' },
+                  },
+                };
+              default:
+                return { done: true, value: undefined };
+            }
+          },
+        };
+      },
+    };
+
+    const onStartMock = vi.fn();
+    const onTextMock = vi.fn();
+    const onTokenMock = vi.fn();
+    const onCompletionMock = vi.fn();
+
+    const protocolStream = AnthropicStream(mockAnthropicStream, {
+      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: message_1\n',
+      'event: data\n',
+      `data: {"id":"message_1","metadata":{}}\n\n`,
+      'id: message_1\n',
+      'event: text\n',
+      `data: "Hello"\n\n`,
+      'id: message_1\n',
+      'event: text\n',
+      `data: " world!"\n\n`,
+      'id: message_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 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 mockReadableStream = new ReadableStream({
+      start(controller) {
+        controller.enqueue({
+          type: 'message_start',
+          message: { id: 'message_1', metadata: {} },
+        });
+        controller.enqueue(toolUseEvent);
+        controller.enqueue({
+          type: 'message_stop',
+        });
+        controller.close();
+      },
+    });
+
+    const onToolCallMock = vi.fn();
+
+    const protocolStream = AnthropicStream(mockReadableStream, {
+      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: 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(onToolCallMock).toHaveBeenCalledTimes(1);
+  });
+
+  it('should handle ReadableStream input', async () => {
+    const mockReadableStream = new ReadableStream({
+      start(controller) {
+        controller.enqueue({
+          type: 'message_start',
+          message: { id: 'message_1', metadata: {} },
+        });
+        controller.enqueue({
+          type: 'content_block_delta',
+          delta: { type: 'text_delta', text: 'Hello' },
+        });
+        controller.enqueue({
+          type: 'message_stop',
+        });
+        controller.close();
+      },
+    });
+
+    const protocolStream = AnthropicStream(mockReadableStream);
+
+    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: message_1\n',
+      'event: data\n',
+      `data: {"id":"message_1","metadata":{}}\n\n`,
+      'id: message_1\n',
+      'event: text\n',
+      `data: "Hello"\n\n`,
+      'id: message_1\n',
+      'event: stop\n',
+      `data: "message_stop"\n\n`,
+    ]);
+  });
+});

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

@@ -0,0 +1,91 @@
+import Anthropic from '@anthropic-ai/sdk';
+import type { Stream } from '@anthropic-ai/sdk/streaming';
+import { readableFromAsyncIterable } from 'ai';
+
+import { ChatStreamCallbacks } from '../../types';
+import {
+  StreamProtocolChunk,
+  StreamProtocolToolCallChunk,
+  StreamStack,
+  StreamToolCallChunkData,
+  createCallbacksTransformer,
+  createSSEProtocolTransformer,
+} from './protocol';
+
+export const transformAnthropicStream = (
+  chunk: Anthropic.MessageStreamEvent,
+  stack: StreamStack,
+): StreamProtocolChunk => {
+  // maybe need another structure to add support for multiple choices
+  switch (chunk.type) {
+    case 'message_start': {
+      stack.id = chunk.message.id;
+      return { data: chunk.message, id: chunk.message.id, type: 'data' };
+    }
+
+    // case 'content_block_start': {
+    //   return { data: chunk.content_block.text, id: stack.id, type: 'data' };
+    // }
+
+    case 'content_block_delta': {
+      switch (chunk.delta.type as string) {
+        default:
+        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;
+
+          const toolCall: StreamToolCallChunkData = {
+            function: { arguments: JSON.stringify(delta.input), name: delta.name },
+            id: delta.id,
+            index: 0,
+            type: 'function',
+          };
+
+          return {
+            data: [toolCall],
+            id: stack.id,
+            type: 'tool_calls',
+          } as StreamProtocolToolCallChunk;
+        }
+      }
+    }
+
+    case 'message_delta': {
+      return { data: chunk.delta.stop_reason, id: stack.id, type: 'stop' };
+    }
+
+    case 'message_stop': {
+      return { data: 'message_stop', id: stack.id, type: 'stop' };
+    }
+
+    default: {
+      return { data: chunk, id: stack.id, type: 'data' };
+    }
+  }
+};
+
+const chatStreamable = async function* (stream: AsyncIterable<Anthropic.MessageStreamEvent>) {
+  for await (const response of stream) {
+    yield response;
+  }
+};
+
+export const AnthropicStream = (
+  stream: Stream<Anthropic.MessageStreamEvent> | ReadableStream,
+  callbacks?: ChatStreamCallbacks,
+) => {
+  const streamStack: StreamStack = { id: '' };
+
+  const readableStream =
+    stream instanceof ReadableStream ? stream : readableFromAsyncIterable(chatStreamable(stream));
+
+  return readableStream
+    .pipeThrough(createSSEProtocolTransformer(transformAnthropicStream, streamStack))
+    .pipeThrough(createCallbacksTransformer(callbacks));
+};

package/src/libs/agent-runtime/utils/streams/bedrock/claude.ts

@@ -0,0 +1,21 @@
+import { InvokeModelWithResponseStreamResponse } from '@aws-sdk/client-bedrock-runtime';
+
+import { nanoid } from '@/utils/uuid';
+
+import { ChatStreamCallbacks } from '../../../types';
+import { transformAnthropicStream } from '../anthropic';
+import { StreamStack, createCallbacksTransformer, createSSEProtocolTransformer } from '../protocol';
+import { createBedrockStream } from './common';
+
+export const AWSBedrockClaudeStream = (
+  res: InvokeModelWithResponseStreamResponse | ReadableStream,
+  cb?: ChatStreamCallbacks,
+): ReadableStream<string> => {
+  const streamStack: StreamStack = { id: 'chat_' + nanoid() };
+
+  const stream = res instanceof ReadableStream ? res : createBedrockStream(res);
+
+  return stream
+    .pipeThrough(createSSEProtocolTransformer(transformAnthropicStream, streamStack))
+    .pipeThrough(createCallbacksTransformer(cb));
+};

package/src/libs/agent-runtime/utils/streams/bedrock/common.ts

@@ -0,0 +1,32 @@
+import {
+  InvokeModelWithResponseStreamResponse,
+  ResponseStream,
+} from '@aws-sdk/client-bedrock-runtime';
+import { readableFromAsyncIterable } from 'ai';
+
+const chatStreamable = async function* (stream: AsyncIterable<ResponseStream>) {
+  for await (const response of stream) {
+    if (response.chunk) {
+      const decoder = new TextDecoder();
+
+      const value = decoder.decode(response.chunk.bytes, { stream: true });
+      try {
+        const chunk = JSON.parse(value);
+
+        yield chunk;
+      } catch (e) {
+        console.log('bedrock stream parser error:', e);
+
+        yield value;
+      }
+    } else {
+      yield response;
+    }
+  }
+};
+
+/**
+ * covert the bedrock response to a readable stream
+ */
+export const createBedrockStream = (res: InvokeModelWithResponseStreamResponse) =>
+  readableFromAsyncIterable(chatStreamable(res.body!));

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

@@ -0,0 +1,3 @@
+export * from './claude';
+export * from './common';
+export * from './llama';

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

@@ -0,0 +1,196 @@
+import { InvokeModelWithResponseStreamResponse } from '@aws-sdk/client-bedrock-runtime';
+import { Readable } from 'stream';
+import { describe, expect, it, vi } from 'vitest';
+
+import * as uuidModule from '@/utils/uuid';
+
+import { AWSBedrockLlamaStream } from './llama';
+
+describe('AWSBedrockLlamaStream', () => {
+  it('should transform Bedrock Llama stream to protocol stream', async () => {
+    vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1');
+    const mockBedrockStream = new ReadableStream({
+      start(controller) {
+        controller.enqueue({ generation: 'Hello', generation_token_count: 1 });
+        controller.enqueue({ generation: ' world!', generation_token_count: 2 });
+        controller.enqueue({ stop_reason: 'stop' });
+        controller.close();
+      },
+    });
+
+    const onStartMock = vi.fn();
+    const onTextMock = vi.fn();
+    const onTokenMock = vi.fn();
+    const onCompletionMock = vi.fn();
+
+    const protocolStream = AWSBedrockLlamaStream(mockBedrockStream, {
+      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 transform Bedrock Llama AsyncIterator to protocol stream', async () => {
+    vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1');
+
+    const mockBedrockStream: InvokeModelWithResponseStreamResponse = {
+      body: {
+        // @ts-ignore
+        async *[Symbol.asyncIterator]() {
+          yield { generation: 'Hello', generation_token_count: 1 };
+          yield { generation: ' world!', generation_token_count: 2 };
+          yield { stop_reason: 'stop' };
+        },
+      },
+    };
+
+    const onStartMock = vi.fn();
+    const onTextMock = vi.fn();
+    const onTokenMock = vi.fn();
+    const onCompletionMock = vi.fn();
+
+    const protocolStream = AWSBedrockLlamaStream(mockBedrockStream, {
+      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 Bedrock response with chunk property', async () => {
+    vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('2');
+
+    const mockBedrockStream: InvokeModelWithResponseStreamResponse = {
+      contentType: 'any',
+      body: {
+        // @ts-ignore
+        async *[Symbol.asyncIterator]() {
+          yield {
+            chunk: {
+              bytes: new TextEncoder().encode('{"generation":"Hello","generation_token_count":1}'),
+            },
+          };
+          yield {
+            chunk: {
+              bytes: new TextEncoder().encode(
+                '{"generation":" world!","generation_token_count":2}',
+              ),
+            },
+          };
+          yield { chunk: { bytes: new TextEncoder().encode('{"stop_reason":"stop"}') } };
+        },
+      },
+    };
+
+    const onStartMock = vi.fn();
+    const onTextMock = vi.fn();
+    const onTokenMock = vi.fn();
+    const onCompletionMock = vi.fn();
+
+    const protocolStream = AWSBedrockLlamaStream(mockBedrockStream, {
+      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_2\n',
+      'event: text\n',
+      `data: "Hello"\n\n`,
+      'id: chat_2\n',
+      'event: text\n',
+      `data: " world!"\n\n`,
+      'id: chat_2\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 mockBedrockStream = new ReadableStream({
+      start(controller) {
+        controller.close();
+      },
+    });
+
+    const protocolStream = AWSBedrockLlamaStream(mockBedrockStream);
+
+    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/bedrock/llama.ts

@@ -0,0 +1,51 @@
+import { InvokeModelWithResponseStreamResponse } from '@aws-sdk/client-bedrock-runtime';
+
+import { nanoid } from '@/utils/uuid';
+
+import { ChatStreamCallbacks } from '../../../types';
+import {
+  StreamProtocolChunk,
+  StreamStack,
+  createCallbacksTransformer,
+  createSSEProtocolTransformer,
+} from '../protocol';
+import { createBedrockStream } from './common';
+
+interface AmazonBedrockInvocationMetrics {
+  firstByteLatency: number;
+  inputTokenCount: number;
+  invocationLatency: number;
+  outputTokenCount: number;
+}
+interface BedrockLlamaStreamChunk {
+  'amazon-bedrock-invocationMetrics'?: AmazonBedrockInvocationMetrics;
+  'generation': string;
+  'generation_token_count': number;
+  'prompt_token_count'?: number | null;
+  'stop_reason'?: null | 'stop' | string;
+}
+
+export const transformLlamaStream = (
+  chunk: BedrockLlamaStreamChunk,
+  stack: StreamStack,
+): StreamProtocolChunk => {
+  // maybe need another structure to add support for multiple choices
+  if (chunk.stop_reason) {
+    return { data: 'finished', id: stack.id, type: 'stop' };
+  }
+
+  return { data: chunk.generation, id: stack.id, type: 'text' };
+};
+
+export const AWSBedrockLlamaStream = (
+  res: InvokeModelWithResponseStreamResponse | ReadableStream,
+  cb?: ChatStreamCallbacks,
+): ReadableStream<string> => {
+  const streamStack: StreamStack = { id: 'chat_' + nanoid() };
+
+  const stream = res instanceof ReadableStream ? res : createBedrockStream(res);
+
+  return stream
+    .pipeThrough(createSSEProtocolTransformer(transformLlamaStream, streamStack))
+    .pipeThrough(createCallbacksTransformer(cb));
+};

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

@@ -0,0 +1,97 @@
+import { EnhancedGenerateContentResponse } from '@google/generative-ai';
+import { describe, expect, it, vi } from 'vitest';
+
+import * as uuidModule from '@/utils/uuid';
+
+import { GoogleGenerativeAIStream } from './google-ai';
+
+describe('GoogleGenerativeAIStream', () => {
+  it('should transform Google Generative AI stream to protocol stream', async () => {
+    vi.spyOn(uuidModule, 'nanoid').mockReturnValueOnce('1');
+
+    const mockGenerateContentResponse = (text: string, functionCalls?: any[]) =>
+      ({
+        text: () => text,
+        functionCall: () => functionCalls?.[0],
+        functionCalls: () => functionCalls,
+      }) as EnhancedGenerateContentResponse;
+
+    const mockGoogleStream = new ReadableStream({
+      start(controller) {
+        controller.enqueue(mockGenerateContentResponse('Hello'));
+
+        controller.enqueue(
+          mockGenerateContentResponse('', [{ name: 'testFunction', args: { arg1: 'value1' } }]),
+        );
+        controller.enqueue(mockGenerateContentResponse(' world!'));
+        controller.close();
+      },
+    });
+
+    const onStartMock = vi.fn();
+    const onTextMock = vi.fn();
+    const onTokenMock = vi.fn();
+    const onToolCallMock = vi.fn();
+    const onCompletionMock = vi.fn();
+
+    const protocolStream = GoogleGenerativeAIStream(mockGoogleStream, {
+      onStart: onStartMock,
+      onText: onTextMock,
+      onToken: onTokenMock,
+      onToolCall: onToolCallMock,
+      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([
+      // text
+      'id: chat_1\n',
+      'event: text\n',
+      `data: "Hello"\n\n`,
+
+      // tool call
+      'id: chat_1\n',
+      'event: tool_calls\n',
+      `data: [{"function":{"arguments":"{\\"arg1\\":\\"value1\\"}","name":"testFunction"},"id":"testFunction_0","index":0,"type":"function"}]\n\n`,
+
+      // text
+      'id: chat_1\n',
+      'event: text\n',
+      `data: " world!"\n\n`,
+    ]);
+
+    expect(onStartMock).toHaveBeenCalledTimes(1);
+    expect(onTextMock).toHaveBeenNthCalledWith(1, '"Hello"');
+    expect(onTextMock).toHaveBeenNthCalledWith(2, '" world!"');
+    expect(onTokenMock).toHaveBeenCalledTimes(2);
+    expect(onToolCallMock).toHaveBeenCalledTimes(1);
+    expect(onCompletionMock).toHaveBeenCalledTimes(1);
+  });
+
+  it('should handle empty stream', async () => {
+    const mockGoogleStream = new ReadableStream({
+      start(controller) {
+        controller.close();
+      },
+    });
+
+    const protocolStream = GoogleGenerativeAIStream(mockGoogleStream);
+
+    const decoder = new TextDecoder();
+    const chunks = [];
+
+    // @ts-ignore
+    for await (const chunk of protocolStream) {
+      chunks.push(decoder.decode(chunk, { stream: true }));
+    }
+
+    expect(chunks).toEqual([]);
+  });
+});