@lobehub/chat 1.18.2 → 1.19.1

package/src/features/Conversation/components/MarkdownElements/LobeThinking/Render.tsx

@@ -0,0 +1,86 @@
+import { Icon } from '@lobehub/ui';
+import { createStyles } from 'antd-style';
+import { BringToFrontIcon, ChevronDown, ChevronRight, Loader2Icon } from 'lucide-react';
+import { memo, useState } from 'react';
+import { useTranslation } from 'react-i18next';
+import { Flexbox } from 'react-layout-kit';
+
+import { ARTIFACT_THINKING_TAG } from '@/const/plugin';
+import { useChatStore } from '@/store/chat';
+import { chatSelectors } from '@/store/chat/selectors';
+import { dotLoading } from '@/styles/loading';
+
+import { MarkdownElementProps } from '../type';
+
+/**
+ * Replace all line breaks in the matched `lobeArtifact` tag with an empty string
+ */
+export const isLobeThinkingClosed = (input: string = '') => {
+  const openTag = `<${ARTIFACT_THINKING_TAG}>`;
+  const closeTag = `</${ARTIFACT_THINKING_TAG}>`;
+
+  return input.includes(openTag) && input.includes(closeTag);
+};
+
+const useStyles = createStyles(({ css, token }) => ({
+  container: css`
+    cursor: pointer;
+
+    padding-block: 8px;
+    padding-inline: 12px;
+    padding-inline-end: 12px;
+
+    color: ${token.colorText};
+
+    background: ${token.colorFillQuaternary};
+    border-radius: 8px;
+  `,
+  title: css`
+    overflow: hidden;
+    display: -webkit-box;
+    -webkit-box-orient: vertical;
+    -webkit-line-clamp: 1;
+
+    font-size: 12px;
+    text-overflow: ellipsis;
+  `,
+}));
+
+const Render = memo<MarkdownElementProps>(({ children, id }) => {
+  const { t } = useTranslation('chat');
+  const { styles, cx } = useStyles();
+
+  const [isGenerating] = useChatStore((s) => {
+    const message = chatSelectors.getMessageById(id)(s);
+    return [!isLobeThinkingClosed(message?.content)];
+  });
+
+  const [showDetail, setShowDetail] = useState(false);
+
+  const expand = showDetail || isGenerating;
+  return (
+    <Flexbox
+      className={styles.container}
+      gap={16}
+      onClick={() => {
+        setShowDetail(!showDetail);
+      }}
+      width={'100%'}
+    >
+      <Flexbox distribution={'space-between'} flex={1} horizontal>
+        <Flexbox gap={8} horizontal>
+          <Icon icon={isGenerating ? Loader2Icon : BringToFrontIcon} spin={isGenerating} />
+          {isGenerating ? (
+            <span className={cx(dotLoading)}>{t('artifact.thinking')}</span>
+          ) : (
+            t('artifact.thought')
+          )}
+        </Flexbox>
+        <Icon icon={expand ? ChevronDown : ChevronRight} />
+      </Flexbox>
+      {expand && children}
+    </Flexbox>
+  );
+});
+
+export default Render;

package/src/features/Conversation/components/MarkdownElements/LobeThinking/index.ts

@@ -0,0 +1,12 @@
+import { ARTIFACT_THINKING_TAG } from '@/const/plugin';
+
+import Component from './Render';
+import rehypePlugin from './rehypePlugin';
+
+const AntThinkingElement = {
+  Component,
+  rehypePlugin,
+  tag: ARTIFACT_THINKING_TAG,
+};
+
+export default AntThinkingElement;

package/src/features/Conversation/components/MarkdownElements/LobeThinking/rehypePlugin.test.ts

