Bump the dependencies group across 1 directory with 12 updates (#1331)

Signed-off-by: dependabot[bot] <[email protected]> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Co-authored-by: Collin Dutter <[email protected]>
griptape-ai · Nov 11, 2024 · dcbde3c · dcbde3c
1 parent c720288
commit dcbde3c
Show file tree

Hide file tree

Showing 4 changed files with 72 additions and 50 deletions.
diff --git a/griptape/tokenizers/anthropic_tokenizer.py b/griptape/tokenizers/anthropic_tokenizer.py
@@ -9,6 +9,7 @@
 
 if TYPE_CHECKING:
     from anthropic import Anthropic
+    from anthropic.types.beta import BetaMessageParam
 
 
 @define()
@@ -21,5 +22,15 @@ class AnthropicTokenizer(BaseTokenizer):
         kw_only=True,
     )
 
-    def count_tokens(self, text: str) -> int:
-        return self.client.count_tokens(text)
+    def count_tokens(self, text: str | list[BetaMessageParam]) -> int:
+        types = import_optional_dependency("anthropic.types.beta")
+
+        # TODO: Refactor all Tokenizers to support Prompt Stack as an input.
+        messages = [types.BetaMessageParam(role="user", content=text)] if isinstance(text, str) else text
+
+        usage = self.client.beta.messages.count_tokens(
+            model=self.model,
+            messages=messages,
+        )
+
+        return usage.input_tokens
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -29,7 +29,7 @@ filetype = "^1.2"
 
 # drivers
 cohere = { version = "^5.11.2", optional = true }
-anthropic = { version = "^0.37.1", optional = true }
+anthropic = { version = "^0.39.0", optional = true }
 transformers = { version = "^4.41.1", optional = true}
 huggingface-hub = { version = "^0.26.2", optional = true }
 boto3 = { version = "^1.34.119", optional = true }

diff --git a/tests/unit/tokenizers/test_anthropic_tokenizer.py b/tests/unit/tokenizers/test_anthropic_tokenizer.py
@@ -1,3 +1,5 @@
+from unittest.mock import Mock
+
 import pytest
 
 from griptape.tokenizers import AnthropicTokenizer
@@ -8,6 +10,16 @@ class TestAnthropicTokenizer:
     def tokenizer(self, request):
         return AnthropicTokenizer(model=request.param)
 
+    @pytest.fixture(autouse=True)
+    def mock_client(self, mocker):
+        mock_client = mocker.patch("anthropic.Anthropic")
+
+        mock_client.return_value = Mock(
+            beta=Mock(messages=Mock(count_tokens=Mock(return_value=Mock(input_tokens=5, output_tokens=10))))
+        )
+
+        return mock_client
+
     @pytest.mark.parametrize(
         ("tokenizer", "expected"),
         [("claude-2.1", 5), ("claude-2.0", 5), ("claude-3-opus", 5), ("claude-3-sonnet", 5), ("claude-3-haiku", 5)],