Skip to content

Commit

Permalink
[Log] Add setLogLevel to Engine Interface
Browse files Browse the repository at this point in the history
  • Loading branch information
Neet-Nestor committed May 27, 2024
1 parent b4bc112 commit 016563e
Show file tree
Hide file tree
Showing 6 changed files with 44 additions and 21 deletions.
11 changes: 10 additions & 1 deletion src/engine.ts
Original file line number Diff line number Diff line change
Expand Up @@ -63,8 +63,8 @@ export async function CreateMLCEngine(
modelId: string,
engineConfig?: MLCEngineConfig,
): Promise<MLCEngine> {
log.setLevel(engineConfig?.logLevel || DefaultLogLevel);
const engine = new MLCEngine();
engine.setLogLevel(engineConfig?.logLevel || DefaultLogLevel);
engine.setInitProgressCallback(engineConfig?.initProgressCallback);
engine.setLogitProcessorRegistry(engineConfig?.logitProcessorRegistry);
await engine.reload(modelId, engineConfig?.chatOpts, engineConfig?.appConfig);
Expand Down Expand Up @@ -640,6 +640,15 @@ export class MLCEngine implements MLCEngineInterface {
return this.getPipeline().getMessage();
}

/**
* Set MLCEngine logging output level
*
* @param logLevel The new log level
*/
setLogLevel(logLevel: LogLevel) {
log.setLevel(logLevel);
}

/**
* Get a new Conversation object based on the chat completion request.
*
Expand Down
12 changes: 6 additions & 6 deletions src/extension_service_worker.ts
Original file line number Diff line number Diff line change
Expand Up @@ -140,10 +140,10 @@ export async function CreateServiceWorkerMLCEngine(
engineConfig?: MLCEngineConfig,
keepAliveMs = 10000,
): Promise<ServiceWorkerMLCEngine> {
const serviceWorkerMLCEngine = new ServiceWorkerMLCEngine(
keepAliveMs,
engineConfig?.logLevel,
);
const serviceWorkerMLCEngine = new ServiceWorkerMLCEngine(keepAliveMs);
if (engineConfig?.logLevel) {
serviceWorkerMLCEngine.setLogLevel(engineConfig.logLevel);
}
serviceWorkerMLCEngine.setInitProgressCallback(
engineConfig?.initProgressCallback,
);
Expand Down Expand Up @@ -192,10 +192,10 @@ class PortAdapter implements ChatWorker {
export class ServiceWorkerMLCEngine extends WebWorkerMLCEngine {
port: chrome.runtime.Port;

constructor(keepAliveMs = 10000, logLevel: LogLevel = "WARN") {
constructor(keepAliveMs = 10000) {
const port = chrome.runtime.connect({ name: "web_llm_service_worker" });
const chatWorker = new PortAdapter(port);
super(chatWorker, logLevel);
super(chatWorker);
this.port = port;
setInterval(() => {
this.keepAlive();
Expand Down
1 change: 1 addition & 0 deletions src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ export {
InitProgressReport,
MLCEngineInterface,
LogitProcessor,
LogLevel,
} from "./types";

export { MLCEngine, CreateMLCEngine } from "./engine";
Expand Down
17 changes: 6 additions & 11 deletions src/service_worker.ts
Original file line number Diff line number Diff line change
Expand Up @@ -206,11 +206,10 @@ export async function CreateServiceWorkerMLCEngine(
"Please refresh the page to retry initializing the service worker.",
);
}
const serviceWorkerMLCEngine = new ServiceWorkerMLCEngine(
serviceWorker,
undefined,
engineConfig?.logLevel,
);
const serviceWorkerMLCEngine = new ServiceWorkerMLCEngine(serviceWorker);
if (engineConfig?.logLevel) {
serviceWorkerMLCEngine.setLogLevel(engineConfig.logLevel);
}
serviceWorkerMLCEngine.setInitProgressCallback(
engineConfig?.initProgressCallback,
);
Expand All @@ -228,15 +227,11 @@ export async function CreateServiceWorkerMLCEngine(
export class ServiceWorkerMLCEngine extends WebWorkerMLCEngine {
missedHeatbeat = 0;

constructor(
worker: IServiceWorker,
keepAliveMs = 10000,
logLevel: LogLevel = "WARN",
) {
constructor(worker: IServiceWorker, keepAliveMs = 10000) {
if (!("serviceWorker" in navigator)) {
throw new Error("Service worker API is not available");
}
super(new ServiceWorker(worker), logLevel);
super(new ServiceWorker(worker));
const onmessage = this.onmessage.bind(this);

(navigator.serviceWorker as ServiceWorkerContainer).addEventListener(
Expand Down
17 changes: 16 additions & 1 deletion src/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -194,6 +194,21 @@ export interface MLCEngineInterface {
inputIds: Array<number>,
isPrefill: boolean,
): Promise<number>;

/**
* Set MLCEngine logging output level
*
* @param logLevel The new log level
*/
setLogLevel(logLevel: LogLevel): void;
}

export type LogLevel = "TRACE" | "DEBUG" | "INFO" | "WARN" | "ERROR" | "SILENT";
export const LOG_LEVELS = {
TRACE: 0,
DEBUG: 1,
INFO: 2,
WARN: 3,
ERROR: 4,
SILENT: 5,
};
export type LogLevel = keyof typeof LOG_LEVELS;
7 changes: 5 additions & 2 deletions src/web_worker.ts
Original file line number Diff line number Diff line change
Expand Up @@ -350,8 +350,7 @@ export class WebWorkerMLCEngine implements MLCEngineInterface {
>();
private pendingPromise = new Map<string, (msg: WorkerResponse) => void>();

constructor(worker: ChatWorker, logLevel: LogLevel = "WARN") {
log.setLevel(logLevel);
constructor(worker: ChatWorker) {
this.worker = worker;
worker.onmessage = (event: any) => {
this.onmessage.bind(this)(event);
Expand Down Expand Up @@ -627,4 +626,8 @@ export class WebWorkerMLCEngine implements MLCEngineInterface {
}
}
}

setLogLevel(logLevel: LogLevel) {
log.setLevel(logLevel);
}
}

0 comments on commit 016563e

Please sign in to comment.