anthropic[patch]: fix input_tokens when cached (#27125)

langchain-ai · Oct 4, 2024 · 0b8416b · 0b8416b
1 parent 64a16f2
commit 0b8416b
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 8 deletions.
diff --git a/libs/partners/anthropic/langchain_anthropic/chat_models.py b/libs/partners/anthropic/langchain_anthropic/chat_models.py
@@ -1253,7 +1253,12 @@ def _create_usage_metadata(anthropic_usage: BaseModel) -> UsageMetadata:
         "cache_creation": getattr(anthropic_usage, "cache_creation_input_tokens", None),
     }
 
-    input_tokens = getattr(anthropic_usage, "input_tokens", 0)
+    # Anthropic input_tokens exclude cached token counts.
+    input_tokens = (
+        getattr(anthropic_usage, "input_tokens", 0)
+        + (input_token_details["cache_read"] or 0)
+        + (input_token_details["cache_creation"] or 0)
+    )
     output_tokens = getattr(anthropic_usage, "output_tokens", 0)
     return UsageMetadata(
         input_tokens=input_tokens,

diff --git a/libs/partners/anthropic/tests/unit_tests/test_chat_models.py b/libs/partners/anthropic/tests/unit_tests/test_chat_models.py
@@ -128,9 +128,9 @@ def test__format_output_cached() -> None:
     expected = AIMessage(  # type: ignore[misc]
         "bar",
         usage_metadata={
-            "input_tokens": 2,
+            "input_tokens": 9,
             "output_tokens": 1,
-            "total_tokens": 3,
+            "total_tokens": 10,
             "input_token_details": {"cache_creation": 3, "cache_read": 4},
         },
     )

diff --git a/libs/standard-tests/langchain_standard_tests/integration_tests/chat_models.py b/libs/standard-tests/langchain_standard_tests/integration_tests/chat_models.py
@@ -153,28 +153,58 @@ def test_usage_metadata(self, model: BaseChatModel) -> None:
 
         if "audio_input" in self.supported_usage_metadata_details["invoke"]:
             msg = self.invoke_with_audio_input()
-            assert isinstance(msg.usage_metadata["input_token_details"]["audio"], int)  # type: ignore[index]
+            assert msg.usage_metadata is not None
+            assert msg.usage_metadata["input_token_details"] is not None
+            assert isinstance(msg.usage_metadata["input_token_details"]["audio"], int)
+            assert msg.usage_metadata["input_tokens"] >= sum(
+                (v or 0)  # type: ignore[misc]
+                for v in msg.usage_metadata["input_token_details"].values()
+            )
         if "audio_output" in self.supported_usage_metadata_details["invoke"]:
             msg = self.invoke_with_audio_output()
-            assert isinstance(msg.usage_metadata["output_token_details"]["audio"], int)  # type: ignore[index]
+            assert msg.usage_metadata is not None
+            assert msg.usage_metadata["output_token_details"] is not None
+            assert isinstance(msg.usage_metadata["output_token_details"]["audio"], int)
+            assert int(msg.usage_metadata["output_tokens"]) >= sum(
+                (v or 0)  # type: ignore[misc]
+                for v in msg.usage_metadata["output_token_details"].values()
+            )
         if "reasoning_output" in self.supported_usage_metadata_details["invoke"]:
             msg = self.invoke_with_reasoning_output()
+            assert msg.usage_metadata is not None
+            assert msg.usage_metadata["output_token_details"] is not None
             assert isinstance(
-                msg.usage_metadata["output_token_details"]["reasoning"],  # type: ignore[index]
+                msg.usage_metadata["output_token_details"]["reasoning"],
                 int,
             )
+            assert msg.usage_metadata["output_tokens"] >= sum(
+                (v or 0)  # type: ignore[misc]
+                for v in msg.usage_metadata["output_token_details"].values()
+            )
         if "cache_read_input" in self.supported_usage_metadata_details["invoke"]:
             msg = self.invoke_with_cache_read_input()
+            assert msg.usage_metadata is not None
+            assert msg.usage_metadata["input_token_details"] is not None
             assert isinstance(
-                msg.usage_metadata["input_token_details"]["cache_read"],  # type: ignore[index]
+                msg.usage_metadata["input_token_details"]["cache_read"],
                 int,
             )
+            assert msg.usage_metadata["input_tokens"] >= sum(
+                (v or 0)  # type: ignore[misc]
+                for v in msg.usage_metadata["input_token_details"].values()
+            )
         if "cache_creation_input" in self.supported_usage_metadata_details["invoke"]:
             msg = self.invoke_with_cache_creation_input()
+            assert msg.usage_metadata is not None
+            assert msg.usage_metadata["input_token_details"] is not None
             assert isinstance(
-                msg.usage_metadata["input_token_details"]["cache_creation"],  # type: ignore[index]
+                msg.usage_metadata["input_token_details"]["cache_creation"],
                 int,
             )
+            assert msg.usage_metadata["input_tokens"] >= sum(
+                (v or 0)  # type: ignore[misc]
+                for v in msg.usage_metadata["input_token_details"].values()
+            )
 
     def test_usage_metadata_streaming(self, model: BaseChatModel) -> None:
         if not self.returns_usage_metadata: