estel/dify
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

fix: use Gemini response metadata for token counting (#11226)

Browse Source
This commit is contained in:
Shota Totsuka
2024-11-30 18:30:55 +09:00
committed by GitHub
parent e80f41a701
commit 594666eb61
1 changed files with 6 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
api/core/model_runtime/model_providers/google/llm/llm.py
View File

@@ -254,6 +254,10 @@ class GoogleLargeLanguageModel(LargeLanguageModel):
assistant_prompt_message = AssistantPromptMessage(content=response.text) assistant_prompt_message = AssistantPromptMessage(content=response.text)
# calculate num tokens # calculate num tokens
if response.usage_metadata:
prompt_tokens = response.usage_metadata.prompt_token_count
completion_tokens = response.usage_metadata.candidates_token_count
else:
prompt_tokens = self.get_num_tokens(model, credentials, prompt_messages) prompt_tokens = self.get_num_tokens(model, credentials, prompt_messages)
completion_tokens = self.get_num_tokens(model, credentials, [assistant_prompt_message]) completion_tokens = self.get_num_tokens(model, credentials, [assistant_prompt_message])