@lobehub/chat 0.156.1 → 0.157.0

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

@@ -0,0 +1,521 @@
+// @vitest-environment node
+import OpenAI from 'openai';
+import { Mock, afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import {
+  AgentRuntimeErrorType,
+  ChatStreamCallbacks,
+  LobeOpenAICompatibleRuntime,
+  ModelProvider,
+} from '@/libs/agent-runtime';
+
+import * as debugStreamModule from '../debugStream';
+import { LobeOpenAICompatibleFactory } from './index';
+
+const provider = 'groq';
+const defaultBaseURL = 'https://api.groq.com/openai/v1';
+const bizErrorType = 'GroqBizError';
+const invalidErrorType = 'InvalidGroqAPIKey';
+
+// Mock the console.error to avoid polluting test output
+vi.spyOn(console, 'error').mockImplementation(() => {});
+
+let instance: LobeOpenAICompatibleRuntime;
+
+const LobeMockProvider = LobeOpenAICompatibleFactory({
+  baseURL: defaultBaseURL,
+  chatCompletion: {
+    handleError: (error) => {
+      // 403 means the location is not supporteds
+      if (error.status === 403)
+        return { error, errorType: AgentRuntimeErrorType.LocationNotSupportError };
+    },
+  },
+  debug: {
+    chatCompletion: () => process.env.DEBUG_MOCKPROVIDER_CHAT_COMPLETION === '1',
+  },
+  errorType: {
+    bizError: AgentRuntimeErrorType.GroqBizError,
+    invalidAPIKey: AgentRuntimeErrorType.InvalidGroqAPIKey,
+  },
+  provider: ModelProvider.Groq,
+});
+
+beforeEach(() => {
+  instance = new LobeMockProvider({ apiKey: 'test' });
+
+  // 使用 vi.spyOn 来模拟 chat.completions.create 方法
+  vi.spyOn(instance['client'].chat.completions, 'create').mockResolvedValue(
+    new ReadableStream() as any,
+  );
+});
+
+afterEach(() => {
+  vi.clearAllMocks();
+});
+
+describe('LobeOpenAICompatibleFactory', () => {
+  describe('init', () => {
+    it('should correctly initialize with an API key', async () => {
+      const instance = new LobeMockProvider({ apiKey: 'test_api_key' });
+      expect(instance).toBeInstanceOf(LobeMockProvider);
+      expect(instance.baseURL).toEqual(defaultBaseURL);
+    });
+  });
+
+  describe('chat', () => {
+    it('should return a Response on successful API call', async () => {
+      // Arrange
+      const mockStream = new ReadableStream();
+      const mockResponse = Promise.resolve(mockStream);
+
+      (instance['client'].chat.completions.create as Mock).mockResolvedValue(mockResponse);
+
+      // Act
+      const result = await instance.chat({
+        messages: [{ content: 'Hello', role: 'user' }],
+        model: 'mistralai/mistral-7b-instruct:free',
+        temperature: 0,
+      });
+
+      // Assert
+      expect(result).toBeInstanceOf(Response);
+    });
+
+    it('should call chat API with corresponding options', async () => {
+      // Arrange
+      const mockStream = new ReadableStream();
+      const mockResponse = Promise.resolve(mockStream);
+
+      (instance['client'].chat.completions.create as Mock).mockResolvedValue(mockResponse);
+
+      // Act
+      const result = await instance.chat({
+        max_tokens: 1024,
+        messages: [{ content: 'Hello', role: 'user' }],
+        model: 'mistralai/mistral-7b-instruct:free',
+        temperature: 0.7,
+        top_p: 1,
+      });
+
+      // Assert
+      expect(instance['client'].chat.completions.create).toHaveBeenCalledWith(
+        {
+          max_tokens: 1024,
+          messages: [{ content: 'Hello', role: 'user' }],
+          model: 'mistralai/mistral-7b-instruct:free',
+          temperature: 0.7,
+          stream: true,
+          top_p: 1,
+        },
+        { headers: { Accept: '*/*' } },
+      );
+      expect(result).toBeInstanceOf(Response);
+    });
+
+    describe('streaming response', () => {
+      it('should handle multiple data chunks correctly', async () => {
+        const mockStream = new ReadableStream({
+          start(controller) {
+            controller.enqueue({
+              id: 'a',
+              object: 'chat.completion.chunk',
+              created: 1709125675,
+              model: 'mistralai/mistral-7b-instruct:free',
+              system_fingerprint: 'fp_86156a94a0',
+              choices: [
+                { index: 0, delta: { content: 'hello' }, logprobs: null, finish_reason: null },
+              ],
+            });
+            controller.close();
+          },
+        });
+        vi.spyOn(instance['client'].chat.completions, 'create').mockResolvedValue(
+          mockStream as any,
+        );
+
+        const result = await instance.chat({
+          messages: [{ content: 'Hello', role: 'user' }],
+          model: 'mistralai/mistral-7b-instruct:free',
+          temperature: 0,
+        });
+
+        const decoder = new TextDecoder();
+        const reader = result.body!.getReader();
+        expect(decoder.decode((await reader.read()).value)).toEqual('id: a\n');
+        expect(decoder.decode((await reader.read()).value)).toEqual('event: text\n');
+        expect(decoder.decode((await reader.read()).value)).toEqual('data: "hello"\n\n');
+        expect((await reader.read()).done).toBe(true);
+      });
+
+      it('should transform non-streaming response to stream correctly', async () => {
+        const mockResponse: OpenAI.ChatCompletion = {
+          id: 'a',
+          object: 'chat.completion',
+          created: 123,
+          model: 'mistralai/mistral-7b-instruct:free',
+          choices: [
+            {
+              index: 0,
+              message: { role: 'assistant', content: 'Hello' },
+              finish_reason: 'stop',
+              logprobs: null,
+            },
+          ],
+          usage: {
+            prompt_tokens: 5,
+            completion_tokens: 5,
+            total_tokens: 10,
+          },
+        };
+        vi.spyOn(instance['client'].chat.completions, 'create').mockResolvedValue(
+          mockResponse as any,
+        );
+
+        const result = await instance.chat({
+          messages: [{ content: 'Hello', role: 'user' }],
+          model: 'mistralai/mistral-7b-instruct:free',
+          temperature: 0,
+          stream: false,
+        });
+
+        const decoder = new TextDecoder();
+
+        const reader = result.body!.getReader();
+        expect(decoder.decode((await reader.read()).value)).toContain('id: a\n');
+        expect(decoder.decode((await reader.read()).value)).toContain('event: text\n');
+        expect(decoder.decode((await reader.read()).value)).toContain('data: "Hello"\n\n');
+
+        expect(decoder.decode((await reader.read()).value)).toContain('id: a\n');
+        expect(decoder.decode((await reader.read()).value)).toContain('event: text\n');
+        expect(decoder.decode((await reader.read()).value)).toContain('');
+
+        expect((await reader.read()).done).toBe(true);
+      });
+    });
+
+    describe('handlePayload option', () => {
+      it('should modify request payload correctly', async () => {
+        const mockCreateMethod = vi.spyOn(instance['client'].chat.completions, 'create');
+
+        await instance.chat({
+          messages: [{ content: 'Hello', role: 'user' }],
+          model: 'mistralai/mistral-7b-instruct:free',
+          temperature: 0,
+        });
+
+        expect(mockCreateMethod).toHaveBeenCalledWith(
+          expect.objectContaining({
+            // 根据实际的 handlePayload 函数,添加断言
+          }),
+          expect.anything(),
+        );
+      });
+    });
+
+    describe('cancel request', () => {
+      it('should cancel ongoing request correctly', async () => {
+        const controller = new AbortController();
+        const mockCreateMethod = vi.spyOn(instance['client'].chat.completions, 'create');
+
+        instance.chat(
+          {
+            messages: [{ content: 'Hello', role: 'user' }],
+            model: 'mistralai/mistral-7b-instruct:free',
+            temperature: 0,
+          },
+          { signal: controller.signal },
+        );
+
+        controller.abort();
+
+        expect(mockCreateMethod).toHaveBeenCalledWith(
+          expect.anything(),
+          expect.objectContaining({
+            signal: controller.signal,
+          }),
+        );
+      });
+    });
+
+    describe('Error', () => {
+      it('should return bizErrorType with an openai error response when OpenAI.APIError is thrown', async () => {
+        // Arrange
+        const apiError = new OpenAI.APIError(
+          400,
+          {
+            status: 400,
+            error: {
+              message: 'Bad Request',
+            },
+          },
+          'Error message',
+          {},
+        );
+
+        vi.spyOn(instance['client'].chat.completions, 'create').mockRejectedValue(apiError);
+
+        // Act
+        try {
+          await instance.chat({
+            messages: [{ content: 'Hello', role: 'user' }],
+            model: 'mistralai/mistral-7b-instruct:free',
+            temperature: 0,
+          });
+        } catch (e) {
+          expect(e).toEqual({
+            endpoint: defaultBaseURL,
+            error: {
+              error: { message: 'Bad Request' },
+              status: 400,
+            },
+            errorType: bizErrorType,
+            provider,
+          });
+        }
+      });
+
+      it('should throw AgentRuntimeError with invalidErrorType if no apiKey is provided', async () => {
+        try {
+          new LobeMockProvider({});
+        } catch (e) {
+          expect(e).toEqual({ errorType: invalidErrorType });
+        }
+      });
+
+      it('should return bizErrorType with the cause when OpenAI.APIError is thrown with cause', async () => {
+        // Arrange
+        const errorInfo = {
+          stack: 'abc',
+          cause: {
+            message: 'api is undefined',
+          },
+        };
+        const apiError = new OpenAI.APIError(400, errorInfo, 'module error', {});
+
+        vi.spyOn(instance['client'].chat.completions, 'create').mockRejectedValue(apiError);
+
+        // Act
+        try {
+          await instance.chat({
+            messages: [{ content: 'Hello', role: 'user' }],
+            model: 'mistralai/mistral-7b-instruct:free',
+            temperature: 0,
+          });
+        } catch (e) {
+          expect(e).toEqual({
+            endpoint: defaultBaseURL,
+            error: {
+              cause: { message: 'api is undefined' },
+              stack: 'abc',
+            },
+            errorType: bizErrorType,
+            provider,
+          });
+        }
+      });
+
+      it('should return bizErrorType with an cause response with desensitize Url', async () => {
+        // Arrange
+        const errorInfo = {
+          stack: 'abc',
+          cause: { message: 'api is undefined' },
+        };
+        const apiError = new OpenAI.APIError(400, errorInfo, 'module error', {});
+
+        instance = new LobeMockProvider({
+          apiKey: 'test',
+
+          baseURL: 'https://api.abc.com/v1',
+        });
+
+        vi.spyOn(instance['client'].chat.completions, 'create').mockRejectedValue(apiError);
+
+        // Act
+        try {
+          await instance.chat({
+            messages: [{ content: 'Hello', role: 'user' }],
+            model: 'mistralai/mistral-7b-instruct:free',
+            temperature: 0,
+          });
+        } catch (e) {
+          expect(e).toEqual({
+            endpoint: 'https://api.***.com/v1',
+            error: {
+              cause: { message: 'api is undefined' },
+              stack: 'abc',
+            },
+            errorType: bizErrorType,
+            provider,
+          });
+        }
+      });
+
+      describe('handleError option', () => {
+        it('should return correct error type for 403 status code', async () => {
+          const error = { status: 403 };
+          vi.spyOn(instance['client'].chat.completions, 'create').mockRejectedValue(error);
+
+          try {
+            await instance.chat({
+              messages: [{ content: 'Hello', role: 'user' }],
+              model: 'mistralai/mistral-7b-instruct:free',
+              temperature: 0,
+            });
+          } catch (e) {
+            expect(e).toEqual({
+              error,
+              errorType: AgentRuntimeErrorType.LocationNotSupportError,
+              provider,
+            });
+          }
+        });
+      });
+
+      it('should throw an InvalidOpenRouterAPIKey error type on 401 status code', async () => {
+        // Mock the API call to simulate a 401 error
+        const error = new Error('Unauthorized') as any;
+        error.status = 401;
+        vi.mocked(instance['client'].chat.completions.create).mockRejectedValue(error);
+
+        try {
+          await instance.chat({
+            messages: [{ content: 'Hello', role: 'user' }],
+            model: 'mistralai/mistral-7b-instruct:free',
+            temperature: 0,
+          });
+        } catch (e) {
+          // Expect the chat method to throw an error with InvalidMoonshotAPIKey
+          expect(e).toEqual({
+            endpoint: defaultBaseURL,
+            error: new Error('Unauthorized'),
+            errorType: invalidErrorType,
+            provider,
+          });
+        }
+      });
+
+      it('should return AgentRuntimeError for non-OpenAI errors', async () => {
+        // Arrange
+        const genericError = new Error('Generic Error');
+
+        vi.spyOn(instance['client'].chat.completions, 'create').mockRejectedValue(genericError);
+
+        // Act
+        try {
+          await instance.chat({
+            messages: [{ content: 'Hello', role: 'user' }],
+            model: 'mistralai/mistral-7b-instruct:free',
+            temperature: 0,
+          });
+        } catch (e) {
+          expect(e).toEqual({
+            endpoint: defaultBaseURL,
+            errorType: 'AgentRuntimeError',
+            provider,
+            error: {
+              name: genericError.name,
+              cause: genericError.cause,
+              message: genericError.message,
+              stack: genericError.stack,
+            },
+          });
+        }
+      });
+    });
+
+    describe('chat with callback and headers', () => {
+      it('should handle callback and headers correctly', async () => {
+        // 模拟 chat.completions.create 方法返回一个可读流
+        const mockCreateMethod = vi
+          .spyOn(instance['client'].chat.completions, 'create')
+          .mockResolvedValue(
+            new ReadableStream({
+              start(controller) {
+                controller.enqueue({
+                  id: 'chatcmpl-8xDx5AETP8mESQN7UB30GxTN2H1SO',
+                  object: 'chat.completion.chunk',
+                  created: 1709125675,
+                  model: 'mistralai/mistral-7b-instruct:free',
+                  system_fingerprint: 'fp_86156a94a0',
+                  choices: [
+                    { index: 0, delta: { content: 'hello' }, logprobs: null, finish_reason: null },
+                  ],
+                });
+                controller.close();
+              },
+            }) as any,
+          );
+
+        // 准备 callback 和 headers
+        const mockCallback: ChatStreamCallbacks = {
+          onStart: vi.fn(),
+          onToken: vi.fn(),
+        };
+        const mockHeaders = { 'Custom-Header': 'TestValue' };
+
+        // 执行测试
+        const result = await instance.chat(
+          {
+            messages: [{ content: 'Hello', role: 'user' }],
+            model: 'mistralai/mistral-7b-instruct:free',
+            temperature: 0,
+          },
+          { callback: mockCallback, headers: mockHeaders },
+        );
+
+        // 验证 callback 被调用
+        await result.text(); // 确保流被消费
+        expect(mockCallback.onStart).toHaveBeenCalled();
+        expect(mockCallback.onToken).toHaveBeenCalledWith('hello');
+
+        // 验证 headers 被正确传递
+        expect(result.headers.get('Custom-Header')).toEqual('TestValue');
+
+        // 清理
+        mockCreateMethod.mockRestore();
+      });
+    });
+
+    describe('DEBUG', () => {
+      it('should call debugStream and return StreamingTextResponse when DEBUG_OPENROUTER_CHAT_COMPLETION is 1', async () => {
+        // Arrange
+        const mockProdStream = new ReadableStream() as any; // 模拟的 prod 流
+        const mockDebugStream = new ReadableStream({
+          start(controller) {
+            controller.enqueue('Debug stream content');
+            controller.close();
+          },
+        }) as any;
+        mockDebugStream.toReadableStream = () => mockDebugStream; // 添加 toReadableStream 方法
+
+        // 模拟 chat.completions.create 返回值，包括模拟的 tee 方法
+        (instance['client'].chat.completions.create as Mock).mockResolvedValue({
+          tee: () => [mockProdStream, { toReadableStream: () => mockDebugStream }],
+        });
+
+        // 保存原始环境变量值
+        const originalDebugValue = process.env.DEBUG_MOCKPROVIDER_CHAT_COMPLETION;
+
+        // 模拟环境变量
+        process.env.DEBUG_MOCKPROVIDER_CHAT_COMPLETION = '1';
+        vi.spyOn(debugStreamModule, 'debugStream').mockImplementation(() => Promise.resolve());
+
+        // 执行测试
+        // 运行你的测试函数，确保它会在条件满足时调用 debugStream
+        // 假设的测试函数调用，你可能需要根据实际情况调整
+        await instance.chat({
+          messages: [{ content: 'Hello', role: 'user' }],
+          model: 'mistralai/mistral-7b-instruct:free',
+          temperature: 0,
+        });
+
+        // 验证 debugStream 被调用
+        expect(debugStreamModule.debugStream).toHaveBeenCalled();
+
+        // 恢复原始环境变量值
+        process.env.DEBUG_MOCKPROVIDER_CHAT_COMPLETION = originalDebugValue;
+      });
+    });
+  });
+});

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

@@ -1,4 +1,3 @@
-import { OpenAIStream, StreamingTextResponse } from 'ai';
 import OpenAI, { ClientOptions } from 'openai';
 
 import { LOBE_DEFAULT_MODEL_LIST } from '@/config/modelProviders';
@@ -11,6 +10,8 @@ import { AgentRuntimeError } from '../createError';
 import { debugStream } from '../debugStream';
 import { desensitizeUrl } from '../desensitizeUrl';
 import { handleOpenAIError } from '../handleOpenAIError';
+import { StreamingResponse } from '../response';
+import { OpenAIStream } from '../streams';
 
 // the model contains the following keywords is not a chat model, so we should filter them out
 const CHAT_MODELS_BLOCK_LIST = [
@@ -72,7 +73,10 @@ export const LobeOpenAICompatibleFactory = ({
       try {
         const postPayload = chatCompletion?.handlePayload
           ? chatCompletion.handlePayload(payload)
-          : (payload as unknown as OpenAI.ChatCompletionCreateParamsStreaming);
+          : ({
+              ...payload,
+              stream: payload.stream ?? true,
+            } as OpenAI.ChatCompletionCreateParamsStreaming);
 
         const response = await this.client.chat.completions.create(postPayload, {
           // https://github.com/lobehub/lobe-chat/pull/318
@@ -80,13 +84,26 @@ export const LobeOpenAICompatibleFactory = ({
           signal: options?.signal,
         });
 
-        const [prod, useForDebug] = response.tee();
+        if (postPayload.stream) {
+          const [prod, useForDebug] = response.tee();
+
+          if (debug?.chatCompletion?.()) {
+            debugStream(useForDebug.toReadableStream()).catch(console.error);
+          }
+
+          return StreamingResponse(OpenAIStream(prod, options?.callback), {
+            headers: options?.headers,
+          });
+        }
 
         if (debug?.chatCompletion?.()) {
-          debugStream(useForDebug.toReadableStream()).catch(console.error);
+          console.log('\n[no stream response]\n');
+          console.log(JSON.stringify(response) + '\n');
         }
 
-        return new StreamingTextResponse(OpenAIStream(prod, options?.callback), {
+        const stream = this.transformResponseToStream(response as unknown as OpenAI.ChatCompletion);
+
+        return StreamingResponse(OpenAIStream(stream, options?.callback), {
           headers: options?.headers,
         });
       } catch (error) {
@@ -160,4 +177,58 @@ export const LobeOpenAICompatibleFactory = ({
 
         .filter(Boolean) as ChatModelCard[];
     }
+
+    /**
+     * make the OpenAI response data as a stream
+     * @private
+     */
+    private transformResponseToStream(data: OpenAI.ChatCompletion) {
+      return new ReadableStream({
+        start(controller) {
+          const chunk: OpenAI.ChatCompletionChunk = {
+            choices: data.choices.map((choice: OpenAI.ChatCompletion.Choice) => ({
+              delta: {
+                content: choice.message.content,
+                role: choice.message.role,
+                tool_calls: choice.message.tool_calls?.map(
+                  (tool, index): OpenAI.ChatCompletionChunk.Choice.Delta.ToolCall => ({
+                    function: tool.function,
+                    id: tool.id,
+                    index,
+                    type: tool.type,
+                  }),
+                ),
+              },
+              finish_reason: null,
+              index: choice.index,
+              logprobs: choice.logprobs,
+            })),
+            created: data.created,
+            id: data.id,
+            model: data.model,
+            object: 'chat.completion.chunk',
+          };
+
+          controller.enqueue(chunk);
+
+          controller.enqueue({
+            choices: data.choices.map((choice: OpenAI.ChatCompletion.Choice) => ({
+              delta: {
+                content: choice.message.content,
+                role: choice.message.role,
+              },
+              finish_reason: choice.finish_reason,
+              index: choice.index,
+              logprobs: choice.logprobs,
+            })),
+            created: data.created,
+            id: data.id,
+            model: data.model,
+            object: 'chat.completion.chunk',
+            system_fingerprint: data.system_fingerprint,
+          } as OpenAI.ChatCompletionChunk);
+          controller.close();
+        },
+      });
+    }
   };

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

@@ -0,0 +1,12 @@
+export const StreamingResponse = (
+  stream: ReadableStream,
+  options?: { headers?: Record<string, string> },
+) => {
+  return new Response(stream, {
+    headers: {
+      'Cache-Control': 'no-cache',
+      'Content-Type': 'text/event-stream',
+      ...options?.headers,
+    },
+  });
+};