@@ -0,0 +1,124 @@
+import { describe, expect, it } from 'vitest';
+
+import rehypePlugin from './rehypePlugin';
+
+describe('rehypePlugin', () => {
+  it('should transform <lobeThinking> tags within paragraphs', () => {
+    const tree = {
+      type: 'root',
+      children: [
+        {
+          type: 'element',
+          tagName: 'p',
+          children: [
+            { type: 'text', value: 'Before ' },
+            { type: 'raw', value: '<lobeThinking>' },
+            { type: 'text', value: 'Thinking content' },
+            { type: 'raw', value: '</lobeThinking>' },
+            { type: 'text', value: ' After' },
+          ],
+        },
+      ],
+    };
+
+    const expectedTree = {
+      type: 'root',
+      children: [
+        {
+          type: 'element',
+          tagName: 'lobeThinking',
+          properties: {},
+          children: [{ type: 'text', value: 'Thinking content' }],
+        },
+      ],
+    };
+
+    const plugin = rehypePlugin();
+    plugin(tree);
+
+    expect(tree).toEqual(expectedTree);
+  });
+
+  it('should not transform when only closing tag is present', () => {
+    const tree = {
+      type: 'root',
+      children: [
+        {
+          type: 'element',
+          tagName: 'p',
+          children: [
+            { type: 'text', value: 'Thinking content' },
+            { type: 'raw', value: '</lobeThinking>' },
+            { type: 'text', value: ' After' },
+          ],
+        },
+      ],
+    };
+
+    const originalTree = JSON.parse(JSON.stringify(tree));
+
+    const plugin = rehypePlugin();
+    plugin(tree);
+
+    expect(tree).toEqual(originalTree);
+  });
+
+  it('should handle multiple paragraphs and transformations', () => {
+    const tree = {
+      type: 'root',
+      children: [
+        {
+          type: 'element',
+          tagName: 'p',
+          children: [{ type: 'text', value: 'Normal paragraph' }],
+        },
+        {
+          type: 'element',
+          tagName: 'p',
+          children: [
+            { type: 'raw', value: '<lobeThinking>' },
+            { type: 'text', value: 'First thinking' },
+            { type: 'raw', value: '</lobeThinking>' },
+          ],
+        },
+        {
+          type: 'element',
+          tagName: 'p',
+          children: [
+            { type: 'raw', value: '<lobeThinking>' },
+            { type: 'text', value: 'Second thinking' },
+            { type: 'raw', value: '</lobeThinking>' },
+          ],
+        },
+      ],
+    };
+
+    const expectedTree = {
+      type: 'root',
+      children: [
+        {
+          type: 'element',
+          tagName: 'p',
+          children: [{ type: 'text', value: 'Normal paragraph' }],
+        },
+        {
+          type: 'element',
+          tagName: 'lobeThinking',
+          properties: {},
+          children: [{ type: 'text', value: 'First thinking' }],
+        },
+        {
+          type: 'element',
+          tagName: 'lobeThinking',
+          properties: {},
+          children: [{ type: 'text', value: 'Second thinking' }],
+        },
+      ],
+    };
+
+    const plugin = rehypePlugin();
+    plugin(tree);
+
+    expect(tree).toEqual(expectedTree);
+  });
+});

package/src/features/Conversation/components/MarkdownElements/LobeThinking/rehypePlugin.ts

@@ -0,0 +1,51 @@
+import type { Node } from 'unist';
+import { visit } from 'unist-util-visit';
+
+import { ARTIFACT_THINKING_TAG } from '@/const/plugin';
+
+// eslint-disable-next-line unicorn/consistent-function-scoping
+const rehypePlugin = () => (tree: Node) => {
+  visit(tree, 'element', (node: any, index, parent) => {
+    if (node.type === 'element' && node.tagName === 'p') {
+      const children = node.children || [];
+      const openTagIndex = children.findIndex(
+        (child: any) => child.type === 'raw' && child.value === `<${ARTIFACT_THINKING_TAG}>`,
+      );
+      const closeTagIndex = children.findIndex(
+        (child: any) => child.type === 'raw' && child.value === `</${ARTIFACT_THINKING_TAG}>`,
+      );
+
+      if (openTagIndex !== -1) {
+        // 有闭合标签的情况
+        if (closeTagIndex !== -1 && closeTagIndex > openTagIndex) {
+          const content = children.slice(openTagIndex + 1, closeTagIndex);
+          const lobeThinkingNode = {
+            children: content,
+            properties: {},
+            tagName: ARTIFACT_THINKING_TAG,
+            type: 'element',
+          };
+
+          // Replace the entire paragraph with our new lobeThinking node
+          parent.children.splice(index, 1, lobeThinkingNode);
+          return index; // Skip processing the newly inserted node
+        } else {
+          // 无闭合标签的情况
+          const content = children.slice(openTagIndex + 1);
+          const lobeThinkingNode = {
+            children: content,
+            properties: {},
+            tagName: ARTIFACT_THINKING_TAG,
+            type: 'element',
+          };
+
+          // Replace the entire paragraph with our new lobeThinking node
+          parent.children.splice(index, 1, lobeThinkingNode);
+          return index; // Skip processing the newly inserted node
+        }
+      }
+    }
+  });
+};
+
+export default rehypePlugin;

