Browse Source

Yields the usage tokens

Vignesh Subbiah 11 months ago
parent
commit
5329c712de
1 changed files with 9 additions and 0 deletions
  1. 9 0
      src/api/providers/unbound.ts

+ 9 - 0
src/api/providers/unbound.ts

@@ -96,6 +96,7 @@ export class UnboundHandler implements ApiHandler, SingleCompletionHandler {
 
 		for await (const chunk of completion) {
 			const delta = chunk.choices[0]?.delta
+			const usage = chunk.usage
 
 			if (delta?.content) {
 				yield {
@@ -103,6 +104,14 @@ export class UnboundHandler implements ApiHandler, SingleCompletionHandler {
 					text: delta.content,
 				}
 			}
+
+			if (usage) {
+				yield {
+					type: "usage",
+					inputTokens: usage?.prompt_tokens || 0,
+					outputTokens: usage?.completion_tokens || 0,
+				}
+			}
 		}
 	}