hz/ebiz-dify-ai
1
0
Fork 0
mirror of http://112.124.100.131/huang.ze/ebiz-dify-ai.git synced 2025-12-10 11:26:52 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: use Gemini response metadata for token counting (#11743)

Browse Source
This commit is contained in:
Shota Totsuka
2024-12-17 18:42:05 +09:00
committed by GitHub
parent 900e93f758
commit 7d5a385811
1 changed files with 6 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -292,8 +292,12 @@ class GoogleLargeLanguageModel(LargeLanguageModel):
) )
else: else:
# calculate num tokens # calculate num tokens
prompt_tokens = self.get_num_tokens(model, credentials, prompt_messages) if hasattr(response, "usage_metadata") and response.usage_metadata:
completion_tokens = self.get_num_tokens(model, credentials, [assistant_prompt_message]) 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)
completion_tokens = self.get_num_tokens(model, credentials, [assistant_prompt_message])
# transform usage # transform usage
usage = self._calc_response_usage(model, credentials, prompt_tokens, completion_tokens) usage = self._calc_response_usage(model, credentials, prompt_tokens, completion_tokens)