@lobehub/chat 1.29.6 → 1.31.0

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

@@ -0,0 +1,648 @@
+// @vitest-environment node
+import { Mock, afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { ChatCompletionTool } from '@/libs/agent-runtime';
+
+import * as debugStreamModule from '../utils/debugStream';
+import { LobeCloudflareAI } from './index';
+
+const provider = 'cloudflare';
+
+const bizErrorType = 'ProviderBizError';
+const invalidErrorType = 'InvalidProviderAPIKey';
+
+// Mock the console.error to avoid polluting test output
+vi.spyOn(console, 'error').mockImplementation(() => {});
+
+let instance: LobeCloudflareAI;
+const textEncoder = new TextEncoder();
+
+afterEach(() => {
+  vi.restoreAllMocks();
+});
+
+describe('LobeCloudflareAI', () => {
+  const accountID = '80009000a000b000c000d000e000f000';
+  describe('init', () => {
+    it('should correctly initialize with API key and Account ID', async () => {
+      const instance = new LobeCloudflareAI({
+        apiKey: 'test_api_key',
+        baseURLOrAccountID: accountID,
+      });
+      expect(instance).toBeInstanceOf(LobeCloudflareAI);
+      expect(instance.baseURL).toBe(
+        `https://api.cloudflare.com/client/v4/accounts/${accountID}/ai/run/`,
+      );
+      expect(instance.accountID).toBe(accountID);
+    });
+
+    it('should correctly initialize with API key and Gateway URL', async () => {
+      const baseURL = `https://gateway.ai.cloudflare.com/v1/${accountID}/test-gateway/workers-ai`;
+      const instance = new LobeCloudflareAI({
+        apiKey: 'test_api_key',
+        baseURLOrAccountID: baseURL,
+      });
+      expect(instance).toBeInstanceOf(LobeCloudflareAI);
+      expect(instance.baseURL).toBe(baseURL + '/'); // baseURL MUST end with '/'.
+      expect(instance.accountID).toBe(accountID);
+    });
+  });
+
+  describe('chat', () => {
+    beforeEach(() => {
+      instance = new LobeCloudflareAI({
+        apiKey: 'test_api_key',
+        baseURLOrAccountID: accountID,
+      });
+
+      // Mock fetch
+      vi.spyOn(globalThis, 'fetch').mockResolvedValue(
+        new Response(
+          new ReadableStream<Uint8Array>({
+            start(controller) {
+              controller.enqueue(textEncoder.encode('data: {"response": "Hello, world!"}\n\n'));
+              controller.close();
+            },
+          }),
+        ),
+      );
+    });
+
+    it('should return a Response on successful API call', async () => {
+      const result = await instance.chat({
+        messages: [{ content: 'Hello', role: 'user' }],
+        model: '@hf/meta-llama/meta-llama-3-8b-instruct',
+        temperature: 0,
+      });
+
+      // Assert
+      expect(result).toBeInstanceOf(Response);
+    });
+
+    it('should handle text messages correctly', async () => {
+      // Arrange
+      const textEncoder = new TextEncoder();
+      const mockResponse = new Response(
+        new ReadableStream<Uint8Array>({
+          start(controller) {
+            controller.enqueue(textEncoder.encode('data: {"response": "Hello, world!"}\n\n'));
+            controller.close();
+          },
+        }),
+      );
+      (globalThis.fetch as Mock).mockResolvedValue(mockResponse);
+
+      // Act
+      const result = await instance.chat({
+        messages: [{ content: 'Hello', role: 'user' }],
+        model: '@hf/meta-llama/meta-llama-3-8b-instruct',
+        temperature: 0,
+        top_p: 1,
+      });
+
+      // Assert
+      expect(globalThis.fetch).toHaveBeenCalledWith(
+        // url
+        expect.objectContaining({
+          pathname: `/client/v4/accounts/${accountID}/ai/run/@hf/meta-llama/meta-llama-3-8b-instruct`,
+        }),
+        // body
+        expect.objectContaining({
+          body: expect.any(String),
+          method: 'POST',
+        }),
+      );
+
+      const fetchCallArgs = (globalThis.fetch as Mock).mock.calls[0];
+      const body = JSON.parse(fetchCallArgs[1].body);
+      expect(body).toEqual(
+        expect.objectContaining({
+          //max_tokens: 4096,
+          messages: [{ content: 'Hello', role: 'user' }],
+          //stream: true,
+          temperature: 0,
+          top_p: 1,
+        }),
+      );
+
+      expect(result).toBeInstanceOf(Response);
+    });
+
+    it('should handle system prompt correctly', async () => {
+      // Arrange
+      const textEncoder = new TextEncoder();
+      const mockResponse = new Response(
+        new ReadableStream<Uint8Array>({
+          start(controller) {
+            controller.enqueue(textEncoder.encode('data: {"response": "Hello, world!"}\n\n'));
+            controller.close();
+          },
+        }),
+      );
+      (globalThis.fetch as Mock).mockResolvedValue(mockResponse);
+
+      // Act
+      const result = await instance.chat({
+        messages: [
+          { content: 'You are an awesome greeter', role: 'system' },
+          { content: 'Hello', role: 'user' },
+        ],
+        model: '@hf/meta-llama/meta-llama-3-8b-instruct',
+        temperature: 0,
+      });
+
+      // Assert
+      expect(globalThis.fetch).toHaveBeenCalledWith(
+        // url
+        expect.objectContaining({
+          pathname: `/client/v4/accounts/${accountID}/ai/run/@hf/meta-llama/meta-llama-3-8b-instruct`,
+        }),
+        // body
+        expect.objectContaining({
+          body: expect.any(String),
+          method: 'POST',
+        }),
+      );
+
+      const fetchCallArgs = (globalThis.fetch as Mock).mock.calls[0];
+      const body = JSON.parse(fetchCallArgs[1].body);
+      expect(body).toEqual(
+        expect.objectContaining({
+          //max_tokens: 4096,
+          messages: [
+            { content: 'You are an awesome greeter', role: 'system' },
+            { content: 'Hello', role: 'user' },
+          ],
+          //stream: true,
+          temperature: 0,
+        }),
+      );
+
+      expect(result).toBeInstanceOf(Response);
+    });
+
+    it('should call Cloudflare API with supported opions', async () => {
+      // Arrange
+      const mockResponse = new Response(
+        new ReadableStream<Uint8Array>({
+          start(controller) {
+            controller.enqueue(textEncoder.encode('data: {"response": "Hello, world!"}\n\n'));
+            controller.close();
+          },
+        }),
+      );
+      (globalThis.fetch as Mock).mockResolvedValue(mockResponse);
+
+      // Act
+      const result = await instance.chat({
+        max_tokens: 2048,
+        messages: [{ content: 'Hello', role: 'user' }],
+        model: '@hf/meta-llama/meta-llama-3-8b-instruct',
+        temperature: 0.5,
+        top_p: 1,
+      });
+
+      // Assert
+      expect(globalThis.fetch).toHaveBeenCalledWith(
+        // url
+        expect.objectContaining({
+          pathname: `/client/v4/accounts/${accountID}/ai/run/@hf/meta-llama/meta-llama-3-8b-instruct`,
+        }),
+        // body
+        expect.objectContaining({
+          body: expect.any(String),
+          method: 'POST',
+        }),
+      );
+
+      const fetchCallArgs = (globalThis.fetch as Mock).mock.calls[0];
+      const body = JSON.parse(fetchCallArgs[1].body);
+      expect(body).toEqual(
+        expect.objectContaining({
+          max_tokens: 2048,
+          messages: [{ content: 'Hello', role: 'user' }],
+          //stream: true,
+          temperature: 0.5,
+          top_p: 1,
+        }),
+      );
+
+      expect(result).toBeInstanceOf(Response);
+    });
+
+    it('should call debugStream in DEBUG mode', async () => {
+      // Arrange
+      const mockProdStream = new ReadableStream({
+        start(controller) {
+          controller.enqueue('Hello, world!');
+          controller.close();
+        },
+      }) as any;
+      const mockDebugStream = new ReadableStream({
+        start(controller) {
+          controller.enqueue('Debug stream content');
+          controller.close();
+        },
+      }) as any;
+      mockDebugStream.toReadableStream = () => mockDebugStream;
+
+      (globalThis.fetch as Mock).mockResolvedValue({
+        body: {
+          tee: () => [mockProdStream, { toReadableStream: () => mockDebugStream }],
+        },
+      });
+
+      const originalDebugValue = process.env.DEBUG_CLOUDFLARE_CHAT_COMPLETION;
+
+      process.env.DEBUG_CLOUDFLARE_CHAT_COMPLETION = '1';
+      vi.spyOn(debugStreamModule, 'debugStream').mockImplementation(() => Promise.resolve());
+
+      // Act
+      await instance.chat({
+        messages: [{ content: 'Hello', role: 'user' }],
+        model: '@hf/meta-llama/meta-llama-3-8b-instruct',
+        temperature: 0,
+      });
+
+      // Assert
+      expect(debugStreamModule.debugStream).toHaveBeenCalled();
+
+      // Cleanup
+      process.env.DEBUG_CLOUDFLARE_CHAT_COMPLETION = originalDebugValue;
+    });
+
+    describe('chat with tools', () => {
+      it('should call client.beta.tools.messages.create when tools are provided', async () => {
+        // Arrange
+        const tools: ChatCompletionTool[] = [
+          { function: { name: 'tool1', description: 'desc1' }, type: 'function' },
+        ];
+
+        // Act
+        await instance.chat({
+          messages: [{ content: 'Hello', role: 'user' }],
+          model: '@hf/meta-llama/meta-llama-3-8b-instruct',
+          temperature: 1,
+          tools,
+        });
+
+        // Assert
+        expect(globalThis.fetch).toHaveBeenCalled();
+
+        const fetchCallArgs = (globalThis.fetch as Mock).mock.calls[0];
+        const body = JSON.parse(fetchCallArgs[1].body);
+        expect(body).toEqual(
+          expect.objectContaining({
+            tools: tools.map((t) => t.function),
+          }),
+        );
+      });
+    });
+
+    describe('Error', () => {
+      it('should throw ProviderBizError error on 400 error', async () => {
+        // Arrange
+        const apiError = {
+          status: 400,
+          error: {
+            type: 'error',
+            error: {
+              type: 'authentication_error',
+              message: 'invalid x-api-key',
+            },
+          },
+        };
+        (globalThis.fetch as Mock).mockRejectedValue(apiError);
+
+        try {
+          // Act
+          await instance.chat({
+            messages: [{ content: 'Hello', role: 'user' }],
+            model: '@hf/meta-llama/meta-llama-3-8b-instruct',
+            temperature: 0,
+          });
+        } catch (e) {
+          // Assert
+          expect(e).toEqual({
+            endpoint: expect.stringMatching(/https:\/\/.+/),
+            error: apiError,
+            errorType: bizErrorType,
+            provider,
+          });
+        }
+      });
+
+      it('should throw InvalidProviderAPIKey if no accountID is provided', async () => {
+        try {
+          new LobeCloudflareAI({
+            apiKey: 'test',
+          });
+        } catch (e) {
+          expect(e).toEqual({ errorType: invalidErrorType });
+        }
+      });
+
+      it('should throw InvalidProviderAPIKey if no apiKey is provided', async () => {
+        try {
+          new LobeCloudflareAI({
+            baseURLOrAccountID: accountID,
+          });
+        } catch (e) {
+          expect(e).toEqual({ errorType: invalidErrorType });
+        }
+      });
+
+      it('should not throw Error when apiKey is not provided but baseURL is provided', async () => {
+        const customInstance = new LobeCloudflareAI({
+          baseURLOrAccountID: 'https://custom.cloudflare.url/',
+        });
+        expect(customInstance).toBeInstanceOf(LobeCloudflareAI);
+        expect(customInstance.apiKey).toBeUndefined();
+        expect(customInstance.baseURL).toBe('https://custom.cloudflare.url/');
+      });
+    });
+
+    describe('Error handling', () => {
+      it('should throw ProviderBizError on other error status codes', async () => {
+        // Arrange
+        const apiError = { status: 400 };
+        (globalThis.fetch as Mock).mockRejectedValue(apiError);
+
+        // Act & Assert
+        await expect(
+          instance.chat({
+            messages: [{ content: 'Hello', role: 'user' }],
+            model: '@hf/meta-llama/meta-llama-3-8b-instruct',
+            temperature: 1,
+          }),
+        ).rejects.toEqual({
+          endpoint: expect.stringMatching(/https:\/\/.+/),
+          error: apiError,
+          errorType: bizErrorType,
+          provider,
+        });
+      });
+
+      it('should desensitize accountID in error message', async () => {
+        // Arrange
+        const apiError = { status: 400 };
+        const customInstance = new LobeCloudflareAI({
+          apiKey: 'test',
+          baseURLOrAccountID: accountID,
+        });
+        (globalThis.fetch as Mock).mockRejectedValue(apiError);
+
+        // Act & Assert
+        await expect(
+          customInstance.chat({
+            messages: [{ content: 'Hello', role: 'user' }],
+            model: '@hf/meta-llama/meta-llama-3-8b-instruct',
+            temperature: 0,
+          }),
+        ).rejects.toEqual({
+          endpoint: expect.not.stringContaining(accountID),
+          error: apiError,
+          errorType: bizErrorType,
+          provider,
+        });
+      });
+    });
+
+    describe('Options', () => {
+      it('should pass signal to API call', async () => {
+        // Arrange
+        const controller = new AbortController();
+
+        // Act
+        await instance.chat(
+          {
+            messages: [{ content: 'Hello', role: 'user' }],
+            model: '@hf/meta-llama/meta-llama-3-8b-instruct',
+            temperature: 1,
+          },
+          { signal: controller.signal },
+        );
+
+        // Assert
+        expect(globalThis.fetch).toHaveBeenCalledWith(
+          expect.any(URL),
+          expect.objectContaining({ signal: controller.signal }),
+        );
+      });
+
+      it('should apply callback to the returned stream', async () => {
+        // Arrange
+        const callback = vi.fn();
+
+        // Act
+        await instance.chat(
+          {
+            messages: [{ content: 'Hello', role: 'user' }],
+            model: '@hf/meta-llama/meta-llama-3-8b-instruct',
+            temperature: 0,
+          },
+          {
+            callback: { onStart: callback },
+          },
+        );
+
+        // Assert
+        expect(callback).toHaveBeenCalled();
+      });
+
+      it('should set headers on the response', async () => {
+        // Arrange
+        const headers = { 'X-Test-Header': 'test' };
+
+        // Act
+        const result = await instance.chat(
+          {
+            messages: [{ content: 'Hello', role: 'user' }],
+            model: '@hf/meta-llama/meta-llama-3-8b-instruct',
+            temperature: 1,
+          },
+          { headers },
+        );
+
+        // Assert
+        expect(result.headers.get('X-Test-Header')).toBe('test');
+      });
+    });
+
+    describe('Edge cases', () => {
+      it('should handle empty messages array', async () => {
+        // Act & Assert
+        await expect(
+          instance.chat({
+            messages: [],
+            model: '@hf/meta-llama/meta-llama-3-8b-instruct',
+            temperature: 1,
+          }),
+        ).resolves.toBeInstanceOf(Response);
+      });
+    });
+  });
+
+  describe('models', () => {
+    it('should send request', async () => {
+      // Arrange
+      const apiKey = 'test_api_key';
+      const instance = new LobeCloudflareAI({ apiKey, baseURLOrAccountID: accountID });
+
+      vi.spyOn(globalThis, 'fetch').mockResolvedValue(
+        new Response(
+          JSON.stringify({
+            result: [
+              {
+                description: 'Model 1',
+                name: 'model1',
+                task: { name: 'Text Generation' },
+                properties: [{ property_id: 'beta', value: 'false' }],
+              },
+              {
+                description: 'Model 2',
+                name: 'model2',
+                task: { name: 'Text Generation' },
+                properties: [{ property_id: 'beta', value: 'true' }],
+              },
+            ],
+          }),
+        ),
+      );
+
+      // Act
+      const result = await instance.models();
+
+      // Assert
+      expect(globalThis.fetch).toHaveBeenCalledWith(
+        `https://api.cloudflare.com/client/v4/accounts/${accountID}/ai/models/search`,
+        {
+          headers: {
+            'Authorization': `Bearer ${apiKey}`,
+            'Content-Type': 'application/json',
+          },
+          method: 'GET',
+        },
+      );
+
+      expect(result).toHaveLength(2);
+    });
+
+    it('should set id to name', async () => {
+      // Arrange
+      const instance = new LobeCloudflareAI({
+        apiKey: 'test_api_key',
+        baseURLOrAccountID: accountID,
+      });
+
+      vi.spyOn(globalThis, 'fetch').mockResolvedValue(
+        new Response(
+          JSON.stringify({
+            result: [
+              {
+                id: 'id1',
+                name: 'name1',
+                task: { name: 'Text Generation' },
+              },
+            ],
+          }),
+        ),
+      );
+
+      // Act
+      const result = await instance.models();
+
+      // Assert
+      expect(result).toEqual([
+        expect.objectContaining({
+          displayName: 'name1',
+          id: 'name1',
+        }),
+      ]);
+    });
+
+    it('should filter text generation models', async () => {
+      // Arrange
+      const instance = new LobeCloudflareAI({
+        apiKey: 'test_api_key',
+        baseURLOrAccountID: accountID,
+      });
+
+      vi.spyOn(globalThis, 'fetch').mockResolvedValue(
+        new Response(
+          JSON.stringify({
+            result: [
+              {
+                id: '1',
+                name: 'model1',
+                task: { name: 'Text Generation' },
+              },
+              {
+                id: '2',
+                name: 'model2',
+                task: { name: 'Text Classification' },
+              },
+            ],
+          }),
+        ),
+      );
+
+      // Act
+      const result = await instance.models();
+
+      // Assert
+      expect(result).toEqual([
+        expect.objectContaining({
+          displayName: 'model1',
+          id: 'model1',
+        }),
+      ]);
+    });
+
+    it('should enable non-beta models and mark beta models', async () => {
+      // Arrange
+      const instance = new LobeCloudflareAI({
+        apiKey: 'test_api_key',
+        baseURLOrAccountID: accountID,
+      });
+
+      vi.spyOn(globalThis, 'fetch').mockResolvedValue(
+        new Response(
+          JSON.stringify({
+            result: [
+              {
+                id: '1',
+                name: 'model1',
+                task: { name: 'Text Generation' },
+                properties: [{ property_id: 'beta', value: 'false' }],
+              },
+              {
+                id: '2',
+                name: 'model2',
+                task: { name: 'Text Generation' },
+                properties: [{ property_id: 'beta', value: 'true' }],
+              },
+            ],
+          }),
+        ),
+      );
+
+      // Act
+      const result = await instance.models();
+
+      // Assert
+      expect(result).toEqual([
+        expect.objectContaining({
+          displayName: 'model1',
+          enabled: true,
+          id: 'model1',
+        }),
+        expect.objectContaining({
+          displayName: 'model2 (Beta)',
+          enabled: false,
+          id: 'model2',
+        }),
+      ]);
+    });
+  });
+});