瀏覽代碼

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

tags/0.14.1
Shota Totsuka 10 月之前
父節點
當前提交
7d5a385811
No account linked to committer's email address
共有 1 個文件被更改,包括 6 次插入2 次删除
  1. 6
    2
      api/core/model_runtime/model_providers/google/llm/llm.py

+ 6
- 2
api/core/model_runtime/model_providers/google/llm/llm.py 查看文件

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

# transform usage
usage = self._calc_response_usage(model, credentials, prompt_tokens, completion_tokens)

Loading…
取消
儲存