pi-llama-cpp 0.3.2 → 0.3.3

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-llama-cpp",
-  "version": "0.3.2",
+  "version": "0.3.3",
   "description": "Pi extension for llama.cpp integration. Supports both router and single modes.",
   "keywords": [
     "pi",

package/src/models/baseModel.ts

@@ -69,6 +69,8 @@ export abstract class BaseModel {
 
   /**
    * Gets the load status of the model
+   *
+   * @returns The current status
    */
   public async getStatus(): Promise<Status> {
     try {
@@ -79,8 +81,10 @@ export abstract class BaseModel {
       if (is_sleeping) return Status.SLEEPING;
       if (!error) return Status.LOADED;
       if (error.code === 503) return Status.LOADING;
+      if (error.code === 400 && error.message === "model is not loaded")
+        return Status.UNLOADED;
 
-      return Status.UNLOADED;
+      return Status.FAILED;
     } catch (err) {
       return Status.FAILED;
     }
@@ -88,6 +92,8 @@ export abstract class BaseModel {
 
   /**
    * Gets the context size of a particular model
+   *
+   * @returns The detected context size
    */
   async getContextSize(): Promise<number> {
     try {
@@ -130,6 +136,7 @@ export abstract class BaseModel {
 
   /**
    * Converts the llama-server model into a configuration object used by Pi
+   *
    * @returns A Pi configuration object
    */
   async toProviderConfig(): Promise<ProviderModelConfig> {
@@ -167,15 +174,21 @@ export abstract class BaseModel {
    * Polls llama-server to check when the model is loaded
    *
    * @param startTime The initial polling timestamp
+   * @param timeout The maximum amount of ms before timeout. Defaults to POLLING_TIMEOUT
+   * @param interval The polling interval. Defaults to POLLING_INTERVAL
    */
-  async pollStatus(startTime = Date.now()): Promise<void> {
+  async pollStatus(
+    startTime: number = Date.now(),
+    timeout: number = POLLING_TIMEOUT,
+    interval: number = POLLING_INTERVAL,
+  ): Promise<void> {
     while ((await this.getStatus()) === Status.LOADING) {
       // Force a timeout if we wasted too much time polling
-      if (Date.now() - startTime > POLLING_TIMEOUT) {
-        const message = `Model loading timed out after ${POLLING_TIMEOUT} ms: ${this.id}`;
+      if (Date.now() - startTime > timeout) {
+        const message = `Model loading timed out after ${timeout} ms: ${this.id}`;
         throw new Error(message);
       }
-      await new Promise((r) => setTimeout(r, POLLING_INTERVAL));
+      await new Promise((r) => setTimeout(r, interval));
     }
   }
 }

package/src/models/routerModel.ts

@@ -1,4 +1,4 @@
-import { DEFAULT_CTX } from "../constants";
+import { DEFAULT_CTX, POLLING_INTERVAL, POLLING_TIMEOUT } from "../constants";
 import { Mode } from "../enums/mode";
 import { Status } from "../enums/status";
 import { ModelsEndpoint } from "../interfaces/endpoints/models";
@@ -22,23 +22,8 @@ export class RouterModel extends BaseModel {
     if (!model) return Status.FAILED;
 
     const status = this.statusMapper[model.status!.value];
-    if (status === Status.UNLOADED) {
-      if (this.model.status!.failed) {
-        /**
-         * Workaround for the currently-bugged /models status detection
-         * (I suspect it was introduced in PR #22683 of llama.cpp)
-         *
-         * This workaround will show an eternal "loading" status when the model's real status
-         * is "failed", which is acceptable, because models in "failed" or "loading" status
-         * shouldn't be used.
-         *
-         * In exchange, it will allow unloaded models to be correctly shown as "unloaded".
-         */
-        // return Status.FAILED;  // <-- Original implementation
-        return await super.getStatus();
-      }
-
-      return Status.UNLOADED;
+    if (status === Status.UNLOADED || status === Status.LOADING) {
+      return super.getStatus();
     }
 
     return status;
@@ -48,22 +33,33 @@ export class RouterModel extends BaseModel {
    * Workaround for the currently-bugged /models status detection
    * (I suspect it was introduced in PR #22683 of llama.cpp)
    *
-   * @returns The detected status
+   * When a model is loaded for the very first time,
+   * this workaround will try to poll to /props instead of /models
+   * for up to 5 seconds to try to detect if the model is really loading,
+   * or if it definitely failed.
+   *
+   * The tradeoff is that we'll have to wait for 5 seconds
+   * while the model is "loading", while not really loading.
+   *
+   * In exchange, it will allow unloaded models to be correctly shown as "unloaded".
    */
-  private async getStatusWorkaround(): Promise<Status> {
-    try {
-      const { is_sleeping, error } = await rpc<PropsEndpoint>(
-        `/props?model=${this.id}`,
-      );
-
-      if (is_sleeping) return Status.SLEEPING;
-      if (!error) return Status.LOADED;
-      if (error.code === 503) return Status.LOADING;
-
-      return Status.UNLOADED;
-    } catch (err) {
-      return Status.FAILED;
+  async pollStatus(startTime = Date.now()): Promise<void> {
+    let elapsed = 0;
+    const limit = 5000;
+
+    // Grab the glitch
+    while (Date.now() - startTime <= limit) {
+      try {
+        await rpc<PropsEndpoint>(`/props?model=${this.id}`);
+        break;
+      } catch {
+        elapsed += POLLING_INTERVAL;
+        await new Promise((r) => setTimeout(r, POLLING_INTERVAL));
+      }
     }
+
+    const timeout = POLLING_TIMEOUT - elapsed;
+    return await super.pollStatus(startTime, timeout);
   }
 
   async getCapabilities(): Promise<["text"] | ["image"]> {

package/src/models/singleModel.ts

@@ -1,8 +1,23 @@
+import { DEFAULT_CTX } from "../constants";
 import { Mode } from "../enums/mode";
+import { PropsEndpoint } from "../interfaces/endpoints/props";
+import { rpc } from "../tools/retriever";
 import { BaseModel } from "./baseModel";
 
 export class SingleModel extends BaseModel {
   get mode(): Mode {
     return Mode.SINGLE;
   }
+
+  async getContextSize(): Promise<number> {
+    try {
+      const { default_generation_settings } = await rpc<PropsEndpoint>(
+        `/props?model=${this.id}`,
+      );
+      const { n_ctx } = default_generation_settings;
+      return n_ctx;
+    } catch {
+      return DEFAULT_CTX;
+    }
+  }
 }