support context caching

fix tests
lobehub · Mar 6, 2025 · 3e8cd2e · 3e8cd2e
1 parent 9f09952
commit 3e8cd2e
Show file tree

Hide file tree

Showing 5 changed files with 109 additions and 13 deletions.
diff --git a/src/libs/agent-runtime/anthropic/index.test.ts b/src/libs/agent-runtime/anthropic/index.test.ts
@@ -81,7 +81,12 @@ describe('LobeAnthropicAI', () => {
       expect(instance['client'].messages.create).toHaveBeenCalledWith(
         {
           max_tokens: 4096,
-          messages: [{ content: 'Hello', role: 'user' }],
+          messages: [
+            {
+              content: [{ cache_control: { type: 'ephemeral' }, text: 'Hello', type: 'text' }],
+              role: 'user',
+            },
+          ],
           model: 'claude-3-haiku-20240307',
           stream: true,
           temperature: 0,
@@ -117,10 +122,21 @@ describe('LobeAnthropicAI', () => {
       expect(instance['client'].messages.create).toHaveBeenCalledWith(
         {
           max_tokens: 4096,
-          messages: [{ content: 'Hello', role: 'user' }],
+          messages: [
+            {
+              content: [{ cache_control: { type: 'ephemeral' }, text: 'Hello', type: 'text' }],
+              role: 'user',
+            },
+          ],
           model: 'claude-3-haiku-20240307',
           stream: true,
-          system: 'You are an awesome greeter',
+          system: [
+            {
+              cache_control: { type: 'ephemeral' },
+              type: 'text',
+              text: 'You are an awesome greeter',
+            },
+          ],
           temperature: 0,
         },
         {},
@@ -152,7 +168,12 @@ describe('LobeAnthropicAI', () => {
       expect(instance['client'].messages.create).toHaveBeenCalledWith(
         {
           max_tokens: 2048,
-          messages: [{ content: 'Hello', role: 'user' }],
+          messages: [
+            {
+              content: [{ cache_control: { type: 'ephemeral' }, text: 'Hello', type: 'text' }],
+              role: 'user',
+            },
+          ],
           model: 'claude-3-haiku-20240307',
           stream: true,
           temperature: 0.25,
@@ -189,7 +210,12 @@ describe('LobeAnthropicAI', () => {
       expect(instance['client'].messages.create).toHaveBeenCalledWith(
         {
           max_tokens: 2048,
-          messages: [{ content: 'Hello', role: 'user' }],
+          messages: [
+            {
+              content: [{ cache_control: { type: 'ephemeral' }, text: 'Hello', type: 'text' }],
+              role: 'user',
+            },
+          ],
           model: 'claude-3-haiku-20240307',
           stream: true,
           temperature: 0.25,

diff --git a/src/libs/agent-runtime/anthropic/index.ts b/src/libs/agent-runtime/anthropic/index.ts
@@ -97,10 +97,33 @@ export class LobeAnthropicAI implements LobeRuntimeAI {
   }
 
   private async buildAnthropicPayload(payload: ChatStreamPayload) {
-    const { messages, model, max_tokens, temperature, top_p, tools, thinking } = payload;
+    const {
+      messages,
+      model,
+      max_tokens,
+      temperature,
+      top_p,
+      tools,
+      thinking,
+      enabledContextCaching = true,
+    } = payload;
     const system_message = messages.find((m) => m.role === 'system');
     const user_messages = messages.filter((m) => m.role !== 'system');
 
+    const systemPrompts = !!system_message?.content
+      ? ([
+          {
+            cache_control: enabledContextCaching ? { type: 'ephemeral' } : undefined,
+            text: system_message?.content as string,
+            type: 'text',
+          },
+        ] as Anthropic.TextBlockParam[])
+      : undefined;
+
+    const postMessages = await buildAnthropicMessages(user_messages, { enabledContextCaching });
+
+    const postTools = buildAnthropicTools(tools);
+
     if (!!thinking) {
       const maxTokens =
         max_tokens ?? (thinking?.budget_tokens ? thinking?.budget_tokens + 4096 : 4096);
@@ -109,22 +132,21 @@ export class LobeAnthropicAI implements LobeRuntimeAI {
       // `top_p` must be unset when thinking is enabled.
       return {
         max_tokens: maxTokens,
-        messages: await buildAnthropicMessages(user_messages),
+        messages: postMessages,
         model,
-        system: system_message?.content as string,
-
+        system: systemPrompts,
         thinking,
-        tools: buildAnthropicTools(tools),
+        tools: postTools,
       } satisfies Anthropic.MessageCreateParams;
     }
 
     return {
       max_tokens: max_tokens ?? 4096,
-      messages: await buildAnthropicMessages(user_messages),
+      messages: postMessages,
       model,
-      system: system_message?.content as string,
+      system: systemPrompts,
       temperature: payload.temperature !== undefined ? temperature / 2 : undefined,
-      tools: buildAnthropicTools(tools),
+      tools: postTools,
       top_p,
     } satisfies Anthropic.MessageCreateParams;
   }

diff --git a/src/libs/agent-runtime/types/chat.ts b/src/libs/agent-runtime/types/chat.ts
@@ -46,6 +46,10 @@ export interface OpenAIChatMessage {
  * @title Chat Stream Payload
  */
 export interface ChatStreamPayload {
+  /**
+   * 开启上下文缓存
+   */
+  enabledContextCaching?: boolean;
   /**
    * 是否开启搜索
    */

diff --git a/src/libs/agent-runtime/utils/anthropicHelpers.test.ts b/src/libs/agent-runtime/utils/anthropicHelpers.test.ts
@@ -619,6 +619,26 @@ describe('anthropicHelpers', () => {
         { content: '继续', role: 'user' },
       ]);
     });
+
+    it('should enable cache control', async () => {
+      const messages: OpenAIChatMessage[] = [
+        { content: 'Hello', role: 'user' },
+        { content: 'Hello', role: 'user' },
+        { content: 'Hi', role: 'assistant' },
+      ];
+
+      const contents = await buildAnthropicMessages(messages, { enabledContextCaching: true });
+
+      expect(contents).toHaveLength(3);
+      expect(contents).toEqual([
+        { content: 'Hello', role: 'user' },
+        { content: 'Hello', role: 'user' },
+        {
+          content: [{ cache_control: { type: 'ephemeral' }, text: 'Hi', type: 'text' }],
+          role: 'assistant',
+        },
+      ]);
+    });
   });
 
   describe('buildAnthropicTools', () => {

diff --git a/src/libs/agent-runtime/utils/anthropicHelpers.ts b/src/libs/agent-runtime/utils/anthropicHelpers.ts
@@ -130,6 +130,7 @@ export const buildAnthropicMessage = async (
 
 export const buildAnthropicMessages = async (
   oaiMessages: OpenAIChatMessage[],
+  options: { enabledContextCaching?: boolean } = {},
 ): Promise<Anthropic.Messages.MessageParam[]> => {
   const messages: Anthropic.Messages.MessageParam[] = [];
   let pendingToolResults: Anthropic.ToolResultBlockParam[] = [];
@@ -180,8 +181,31 @@ export const buildAnthropicMessages = async (
     }
   }
 
+  const lastMessage = messages.at(-1);
+  if (options.enabledContextCaching && !!lastMessage) {
+    if (typeof lastMessage.content === 'string') {
+      lastMessage.content = [
+        {
+          cache_control: { type: 'ephemeral' },
+          text: lastMessage.content as string,
+          type: 'text',
+        },
+      ];
+    } else {
+      const lastContent = lastMessage.content.at(-1);
+
+      if (
+        lastContent &&
+        lastContent.type !== 'thinking' &&
+        lastContent.type !== 'redacted_thinking'
+      ) {
+        lastContent.cache_control = { type: 'ephemeral' };
+      }
+    }
+  }
   return messages;
 };
+
 export const buildAnthropicTools = (tools?: OpenAI.ChatCompletionTool[]) =>
   tools?.map(
     (tool): Anthropic.Tool => ({