@lobehub/chat 1.29.6 → 1.31.0

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

@@ -0,0 +1,123 @@
+import { ChatModelCard } from '@/types/llm';
+
+import { LobeRuntimeAI } from '../BaseAI';
+import { AgentRuntimeErrorType } from '../error';
+import { ChatCompetitionOptions, ChatStreamPayload, ModelProvider } from '../types';
+import {
+  CloudflareStreamTransformer,
+  DEFAULT_BASE_URL_PREFIX,
+  convertModelManifest,
+  desensitizeCloudflareUrl,
+  fillUrl,
+} from '../utils/cloudflareHelpers';
+import { AgentRuntimeError } from '../utils/createError';
+import { debugStream } from '../utils/debugStream';
+import { StreamingResponse } from '../utils/response';
+import { createCallbacksTransformer } from '../utils/streams';
+
+export interface LobeCloudflareParams {
+  apiKey?: string;
+  baseURLOrAccountID?: string;
+}
+
+export class LobeCloudflareAI implements LobeRuntimeAI {
+  baseURL: string;
+  accountID: string;
+  apiKey?: string;
+
+  constructor({ apiKey, baseURLOrAccountID }: LobeCloudflareParams) {
+    if (!baseURLOrAccountID) {
+      throw AgentRuntimeError.createError(AgentRuntimeErrorType.InvalidProviderAPIKey);
+    }
+    if (baseURLOrAccountID.startsWith('http')) {
+      this.baseURL = baseURLOrAccountID.endsWith('/')
+        ? baseURLOrAccountID
+        : baseURLOrAccountID + '/';
+      // Try get accountID from baseURL
+      this.accountID = baseURLOrAccountID.replaceAll(/^.*\/([\dA-Fa-f]{32})\/.*$/g, '$1');
+    } else {
+      if (!apiKey) {
+        throw AgentRuntimeError.createError(AgentRuntimeErrorType.InvalidProviderAPIKey);
+      }
+      this.accountID = baseURLOrAccountID;
+      this.baseURL = fillUrl(baseURLOrAccountID);
+    }
+    this.apiKey = apiKey;
+  }
+
+  async chat(payload: ChatStreamPayload, options?: ChatCompetitionOptions): Promise<Response> {
+    try {
+      const { model, tools, ...restPayload } = payload;
+      const functions = tools?.map((tool) => tool.function);
+      const headers = options?.headers || {};
+      if (this.apiKey) {
+        headers['Authorization'] = `Bearer ${this.apiKey}`;
+      }
+      const url = new URL(model, this.baseURL);
+      const response = await fetch(url, {
+        body: JSON.stringify({ tools: functions, ...restPayload }),
+        headers: { 'Content-Type': 'application/json', ...headers },
+        method: 'POST',
+        signal: options?.signal,
+      });
+
+      const desensitizedEndpoint = desensitizeCloudflareUrl(url.toString());
+
+      switch (response.status) {
+        case 400: {
+          throw AgentRuntimeError.chat({
+            endpoint: desensitizedEndpoint,
+            error: response,
+            errorType: AgentRuntimeErrorType.ProviderBizError,
+            provider: ModelProvider.Cloudflare,
+          });
+        }
+      }
+
+      // Only tee when debugging
+      let responseBody: ReadableStream;
+      if (process.env.DEBUG_CLOUDFLARE_CHAT_COMPLETION === '1') {
+        const [prod, useForDebug] = response.body!.tee();
+        debugStream(useForDebug).catch();
+        responseBody = prod;
+      } else {
+        responseBody = response.body!;
+      }
+
+      return StreamingResponse(
+        responseBody
+          .pipeThrough(new TransformStream(new CloudflareStreamTransformer()))
+          .pipeThrough(createCallbacksTransformer(options?.callback)),
+        { headers: options?.headers },
+      );
+    } catch (error) {
+      const desensitizedEndpoint = desensitizeCloudflareUrl(this.baseURL);
+
+      throw AgentRuntimeError.chat({
+        endpoint: desensitizedEndpoint,
+        error: error as any,
+        errorType: AgentRuntimeErrorType.ProviderBizError,
+        provider: ModelProvider.Cloudflare,
+      });
+    }
+  }
+
+  async models(): Promise<ChatModelCard[]> {
+    const url = `${DEFAULT_BASE_URL_PREFIX}/client/v4/accounts/${this.accountID}/ai/models/search`;
+    const response = await fetch(url, {
+      headers: {
+        'Authorization': `Bearer ${this.apiKey}`,
+        'Content-Type': 'application/json',
+      },
+      method: 'GET',
+    });
+    const j = await response.json();
+    const models: any[] = j['result'].filter(
+      (model: any) => model['task']['name'] === 'Text Generation',
+    );
+    const chatModels: ChatModelCard[] = models
+      .map((model) => convertModelManifest(model))
+      .sort((a, b) => a.displayName.localeCompare(b.displayName));
+    return chatModels;
+  }
+}

