@lobehub/chat 0.146.1 → 0.147.0

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

@@ -6,11 +6,7 @@ import { ClientOptions } from 'openai';
 
 import { LobeRuntimeAI } from '../BaseAI';
 import { AgentRuntimeErrorType } from '../error';
-import {
-  ChatCompetitionOptions,
-  ChatStreamPayload,
-  ModelProvider
-} from '../types';
+import { ChatCompetitionOptions, ChatStreamPayload, ModelProvider } from '../types';
 import { AgentRuntimeError } from '../utils/createError';
 import { debugStream } from '../utils/debugStream';
 import { desensitizeUrl } from '../utils/desensitizeUrl';
@@ -20,12 +16,12 @@ const DEFAULT_BASE_URL = 'https://api.anthropic.com';
 
 export class LobeAnthropicAI implements LobeRuntimeAI {
   private client: Anthropic;
-  
+
   baseURL: string;
 
   constructor({ apiKey, baseURL = DEFAULT_BASE_URL }: ClientOptions) {
     if (!apiKey) throw AgentRuntimeError.createError(AgentRuntimeErrorType.InvalidAnthropicAPIKey);
-    
+
     this.client = new Anthropic({ apiKey, baseURL });
     this.baseURL = this.client.baseURL;
   }

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

@@ -0,0 +1,166 @@
+// @vitest-environment node
+import { AzureKeyCredential, OpenAIClient } from '@azure/openai';
+import { Mock, afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import * as debugStreamModule from '../utils/debugStream';
+import { LobeAzureOpenAI } from './index';
+
+// Mock the console.error to avoid polluting test output
+vi.spyOn(console, 'error').mockImplementation(() => {});
+
+describe('LobeAzureOpenAI', () => {
+  let instance: LobeAzureOpenAI;
+
+  beforeEach(() => {
+    instance = new LobeAzureOpenAI(
+      'https://test.openai.azure.com/',
+      'test_key',
+      '2023-03-15-preview',
+    );
+
+    // 使用 vi.spyOn 来模拟 streamChatCompletions 方法
+    vi.spyOn(instance['client'], 'streamChatCompletions').mockResolvedValue(
+      new ReadableStream() as any,
+    );
+  });
+
+  afterEach(() => {
+    vi.clearAllMocks();
+  });
+
+  describe('constructor', () => {
+    it('should throw InvalidAzureAPIKey error when apikey or endpoint is missing', () => {
+      try {
+        new LobeAzureOpenAI();
+      } catch (e) {
+        expect(e).toEqual({ errorType: 'InvalidAzureAPIKey' });
+      }
+    });
+
+    it('should create an instance of OpenAIClient with correct parameters', () => {
+      const endpoint = 'https://test.openai.azure.com/';
+      const apikey = 'test_key';
+      const apiVersion = '2023-03-15-preview';
+
+      const instance = new LobeAzureOpenAI(endpoint, apikey, apiVersion);
+
+      expect(instance.client).toBeInstanceOf(OpenAIClient);
+      expect(instance.baseURL).toBe(endpoint);
+    });
+  });
+
+  describe('chat', () => {
+    it('should return a StreamingTextResponse on successful API call', async () => {
+      // Arrange
+      const mockStream = new ReadableStream();
+      const mockResponse = Promise.resolve(mockStream);
+
+      (instance['client'].streamChatCompletions as Mock).mockResolvedValue(mockResponse);
+
+      // Act
+      const result = await instance.chat({
+        messages: [{ content: 'Hello', role: 'user' }],
+        model: 'text-davinci-003',
+        temperature: 0,
+      });
+
+      // Assert
+      expect(result).toBeInstanceOf(Response);
+    });
+
+    describe('Error', () => {
+      it('should return AzureBizError with DeploymentNotFound error', async () => {
+        // Arrange
+        const error = {
+          code: 'DeploymentNotFound',
+          message: 'Deployment not found',
+        };
+
+        (instance['client'].streamChatCompletions as Mock).mockRejectedValue(error);
+
+        // Act
+        try {
+          await instance.chat({
+            messages: [{ content: 'Hello', role: 'user' }],
+            model: 'text-davinci-003',
+            temperature: 0,
+          });
+        } catch (e) {
+          // Assert
+          expect(e).toEqual({
+            endpoint: 'https://test.openai.azure.com/',
+            error: {
+              code: 'DeploymentNotFound',
+              message: 'Deployment not found',
+              deployId: 'text-davinci-003',
+            },
+            errorType: 'AzureBizError',
+            provider: 'azure',
+          });
+        }
+      });
+
+      it('should return AgentRuntimeError for non-Azure errors', async () => {
+        // Arrange
+        const genericError = new Error('Generic Error');
+
+        (instance['client'].streamChatCompletions as Mock).mockRejectedValue(genericError);
+
+        // Act
+        try {
+          await instance.chat({
+            messages: [{ content: 'Hello', role: 'user' }],
+            model: 'text-davinci-003',
+            temperature: 0,
+          });
+        } catch (e) {
+          // Assert
+          expect(e).toEqual({
+            endpoint: 'https://test.openai.azure.com/',
+            errorType: 'AgentRuntimeError',
+            provider: 'azure',
+            error: {
+              name: genericError.name,
+              cause: genericError.cause,
+              message: genericError.message,
+            },
+          });
+        }
+      });
+    });
+
+    describe('DEBUG', () => {
+      it('should call debugStream when DEBUG_CHAT_COMPLETION is 1', async () => {
+        // Arrange
+        const mockProdStream = new ReadableStream() as any;
+        const mockDebugStream = new ReadableStream({
+          start(controller) {
+            controller.enqueue('Debug stream content');
+            controller.close();
+          },
+        }) as any;
+        mockDebugStream.toReadableStream = () => mockDebugStream;
+
+        (instance['client'].streamChatCompletions as Mock).mockResolvedValue({
+          tee: () => [mockProdStream, { toReadableStream: () => mockDebugStream }],
+        });
+
+        process.env.DEBUG_AZURE_CHAT_COMPLETION = '1';
+        vi.spyOn(debugStreamModule, 'debugStream').mockImplementation(() => Promise.resolve());
+
+        // Act
+        await instance.chat({
+          messages: [{ content: 'Hello', role: 'user' }],
+          model: 'text-davinci-003',
+          temperature: 0,
+        });
+
+        // Assert
+        expect(debugStreamModule.debugStream).toHaveBeenCalled();
+
+        // Restore
+        delete process.env.DEBUG_AZURE_CHAT_COMPLETION;
+      });
+    });
+  });
+});

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

@@ -11,16 +11,15 @@ import { AgentRuntimeErrorType } from '../error';
 import { ChatStreamPayload, ModelProvider } from '../types';
 import { AgentRuntimeError } from '../utils/createError';
 import { debugStream } from '../utils/debugStream';
-import { DEBUG_CHAT_COMPLETION } from '../utils/env';
 
 export class LobeAzureOpenAI implements LobeRuntimeAI {
-  private _client: OpenAIClient;
+  client: OpenAIClient;
 
   constructor(endpoint?: string, apikey?: string, apiVersion?: string) {
     if (!apikey || !endpoint)
       throw AgentRuntimeError.createError(AgentRuntimeErrorType.InvalidAzureAPIKey);
 
-    this._client = new OpenAIClient(endpoint, new AzureKeyCredential(apikey), { apiVersion });
+    this.client = new OpenAIClient(endpoint, new AzureKeyCredential(apikey), { apiVersion });
 
     this.baseURL = endpoint;
   }
@@ -34,7 +33,7 @@ export class LobeAzureOpenAI implements LobeRuntimeAI {
     // ============  2. send api   ============ //
 
     try {
-      const response = await this._client.streamChatCompletions(
+      const response = await this.client.streamChatCompletions(
         model,
         messages as ChatRequestMessage[],
         params as GetChatCompletionsOptions,
@@ -45,7 +44,7 @@ export class LobeAzureOpenAI implements LobeRuntimeAI {
 
       const [debug, prod] = stream.tee();
 
-      if (DEBUG_CHAT_COMPLETION) {
+      if (process.env.DEBUG_AZURE_CHAT_COMPLETION === '1') {
         debugStream(debug).catch(console.error);
       }
 
@@ -53,10 +52,18 @@ export class LobeAzureOpenAI implements LobeRuntimeAI {
     } catch (e) {
       let error = e as { [key: string]: any; code: string; message: string };
 
-      switch (error.code) {
-        case 'DeploymentNotFound': {
-          error = { ...error, deployId: model };
+      if (error.code) {
+        switch (error.code) {
+          case 'DeploymentNotFound': {
+            error = { ...error, deployId: model };
+          }
         }
+      } else {
+        error = {
+          cause: error.cause,
+          message: error.message,
+          name: error.name,
+        } as any;
       }
 
       const errorType = error.code
@@ -64,6 +71,7 @@ export class LobeAzureOpenAI implements LobeRuntimeAI {
         : AgentRuntimeErrorType.AgentRuntimeError;
 
       throw AgentRuntimeError.chat({
+        endpoint: this.baseURL,
         error,
         errorType,
         provider: ModelProvider.Azure,

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

@@ -1,9 +1,9 @@
 // @vitest-environment node
+import { InvokeModelWithResponseStreamCommand } from '@aws-sdk/client-bedrock-runtime';
 import { Mock, afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
 
-import {
-  InvokeModelWithResponseStreamCommand,
-} from '@aws-sdk/client-bedrock-runtime';
+import { AgentRuntimeErrorType, ModelProvider } from '@/libs/agent-runtime';
+
 import * as debugStreamModule from '../utils/debugStream';
 import { LobeBedrockAI } from './index';
 
@@ -12,13 +12,13 @@ const provider = 'bedrock';
 // Mock the console.error to avoid polluting test output
 vi.spyOn(console, 'error').mockImplementation(() => {});
 
-vi.mock("@aws-sdk/client-bedrock-runtime", async (importOriginal) => {
+vi.mock('@aws-sdk/client-bedrock-runtime', async (importOriginal) => {
   const module = await importOriginal();
   return {
     ...(module as any),
-    InvokeModelWithResponseStreamCommand: vi.fn()
-  }
-})
+    InvokeModelWithResponseStreamCommand: vi.fn(),
+  };
+});
 
 let instance: LobeBedrockAI;
 
@@ -49,16 +49,44 @@ describe('LobeBedrockAI', () => {
   });
 
   describe('chat', () => {
+    it('should call invokeLlamaModel when model starts with "meta"', async () => {
+      // @ts-ignore
+      const spy = vi.spyOn(instance, 'invokeLlamaModel');
 
-    describe('Claude model', () => {
+      // Act
+      await instance.chat({
+        messages: [{ content: 'Hello', role: 'user' }],
+        model: 'meta.llama:1',
+        temperature: 0,
+      });
+
+      // Assert
+      expect(spy).toHaveBeenCalled();
+    });
 
+    it('should call invokeClaudeModel when model does not start with "meta"', async () => {
+      // @ts-ignore
+      const spy = vi.spyOn(instance, 'invokeClaudeModel');
+
+      // Act
+      await instance.chat({
+        messages: [{ content: 'Hello', role: 'user' }],
+        model: 'anthropic.claude-v2:1',
+        temperature: 0,
+      });
+
+      // Assert
+      expect(spy).toHaveBeenCalled();
+    });
+
+    describe('Claude model', () => {
       it('should return a Response on successful API call', async () => {
         const result = await instance.chat({
           messages: [{ content: 'Hello', role: 'user' }],
           model: 'anthropic.claude-v2:1',
           temperature: 0,
         });
-  
+
         // Assert
         expect(result).toBeInstanceOf(Response);
       });
@@ -73,7 +101,7 @@ describe('LobeBedrockAI', () => {
         });
         const mockResponse = Promise.resolve(mockStream);
         (instance['client'].send as Mock).mockResolvedValue(mockResponse);
-  
+
         // Act
         const result = await instance.chat({
           messages: [{ content: 'Hello', role: 'user' }],
@@ -81,12 +109,12 @@ describe('LobeBedrockAI', () => {
           temperature: 0,
           top_p: 1,
         });
-  
+
         // Assert
         expect(InvokeModelWithResponseStreamCommand).toHaveBeenCalledWith({
           accept: 'application/json',
           body: JSON.stringify({
-            anthropic_version: "bedrock-2023-05-31",
+            anthropic_version: 'bedrock-2023-05-31',
             max_tokens: 4096,
             messages: [{ content: 'Hello', role: 'user' }],
             temperature: 0,
@@ -108,7 +136,7 @@ describe('LobeBedrockAI', () => {
         });
         const mockResponse = Promise.resolve(mockStream);
         (instance['client'].send as Mock).mockResolvedValue(mockResponse);
-  
+
         // Act
         const result = await instance.chat({
           messages: [
@@ -119,12 +147,12 @@ describe('LobeBedrockAI', () => {
           temperature: 0,
           top_p: 1,
         });
-  
+
         // Assert
         expect(InvokeModelWithResponseStreamCommand).toHaveBeenCalledWith({
           accept: 'application/json',
           body: JSON.stringify({
-            anthropic_version: "bedrock-2023-05-31",
+            anthropic_version: 'bedrock-2023-05-31',
             max_tokens: 4096,
             messages: [{ content: 'Hello', role: 'user' }],
             system: 'You are an awesome greeter',
@@ -147,7 +175,7 @@ describe('LobeBedrockAI', () => {
         });
         const mockResponse = Promise.resolve(mockStream);
         (instance['client'].send as Mock).mockResolvedValue(mockResponse);
-  
+
         // Act
         const result = await instance.chat({
           max_tokens: 2048,
@@ -156,12 +184,12 @@ describe('LobeBedrockAI', () => {
           temperature: 0.5,
           top_p: 1,
         });
-  
+
         // Assert
         expect(InvokeModelWithResponseStreamCommand).toHaveBeenCalledWith({
           accept: 'application/json',
           body: JSON.stringify({
-            anthropic_version: "bedrock-2023-05-31",
+            anthropic_version: 'bedrock-2023-05-31',
             max_tokens: 2048,
             messages: [{ content: 'Hello', role: 'user' }],
             temperature: 0.5,
@@ -172,7 +200,7 @@ describe('LobeBedrockAI', () => {
         });
         expect(result).toBeInstanceOf(Response);
       });
-  
+
       it('should call Anthropic model without unsupported opions', async () => {
         // Arrange
         const mockStream = new ReadableStream({
@@ -183,7 +211,7 @@ describe('LobeBedrockAI', () => {
         });
         const mockResponse = Promise.resolve(mockStream);
         (instance['client'].send as Mock).mockResolvedValue(mockResponse);
-  
+
         // Act
         const result = await instance.chat({
           frequency_penalty: 0.5, // Unsupported option
@@ -194,12 +222,12 @@ describe('LobeBedrockAI', () => {
           temperature: 0.5,
           top_p: 1,
         });
-  
+
         // Assert
         expect(InvokeModelWithResponseStreamCommand).toHaveBeenCalledWith({
           accept: 'application/json',
           body: JSON.stringify({
-            anthropic_version: "bedrock-2023-05-31",
+            anthropic_version: 'bedrock-2023-05-31',
             max_tokens: 2048,
             messages: [{ content: 'Hello', role: 'user' }],
             temperature: 0.5,
@@ -211,7 +239,156 @@ describe('LobeBedrockAI', () => {
         expect(result).toBeInstanceOf(Response);
       });
 
+      it('should call debugStream when DEBUG_BEDROCK_CHAT_COMPLETION is set to "1"', async () => {
+        // Arrange
+        process.env.DEBUG_BEDROCK_CHAT_COMPLETION = '1';
+        const spy = vi.spyOn(debugStreamModule, 'debugStream');
+
+        // Act
+        await instance.chat({
+          messages: [{ content: 'Hello', role: 'user' }],
+          model: 'anthropic.claude-v2:1',
+          temperature: 0,
+        });
+
+        // Assert
+        expect(spy).toHaveBeenCalled();
+
+        // Clean up
+        delete process.env.DEBUG_BEDROCK_CHAT_COMPLETION;
+      });
+
+      it('should handle errors and throw AgentRuntimeError', async () => {
+        // Arrange
+        const errorMessage = 'An error occurred';
+        const errorMetadata = { statusCode: 500 };
+        const mockError = new Error(errorMessage);
+        (mockError as any).$metadata = errorMetadata;
+        (instance['client'].send as Mock).mockRejectedValue(mockError);
+
+        // Act & Assert
+        await expect(
+          instance.chat({
+            max_tokens: 100,
+            messages: [{ content: 'Hello', role: 'user' }],
+            model: 'anthropic.claude-v2:1',
+            temperature: 0,
+          }),
+        ).rejects.toThrow(
+          expect.objectContaining({
+            error: {
+              body: errorMetadata,
+              message: errorMessage,
+              type: 'Error',
+            },
+            errorType: AgentRuntimeErrorType.BedrockBizError,
+            provider: ModelProvider.Bedrock,
+            region: 'us-west-2',
+          }),
+        );
+      });
     });
 
+    describe('Llama Model', () => {
+      it('should call Llama model with valid payload', async () => {
+        // Arrange
+        const mockStream = new ReadableStream({
+          start(controller) {
+            controller.enqueue('Hello, world!');
+            controller.close();
+          },
+        });
+        const mockResponse = Promise.resolve(mockStream);
+        (instance['client'].send as Mock).mockResolvedValue(mockResponse);
+
+        // Act
+        const result = await instance.chat({
+          temperature: 0,
+          max_tokens: 100,
+          messages: [{ content: 'Hello', role: 'user' }],
+          model: 'meta.llama:1',
+        });
+
+        // Assert
+        expect(InvokeModelWithResponseStreamCommand).toHaveBeenCalledWith({
+          accept: 'application/json',
+          body: JSON.stringify({
+            max_gen_len: 100,
+            prompt: '<s>[INST] Hello [/INST]',
+          }),
+          contentType: 'application/json',
+          modelId: 'meta.llama:1',
+        });
+        expect(result).toBeInstanceOf(Response);
+      });
+
+      it('should handle errors and throw AgentRuntimeError', async () => {
+        // Arrange
+        const errorMessage = 'An error occurred';
+        const errorMetadata = { statusCode: 500 };
+        const mockError = new Error(errorMessage);
+        (mockError as any).$metadata = errorMetadata;
+        (instance['client'].send as Mock).mockRejectedValue(mockError);
+
+        // Act & Assert
+        await expect(
+          instance.chat({
+            max_tokens: 100,
+            messages: [{ content: 'Hello', role: 'user' }],
+            model: 'meta.llama:1',
+            temperature: 0,
+          }),
+        ).rejects.toThrow(
+          expect.objectContaining({
+            error: {
+              body: errorMetadata,
+              message: errorMessage,
+              region: 'us-west-2',
+              type: 'Error',
+            },
+            errorType: AgentRuntimeErrorType.BedrockBizError,
+            provider: ModelProvider.Bedrock,
+            region: 'us-west-2',
+          }),
+        );
+      });
+
+      it('should call debugStream when DEBUG_BEDROCK_CHAT_COMPLETION is set to "1"', async () => {
+        // Arrange
+        process.env.DEBUG_BEDROCK_CHAT_COMPLETION = '1';
+        const spy = vi.spyOn(debugStreamModule, 'debugStream');
+
+        // Act
+        await instance.chat({
+          messages: [{ content: 'Hello', role: 'user' }],
+          model: 'meta.llama:1',
+          temperature: 0,
+        });
+
+        // Assert
+        expect(spy).toHaveBeenCalled();
+
+        // Clean up
+        delete process.env.DEBUG_BEDROCK_CHAT_COMPLETION;
+      });
+    });
+
+    it('should call options.callback when provided', async () => {
+      // Arrange
+      const onStart = vi.fn();
+
+      // Act
+      await instance.chat(
+        {
+          messages: [{ content: 'Hello', role: 'user' }],
+          model: 'anthropic.claude-v2:1',
+          temperature: 0,
+        },
+        { callback: { onStart } },
+      );
+
+      // Assert
+      expect(onStart).toHaveBeenCalled();
+    });
   });
 });

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

@@ -2,23 +2,15 @@ import {
   BedrockRuntimeClient,
   InvokeModelWithResponseStreamCommand,
 } from '@aws-sdk/client-bedrock-runtime';
-import {
-  AWSBedrockLlama2Stream,
-  AWSBedrockStream,
-  StreamingTextResponse
-} from 'ai';
+import { AWSBedrockLlama2Stream, AWSBedrockStream, StreamingTextResponse } from 'ai';
 import { experimental_buildLlama2Prompt } from 'ai/prompts';
 
 import { LobeRuntimeAI } from '../BaseAI';
 import { AgentRuntimeErrorType } from '../error';
-import {
-  ChatCompetitionOptions,
-  ChatStreamPayload,
-  ModelProvider,
-} from '../types';
+import { ChatCompetitionOptions, ChatStreamPayload, ModelProvider } from '../types';
+import { buildAnthropicMessages } from '../utils/anthropicHelpers';
 import { AgentRuntimeError } from '../utils/createError';
 import { debugStream } from '../utils/debugStream';
-import { buildAnthropicMessages } from '../utils/anthropicHelpers';
 
 export interface LobeBedrockAIParams {
   accessKeyId?: string;
@@ -54,7 +46,7 @@ export class LobeBedrockAI implements LobeRuntimeAI {
 
   private invokeClaudeModel = async (
     payload: ChatStreamPayload,
-    options?: ChatCompetitionOptions
+    options?: ChatCompetitionOptions,
   ): Promise<StreamingTextResponse> => {
     const { max_tokens, messages, model, temperature, top_p } = payload;
     const system_message = messages.find((m) => m.role === 'system');
@@ -63,7 +55,7 @@ export class LobeBedrockAI implements LobeRuntimeAI {
     const command = new InvokeModelWithResponseStreamCommand({
       accept: 'application/json',
       body: JSON.stringify({
-        anthropic_version: "bedrock-2023-05-31",
+        anthropic_version: 'bedrock-2023-05-31',
         max_tokens: max_tokens || 4096,
         messages: buildAnthropicMessages(user_messages),
         system: system_message?.content as string,
@@ -79,7 +71,11 @@ export class LobeBedrockAI implements LobeRuntimeAI {
       const bedrockResponse = await this.client.send(command);
 
       // Convert the response into a friendly text-stream
-      const stream = AWSBedrockStream(bedrockResponse, options?.callback, (chunk) => chunk.delta?.text);
+      const stream = AWSBedrockStream(
+        bedrockResponse,
+        options?.callback,
+        (chunk) => chunk.delta?.text,
+      );
 
       const [debug, output] = stream.tee();
 
@@ -105,9 +101,7 @@ export class LobeBedrockAI implements LobeRuntimeAI {
     }
   };
 
-  private invokeLlamaModel = async (
-    payload: ChatStreamPayload
-  ): Promise<StreamingTextResponse> => {
+  private invokeLlamaModel = async (payload: ChatStreamPayload): Promise<StreamingTextResponse> => {
     const { max_tokens, messages, model } = payload;
     const command = new InvokeModelWithResponseStreamCommand({
       accept: 'application/json',
@@ -149,7 +143,6 @@ export class LobeBedrockAI implements LobeRuntimeAI {
       });
     }
   };
-
 }
 
 export default LobeBedrockAI;