diff --git a/langchain/chat_models/anthropic.py b/langchain/chat_models/anthropic.py index daed935b..3ac59507 100644 --- a/langchain/chat_models/anthropic.py +++ b/langchain/chat_models/anthropic.py @@ -141,3 +141,9 @@ class ChatAnthropic(BaseChatModel, _AnthropicCommon): completion = response["completion"] message = AIMessage(content=completion) return ChatResult(generations=[ChatGeneration(message=message)]) + + def get_num_tokens(self, text: str) -> int: + """Calculate number of tokens.""" + if not self.count_tokens: + raise NameError("Please ensure the anthropic package is loaded") + return self.count_tokens(text) diff --git a/langchain/llms/anthropic.py b/langchain/llms/anthropic.py index b71fe682..5c2349f4 100644 --- a/langchain/llms/anthropic.py +++ b/langchain/llms/anthropic.py @@ -97,12 +97,6 @@ class _AnthropicCommon(BaseModel): return stop - def get_num_tokens(self, text: str) -> int: - """Calculate number of tokens.""" - if not self.count_tokens: - raise NameError("Please ensure the anthropic package is loaded") - return self.count_tokens(text) - class Anthropic(LLM, _AnthropicCommon): r"""Wrapper around Anthropic's large language models. @@ -263,3 +257,9 @@ class Anthropic(LLM, _AnthropicCommon): stop_sequences=stop, **self._default_params, ) + + def get_num_tokens(self, text: str) -> int: + """Calculate number of tokens.""" + if not self.count_tokens: + raise NameError("Please ensure the anthropic package is loaded") + return self.count_tokens(text)