package/src/features/Conversation/components/MarkdownElements/index.ts

@@ -0,0 +1,4 @@
+import LobeArtifact from './LobeArtifact';
+import LobeThinking from './LobeThinking';
+
+export const markdownElements = [LobeArtifact, LobeThinking];

package/src/features/Conversation/components/MarkdownElements/type.ts

@@ -0,0 +1,7 @@
+import { ReactNode } from 'react';
+
+export interface MarkdownElementProps {
+  children: ReactNode;
+  id: string;
+  type: string;
+}

package/src/libs/agent-runtime/AgentRuntime.ts

@@ -3,6 +3,7 @@ import { ClientOptions } from 'openai';
 import type { TracePayload } from '@/const/trace';
 
 import { LobeRuntimeAI } from './BaseAI';
+import { LobeAi21AI } from './ai21';
 import { LobeAi360AI } from './ai360';
 import { LobeAnthropicAI } from './anthropic';
 import { LobeAzureOpenAI } from './azureOpenai';
@@ -10,6 +11,7 @@ import { LobeBaichuanAI } from './baichuan';
 import { LobeBedrockAI, LobeBedrockAIParams } from './bedrock';
 import { LobeDeepSeekAI } from './deepseek';
 import { LobeFireworksAI } from './fireworksai';
+import { LobeGithubAI } from './github';
 import { LobeGoogleAI } from './google';
 import { LobeGroq } from './groq';
 import { LobeMinimaxAI } from './minimax';
