@lobehub/chat 0.134.1 → 0.135.1

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

@@ -2,15 +2,23 @@ import {
   BedrockRuntimeClient,
   InvokeModelWithResponseStreamCommand,
 } from '@aws-sdk/client-bedrock-runtime';
-import { AWSBedrockAnthropicStream, AWSBedrockLlama2Stream, StreamingTextResponse } from 'ai';
-import { experimental_buildAnthropicPrompt, experimental_buildLlama2Prompt } from 'ai/prompts';
+import {
+  AWSBedrockLlama2Stream,
+  AWSBedrockStream,
+  StreamingTextResponse
+} from 'ai';
+import { experimental_buildLlama2Prompt } from 'ai/prompts';
 
 import { LobeRuntimeAI } from '../BaseAI';
 import { AgentRuntimeErrorType } from '../error';
-import { ChatStreamPayload, ModelProvider } from '../types';
+import {
+  ChatCompetitionOptions,
+  ChatStreamPayload,
+  ModelProvider,
+} from '../types';
 import { AgentRuntimeError } from '../utils/createError';
 import { debugStream } from '../utils/debugStream';
-import { DEBUG_CHAT_COMPLETION } from '../utils/env';
+import { buildAnthropicMessages } from '../utils/anthropicHelpers';
 
 export interface LobeBedrockAIParams {
   accessKeyId?: string;
@@ -38,23 +46,32 @@ export class LobeBedrockAI implements LobeRuntimeAI {
     });
   }
 
