apunkt/SurfSense
1
0
Fork 0
mirror of https://github.com/MODSetter/SurfSense.git synced 2026-06-04 20:05:16 +02:00
Code Issues Projects Releases Packages Wiki Activity

feat: improve token usage tracking and response handling in chat routes and services

Browse source
This commit is contained in:
Anish Sarkar 2026-04-14 14:28:31 +05:30
parent 55099a20ac
commit 5af6005163
3 changed files with 39 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

7
surfsense_backend/app/services/token_tracking_service.py
View file

@ -87,6 +87,7 @@ def start_turn() -> TurnTokenAccumulator:
"""Create a fresh accumulator for the current async context and return it."""
acc = TurnTokenAccumulator()
_turn_accumulator.set(acc)
logger.info("[TokenTracking] start_turn: new accumulator created (id=%s)", id(acc))
return acc
@ -106,10 +107,12 @@ class TokenTrackingCallback(CustomLogger):
) -> None:
acc = _turn_accumulator.get()
if acc is None:
logger.debug("[TokenTracking] async_log_success_event fired but no accumulator in context")
return
usage = getattr(response_obj, "usage", None)
if not usage:
logger.debug("[TokenTracking] async_log_success_event fired but response has no usage data")
return
prompt_tokens = getattr(usage, "prompt_tokens", 0) or 0
@ -124,6 +127,10 @@ class TokenTrackingCallback(CustomLogger):
completion_tokens=completion_tokens,
total_tokens=total_tokens,
)
logger.info(
"[TokenTracking] Captured: model=%s prompt=%d completion=%d total=%d (accumulator now has %d calls)",
model, prompt_tokens, completion_tokens, total_tokens, len(acc.calls),
)
token_tracker = TokenTrackingCallback()