package/src/libs/agent-runtime/types/type.ts

@@ -28,6 +28,7 @@ export enum ModelProvider {
   Azure = 'azure',
   Baichuan = 'baichuan',
   Bedrock = 'bedrock',
+  Cloudflare = 'cloudflare',
   DeepSeek = 'deepseek',
   FireworksAI = 'fireworksai',
   Github = 'github',
@@ -52,6 +53,7 @@ export enum ModelProvider {
   TogetherAI = 'togetherai',
   Upstage = 'upstage',
   Wenxin = 'wenxin',
+  XAI = 'xai',
   ZeroOne = 'zeroone',
   ZhiPu = 'zhipu',
 }

package/src/libs/agent-runtime/utils/cloudflareHelpers.test.ts

@@ -0,0 +1,339 @@
+// @vitest-environment node
+import { Mock, afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import * as desensitizeTool from '../utils/desensitizeUrl';
+import {
+  CloudflareStreamTransformer,
+  desensitizeCloudflareUrl,
+  fillUrl,
+  getModelBeta,
+  getModelDisplayName,
+  getModelFunctionCalling,
+  getModelTokens,
+} from './cloudflareHelpers';
+
+//const {
+//  getModelBeta,
+//  getModelDisplayName,
+//  getModelFunctionCalling,
+//  getModelTokens,
+//} = require('./cloudflareHelpers');
+
+//const cloudflareHelpers = require('./cloudflareHelpers');
+//const getModelBeta = cloudflareHelpers.__get__('getModelBeta');
+//const getModelDisplayName = cloudflareHelpers.__get__('getModelDisplayName');
+//const getModelFunctionCalling = cloudflareHelpers.__get__('getModelFunctionCalling');
+//const getModelTokens = cloudflareHelpers.__get__('getModelTokens');
+
+afterEach(() => {
+  vi.restoreAllMocks();
+});
+
+describe('cloudflareHelpers', () => {
+  describe('CloudflareStreamTransformer', () => {
+    let transformer: CloudflareStreamTransformer;
+    beforeEach(() => {
+      transformer = new CloudflareStreamTransformer();
+    });
+
+    describe('parseChunk', () => {
+      let chunks: string[];
+      let controller: TransformStreamDefaultController;
+
+      beforeEach(() => {
+        chunks = [];
+        controller = Object.create(TransformStreamDefaultController.prototype);
+        vi.spyOn(controller, 'enqueue').mockImplementation((chunk) => {
+          chunks.push(chunk);
+        });
+      });
+
+      it('should parse chunk', () => {
+        // Arrange
+        const chunk = 'data: {"key": "value", "response": "response1"}';
+        const textDecoder = new TextDecoder();
+
+        // Act
+        transformer['parseChunk'](chunk, controller);
+
+        // Assert
+        expect(chunks.length).toBe(2);
+        expect(chunks[0]).toBe('event: text\n');
+        expect(chunks[1]).toBe('data: "response1"\n\n');
+      });
+
+      it('should not replace `data` in text', () => {
+        // Arrange
+        const chunk = 'data: {"key": "value", "response": "data: a"}';
+        const textDecoder = new TextDecoder();
+
+        // Act
+        transformer['parseChunk'](chunk, controller);
+
+        // Assert
+        expect(chunks.length).toBe(2);
+        expect(chunks[0]).toBe('event: text\n');
+        expect(chunks[1]).toBe('data: "data: a"\n\n');
+      });
+    });
+
+    describe('transform', () => {
+      const textDecoder = new TextDecoder();
+      const textEncoder = new TextEncoder();
+      let chunks: string[];
+
+      beforeEach(() => {
+        chunks = [];
+        vi.spyOn(
+          transformer as any as {
+            parseChunk: (chunk: string, controller: TransformStreamDefaultController) => void;
+          },
+          'parseChunk',
+        ).mockImplementation((chunk: string, _) => {
+          chunks.push(chunk);
+        });
+      });
+
+      it('should split single chunk', async () => {
+        // Arrange
+        const chunk = textEncoder.encode('data: {"key": "value", "response": "response1"}\n\n');
+
+        // Act
+        await transformer.transform(chunk, undefined!);
+
+        // Assert
+        expect(chunks.length).toBe(1);
+        expect(chunks[0]).toBe('data: {"key": "value", "response": "response1"}');
+      });
+
+      it('should split multiple chunks', async () => {
+        // Arrange
+        const chunk = textEncoder.encode(
+          'data: {"key": "value", "response": "response1"}\n\n' +
+            'data: {"key": "value", "response": "response2"}\n\n',
+        );
+
+        // Act
+        await transformer.transform(chunk, undefined!);
+
+        // Assert
+        expect(chunks.length).toBe(2);
+        expect(chunks[0]).toBe('data: {"key": "value", "response": "response1"}');
+        expect(chunks[1]).toBe('data: {"key": "value", "response": "response2"}');
+      });
+
+      it('should ignore empty chunk', async () => {
+        // Arrange
+        const chunk = textEncoder.encode('\n\n');
+
+        // Act
+        await transformer.transform(chunk, undefined!);
+
+        // Assert
+        expect(chunks.join()).toBe('');
+      });
+
+      it('should split and concat delayed chunks', async () => {
+        // Arrange
+        const chunk1 = textEncoder.encode('data: {"key": "value", "respo');
+        const chunk2 = textEncoder.encode('nse": "response1"}\n\ndata: {"key": "val');
+        const chunk3 = textEncoder.encode('ue", "response": "response2"}\n\n');
+
+        // Act & Assert
+        await transformer.transform(chunk1, undefined!);
+        expect(transformer['parseChunk']).not.toHaveBeenCalled();
+        expect(chunks.length).toBe(0);
+        expect(transformer['buffer']).toBe('data: {"key": "value", "respo');
+
+        await transformer.transform(chunk2, undefined!);
+        expect(chunks.length).toBe(1);
+        expect(chunks[0]).toBe('data: {"key": "value", "response": "response1"}');
+        expect(transformer['buffer']).toBe('data: {"key": "val');
+
+        await transformer.transform(chunk3, undefined!);
+        expect(chunks.length).toBe(2);
+        expect(chunks[1]).toBe('data: {"key": "value", "response": "response2"}');
+        expect(transformer['buffer']).toBe('');
+      });
+
+      it('should ignore standalone [DONE]', async () => {
+        // Arrange
+        const chunk = textEncoder.encode('data: [DONE]\n\n');
+
+        // Act
+        await transformer.transform(chunk, undefined!);
+
+        // Assert
+        expect(transformer['parseChunk']).not.toHaveBeenCalled();
+        expect(chunks.length).toBe(0);
+        expect(transformer['buffer']).toBe('');
+      });
+
+      it('should ignore [DONE] in chunk', async () => {
+        // Arrange
+        const chunk = textEncoder.encode(
+          'data: {"key": "value", "response": "response1"}\n\ndata: [DONE]\n\n',
+        );
+
+        // Act
+        await transformer.transform(chunk, undefined!);
+
+        // Assert
+        expect(chunks.length).toBe(1);
+        expect(chunks[0]).toBe('data: {"key": "value", "response": "response1"}');
+        expect(transformer['buffer']).toBe('');
+      });
+    });
+  });
+
+  describe('fillUrl', () => {
+    it('should return URL with account id', () => {
+      const url = fillUrl('80009000a000b000c000d000e000f000');
+      expect(url).toBe(
+        'https://api.cloudflare.com/client/v4/accounts/80009000a000b000c000d000e000f000/ai/run/',
+      );
+    });
+  });
+
+  describe('maskAccountId', () => {
+    describe('desensitizeAccountId', () => {
+      it('should replace account id with **** in official API endpoint', () => {
+        const url =
+          'https://api.cloudflare.com/client/v4/accounts/80009000a000b000c000d000e000f000/ai/run/';
+        const maskedUrl = desensitizeCloudflareUrl(url);
+        expect(maskedUrl).toBe('https://api.cloudflare.com/client/v4/accounts/****/ai/run/');
+      });
+
+      it('should replace account id with **** in custom API endpoint', () => {
+        const url =
+          'https://api.cloudflare.com/custom/prefix/80009000a000b000c000d000e000f000/custom/suffix/';
+        const maskedUrl = desensitizeCloudflareUrl(url);
+        expect(maskedUrl).toBe('https://api.cloudflare.com/custom/prefix/****/custom/suffix/');
+      });
+    });
+
+    describe('desensitizeCloudflareUrl', () => {
+      it('should mask account id in official API endpoint', () => {
+        const url =
+          'https://api.cloudflare.com/client/v4/accounts/80009000a000b000c000d000e000f000/ai/run/';
+        const maskedUrl = desensitizeCloudflareUrl(url);
+        expect(maskedUrl).toBe('https://api.cloudflare.com/client/v4/accounts/****/ai/run/');
+      });
+
+      it('should call desensitizeUrl for custom API endpoint', () => {
+        const url = 'https://custom.url/path';
+        vi.spyOn(desensitizeTool, 'desensitizeUrl').mockImplementation(
+          (_) => 'https://custom.mocked.url',
+        );
+        const maskedUrl = desensitizeCloudflareUrl(url);
+        expect(desensitizeTool.desensitizeUrl).toHaveBeenCalledWith('https://custom.url');
+        expect(maskedUrl).toBe('https://custom.mocked.url/path');
+      });
+
+      it('should mask account id in custom API endpoint', () => {
+        const url =
+          'https://custom.url/custom/prefix/80009000a000b000c000d000e000f000/custom/suffix/';
+        const maskedUrl = desensitizeCloudflareUrl(url);
+        expect(maskedUrl).toBe('https://cu****om.url/custom/prefix/****/custom/suffix/');
+      });
+
+      it('should mask account id in custom API endpoint with query params', () => {
+        const url =
+          'https://custom.url/custom/prefix/80009000a000b000c000d000e000f000/custom/suffix/?query=param';
+        const maskedUrl = desensitizeCloudflareUrl(url);
+        expect(maskedUrl).toBe(
+          'https://cu****om.url/custom/prefix/****/custom/suffix/?query=param',
+        );
+      });
+
+      it('should mask account id in custom API endpoint with port', () => {
+        const url =
+          'https://custom.url:8080/custom/prefix/80009000a000b000c000d000e000f000/custom/suffix/';
+        const maskedUrl = desensitizeCloudflareUrl(url);
+        expect(maskedUrl).toBe('https://cu****om.url:****/custom/prefix/****/custom/suffix/');
+      });
+    });
+  });
+
+  describe('modelManifest', () => {
+    describe('getModelBeta', () => {
+      it('should get beta property', () => {
+        const model = { properties: [{ property_id: 'beta', value: 'true' }] };
+        const beta = getModelBeta(model);
+        expect(beta).toBe(true);
+      });
+
+      it('should return false if beta property is false', () => {
+        const model = { properties: [{ property_id: 'beta', value: 'false' }] };
+        const beta = getModelBeta(model);
+        expect(beta).toBe(false);
+      });
+
+      it('should return false if beta property is not present', () => {
+        const model = { properties: [] };
+        const beta = getModelBeta(model);
+        expect(beta).toBe(false);
+      });
+    });
+
+    describe('getModelDisplayName', () => {
+      it('should return display name with beta suffix', () => {
+        const model = { name: 'model', properties: [{ property_id: 'beta', value: 'true' }] };
+        const name = getModelDisplayName(model, true);
+        expect(name).toBe('model (Beta)');
+      });
+
+      it('should return display name without beta suffix', () => {
+        const model = { name: 'model', properties: [] };
+        const name = getModelDisplayName(model, false);
+        expect(name).toBe('model');
+      });
+
+      it('should return model["name"]', () => {
+        const model = { id: 'modelID', name: 'modelName' };
+        const name = getModelDisplayName(model, false);
+        expect(name).toBe('modelName');
+      });
+
+      it('should return last part of model["name"]', () => {
+        const model = { name: '@provider/modelFamily/modelName' };
+        const name = getModelDisplayName(model, false);
+        expect(name).toBe('modelName');
+      });
+    });
+
+    describe('getModelFunctionCalling', () => {
+      it('should return true if function_calling property is true', () => {
+        const model = { properties: [{ property_id: 'function_calling', value: 'true' }] };
+        const functionCalling = getModelFunctionCalling(model);
+        expect(functionCalling).toBe(true);
+      });
+
+      it('should return false if function_calling property is false', () => {
+        const model = { properties: [{ property_id: 'function_calling', value: 'false' }] };
+        const functionCalling = getModelFunctionCalling(model);
+        expect(functionCalling).toBe(false);
+      });
+
+      it('should return false if function_calling property is not set', () => {
+        const model = { properties: [] };
+        const functionCalling = getModelFunctionCalling(model);
+        expect(functionCalling).toBe(false);
+      });
+    });
+
+    describe('getModelTokens', () => {
+      it('should return tokens property value', () => {
+        const model = { properties: [{ property_id: 'max_total_tokens', value: '100' }] };
+        const tokens = getModelTokens(model);
+        expect(tokens).toBe(100);
+      });
+
+      it('should return undefined if tokens property is not present', () => {
+        const model = { properties: [] };
+        const tokens = getModelTokens(model);
+        expect(tokens).toBeUndefined();
+      });
+    });
+  });
+});

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

@@ -0,0 +1,134 @@
+import { desensitizeUrl } from '../utils/desensitizeUrl';
+
+class CloudflareStreamTransformer {
+  private textDecoder = new TextDecoder();
+  private buffer: string = '';
+
+  private parseChunk(chunk: string, controller: TransformStreamDefaultController) {
+    const dataPrefix = /^data: /;
+    const json = chunk.replace(dataPrefix, '');
+    const parsedChunk = JSON.parse(json);
+    controller.enqueue(`event: text\n`);
+    controller.enqueue(`data: ${JSON.stringify(parsedChunk.response)}\n\n`);
+  }
+
+  public async transform(chunk: Uint8Array, controller: TransformStreamDefaultController) {
+    let textChunk = this.textDecoder.decode(chunk);
+    if (this.buffer.trim() !== '') {
+      textChunk = this.buffer + textChunk;
+      this.buffer = '';
+    }
+    const splits = textChunk.split('\n\n');
+    for (let i = 0; i < splits.length - 1; i++) {
+      if (/\[DONE]/.test(splits[i].trim())) {
+        return;
+      }
+      this.parseChunk(splits[i], controller);
+    }
+    const lastChunk = splits.at(-1)!;
+    if (lastChunk.trim() !== '') {
+      this.buffer += lastChunk; // does not need to be trimmed.
+    } // else drop.
+  }
+}
+
+const CF_PROPERTY_NAME = 'property_id';
+const DEFAULT_BASE_URL_PREFIX = 'https://api.cloudflare.com';
+
+function fillUrl(accountID: string): string {
+  return `${DEFAULT_BASE_URL_PREFIX}/client/v4/accounts/${accountID}/ai/run/`;
+}
+
+function desensitizeAccountId(path: string): string {
+  return path.replace(/\/[\dA-Fa-f]{32}\//, '/****/');
+}
+
+function desensitizeCloudflareUrl(url: string): string {
+  const urlObj = new URL(url);
+  let { protocol, hostname, port, pathname, search } = urlObj;
+  if (url.startsWith(DEFAULT_BASE_URL_PREFIX)) {
+    return `${protocol}//${hostname}${port ? `:${port}` : ''}${desensitizeAccountId(pathname)}${search}`;
+  } else {
+    const desensitizedUrl = desensitizeUrl(`${protocol}//${hostname}${port ? `:${port}` : ''}`);
+    if (desensitizedUrl.endsWith('/') && pathname.startsWith('/')) {
+      pathname = pathname.slice(1);
+    }
+    return `${desensitizedUrl}${desensitizeAccountId(pathname)}${search}`;
+  }
+}
+
+function getModelBeta(model: any): boolean {
+  try {
+    const betaProperty = model['properties'].filter(
+      (property: any) => property[CF_PROPERTY_NAME] === 'beta',
+    );
+    if (betaProperty.length === 1) {
+      return betaProperty[0]['value'] === 'true'; // This is a string now.
+    }
+    return false;
+  } catch {
+    return false;
+  }
+}
+
+function getModelDisplayName(model: any, beta: boolean): string {
+  const modelId = model['name'];
+  let name = modelId.split('/').at(-1)!;
+  if (beta) {
+    name += ' (Beta)';
+  }
+  return name;
+}
+
+// eslint-disable-next-line @typescript-eslint/no-unused-vars, unused-imports/no-unused-vars
+function getModelFunctionCalling(model: any): boolean {
+  try {
+    const fcProperty = model['properties'].filter(
+      (property: any) => property[CF_PROPERTY_NAME] === 'function_calling',
+    );
+    if (fcProperty.length === 1) {
+      return fcProperty[0]['value'] === 'true';
+    }
+    return false;
+  } catch {
+    return false;
+  }
+}
+
+function getModelTokens(model: any): number | undefined {
+  try {
+    const tokensProperty = model['properties'].filter(
+      (property: any) => property[CF_PROPERTY_NAME] === 'max_total_tokens',
+    );
+    if (tokensProperty.length === 1) {
+      return parseInt(tokensProperty[0]['value']);
+    }
+    return undefined;
+  } catch {
+    return undefined;
+  }
+}
+
+function convertModelManifest(model: any) {
+  const modelBeta = getModelBeta(model);
+  return {
+    description: model['description'],
+    displayName: getModelDisplayName(model, modelBeta),
+    enabled: !modelBeta,
+    functionCall: false, //getModelFunctionCalling(model),
+    id: model['name'],
+    tokens: getModelTokens(model),
+  };
+}
+
+export {
+  CloudflareStreamTransformer,
+  convertModelManifest,
+  DEFAULT_BASE_URL_PREFIX,
+  desensitizeCloudflareUrl,
+  fillUrl,
+  getModelBeta,
+  getModelDisplayName,
+  getModelFunctionCalling,
+  getModelTokens,
+};