fix: token tracking

This commit is contained in:
Han Xiao
2025-02-12 10:14:00 +08:00
parent 959595f971
commit ff59d20dce

View File

@@ -77,7 +77,7 @@ async function emitRemainingContent(
requestId: string, requestId: string,
created: number, created: number,
model: string, model: string,
content: string content: string,
) { ) {
if (!content) return; if (!content) return;
@@ -92,7 +92,7 @@ async function emitRemainingContent(
delta: {content}, delta: {content},
logprobs: null, logprobs: null,
finish_reason: null finish_reason: null
}] }],
}; };
res.write(`data: ${JSON.stringify(chunk)}\n\n`); res.write(`data: ${JSON.stringify(chunk)}\n\n`);
} }
@@ -358,6 +358,7 @@ app.post('/v1/chat/completions', (async (req: Request, res: Response) => {
} }
} }
const usage = context.tokenTracker.getUsageDetails();
if (body.stream) { if (body.stream) {
// Complete any ongoing streaming before sending final answer // Complete any ongoing streaming before sending final answer
await completeCurrentStreaming(streamingState, res, requestId, created, body.model); await completeCurrentStreaming(streamingState, res, requestId, created, body.model);
@@ -390,12 +391,13 @@ app.post('/v1/chat/completions', (async (req: Request, res: Response) => {
delta: {content: result.action === 'answer' ? buildMdFromAnswer(result) : result.think}, delta: {content: result.action === 'answer' ? buildMdFromAnswer(result) : result.think},
logprobs: null, logprobs: null,
finish_reason: 'stop' finish_reason: 'stop'
}] }],
usage
}; };
res.write(`data: ${JSON.stringify(answerChunk)}\n\n`); res.write(`data: ${JSON.stringify(answerChunk)}\n\n`);
res.end(); res.end();
} else { } else {
const usage = context.tokenTracker.getUsageDetails();
const response: ChatCompletionResponse = { const response: ChatCompletionResponse = {
id: requestId, id: requestId,
object: 'chat.completion', object: 'chat.completion',
@@ -473,7 +475,8 @@ app.post('/v1/chat/completions', (async (req: Request, res: Response) => {
delta: {content: errorMessage}, delta: {content: errorMessage},
logprobs: null, logprobs: null,
finish_reason: 'stop' finish_reason: 'stop'
}] }],
usage
}; };
res.write(`data: ${JSON.stringify(errorChunk)}\n\n`); res.write(`data: ${JSON.stringify(errorChunk)}\n\n`);
res.end(); res.end();