Skip to content

Commit

Permalink
feat: collect usage for mlcllm response
Browse files Browse the repository at this point in the history
  • Loading branch information
Neet-Nestor committed Jun 23, 2024
1 parent 493503e commit 87bc89e
Showing 1 changed file with 9 additions and 3 deletions.
12 changes: 9 additions & 3 deletions app/client/mlcllm.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import log from "loglevel";
import { ChatOptions, LLMApi } from "./api";
import { ChatCompletionFinishReason } from "@mlc-ai/web-llm";
import { ChatCompletionFinishReason, CompletionUsage } from "@mlc-ai/web-llm";

export class MlcLLMApi implements LLMApi {
private endpoint: string;
Expand All @@ -24,6 +24,7 @@ export class MlcLLMApi implements LLMApi {

let reply: string = "";
let stopReason: ChatCompletionFinishReason | undefined;
let usage: CompletionUsage | undefined;

try {
const response = await fetch(`${this.endpoint}/v1/chat/completions`, {
Expand Down Expand Up @@ -52,6 +53,10 @@ export class MlcLLMApi implements LLMApi {
if (data.choices[0].finish_reason) {
stopReason = data.choices[0].finish_reason;
}

if (data.usage) {
usage = data.usage;
}
}
}

Expand All @@ -60,12 +65,13 @@ export class MlcLLMApi implements LLMApi {
break;
}
}
options.onFinish(reply, stopReason);
options.onFinish(reply, stopReason, usage);
} else {
const data = await response.json();
options.onFinish(
data.choices[0].message.content,
data.choices[0].finish_reason,
data.choices[0].finish_reason || undefined,
data.usage || undefined,
);
}
} catch (error: any) {
Expand Down

0 comments on commit 87bc89e

Please sign in to comment.