andre/crewAI
1
0
Fork 0
mirror of https://github.com/crewAIInc/crewAI.git synced 2026-05-03 08:12:39 +00:00
Code Issues Actions 5 Packages Projects Releases Wiki Activity

more wip.

Browse Source
This commit is contained in:
Brandon Hancock
2024-06-26 20:18:23 -07:00
parent be0a4c2fe5
commit 764234c426
3 changed files with 32 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
src/crewai/utilities/token_counter_callback.py
View File

@@ -35,18 +35,18 @@ class TokenProcess:
class TokenCalcHandler(BaseCallbackHandler):
id = uuid.uuid4() # TODO: REMOVE THIS
model: str = ""
model_name: str = ""
token_cost_process: TokenProcess
def __init__(self, model, token_cost_process):
self.model = model
def __init__(self, model_name, token_cost_process):
self.model_name = model_name
self.token_cost_process = token_cost_process
def on_llm_start(
self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
) -> None:
try:
encoding = tiktoken.encoding_for_model(self.model)
encoding = tiktoken.encoding_for_model(self.model_name)
except KeyError:
encoding = tiktoken.get_encoding("cl100k_base")