From 45a2bde4efa944862a6f277b60aa1616320f008a Mon Sep 17 00:00:00 2001 From: Michael Chi Date: Tue, 19 Aug 2025 00:59:47 +0000 Subject: [PATCH] feat: log Gemini token usage --- backend/smartnpc/src/utils/baseball_streaming.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/backend/smartnpc/src/utils/baseball_streaming.py b/backend/smartnpc/src/utils/baseball_streaming.py index 7e2bd23..cff64f3 100644 --- a/backend/smartnpc/src/utils/baseball_streaming.py +++ b/backend/smartnpc/src/utils/baseball_streaming.py @@ -227,6 +227,20 @@ async def chat_streaming(req:NPCSceneConversationRequest, contents = contents, config = generate_content_config, ): + try: + if chunk.usage_metadata.prompt_token_count: + usage_metadata = chunk.usage_metadata + completion_tokens = usage_metadata.total_token_count - usage_metadata.prompt_token_count # pylint: disable=line-too-long + usage = { + "prompt_tokens": usage_metadata.prompt_token_count, + "total_tokens": usage_metadata.total_token_count, + "completion_tokens": completion_tokens + } + logging.info(usage) + except Exception as e: + logging.error( + "Exception getting token count: %s", e + ) parser.parse_chunk(chunk.text) json_parts = parser.get_parts() for json_part_key in json_parts.keys(): # pylint: disable=consider-using-dict-items, consider-iterating-dictionary