From b939692187d67ff5d03c5b8f2f6e28412d2330e1 Mon Sep 17 00:00:00 2001 From: Louis Date: Fri, 8 Dec 2023 16:38:13 +0700 Subject: [PATCH] chore: stop inference event --- core/src/events.ts | 2 + core/src/extensions/inference.ts | 18 +------- .../inference-nitro-extension/src/index.ts | 43 ++++++------------- .../inference-openai-extension/src/index.ts | 39 +++++------------ web/screens/Chat/MessageToolbar/index.tsx | 5 +-- .../ExploreModelItemHeader/index.tsx | 20 ++------- 6 files changed, 32 insertions(+), 95 deletions(-) diff --git a/core/src/events.ts b/core/src/events.ts index bfaf3ea58..1acbef918 100644 --- a/core/src/events.ts +++ b/core/src/events.ts @@ -18,6 +18,8 @@ export enum EventName { OnModelStop = "OnModelStop", /** The `OnModelStopped` event is emitted when a model stopped ok. */ OnModelStopped = "OnModelStopped", + /** The `OnInferenceStopped` event is emitted when a inference is stopped. */ + OnInferenceStopped = "OnInferenceStopped", } /** diff --git a/core/src/extensions/inference.ts b/core/src/extensions/inference.ts index 483ba1339..9453a06d5 100644 --- a/core/src/extensions/inference.ts +++ b/core/src/extensions/inference.ts @@ -5,26 +5,10 @@ import { BaseExtension } from "../extension"; * Inference extension. Start, stop and inference models. */ export abstract class InferenceExtension extends BaseExtension { - /** - * Initializes the model for the extension. - * @param modelId - The ID of the model to initialize. - */ - abstract initModel(modelId: string, settings?: ModelSettingParams): Promise; - - /** - * Stops the model for the extension. - */ - abstract stopModel(): Promise; - - /** - * Stops the streaming inference. - */ - abstract stopInference(): Promise; - /** * Processes an inference request. * @param data - The data for the inference request. * @returns The result of the inference request. */ - abstract inferenceRequest(data: MessageRequest): Promise; + abstract inference(data: MessageRequest): Promise; } diff --git a/extensions/inference-nitro-extension/src/index.ts b/extensions/inference-nitro-extension/src/index.ts index a1c125ce5..975d94100 100644 --- a/extensions/inference-nitro-extension/src/index.ts +++ b/extensions/inference-nitro-extension/src/index.ts @@ -74,41 +74,17 @@ export default class JanInferenceNitroExtension implements InferenceExtension { events.on(EventName.OnModelStop, (model: Model) => { JanInferenceNitroExtension.handleModelStop(model); }); + + events.on(EventName.OnInferenceStopped, () => { + JanInferenceNitroExtension.handleInferenceStopped(this); + }); } /** * Stops the model inference. */ - onUnload(): void { - this.stopModel(); - } + onUnload(): void {} - /** - * Initializes the model with the specified file name. - * @param {string} modelId - The ID of the model to initialize. - * @returns {Promise} A promise that resolves when the model is initialized. - */ - async initModel( - modelId: string, - settings?: ModelSettingParams - ): Promise {} - - /** - * Stops the model. - * @returns {Promise} A promise that resolves when the model is stopped. - */ - async stopModel(): Promise { - return executeOnMain(MODULE, "killSubprocess"); - } - - /** - * Stops streaming inference. - * @returns {Promise} A promise that resolves when the streaming is stopped. - */ - async stopInference(): Promise { - this.isCancelled = true; - this.controller?.abort(); - } private async writeDefaultEngineSettings() { try { @@ -160,12 +136,19 @@ export default class JanInferenceNitroExtension implements InferenceExtension { } } + private static async handleInferenceStopped( + instance: JanInferenceNitroExtension + ) { + instance.isCancelled = true; + instance.controller?.abort(); + } + /** * Makes a single response inference request. * @param {MessageRequest} data - The data for the inference request. * @returns {Promise} A promise that resolves with the inference response. */ - async inferenceRequest(data: MessageRequest): Promise { + async inference(data: MessageRequest): Promise { const timestamp = Date.now(); const message: ThreadMessage = { thread_id: data.threadId, diff --git a/extensions/inference-openai-extension/src/index.ts b/extensions/inference-openai-extension/src/index.ts index 4eb179da8..06e0f5e04 100644 --- a/extensions/inference-openai-extension/src/index.ts +++ b/extensions/inference-openai-extension/src/index.ts @@ -71,6 +71,9 @@ export default class JanInferenceOpenAIExtension implements InferenceExtension { events.on(EventName.OnModelStop, (model: OpenAIModel) => { JanInferenceOpenAIExtension.handleModelStop(model); }); + events.on(EventName.OnInferenceStopped, () => { + JanInferenceOpenAIExtension.handleInferenceStopped(this); + }); } /** @@ -78,18 +81,6 @@ export default class JanInferenceOpenAIExtension implements InferenceExtension { */ onUnload(): void {} - /** - * Initializes the model with the specified file name. - * @param {string} modelId - The ID of the model to initialize. - * @returns {Promise} A promise that resolves when the model is initialized. - */ - async initModel( - modelId: string, - settings?: ModelSettingParams - ): Promise { - return; - } - static async writeDefaultEngineSettings() { try { const engineFile = join( @@ -110,27 +101,13 @@ export default class JanInferenceOpenAIExtension implements InferenceExtension { console.error(err); } } - /** - * Stops the model. - * @returns {Promise} A promise that resolves when the model is stopped. - */ - async stopModel(): Promise {} - - /** - * Stops streaming inference. - * @returns {Promise} A promise that resolves when the streaming is stopped. - */ - async stopInference(): Promise { - this.isCancelled = true; - this.controller?.abort(); - } /** * Makes a single response inference request. * @param {MessageRequest} data - The data for the inference request. * @returns {Promise} A promise that resolves with the inference response. */ - async inferenceRequest(data: MessageRequest): Promise { + async inference(data: MessageRequest): Promise { const timestamp = Date.now(); const message: ThreadMessage = { thread_id: data.threadId, @@ -168,7 +145,6 @@ export default class JanInferenceOpenAIExtension implements InferenceExtension { JanInferenceOpenAIExtension.writeDefaultEngineSettings(); // Todo: Check model list with API key events.emit(EventName.OnModelReady, model); - // events.emit(EventName.OnModelFail, model) } } @@ -179,6 +155,13 @@ export default class JanInferenceOpenAIExtension implements InferenceExtension { events.emit(EventName.OnModelStopped, model); } + private static async handleInferenceStopped( + instance: JanInferenceOpenAIExtension + ) { + instance.isCancelled = true; + instance.controller?.abort(); + } + /** * Handles a new message request by making an inference request and emitting events. * Function registered in event manager, should be static to avoid binding issues. diff --git a/web/screens/Chat/MessageToolbar/index.tsx b/web/screens/Chat/MessageToolbar/index.tsx index f877e1bdb..fe7cac1f5 100644 --- a/web/screens/Chat/MessageToolbar/index.tsx +++ b/web/screens/Chat/MessageToolbar/index.tsx @@ -30,9 +30,8 @@ const MessageToolbar = ({ message }: { message: ThreadMessage }) => { const { resendChatMessage } = useSendChatMessage() const onStopInferenceClick = async () => { - await extensionManager - .get(ExtensionType.Inference) - ?.stopInference() + events.emit(EventName.OnInferenceStopped, {}) + setTimeout(() => { events.emit(EventName.OnMessageUpdate, { ...message, diff --git a/web/screens/ExploreModels/ExploreModelItemHeader/index.tsx b/web/screens/ExploreModels/ExploreModelItemHeader/index.tsx index 69cd8af3e..ba23056c6 100644 --- a/web/screens/ExploreModels/ExploreModelItemHeader/index.tsx +++ b/web/screens/ExploreModels/ExploreModelItemHeader/index.tsx @@ -55,23 +55,9 @@ const ExploreModelItemHeader: React.FC = ({ model, onClick, open }) => { const isDownloaded = downloadedModels.find((md) => md.id === model.id) != null - let downloadButton; - - if (model.engine === 'openai') { - downloadButton = ( - - ); - } else if (model.engine === 'nitro') { - downloadButton = ( - - ); - } + let downloadButton = ( + + ) const onUseModelClick = () => { startModel(model.id)