fix: fixed bug in claude3 tokenizer (#194)

epam · Dec 2, 2024 · 2e32110 · 2e32110
1 parent 71ee9b3
commit 2e32110
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 9 deletions.
diff --git a/aidial_adapter_bedrock/dial_api/storage.py b/aidial_adapter_bedrock/dial_api/storage.py
@@ -40,8 +40,9 @@ async def _get_bucket(self, session: aiohttp.ClientSession) -> Bucket:
                 headers=self.auth_headers,
             ) as response:
                 response.raise_for_status()
-                self.bucket = await response.json()
+                self.bucket = bucket = await response.json()
                 log.debug(f"bucket: {self.bucket}")
+                return bucket
 
         return self.bucket
 

diff --git a/aidial_adapter_bedrock/llm/model/claude/v3/tokenizer.py b/aidial_adapter_bedrock/llm/model/claude/v3/tokenizer.py
@@ -173,8 +173,8 @@ def _tokenize_tool_system_message(
         ):
             return 294 if tool_choice == "auto" else 261
         case (
-            ChatCompletionDeployment.ANTHROPIC_CLAUDE_V3_5_SONNET_V2,
-            ChatCompletionDeployment.ANTHROPIC_CLAUDE_V3_5_SONNET_V2_US,
+            ChatCompletionDeployment.ANTHROPIC_CLAUDE_V3_5_SONNET_V2
+            | ChatCompletionDeployment.ANTHROPIC_CLAUDE_V3_5_SONNET_V2_US
         ):
             return 346 if tool_choice == "auto" else 313
         case (

diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -35,7 +35,7 @@ pytest-dotenv = "^0.5.2"
 pytest-xdist = "^3.5.0"
 
 [tool.poetry.group.lint.dependencies]
-pyright = "1.1.324"
+pyright = "1.1.389"
 black = "24.3.0"
 isort = "5.12.0"
 autoflake = "2.2.0"