Skip to content

Commit 8a09c96

Browse files
committed
fix token counting for r1
1 parent 45160cd commit 8a09c96

File tree

1 file changed

+21
-3
lines changed

1 file changed

+21
-3
lines changed

src/agent/custom_massage_manager.py

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,11 +9,15 @@
99
from browser_use.agent.views import ActionResult, AgentStepInfo
1010
from browser_use.browser.views import BrowserState
1111
from langchain_core.language_models import BaseChatModel
12+
from langchain_anthropic import ChatAnthropic
13+
from langchain_core.language_models import BaseChatModel
1214
from langchain_core.messages import (
13-
HumanMessage,
14-
AIMessage
15+
AIMessage,
16+
BaseMessage,
17+
HumanMessage,
1518
)
16-
19+
from langchain_openai import ChatOpenAI
20+
from ..utils.llm import DeepSeekR1ChatOpenAI
1721
from .custom_prompts import CustomAgentMessagePrompt
1822

1923
logger = logging.getLogger(__name__)
@@ -108,3 +112,17 @@ def add_state_message(
108112
step_info=step_info,
109113
).get_user_message()
110114
self._add_message_with_tokens(state_message)
115+
116+
def _count_text_tokens(self, text: str) -> int:
117+
if isinstance(self.llm, (ChatOpenAI, ChatAnthropic, DeepSeekR1ChatOpenAI)):
118+
try:
119+
tokens = self.llm.get_num_tokens(text)
120+
except Exception:
121+
tokens = (
122+
len(text) // self.ESTIMATED_TOKENS_PER_CHARACTER
123+
) # Rough estimate if no tokenizer available
124+
else:
125+
tokens = (
126+
len(text) // self.ESTIMATED_TOKENS_PER_CHARACTER
127+
) # Rough estimate if no tokenizer available
128+
return tokens

0 commit comments

Comments
 (0)