@@ -116,6 +118,7 @@ class AgentRuntime {
   static async initializeWithProviderOptions(
     provider: string,
     params: Partial<{
+      ai21: Partial<ClientOptions>;
       ai360: Partial<ClientOptions>;
       anthropic: Partial<ClientOptions>;
       azure: { apiVersion?: string; apikey?: string; endpoint?: string };
@@ -123,6 +126,7 @@ class AgentRuntime {
       bedrock: Partial<LobeBedrockAIParams>;
       deepseek: Partial<ClientOptions>;
       fireworksai: Partial<ClientOptions>;
+      github: Partial<ClientOptions>;
       google: { apiKey?: string; baseURL?: string };
       groq: Partial<ClientOptions>;
       minimax: Partial<ClientOptions>;
@@ -218,6 +222,11 @@ class AgentRuntime {
         break;
       }
 
+      case ModelProvider.Github: {
+        runtimeModel = new LobeGithubAI(params.github);
+        break;
+      }
+
       case ModelProvider.OpenRouter: {
         runtimeModel = new LobeOpenRouterAI(params.openrouter);
         break;
@@ -282,6 +291,11 @@ class AgentRuntime {
         runtimeModel = new LobeSparkAI(params.spark);
         break;
       }
+
+      case ModelProvider.Ai21: {
+        runtimeModel = new LobeAi21AI(params.ai21);
+        break;
+      }
     }
 
     return new AgentRuntime(runtimeModel);

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

@@ -0,0 +1,255 @@
+// @vitest-environment node
+import OpenAI from 'openai';
+import { Mock, afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import {
+  ChatStreamCallbacks,
+  LobeOpenAICompatibleRuntime,
+  ModelProvider,
+} from '@/libs/agent-runtime';
+
+import * as debugStreamModule from '../utils/debugStream';
+import { LobeAi21AI } from './index';
+
+const provider = ModelProvider.Ai21;
+const defaultBaseURL = 'https://api.ai21.com/studio/v1';
+
+const bizErrorType = 'ProviderBizError';
+const invalidErrorType = 'InvalidProviderAPIKey';
+
+// Mock the console.error to avoid polluting test output
+vi.spyOn(console, 'error').mockImplementation(() => {});
+
+let instance: LobeOpenAICompatibleRuntime;
+
+beforeEach(() => {
+  instance = new LobeAi21AI({ apiKey: 'test' });
+
+  // 使用 vi.spyOn 来模拟 chat.completions.create 方法
+  vi.spyOn(instance['client'].chat.completions, 'create').mockResolvedValue(
+    new ReadableStream() as any,
+  );
+});
+
+afterEach(() => {
+  vi.clearAllMocks();
+});
+
+describe('LobeAi21AI', () => {
+  describe('init', () => {
+    it('should correctly initialize with an API key', async () => {
+      const instance = new LobeAi21AI({ apiKey: 'test_api_key' });
+      expect(instance).toBeInstanceOf(LobeAi21AI);
+      expect(instance.baseURL).toEqual(defaultBaseURL);
+    });
+  });
+
+  describe('chat', () => {
+    describe('Error', () => {
+      it('should return OpenAIBizError 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: 'jamba-1.5-mini',
+            temperature: 0,
+          });
+        } catch (e) {
+          expect(e).toEqual({
+            endpoint: defaultBaseURL,
+            error: {
+              error: { message: 'Bad Request' },
+              status: 400,
+            },
+            errorType: bizErrorType,
+            provider,
+          });
+        }
+      });
+
+      it('should throw AgentRuntimeError with NoOpenAIAPIKey if no apiKey is provided', async () => {
+        try {
+          new LobeAi21AI({});
+        } catch (e) {
+          expect(e).toEqual({ errorType: invalidErrorType });
+        }
+      });
+
+      it('should return OpenAIBizError 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: 'jamba-1.5-mini',
+            temperature: 0,
+          });
+        } catch (e) {
+          expect(e).toEqual({
+            endpoint: defaultBaseURL,
+            error: {
+              cause: { message: 'api is undefined' },
+              stack: 'abc',
+            },
+            errorType: bizErrorType,
+            provider,
+          });
+        }
+      });
+
+      it('should return OpenAIBizError 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 LobeAi21AI({
+          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: 'jamba-1.5-mini',
+            temperature: 0,
+          });
+        } catch (e) {
+          expect(e).toEqual({
+            endpoint: 'https://api.***.com/v1',
+            error: {
+              cause: { message: 'api is undefined' },
+              stack: 'abc',
+            },
+            errorType: bizErrorType,
+            provider,
+          });
+        }
+      });
+
+      it('should throw an InvalidAi21APIKey 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: 'jamba-1.5-mini',
+            temperature: 0,
+          });
+        } catch (e) {
+          // Expect the chat method to throw an error with InvalidAi21APIKey
+          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: 'jamba-1.5-mini',
+            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('DEBUG', () => {
+      it('should call debugStream and return StreamingTextResponse when DEBUG_AI21_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_AI21_CHAT_COMPLETION;
+
+        // 模拟环境变量
+        process.env.DEBUG_AI21_CHAT_COMPLETION = '1';
+        vi.spyOn(debugStreamModule, 'debugStream').mockImplementation(() => Promise.resolve());
+
+        // 执行测试
+        // 运行你的测试函数，确保它会在条件满足时调用 debugStream
+        // 假设的测试函数调用，你可能需要根据实际情况调整
+        await instance.chat({
+          messages: [{ content: 'Hello', role: 'user' }],
+          model: 'jamba-1.5-mini',
+          stream: true,
+          temperature: 0,
+        });
+
+        // 验证 debugStream 被调用
+        expect(debugStreamModule.debugStream).toHaveBeenCalled();
+
+        // 恢复原始环境变量值
+        process.env.DEBUG_AI21_CHAT_COMPLETION = originalDebugValue;
+      });
+    });
+  });
+});

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

@@ -0,0 +1,18 @@
+import { ModelProvider } from '../types';
+import { LobeOpenAICompatibleFactory } from '../utils/openaiCompatibleFactory';
+
+export const LobeAi21AI = LobeOpenAICompatibleFactory({
+  baseURL: 'https://api.ai21.com/studio/v1',
+  chatCompletion: {
+    handlePayload: (payload) => {
+      return {
+        ...payload,
+        stream: !payload.tools,
+      } as any;
+    },
+  },
+  debug: {
+    chatCompletion: () => process.env.DEBUG_AI21_CHAT_COMPLETION === '1',
+  },
+  provider: ModelProvider.Ai21,
+});

package/src/libs/agent-runtime/error.ts

@@ -13,6 +13,8 @@ export const AgentRuntimeErrorType = {
   InvalidBedrockCredentials: 'InvalidBedrockCredentials',
   StreamChunkError: 'StreamChunkError',
 
+  InvalidGithubToken: 'InvalidGithubToken',
+
   /**
    * @deprecated
    */