diff options
Diffstat (limited to 'js/src/session/session.ts')
| -rw-r--r-- | js/src/session/session.ts | 12 |
1 files changed, 7 insertions, 5 deletions
diff --git a/js/src/session/session.ts b/js/src/session/session.ts index 1ab09d40d..8b3f990de 100644 --- a/js/src/session/session.ts +++ b/js/src/session/session.ts @@ -162,6 +162,13 @@ export namespace Session { msgs.push(system); state().messages.set(sessionID, msgs); generateText({ + onStepFinish: (step) => { + update(sessionID, (draft) => { + draft.tokens.input += step.usage.inputTokens || 0; + draft.tokens.output += step.usage.outputTokens || 0; + draft.tokens.reasoning += step.usage.reasoningTokens || 0; + }); + }, messages: convertToModelMessages([ { role: "system", @@ -293,11 +300,6 @@ export namespace Session { session.tokens.input += usage.inputTokens || 0; session.tokens.output += usage.outputTokens || 0; session.tokens.reasoning += usage.reasoningTokens || 0; - await update(sessionID, (draft) => { - draft.tokens.input += usage.inputTokens || 0; - draft.tokens.output += usage.outputTokens || 0; - draft.tokens.reasoning += usage.reasoningTokens || 0; - }); return next; } } |