@lobehub/chat 1.40.4 → 1.42.0

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

@@ -1,10 +1,13 @@
 // @vitest-environment node
 import OpenAI from 'openai';
+import type { Stream } from 'openai/streaming';
+
 import { Mock, afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
 
 import {
   AgentRuntimeErrorType,
   ChatStreamCallbacks,
+  ChatStreamPayload,
   LobeOpenAICompatibleRuntime,
   ModelProvider,
 } from '@/libs/agent-runtime';
@@ -797,6 +800,134 @@ describe('LobeOpenAICompatibleFactory', () => {
       });
     });
 
+    it('should use custom stream handler when provided', async () => {
+      // Create a custom stream handler that handles both ReadableStream and OpenAI Stream
+      const customStreamHandler = vi.fn((stream: ReadableStream | Stream<OpenAI.ChatCompletionChunk>) => {
+        const readableStream = stream instanceof ReadableStream ? stream : stream.toReadableStream();
+        return new ReadableStream({
+          start(controller) {
+            const reader = readableStream.getReader();
+            const process = async () => {
+              try {
+                while (true) {
+                  const { done, value } = await reader.read();
+                  if (done) break;
+                  controller.enqueue(value);
+                }
+              } finally {
+                controller.close();
+              }
+            };
+            process();
+          },
+        });
+      });
+
+      const LobeMockProvider = LobeOpenAICompatibleFactory({
+        baseURL: 'https://api.test.com/v1',
+        chatCompletion: {
+          handleStream: customStreamHandler,
+        },
+        provider: ModelProvider.OpenAI,
+      });
+
+      const instance = new LobeMockProvider({ apiKey: 'test' });
+
+      // Create a mock stream
+      const mockStream = new ReadableStream({
+        start(controller) {
+          controller.enqueue({
+            id: 'test-id',
+            choices: [{ delta: { content: 'Hello' }, index: 0 }],
+            created: Date.now(),
+            model: 'test-model',
+            object: 'chat.completion.chunk',
+          });
+          controller.close();
+        },
+      });
+
+      vi.spyOn(instance['client'].chat.completions, 'create').mockResolvedValue({
+        tee: () => [mockStream, mockStream],
+      } as any);
+
+      const payload: ChatStreamPayload = {
+        messages: [{ content: 'Test', role: 'user' }],
+        model: 'test-model',
+        temperature: 0.7,
+      };
+
+      await instance.chat(payload);
+
+      expect(customStreamHandler).toHaveBeenCalled();
+    });
+
+    it('should use custom transform handler for non-streaming response', async () => {
+      const customTransformHandler = vi.fn((data: OpenAI.ChatCompletion): ReadableStream => {
+        return new ReadableStream({
+          start(controller) {
+            // Transform the completion to chunk format
+            controller.enqueue({
+              id: data.id,
+              choices: data.choices.map((choice) => ({
+                delta: { content: choice.message.content },
+                index: choice.index,
+              })),
+              created: data.created,
+              model: data.model,
+              object: 'chat.completion.chunk',
+            });
+            controller.close();
+          },
+        });
+      });
+
+      const LobeMockProvider = LobeOpenAICompatibleFactory({
+        baseURL: 'https://api.test.com/v1',
+        chatCompletion: {
+          handleTransformResponseToStream: customTransformHandler,
+        },
+        provider: ModelProvider.OpenAI,
+      });
+
+      const instance = new LobeMockProvider({ apiKey: 'test' });
+
+      const mockResponse: OpenAI.ChatCompletion = {
+        id: 'test-id',
+        choices: [
+          {
+            index: 0,
+            message: { 
+              role: 'assistant', 
+              content: 'Test response',
+              refusal: null
+            },
+            logprobs: null,
+            finish_reason: 'stop',
+          },
+        ],
+        created: Date.now(),
+        model: 'test-model',
+        object: 'chat.completion',
+        usage: { completion_tokens: 2, prompt_tokens: 1, total_tokens: 3 },
+      };
+
+      vi.spyOn(instance['client'].chat.completions, 'create').mockResolvedValue(
+        mockResponse as any,
+      );
+
+      const payload: ChatStreamPayload = {
+        messages: [{ content: 'Test', role: 'user' }],
+        model: 'test-model',
+        temperature: 0.7,
+        stream: false,
+      };
+
+      await instance.chat(payload);
+
+      expect(customTransformHandler).toHaveBeenCalledWith(mockResponse);
+    });
+
     describe('DEBUG', () => {
       it('should call debugStream and return StreamingTextResponse when DEBUG_OPENROUTER_CHAT_COMPLETION is 1', async () => {
         // Arrange

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

@@ -25,6 +25,7 @@ import { handleOpenAIError } from '../handleOpenAIError';
 import { convertOpenAIMessages } from '../openaiHelpers';
 import { StreamingResponse } from '../response';
 import { OpenAIStream, OpenAIStreamOptions } from '../streams';
+import { ChatStreamCallbacks } from '../../types';
 
 // the model contains the following keywords is not a chat model, so we should filter them out
 export const CHAT_MODELS_BLOCK_LIST = [
@@ -62,10 +63,17 @@ interface OpenAICompatibleFactoryOptions<T extends Record<string, any> = any> {
       payload: ChatStreamPayload,
       options: ConstructorOptions<T>,
     ) => OpenAI.ChatCompletionCreateParamsStreaming;
+    handleStream?: (
+      stream: Stream<OpenAI.ChatCompletionChunk> | ReadableStream,
+      callbacks?: ChatStreamCallbacks,
+    ) => ReadableStream;
     handleStreamBizErrorType?: (error: {
       message: string;
       name: string;
     }) => ILobeAgentRuntimeErrorType | undefined;
+    handleTransformResponseToStream?: (
+      data: OpenAI.ChatCompletion,
+    ) => ReadableStream<OpenAI.ChatCompletionChunk>;
     noUserId?: boolean;
   };
   constructorOptions?: ConstructorOptions<T>;
@@ -228,7 +236,8 @@ export const LobeOpenAICompatibleFactory = <T extends Record<string, any> = any>
             debugStream(useForDebugStream).catch(console.error);
           }
 
-          return StreamingResponse(OpenAIStream(prod, streamOptions), {
+          const streamHandler = chatCompletion?.handleStream || OpenAIStream;
+          return StreamingResponse(streamHandler(prod, streamOptions), {
             headers: options?.headers,
           });
         }
@@ -239,9 +248,11 @@ export const LobeOpenAICompatibleFactory = <T extends Record<string, any> = any>
 
         if (responseMode === 'json') return Response.json(response);
 
-        const stream = transformResponseToStream(response as unknown as OpenAI.ChatCompletion);
+        const transformHandler = chatCompletion?.handleTransformResponseToStream || transformResponseToStream;
+        const stream = transformHandler(response as unknown as OpenAI.ChatCompletion);
 
-        return StreamingResponse(OpenAIStream(stream, streamOptions), {
+        const streamHandler = chatCompletion?.handleStream || OpenAIStream;
+        return StreamingResponse(streamHandler(stream, streamOptions), {
           headers: options?.headers,
         });
       } catch (error) {

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

@@ -7,3 +7,4 @@ export * from './ollama';
 export * from './openai';
 export * from './protocol';
 export * from './qwen';
+export * from './spark';

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

@@ -0,0 +1,199 @@
+import { beforeAll, describe, expect, it, vi } from 'vitest';
+import { SparkAIStream, transformSparkResponseToStream } from './spark';
+import type OpenAI from 'openai';
+
+describe('SparkAIStream', () => {
+  beforeAll(() => {});
+
+  it('should transform non-streaming response to stream', async () => {
+    const mockResponse = {
+      id: "cha000ceba6@dx193d200b580b8f3532",
+      object: "chat.completion",
+      created: 1734395014,
+      model: "max-32k",
+      choices: [
+        {
+          message: {
+            role: "assistant",
+            content: "",
+            refusal: null,
+            tool_calls: {
+              type: "function",
+              function: {
+                arguments: '{"city":"Shanghai"}',
+                name: "realtime-weather____fetchCurrentWeather"
+              },
+              id: "call_1"
+            }
+          },
+          index: 0,
+          logprobs: null,
+          finish_reason: "tool_calls"
+        }
+      ],
+      usage: {
+        prompt_tokens: 8,
+        completion_tokens: 0,
+        total_tokens: 8
+      }
+    } as unknown as OpenAI.ChatCompletion;
+
+    const stream = transformSparkResponseToStream(mockResponse);
+    const decoder = new TextDecoder();
+    const chunks = [];
+
+    // @ts-ignore
+    for await (const chunk of stream) {
+      chunks.push(chunk);
+    }
+
+    expect(chunks).toHaveLength(2);
+    expect(chunks[0].choices[0].delta.tool_calls).toEqual([{
+      function: {
+        arguments: '{"city":"Shanghai"}',
+        name: "realtime-weather____fetchCurrentWeather"
+      },
+      id: "call_1",
+      index: 0,
+      type: "function"
+    }]);
+    expect(chunks[1].choices[0].finish_reason).toBeDefined();
+  });
+
+  it('should transform streaming response with tool calls', async () => {
+    const mockStream = new ReadableStream({
+      start(controller) {
+        controller.enqueue({
+          id: "cha000b0bf9@dx193d1ffa61cb894532",
+          object: "chat.completion.chunk",
+          created: 1734395014,
+          model: "max-32k",
+          choices: [
+            {
+              delta: {
+                role: "assistant",
+                content: "",
+                tool_calls: {
+                  type: "function",
+                  function: {
+                    arguments: '{"city":"Shanghai"}',
+                    name: "realtime-weather____fetchCurrentWeather"
+                  },
+                  id: "call_1"
+                }
+              },
+              index: 0
+            }
+          ]
+        } as unknown as OpenAI.ChatCompletionChunk);
+        controller.close();
+      }
+    });
+
+    const onToolCallMock = vi.fn();
+
+    const protocolStream = SparkAIStream(mockStream, {
+      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: cha000b0bf9@dx193d1ffa61cb894532\n',
+      'event: tool_calls\n',
+      `data: [{\"function\":{\"arguments\":\"{\\\"city\\\":\\\"Shanghai\\\"}\",\"name\":\"realtime-weather____fetchCurrentWeather\"},\"id\":\"call_1\",\"index\":0,\"type\":\"function\"}]\n\n`
+    ]);
+
+    expect(onToolCallMock).toHaveBeenCalledTimes(1);
+  });
+
+  it('should handle text content in stream', async () => {
+    const mockStream = new ReadableStream({
+      start(controller) {
+        controller.enqueue({
+          id: "test-id",
+          object: "chat.completion.chunk",
+          created: 1734395014,
+          model: "max-32k",
+          choices: [
+            {
+              delta: {
+                content: "Hello",
+                role: "assistant"
+              },
+              index: 0
+            }
+          ]
+        } as OpenAI.ChatCompletionChunk);
+        controller.enqueue({
+          id: "test-id",
+          object: "chat.completion.chunk",
+          created: 1734395014,
+          model: "max-32k",
+          choices: [
+            {
+              delta: {
+                content: " World",
+                role: "assistant"
+              },
+              index: 0
+            }
+          ]
+        } as OpenAI.ChatCompletionChunk);
+        controller.close();
+      }
+    });
+
+    const onTextMock = vi.fn();
+    
+    const protocolStream = SparkAIStream(mockStream, {
+      onText: onTextMock
+    });
+
+    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: test-id\n',
+      'event: text\n',
+      'data: "Hello"\n\n',
+      'id: test-id\n',
+      'event: text\n',
+      'data: " World"\n\n'
+    ]);
+
+    expect(onTextMock).toHaveBeenNthCalledWith(1, '"Hello"');
+    expect(onTextMock).toHaveBeenNthCalledWith(2, '" World"');
+  });
+
+  it('should handle empty stream', async () => {
+    const mockStream = new ReadableStream({
+      start(controller) {
+        controller.close();
+      }
+    });
+
+    const protocolStream = SparkAIStream(mockStream);
+
+    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/spark.ts

@@ -0,0 +1,134 @@
+import OpenAI from 'openai';
+import type { Stream } from 'openai/streaming';
+
+import { ChatStreamCallbacks } from '../../types';
+import {
+  StreamProtocolChunk,
+  StreamProtocolToolCallChunk,
+  convertIterableToStream,
+  createCallbacksTransformer,
+  createSSEProtocolTransformer,
+  generateToolCallId,
+} from './protocol';
+
+export function transformSparkResponseToStream(data: OpenAI.ChatCompletion) {
+  return new ReadableStream({
+    start(controller) {
+      const chunk: OpenAI.ChatCompletionChunk = {
+        choices: data.choices.map((choice: OpenAI.ChatCompletion.Choice) => {
+          const toolCallsArray = choice.message.tool_calls
+            ? Array.isArray(choice.message.tool_calls)
+              ? choice.message.tool_calls
+              : [choice.message.tool_calls]
+            : []; // 如果不是数组，包装成数组
+
+          return {
+            delta: {
+              content: choice.message.content,
+              role: choice.message.role,
+              tool_calls: toolCallsArray.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: null,
+            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();
+    },
+  });
+}
+
+export const transformSparkStream = (chunk: OpenAI.ChatCompletionChunk): StreamProtocolChunk => {
+  const item = chunk.choices[0];
+
+  if (!item) {
+    return { data: chunk, id: chunk.id, type: 'data' };
+  }
+
+  if (item.delta?.tool_calls) {
+    const toolCallsArray = Array.isArray(item.delta.tool_calls)
+      ? item.delta.tool_calls
+      : [item.delta.tool_calls]; // 如果不是数组，包装成数组
+
+    if (toolCallsArray.length > 0) {
+      return {
+        data: toolCallsArray.map((toolCall, index) => ({
+          function: toolCall.function,
+          id: toolCall.id || generateToolCallId(index, toolCall.function?.name),
+          index: typeof toolCall.index !== 'undefined' ? toolCall.index : index,
+          type: toolCall.type || 'function',
+        })),
+        id: chunk.id,
+        type: 'tool_calls',
+      } as StreamProtocolToolCallChunk;
+    }
+  }
+
+  if (item.finish_reason) {
+    // one-api 的流式接口，会出现既有 finish_reason ，也有 content 的情况
+    //  {"id":"demo","model":"deepl-en","choices":[{"index":0,"delta":{"role":"assistant","content":"Introduce yourself."},"finish_reason":"stop"}]}
+
+    if (typeof item.delta?.content === 'string' && !!item.delta.content) {
+      return { data: item.delta.content, id: chunk.id, type: 'text' };
+    }
+
+    return { data: item.finish_reason, id: chunk.id, type: 'stop' };
+  }
+
+  if (typeof item.delta?.content === 'string') {
+    return { data: item.delta.content, id: chunk.id, type: 'text' };
+  }
+
+  if (item.delta?.content === null) {
+    return { data: item.delta, id: chunk.id, type: 'data' };
+  }
+
+  return {
+    data: { delta: item.delta, id: chunk.id, index: item.index },
+    id: chunk.id,
+    type: 'data',
+  };
+};
+
+export const SparkAIStream = (
+  stream: Stream<OpenAI.ChatCompletionChunk> | ReadableStream,
+  callbacks?: ChatStreamCallbacks,
+) => {
+  const readableStream =
+    stream instanceof ReadableStream ? stream : convertIterableToStream(stream);
+
+  return readableStream
+    .pipeThrough(createSSEProtocolTransformer(transformSparkStream))
+    .pipeThrough(createCallbacksTransformer(callbacks));
+};

package/src/libs/next-auth/sso-providers/index.ts

@@ -8,6 +8,7 @@ import GenericOIDC from './generic-oidc';
 import Github from './github';
 import Logto from './logto';
 import MicrosoftEntraID from './microsoft-entra-id';
+import WeChat from './wechat';
 import Zitadel from './zitadel';
 
 export const ssoProviders = [
@@ -22,4 +23,5 @@ export const ssoProviders = [
   CloudflareZeroTrust,
   Casdoor,
   MicrosoftEntraID,
+  WeChat,
 ];

package/src/libs/next-auth/sso-providers/wechat.ts

@@ -0,0 +1,24 @@
+import WeChat from '@auth/core/providers/wechat';
+
+import { CommonProviderConfig } from './sso.config';
+
+const provider = {
+  id: 'wechat',
+  provider: WeChat({
+    ...CommonProviderConfig,
+    clientId: process.env.AUTH_WECHAT_ID,
+    clientSecret: process.env.AUTH_WECHAT_SECRET,
+    platformType: 'WebsiteApp',
+    profile: (profile) => {
+      return {
+        email: null,
+        id: profile.unionid,
+        image: profile.headimgurl,
+        name: profile.nickname,
+        providerAccountId: profile.unionid,
+      };
+    },
+  }),
+};
+
+export default provider;

package/src/server/modules/AssistantStore/index.test.ts

@@ -7,19 +7,19 @@ const baseURL = 'https://registry.npmmirror.com/@lobehub/agents-index/v1/files/p
 describe('AssistantStore', () => {
   it('should return the default index URL when no language is provided', () => {
     const agentMarket = new AssistantStore();
-    const url = agentMarket.getAgentIndexUrl();
+    const url = agentMarket['getAgentIndexUrl']();
     expect(url).toBe(`${baseURL}/index.en-US.json`);
   });
 
   it('should return the index URL for a not supported language', () => {
     const agentMarket = new AssistantStore();
-    const url = agentMarket.getAgentIndexUrl('xxx' as any);
+    const url = agentMarket['getAgentIndexUrl']('xxx' as any);
     expect(url).toBe('https://registry.npmmirror.com/@lobehub/agents-index/v1/files/public');
   });
 
   it('should return the zh-CN URL for zh locale', () => {
     const agentMarket = new AssistantStore();
-    const url = agentMarket.getAgentIndexUrl('zh' as any);
+    const url = agentMarket['getAgentIndexUrl']('zh' as any);
     expect(url).toBe(
       'https://registry.npmmirror.com/@lobehub/agents-index/v1/files/public/index.zh-CN.json',
     );
@@ -27,7 +27,7 @@ describe('AssistantStore', () => {
 
   it('should return the default URL for en locale', () => {
     const agentMarket = new AssistantStore();
-    const url = agentMarket.getAgentIndexUrl('en' as any);
+    const url = agentMarket['getAgentIndexUrl']('en' as any);
     expect(url).toBe(
       'https://registry.npmmirror.com/@lobehub/agents-index/v1/files/public/index.en-US.json',
     );
@@ -35,7 +35,7 @@ describe('AssistantStore', () => {
 
   it('should return the base URL if the provided language is not supported', () => {
     const agentMarket = new AssistantStore();
-    const url = agentMarket.getAgentIndexUrl('fr' as any);
+    const url = agentMarket['getAgentIndexUrl']('fr' as any);
     expect(url).toBe(baseURL);
   });
 

package/src/server/modules/AssistantStore/index.ts

@@ -3,6 +3,8 @@ import urlJoin from 'url-join';
 import { appEnv } from '@/config/app';
 import { DEFAULT_LANG, isLocaleNotSupport } from '@/const/locale';
 import { Locales, normalizeLocale } from '@/locales/resources';
+import { EdgeConfig } from '@/server/modules/EdgeConfig';
+import { AgentStoreIndex } from '@/types/discover';
 
 export class AssistantStore {
   private readonly baseUrl: string;
@@ -11,7 +13,7 @@ export class AssistantStore {
     this.baseUrl = baseUrl || appEnv.AGENTS_INDEX_URL;
   }
 
-  getAgentIndexUrl = (lang: Locales = DEFAULT_LANG) => {
+  private getAgentIndexUrl = (lang: Locales = DEFAULT_LANG) => {
     if (isLocaleNotSupport(lang)) return this.baseUrl;
 
     return urlJoin(this.baseUrl, `index.${normalizeLocale(lang)}.json`);
@@ -22,4 +24,40 @@ export class AssistantStore {
 
     return urlJoin(this.baseUrl, `${identifier}.${normalizeLocale(lang)}.json`);
   };
+
+  getAgentIndex = async (locale: Locales = DEFAULT_LANG, revalidate?: number) => {
+    try {
+      let res: Response;
+
+      res = await fetch(this.getAgentIndexUrl(locale as any), { next: { revalidate } });
+
+      if (res.status === 404) {
+        res = await fetch(this.getAgentIndexUrl(DEFAULT_LANG), { next: { revalidate } });
+      }
+
+      if (!res.ok) {
+        console.error('fetch agent index error:', await res.text());
+        return [];
+      }
+
+      const data: AgentStoreIndex = await res.json();
+
+      // Get the assistant whitelist from Edge Config
+      const edgeConfig = new EdgeConfig();
+
+      if (!!appEnv.VERCEL_EDGE_CONFIG) {
+        const assistantWhitelist = await edgeConfig.getAgentWhitelist();
+
+        if (assistantWhitelist && assistantWhitelist?.length > 0) {
+          data.agents = data.agents.filter((item) => assistantWhitelist.includes(item.identifier));
+        }
+      }
+
+      return data;
+    } catch (e) {
+      console.error('fetch agent index error:', e);
+
+      throw e;
+    }
+  };
 }