소스 검색

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

tags/0.13.0
Shota Totsuka 11 달 전
부모
커밋
594666eb61
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 파일 보기

@@ -254,8 +254,12 @@ class GoogleLargeLanguageModel(LargeLanguageModel):
assistant_prompt_message = AssistantPromptMessage(content=response.text)

# 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 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…
취소
저장