From 58e12f35c978cc7ce40839f277f4ae3cbf7bc234 Mon Sep 17 00:00:00 2001
From: Louis <louis@jan.ai>
Date: Thu, 14 Mar 2024 23:59:42 +0700
Subject: [PATCH] fix: wrong engine handling (#2363)

---
 extensions/tensorrt-llm-extension/src/index.ts | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/extensions/tensorrt-llm-extension/src/index.ts b/extensions/tensorrt-llm-extension/src/index.ts
index cd85601dd..02c676841 100644
--- a/extensions/tensorrt-llm-extension/src/index.ts
+++ b/extensions/tensorrt-llm-extension/src/index.ts
@@ -129,6 +129,8 @@ export default class TensorRTLLMExtension extends LocalOAIEngine {
   }
 
   async onModelInit(model: Model): Promise<void> {
+    if (model.engine !== this.provider) return
+
     if ((await this.installationState()) === 'Installed')
       return super.onModelInit(model)
     else {
@@ -138,7 +140,6 @@ export default class TensorRTLLMExtension extends LocalOAIEngine {
           message: 'EXTENSION_IS_NOT_INSTALLED::TensorRT-LLM extension',
         },
       })
-      return
     }
   }
 
@@ -162,6 +163,7 @@ export default class TensorRTLLMExtension extends LocalOAIEngine {
   }
 
   inference(data: MessageRequest): void {
+    if (!this.isRunning) return
     // TensorRT LLM Extension supports streaming only
     if (data.model) data.model.parameters.stream = true
     super.inference(data)