-  async chat(payload: ChatStreamPayload) {
+  async chat(payload: ChatStreamPayload, options?: ChatCompetitionOptions) {
     if (payload.model.startsWith('meta')) return this.invokeLlamaModel(payload);
 
-    return this.invokeClaudeModel(payload);
+    return this.invokeClaudeModel(payload, options);
   }
 
   private invokeClaudeModel = async (
     payload: ChatStreamPayload,
+    options?: ChatCompetitionOptions
   ): Promise<StreamingTextResponse> => {
+    const { max_tokens, messages, model, temperature, top_p } = payload;
+    const system_message = messages.find((m) => m.role === 'system');
+    const user_messages = messages.filter((m) => m.role !== 'system');
+
     const command = new InvokeModelWithResponseStreamCommand({
       accept: 'application/json',
       body: JSON.stringify({
-        max_tokens_to_sample: payload.max_tokens || 400,
-        prompt: experimental_buildAnthropicPrompt(payload.messages as any),
+        anthropic_version: "bedrock-2023-05-31",
+        max_tokens: max_tokens || 4096,
+        messages: buildAnthropicMessages(user_messages),
+        system: system_message?.content as string,
+        temperature: temperature,
+        top_p: top_p,
       }),
       contentType: 'application/json',
-      modelId: payload.model,
+      modelId: model,
     });
 
     try {
@@ -62,11 +79,11 @@ export class LobeBedrockAI implements LobeRuntimeAI {
       const bedrockResponse = await this.client.send(command);
 
       // Convert the response into a friendly text-stream
-      const stream = AWSBedrockAnthropicStream(bedrockResponse);
+      const stream = AWSBedrockStream(bedrockResponse, options?.callback, (chunk) => chunk.delta?.text);
 
       const [debug, output] = stream.tee();
 
-      if (DEBUG_CHAT_COMPLETION) {
+      if (process.env.DEBUG_BEDROCK_CHAT_COMPLETION === '1') {
         debugStream(debug).catch(console.error);
       }
 
@@ -88,15 +105,18 @@ export class LobeBedrockAI implements LobeRuntimeAI {
     }
   };
 
-  private invokeLlamaModel = async (payload: ChatStreamPayload) => {
+  private invokeLlamaModel = async (
+    payload: ChatStreamPayload
+  ): Promise<StreamingTextResponse> => {
+    const { max_tokens, messages, model } = payload;
     const command = new InvokeModelWithResponseStreamCommand({
       accept: 'application/json',
       body: JSON.stringify({
-        max_gen_len: payload.max_tokens || 400,
-        prompt: experimental_buildLlama2Prompt(payload.messages as any),
+        max_gen_len: max_tokens || 400,
+        prompt: experimental_buildLlama2Prompt(messages as any),
       }),
       contentType: 'application/json',
-      modelId: payload.model,
+      modelId: model,
     });
 
     try {
@@ -108,7 +128,7 @@ export class LobeBedrockAI implements LobeRuntimeAI {
 
       const [debug, output] = stream.tee();
 
-      if (DEBUG_CHAT_COMPLETION) {
+      if (process.env.DEBUG_BEDROCK_CHAT_COMPLETION === '1') {
         debugStream(debug).catch(console.error);
       }
       // Respond with the stream
@@ -129,6 +149,7 @@ export class LobeBedrockAI implements LobeRuntimeAI {
       });
     }
   };
+
 }
 
 export default LobeBedrockAI;

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

@@ -0,0 +1,59 @@
+import { describe, expect, it } from 'vitest';
+
+import {
+  buildAnthropicMessage,
+  buildAnthropicBlock,
+} from './anthropicHelpers';
+
+import { parseDataUri } from './uriParser';
+import {
+  OpenAIChatMessage,
+  UserMessageContentPart,
+} from '../types/chat';
+
+describe('anthropicHelpers', () => {
+
+  // Mock the parseDataUri function since it's an implementation detail
+  vi.mock('./uriParser', () => ({
+    parseDataUri: vi.fn().mockReturnValue({
+      mimeType: 'image/jpeg',
+      base64: 'base64EncodedString',
+    }),
+  }));
+
+  describe('buildAnthropicBlock', () => {
+    it('should return the content as is for text type', () => {
+      const content: UserMessageContentPart =
+        { type: 'text', text: 'Hello!' };
+      const result = buildAnthropicBlock(content);
+      expect(result).toEqual(content);
+    });
+  
+    it('should transform an image URL into an Anthropic.ImageBlockParam', () => {
+      const content: UserMessageContentPart =
+        { type: 'image_url', image_url: { url: 'data:image/jpeg;base64,base64EncodedString' } };
+      const result = buildAnthropicBlock(content);
+      expect(parseDataUri).toHaveBeenCalledWith(content.image_url.url);
+      expect(result).toEqual({
+        source: {
+          data: 'base64EncodedString',
+          media_type: 'image/jpeg',
+          type: 'base64',
+        },
+        type: 'image',
+      });
+    });
+  });
+
+  describe('buildAnthropicMessage', () => {
+    it('should correctly convert system message to assistant message', () => {
+      const message: OpenAIChatMessage = 
+        { content: [{ type: 'text', text: 'Hello!' }], role: 'system' };
+      const result = buildAnthropicMessage(message);
+      expect(result).toEqual(
+        { content: [{ type: 'text', text: 'Hello!' }], role: 'assistant' }
+      );
+    });
+  });
+
+});

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

@@ -0,0 +1,47 @@
+import Anthropic from '@anthropic-ai/sdk';
+
+import {
+  OpenAIChatMessage,
+  UserMessageContentPart,
+} from '../types';
+
+import { parseDataUri } from './uriParser';
+
+export const buildAnthropicBlock = (
+  content: UserMessageContentPart,
+): Anthropic.ContentBlock | Anthropic.ImageBlockParam => {
+  switch (content.type) {
+    case 'text': {
+      return content;
+    }
+
+    case 'image_url': {
+      const { mimeType, base64 } = parseDataUri(content.image_url.url);
+
+      return {
+        source: {
+          data: base64 as string,
+          media_type: mimeType as Anthropic.ImageBlockParam.Source['media_type'],
+          type: 'base64',
+        },
+        type: 'image',
+      };
+    }
+  }
+}
+
+export const buildAnthropicMessage = (
+  message: OpenAIChatMessage,
+): Anthropic.Messages.MessageParam => {
+  const content = message.content as string | UserMessageContentPart[];
+  return {
+    content:
+      typeof content === 'string' ? content : content.map((c) => buildAnthropicBlock(c)),
+    role: message.role === 'function' || message.role === 'system' ? 'assistant' : message.role,
+  };
+};
+
+export const buildAnthropicMessages = (
+  messages: OpenAIChatMessage[],
+): Anthropic.Messages.MessageParam[] =>
+  messages.map((message) => buildAnthropicMessage(message));

package/src/libs/swr/index.ts

@@ -0,0 +1,18 @@
+import useSWR, { SWRHook } from 'swr';
+
+/**
+ * 这一类请求方法是比较「死」的请求模式，只会在第一次请求时触发。不会自动刷新，刷新需要搭配 refreshXXX 这样的方法实现，
+ * 适用于 messages、topics、sessions 等由用户在客户端交互产生的数据。
+ */
+// @ts-ignore
+export const useClientDataSWR: SWRHook = (key, fetch, config) =>
+  useSWR(key, fetch, {
+    // default is 2000ms ,it makes the user's quick switch don't work correctly.
+    // Cause issue like this: https://github.com/lobehub/lobe-chat/issues/532
+    // we need to set it to 0.
+    dedupingInterval: 0,
+    refreshWhenOffline: false,
+    revalidateOnFocus: false,
+    revalidateOnReconnect: false,
+    ...config,
+  });

package/src/store/chat/slices/message/action.test.ts

@@ -559,7 +559,7 @@ describe('chatMessage actions', () => {
       });
 
       // 确保 mutate 调用了正确的参数
-      expect(mutate).toHaveBeenCalledWith([activeId, activeTopicId]);
+      expect(mutate).toHaveBeenCalledWith(['SWR_USE_FETCH_MESSAGES', activeId, activeTopicId]);
     });
     it('should handle errors during refreshing messages', async () => {
       useChatStore.setState({ refreshMessages: realRefreshMessages });

package/src/store/chat/slices/message/action.ts

@@ -2,12 +2,13 @@
 // Disable the auto sort key eslint rule to make the code more logic and readable
 import { copyToClipboard } from '@lobehub/ui';
 import { template } from 'lodash-es';
-import useSWR, { SWRResponse, mutate } from 'swr';
+import { SWRResponse, mutate } from 'swr';
 import { StateCreator } from 'zustand/vanilla';
 
 import { LOADING_FLAT, isFunctionMessageAtStart, testFunctionMessageAtEnd } from '@/const/message';
 import { TraceEventType, TraceNameMap } from '@/const/trace';
 import { CreateMessageParams } from '@/database/models/message';
+import { useClientDataSWR } from '@/libs/swr';
 import { chatService } from '@/services/chat';
 import { messageService } from '@/services/message';
 import { topicService } from '@/services/topic';
@@ -25,6 +26,8 @@ import { MessageDispatch, messagesReducer } from './reducer';
 
 const n = setNamespace('message');
 
+const SWR_USE_FETCH_MESSAGES = 'SWR_USE_FETCH_MESSAGES';
+
 interface SendMessageParams {
   message: string;
   files?: { id: string; url: string }[];
@@ -274,9 +277,9 @@ export const chatMessage: StateCreator<
     await get().internalUpdateMessageContent(id, content);
   },
   useFetchMessages: (sessionId, activeTopicId) =>
-    useSWR<ChatMessage[]>(
-      [sessionId, activeTopicId],
-      async ([sessionId, topicId]: [string, string | undefined]) =>
+    useClientDataSWR<ChatMessage[]>(
+      [SWR_USE_FETCH_MESSAGES, sessionId, activeTopicId],
+      async ([, sessionId, topicId]: [string, string, string | undefined]) =>
         messageService.getMessages(sessionId, topicId),
       {
         onSuccess: (messages, key) => {
@@ -289,14 +292,10 @@ export const chatMessage: StateCreator<
             }),
           );
         },
-        // default is 2000ms ,it makes the user's quick switch don't work correctly.
-        // Cause issue like this: https://github.com/lobehub/lobe-chat/issues/532
-        // we need to set it to 0.
-        dedupingInterval: 0,
       },
     ),
   refreshMessages: async () => {
-    await mutate([get().activeId, get().activeTopicId]);
+    await mutate([SWR_USE_FETCH_MESSAGES, get().activeId, get().activeTopicId]);
   },
 
   // the internal process method of the AI message

package/src/store/session/slices/session/action.ts

@@ -6,6 +6,7 @@ import { StateCreator } from 'zustand/vanilla';
 
 import { INBOX_SESSION_ID } from '@/const/session';
 import { SESSION_CHAT_URL } from '@/const/url';
+import { useClientDataSWR } from '@/libs/swr';
 import { sessionService } from '@/services/session';
 import { useGlobalStore } from '@/store/global';
 import { settingsSelectors } from '@/store/global/selectors';
@@ -154,7 +155,7 @@ export const createSessionSlice: StateCreator<
   },
 
   useFetchSessions: () =>
-    useSWR<ChatSessionList>(FETCH_SESSIONS_KEY, sessionService.getSessionsWithGroup, {
+    useClientDataSWR<ChatSessionList>(FETCH_SESSIONS_KEY, sessionService.getSessionsWithGroup, {
       onSuccess: (data) => {
         // 由于 https://github.com/lobehub/lobe-chat/pull/541 的关系
         // 只有触发了 refreshSessions 才会更新 sessions，进而触发页面 rerender