feat: better token counter

This commit is contained in:
Alex
2024-08-31 17:07:40 +01:00
parent c49b7613e0
commit d9309ebc6e
6 changed files with 34 additions and 40 deletions

View File

@@ -2,7 +2,7 @@ import json
from application.retriever.base import BaseRetriever
from application.core.settings import settings
from application.llm.llm_creator import LLMCreator
from application.utils import count_tokens
from application.utils import num_tokens_from_string
from langchain_community.tools import BraveSearch
@@ -78,7 +78,7 @@ class BraveRetSearch(BaseRetriever):
self.chat_history.reverse()
for i in self.chat_history:
if "prompt" in i and "response" in i:
tokens_batch = count_tokens(i["prompt"]) + count_tokens(
tokens_batch = num_tokens_from_string(i["prompt"]) + num_tokens_from_string(
i["response"]
)
if tokens_current_history + tokens_batch < self.token_limit: