Skip to content

Commit 6b8b9c8

Browse files
committed
Update token_counter.py
1 parent ea13d7a commit 6b8b9c8

File tree

1 file changed

+9
-0
lines changed

1 file changed

+9
-0
lines changed

metagpt/utils/token_counter.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
ref4: https://github.com/hwchase17/langchain/blob/master/langchain/chat_models/openai.py
1111
ref5: https://ai.google.dev/models/gemini
1212
"""
13+
import anthropic
1314
import tiktoken
1415
from openai.types import CompletionUsage
1516
from openai.types.chat import ChatCompletionChunk
@@ -377,6 +378,10 @@
377378

378379
def count_input_tokens(messages, model="gpt-3.5-turbo-0125"):
379380
"""Return the number of tokens used by a list of messages."""
381+
if "claude" in model:
382+
vo = anthropic.Client()
383+
num_tokens = vo.count_tokens(str(messages))
384+
return num_tokens
380385
try:
381386
encoding = tiktoken.encoding_for_model(model)
382387
except KeyError:
@@ -463,6 +468,10 @@ def count_output_tokens(string: str, model: str) -> int:
463468
Returns:
464469
int: The number of tokens in the text string.
465470
"""
471+
if "claude" in model:
472+
vo = anthropic.Client()
473+
num_tokens = vo.count_tokens(string)
474+
return num_tokens
466475
try:
467476
encoding = tiktoken.encoding_for_model(model)
468477
except KeyError:

0 commit comments

Comments
 (0)