make tokens count working for non-streaming as well

src/lib/components/InferencePlayground/InferencePlayground.svelte

@@ -37,7 +37,7 @@
 	let showTokenModal = false;
 	let loading = false;
 	let latency = 0;
-	let tokensCount = 0;
+	let generatedTokensCount = 0;
 	let abortController: AbortController | undefined = undefined;
 	let waitForNonStreaming = true;
 
@@ -96,17 +96,21 @@
 						if (streamingMessage) {
 							streamingMessage.content = content;
 							conversation.messages = [...conversation.messages];
-							tokensCount += 1;
+							generatedTokensCount += 1;
 						}
 					},
 					abortController
 				);
 			} else {
 				waitForNonStreaming = true;
-				const newMessage = await handleNonStreamingResponse(hf, conversation);
+				const { message: newMessage, completion_tokens: newTokensCount } = await handleNonStreamingResponse(
+					hf,
+					conversation
+				);
 				// check if the user did not abort the request
 				if (waitForNonStreaming) {
 					conversation.messages = [...conversation.messages, newMessage];
+					generatedTokensCount += newTokensCount;
 				}
 			}
 
@@ -206,7 +210,7 @@
 				<IconDelete />
 			</button>
 			<div class="flex-1 items-center justify-center text-center text-sm text-gray-500">
-				<span class="max-xl:hidden">{tokensCount} tokens · Latency {latency}ms</span>
+				<span class="max-xl:hidden">{generatedTokensCount} tokens · Latency {latency}ms</span>
 			</div>
 			<button
 				type="button"

src/lib/components/InferencePlayground/inferencePlaygroundUtils.ts

@@ -38,7 +38,7 @@ export async function handleStreamingResponse(
 export async function handleNonStreamingResponse(
 	hf: HfInference,
 	conversation: Conversation
-): Promise<ChatCompletionInputMessage> {
+): Promise<{ message: ChatCompletionInputMessage; completion_tokens: number }> {
 	const { model, systemMessage } = conversation;
 	const messages = [
 		...(isSystemPromptSupported(model) && systemMessage.content?.length ? [systemMessage] : []),
@@ -53,7 +53,9 @@ export async function handleNonStreamingResponse(
 	});
 
 	if (response.choices && response.choices.length > 0) {
-		return response.choices[0].message;
+		const { message } = response.choices[0];
+		const { completion_tokens } = response.usage;
+		return { message, completion_tokens };
 	}
 	throw new Error("No response from the model");
 }