File tree 1 file changed +6
-2
lines changed
api/core/model_runtime/model_providers/google/llm
1 file changed +6
-2
lines changed Original file line number Diff line number Diff line change @@ -254,8 +254,12 @@ def _handle_generate_response(
254
254
assistant_prompt_message = AssistantPromptMessage (content = response .text )
255
255
256
256
# calculate num tokens
257
- prompt_tokens = self .get_num_tokens (model , credentials , prompt_messages )
258
- completion_tokens = self .get_num_tokens (model , credentials , [assistant_prompt_message ])
257
+ if response .usage_metadata :
258
+ prompt_tokens = response .usage_metadata .prompt_token_count
259
+ completion_tokens = response .usage_metadata .candidates_token_count
260
+ else :
261
+ prompt_tokens = self .get_num_tokens (model , credentials , prompt_messages )
262
+ completion_tokens = self .get_num_tokens (model , credentials , [assistant_prompt_message ])
259
263
260
264
# transform usage
261
265
usage = self ._calc_response_usage (model , credentials , prompt_tokens , completion_tokens )
You can’t perform that action at this time.
0 commit comments