anthropic[patch]: Fix passing cache control through in messages (#6711)

langchain-ai · Sep 6, 2024 · 1e76bf0 · 1e76bf0
1 parent 1096dee
commit 1e76bf0
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 2 deletions.
diff --git a/libs/langchain-anthropic/src/tests/chat_models.int.test.ts b/libs/langchain-anthropic/src/tests/chat_models.int.test.ts
@@ -393,8 +393,7 @@ test("id is supplied when streaming", async () => {
   expect(finalChunk.id).not.toEqual("");
 });
 
-test("system prompt caching", async () => {
-  const CACHED_TEXT = `## Components
+const CACHED_TEXT = `## Components
 
 LangChain provides standard, extendable interfaces and external integrations for various components useful for building with LLMs.
 Some components LangChain implements, some components we rely on third-party integrations for, and others are a mix.
@@ -655,6 +654,7 @@ LangChain has many different types of output parsers. This is a list of output p
 
 The current date is ${new Date().toISOString()}`;
 
+test("system prompt caching", async () => {
   const model = new ChatAnthropic({
     model: "claude-3-haiku-20240307",
     clientOptions: {
@@ -749,3 +749,45 @@ test.skip("Test ChatAnthropic with custom client", async () => {
   // console.log({ res });
   expect(res.response_metadata.usage).toBeDefined();
 });
+
+test("human message caching", async () => {
+  const model = new ChatAnthropic({
+    model: "claude-3-haiku-20240307",
+    clientOptions: {
+      defaultHeaders: {
+        "anthropic-beta": "prompt-caching-2024-07-31",
+      },
+    },
+  });
+
+  const messages = [
+    new SystemMessage({
+      content: [
+        {
+          type: "text",
+          text: `You are a pirate. Always respond in pirate dialect.\nUse the following as context when answering questions: ${CACHED_TEXT}`,
+        },
+      ],
+    }),
+    new HumanMessage({
+      content: [
+        {
+          type: "text",
+          text: "What types of messages are supported in LangChain?",
+          cache_control: { type: "ephemeral" },
+        },
+      ],
+    }),
+  ];
+
+  const res = await model.invoke(messages);
+  expect(
+    res.response_metadata.usage.cache_creation_input_tokens
+  ).toBeGreaterThan(0);
+  expect(res.response_metadata.usage.cache_read_input_tokens).toBe(0);
+  const res2 = await model.invoke(messages);
+  expect(res2.response_metadata.usage.cache_creation_input_tokens).toBe(0);
+  expect(res2.response_metadata.usage.cache_read_input_tokens).toBeGreaterThan(
+    0
+  );
+});
diff --git a/libs/langchain-anthropic/src/utils/message_inputs.ts b/libs/langchain-anthropic/src/utils/message_inputs.ts
@@ -122,6 +122,9 @@ function _formatContent(content: MessageContent) {
     return content;
   } else {
     const contentBlocks = content.map((contentPart) => {
+      const cacheControl =
+        "cache_control" in contentPart ? contentPart.cache_control : undefined;
+
       if (contentPart.type === "image_url") {
         let source;
         if (typeof contentPart.image_url === "string") {
@@ -132,6 +135,7 @@ function _formatContent(content: MessageContent) {
         return {
           type: "image" as const, // Explicitly setting the type as "image"
           source,
+          ...(cacheControl ? { cache_control: cacheControl } : {}),
         };
       } else if (
         textTypes.find((t) => t === contentPart.type) &&
@@ -141,6 +145,7 @@ function _formatContent(content: MessageContent) {
         return {
           type: "text" as const, // Explicitly setting the type as "text"
           text: contentPart.text,
+          ...(cacheControl ? { cache_control: cacheControl } : {}),
         };
       } else if (toolTypes.find((t) => t === contentPart.type)) {
         const contentPartCopy = { ...contentPart };
@@ -167,6 +172,7 @@ function _formatContent(content: MessageContent) {
         // TODO: Fix when SDK types are fixed
         return {
           ...contentPartCopy,
+          ...(cacheControl ? { cache_control: cacheControl } : {}),
           // eslint-disable-next-line @typescript-eslint/no-explicit-any
         } as any;
       } else {