pi-llama-cpp 0.5.1 → 0.6.0

package/src/server.ts

@@ -0,0 +1,171 @@
+import { PROVIDER_NAME, PROVIDER_PREFIX } from "./constants";
+import { Mode } from "./enums/mode";
+import { HealthEndpoint } from "./interfaces/endpoints/health";
+import { ModelsEndpoint } from "./interfaces/endpoints/models";
+import { PropsEndpoint } from "./interfaces/endpoints/props";
+import { BaseModel } from "./models/baseModel";
+import { LegacyModel } from "./models/legacyModel";
+import { RouterModel } from "./models/routerModel";
+import { SingleModel } from "./models/singleModel";
+import { ConfigResolver } from "./resolver";
+
+export class Server {
+  readonly models: BaseModel[] = [];
+
+  constructor(readonly baseUrl: string) {}
+
+  /**
+   * Generates a unique provider ID from a server URL.
+   */
+  get providerId(): string {
+    return `${PROVIDER_PREFIX}=${this.baseUrl}`;
+  }
+
+  /**
+   * Generates a human-readable provider name from a server URL.
+   */
+  get providerName(): string {
+    return `${PROVIDER_NAME} (${this.baseUrl})`;
+  }
+
+  /**
+   * Retrieves the API key from the resolver
+   * @returns The API key
+   */
+  async getApiKey(): Promise<string> {
+    return await new ConfigResolver().resolveApiKey(this.providerId);
+  }
+
+  /**
+   * Fetches models from the server and populates {@link models}
+   */
+  async initialize() {
+    const { data } = await this.fetchModels();
+    const mode = await this.detectServerMode();
+
+    // Setup models
+    const modelCtor = {
+      [Mode.ROUTER]: RouterModel,
+      [Mode.LEGACY]: LegacyModel,
+      [Mode.SINGLE]: SingleModel,
+    }[mode];
+
+    const models: BaseModel[] = data
+      .map((m) => new modelCtor(m, this))
+      .sort((a, b) => (a.id > b.id ? 1 : a.id === b.id ? 0 : -1));
+
+    this.models.length = 0;
+    this.models.push(...models);
+  }
+
+  /**
+   * Detects the mode of the server
+   *
+   * @returns The detected mode
+   */
+  private async detectServerMode(): Promise<Mode> {
+    const { role } = await this.fetchServerProps();
+    const { data } = await this.fetchModels();
+
+    if (role === "router") return Mode.ROUTER;
+    if ("max_model_len" in data[0]) return Mode.LEGACY;
+    return Mode.SINGLE;
+  }
+
+  /**
+   * Detects if the server is ready
+   * @returns True if it's ready to work
+   */
+  async isReady(): Promise<boolean> {
+    try {
+      const { status } = await this.fetchServerHealth();
+      return status === "ok";
+    } catch {
+      return false;
+    }
+  }
+
+  /**
+   * Retrieves the health status of the server
+   *
+   * @returns The health status
+   */
+  async fetchServerHealth(): Promise<HealthEndpoint> {
+    return await this.rpc<HealthEndpoint>("/health");
+  }
+
+  /**
+   * Fetches models from the server
+   *
+   * @return The models from the server
+   */
+  async fetchModels(): Promise<ModelsEndpoint> {
+    return await this.rpc<ModelsEndpoint>("/v1/models");
+  }
+
+  /**
+   * Fetches general properties of the server
+   *
+   * @return The properties of the server
+   */
+  async fetchServerProps(): Promise<PropsEndpoint> {
+    return await this.rpc<PropsEndpoint>("/props?autoload=false");
+  }
+
+  /**
+   * Fetches properties of a specific model from the server
+   *
+   * @param modelId The ID of the model
+   * @return The properties of the specified model
+   */
+  async fetchModelProps(modelId: string): Promise<PropsEndpoint> {
+    return await this.rpc<PropsEndpoint>(
+      `/props?model=${modelId}&autoload=false`,
+    );
+  }
+
+  /**
+   * Sends a request associated to a specific model from the server
+   *
+   * @param resource The specified resource ("load" | "unload")
+   * @param model The targeted model
+   */
+  async postRequest(
+    resource: "load" | "unload",
+    model: string,
+  ): Promise<ModelsEndpoint> {
+    return await this.rpc<ModelsEndpoint>(`/models/${resource}`, { model });
+  }
+
+  /**
+   * Makes an HTTP request to the llama-server and returns the parsed JSON response
+   *
+   * @param endpoint The endpoint path to fetch (e.g. "/health")
+   * @param body The optional request body for POST requests
+   * @returns The parsed JSON response from the server
+   */
+  private async rpc<T>(
+    endpoint: string,
+    body?: Record<string, unknown>,
+  ): Promise<T> {
+    const url = `${this.baseUrl}${endpoint}`;
+    const apiKey = await this.getApiKey();
+
+    const data = {
+      method: body ? "POST" : "GET",
+      headers: body ? { "Content-Type": "application/json" } : undefined,
+      body: body ? JSON.stringify(body) : undefined,
+    };
+
+    const res = await fetch(url, {
+      ...data,
+      headers: {
+        ...data.headers,
+        ...(apiKey ? { Authorization: `Bearer ${apiKey}` } : {}),
+      },
+    });
+
+    const response: T = await res.json();
+    return response;
+  }
+}

package/tests/commandManager.test.ts

@@ -1,153 +1,202 @@
 import { beforeEach, describe, expect, it, vi } from "vitest";
-import { PROVIDER_ID, PROVIDER_NAME } from "../src/constants";
-import { CommandManager } from "../src/manager";
-
-// Mock modules at top level (vi.mock is hoisted)
-vi.mock("../src/tools/retriever", () => ({
-  isServerReady: vi.fn(),
-  listModels: vi.fn(),
-}));
-
-vi.mock("../src/tools/resolver", () => ({
-  resolveUrl: vi.fn(),
-  resolveApiKey: vi.fn(),
-}));
-
-// Import mocked functions after vi.mock
-import { resolveApiKey, resolveUrl } from "../src/tools/resolver";
-import { isServerReady, listModels } from "../src/tools/retriever";
-
-const mockPi = {
-  registerProvider: vi.fn(),
-};
+import { Action } from "../src/enums/action";
+import { CommandManager } from "../src/managers/command";
+import { ServerManager } from "../src/managers/server";
+import {
+  createMockCtx,
+  createMockModel,
+  createMockPi,
+  createMockServer,
+  mockRpc,
+} from "./mocks";
 
 beforeEach(() => {
   vi.clearAllMocks();
-  (resolveUrl as any).mockResolvedValue("http://127.0.0.1:8080");
-  (resolveApiKey as any).mockResolvedValue("test-key");
+  mockRpc.mockResolvedValue({ data: [] });
 });
 
 describe("CommandManager", () => {
-  it("should register empty models when server is not ready", async () => {
-    (isServerReady as any).mockResolvedValue(false);
-
-    const manager = new CommandManager(mockPi as any);
-    await manager.initialize();
-
-    expect(mockPi.registerProvider).toHaveBeenCalledWith(PROVIDER_ID, {
-      name: PROVIDER_NAME,
-      baseUrl: "http://127.0.0.1:8080",
-      api: "openai-completions",
-      apiKey: "test-key",
-      models: [],
-    });
+  let serverManager: ServerManager;
+  let commandManager: CommandManager;
+  let mockPi: ReturnType<typeof createMockPi>;
+
+  beforeEach(() => {
+    mockPi = createMockPi();
+    serverManager = new ServerManager([]);
+    commandManager = new CommandManager(serverManager);
   });
 
-  it("should update and register models when server is ready", async () => {
-    const mockModel = {
-      name: "test-model",
-      id: "test-model",
-      toProviderConfig: vi
-        .fn()
-        .mockResolvedValue({ id: "test-model", maxTokens: 32000 }),
-    };
-    (isServerReady as any).mockResolvedValue(true);
-    (listModels as any).mockResolvedValue([mockModel]);
-
-    const manager = new CommandManager(mockPi as any);
-    await manager.initialize();
-
-    expect(resolveUrl).toHaveBeenCalledWith(expect.any(String));
-    expect(listModels).toHaveBeenCalled();
-    expect(mockPi.registerProvider).toHaveBeenCalledWith(PROVIDER_ID, {
-      name: PROVIDER_NAME,
-      baseUrl: "http://127.0.0.1:8080",
-      api: "openai-completions",
-      apiKey: "test-key",
-      models: [{ id: "test-model", maxTokens: 32000 }],
+  describe("getArgumentCompletions", () => {
+    it("should provide completions for /models", () => {
+      const completions = commandManager.getArgumentCompletions("");
+      expect(completions).toHaveLength(2);
+      expect(completions?.map((c) => c.value)).toEqual(["info", "unload"]);
     });
-  });
 
-  it("should call notFoundCommand when server is not ready in run()", async () => {
-    (isServerReady as any).mockResolvedValue(false);
-
-    const manager = new CommandManager(mockPi as any);
-    await manager.run("", { ui: { notify: vi.fn() } } as any);
+    it("should filter completions by prefix", () => {
+      const completions = commandManager.getArgumentCompletions("u");
+      expect(completions).toHaveLength(1);
+      expect(completions?.[0].value).toBe("unload");
+    });
 
-    expect(mockPi.registerProvider).not.toHaveBeenCalled();
+    it("should return null when no completions match", () => {
+      const completions = commandManager.getArgumentCompletions("zzz");
+      expect(completions).toBeNull();
+    });
   });
 
-  it("should show info for all models when args is 'info'", async () => {
-    const mockModel = {
-      name: "test-model",
-      id: "test-model",
-      getInfo: vi.fn().mockResolvedValue("Model info for test-model"),
-      toProviderConfig: vi.fn().mockResolvedValue({ id: "test-model" }),
-    };
-    (isServerReady as any).mockResolvedValue(true);
-    (listModels as any).mockResolvedValue([mockModel]);
-
-    const notifyFn = vi.fn();
-    const manager = new CommandManager(mockPi as any);
-    await manager.initialize();
-    await manager.run("info", {
-      ui: { notify: notifyFn, theme: { fg: (_c: string, t: string) => t } },
-    } as any);
-
-    expect(notifyFn).toHaveBeenCalledWith("Model info for test-model", "info");
-    // Called once in initialize() and once in run() to refresh the model list
-    expect(listModels).toHaveBeenCalledTimes(2);
-  });
+  describe("handleCommand", () => {
+    it("should unload all models when args is 'unload'", async () => {
+      const model1 = createMockModel("model-1");
+      const model2 = createMockModel("model-2");
+      const server = createMockServer({
+        baseUrl: "http://127.0.0.1:8080",
+        models: [model1, model2],
+      });
+      serverManager = new ServerManager([server] as any);
+      commandManager = new CommandManager(serverManager);
+
+      const ctx = {
+        ui: {
+          notify: vi.fn(),
+          theme: { fg: (_: string, text: string) => text },
+        },
+      } as any;
+
+      await commandManager.handleCommand("unload", ctx, mockPi as any);
+
+      expect(model1.unload).toHaveBeenCalled();
+      expect(model2.unload).toHaveBeenCalled();
+      expect(ctx.ui.notify).toHaveBeenCalledWith(
+        "Unloaded all Llama.cpp models",
+        "info",
+      );
+    });
 
-  it("should unload all models when args is 'unload'", async () => {
-    const mockModel1 = {
-      name: "model-1",
-      id: "model-1",
-      unload: vi.fn().mockResolvedValue(undefined),
-      toProviderConfig: vi.fn().mockResolvedValue({ id: "model-1" }),
-    };
-    const mockModel2 = {
-      name: "model-2",
-      id: "model-2",
-      unload: vi.fn().mockResolvedValue(undefined),
-      toProviderConfig: vi.fn().mockResolvedValue({ id: "model-2" }),
-    };
-    (isServerReady as any).mockResolvedValue(true);
-    (listModels as any).mockResolvedValue([mockModel1, mockModel2]);
-
-    const notifyFn = vi.fn();
-    const manager = new CommandManager(mockPi as any);
-    await manager.initialize();
-    await manager.run("unload", {
-      ui: { notify: notifyFn },
-    } as any);
-
-    expect(mockModel1.unload).toHaveBeenCalled();
-    expect(mockModel2.unload).toHaveBeenCalled();
-    expect(notifyFn).toHaveBeenCalledWith(
-      "Unloaded all Llama.cpp models",
-      "info",
-    );
+    it("should show model info when args is 'info'", async () => {
+      const model1 = createMockModel("model-1");
+      const model2 = createMockModel("model-2");
+      const server = createMockServer({
+        baseUrl: "http://127.0.0.1:8080",
+        models: [model1, model2],
+      });
+      serverManager = new ServerManager([server] as any);
+      commandManager = new CommandManager(serverManager);
+
+      const ctx = {
+        ui: {
+          notify: vi.fn(),
+          theme: { fg: (_: string, text: string) => text },
+        },
+      } as any;
+
+      await commandManager.handleCommand("info", ctx, mockPi as any);
+
+      expect(model1.getInfo).toHaveBeenCalled();
+      expect(model2.getInfo).toHaveBeenCalled();
+    });
   });
 
-  it("should dispatch modelsCommand when args is empty", async () => {
-    const mockModel = {
-      name: "test-model",
-      id: "test-model",
-      getLabel: vi.fn().mockResolvedValue("test-model"),
-      toProviderConfig: vi.fn().mockResolvedValue({ id: "test-model" }),
+  describe("/models interactive menu", () => {
+    const CHOICE = "model-a   [Server: http://127.0.0.1:8080]";
+
+    /**
+     * Helper to create a CommandManager with mock servers and models.
+     */
+    const createCommandManager = (
+      models: ReturnType<typeof createMockModel>[],
+    ) => {
+      const mockPi = createMockPi();
+      const servers = models.map((model) =>
+        createMockServer({
+          baseUrl: model.serverUrl,
+          models: [model],
+        }),
+      );
+      const serverManager = new ServerManager(servers as any);
+      return {
+        commandManager: new CommandManager(serverManager),
+        serverManager,
+        mockPi,
+      };
     };
-    (isServerReady as any).mockResolvedValue(true);
-    (listModels as any).mockResolvedValue([mockModel]);
-
-    const selectFn = vi.fn().mockReturnValue(null); // cancel immediately
-    const manager = new CommandManager(mockPi as any);
-    await manager.initialize();
-    await manager.run("", {
-      ui: { notify: vi.fn(), select: selectFn },
-    } as any);
-
-    // modelsCommand was called (select is invoked for model picking)
-    expect(selectFn).toHaveBeenCalled();
+
+    it("should return early on cancel (null model selection)", async () => {
+      const models = [createMockModel("model-a")];
+      const { commandManager, mockPi } = createCommandManager(models);
+      const ctx = createMockCtx(() => null);
+
+      await commandManager.handleCommand("", ctx as any, mockPi as any);
+
+      expect(ctx.ui.notify).not.toHaveBeenCalled();
+    });
+
+    it("should show info when INFO action is selected", async () => {
+      const model = createMockModel("model-a");
+      const { commandManager, mockPi } = createCommandManager([model]);
+      let selectCallCount = 0;
+      const ctx = createMockCtx(() => {
+        selectCallCount++;
+        if (selectCallCount === 1) return CHOICE;
+        return Action.INFO;
+      });
+
+      await commandManager.handleCommand("", ctx as any, mockPi as any);
+
+      expect(ctx.ui.notify).toHaveBeenCalledWith(
+        "Model: model-a\nID: model-a",
+        "info",
+      );
+    });
+
+    it("should unload model when UNLOAD action is selected", async () => {
+      const model = createMockModel("model-a");
+      const { commandManager, mockPi } = createCommandManager([model]);
+      let selectCallCount = 0;
+      const ctx = createMockCtx(() => {
+        selectCallCount++;
+        if (selectCallCount === 1) return CHOICE;
+        return Action.UNLOAD;
+      });
+
+      await commandManager.handleCommand("", ctx as any, mockPi as any);
+
+      expect(model.unload).toHaveBeenCalled();
+      expect(ctx.ui.notify).toHaveBeenCalledWith("Unloaded model-a", "info");
+    });
+
+    it("should switch model when SWITCH action is selected", async () => {
+      const model = createMockModel("model-a");
+      const { commandManager, mockPi } = createCommandManager([model]);
+      let selectCallCount = 0;
+      const ctx = createMockCtx(() => {
+        selectCallCount++;
+        if (selectCallCount === 1) return CHOICE;
+        return Action.SWITCH;
+      });
+
+      await commandManager.handleCommand("", ctx as any, mockPi as any);
+
+      expect(mockPi.setModel).toHaveBeenCalled();
+      expect(ctx.ui.notify).toHaveBeenCalledWith("Model model-a ready", "info");
+    });
+
+    it("should loop back to model selection when action is cancelled", async () => {
+      const model = createMockModel("model-a");
+      const { commandManager, mockPi } = createCommandManager([model]);
+
+      let selectCallCount = 0;
+      const ctx = createMockCtx(() => {
+        selectCallCount++;
+        // 1st: select model-a, 2nd: cancel action, 3rd: cancel model => exit
+        if (selectCallCount === 1) return CHOICE;
+        return null;
+      });
+
+      await commandManager.handleCommand("", ctx as any, mockPi as any);
+
+      expect(ctx.ui.select).toHaveBeenCalledTimes(3);
+      expect(ctx.ui.notify).not.toHaveBeenCalled();
+    });
   });
 });

package/tests/legacyModel.test.ts

@@ -0,0 +1,112 @@
+import { beforeEach, describe, expect, it } from "vitest";
+import { Mode } from "../src/enums/mode";
+import { Status } from "../src/enums/status";
+import { DataProperty } from "../src/interfaces/endpoints/models";
+import { LegacyModel } from "../src/models/legacyModel";
+import { createMockServer, mockRpc } from "./mocks";
+
+beforeEach(() => {
+  mockRpc.mockReset();
+});
+
+const createModel = (extra: Partial<DataProperty> = {}): LegacyModel =>
+  new LegacyModel(
+    {
+      id: "test",
+      tags: [],
+      object: "model",
+      owned_by: "test",
+      created: Date.now(),
+      ...extra,
+    },
+    createMockServer(),
+  );
+
+describe("LegacyModel mode", () => {
+  it("should always return LEGACY mode", () => {
+    const model = createModel();
+    expect(model.mode).toBe(Mode.LEGACY);
+  });
+});
+
+describe("LegacyModel capabilities", () => {
+  it("should detect image capability when multimodal is in capabilities", async () => {
+    mockRpc.mockResolvedValueOnce({ modalities: { vision: true } });
+
+    const model = createModel();
+    const capabilities = await model.getCapabilities();
+
+    expect(capabilities).toEqual(["text", "image"]);
+  });
+
+  it("should detect text-only capability when multimodal is not in capabilities", async () => {
+    mockRpc.mockResolvedValueOnce({ modalities: { vision: false } });
+
+    const model = createModel();
+    const capabilities = await model.getCapabilities();
+
+    expect(capabilities).toEqual(["text"]);
+  });
+});
+
+describe("LegacyModel getStatus", () => {
+  it("should return LOADED when not sleeping", async () => {
+    mockRpc.mockResolvedValueOnce({ is_sleeping: false });
+
+    const model = createModel();
+    const status = await model.getStatus();
+
+    expect(status).toBe(Status.LOADED);
+    expect(mockRpc).toHaveBeenCalledWith(
+      `/props?model=${model.id}&autoload=false`,
+    );
+  });
+
+  it("should return SLEEPING when is_sleeping is true", async () => {
+    mockRpc.mockResolvedValueOnce({ is_sleeping: true });
+
+    const model = createModel();
+    const status = await model.getStatus();
+
+    expect(status).toBe(Status.SLEEPING);
+  });
+});
+
+describe("LegacyModel getContextSize", () => {
+  it("should use max_model_len when it is non-zero", async () => {
+    mockRpc.mockResolvedValueOnce({ n_ctx: 4096 });
+    mockRpc.mockResolvedValueOnce({
+      data: [{ max_model_len: 8192 }],
+    });
+
+    const model = createModel();
+    const ctxSize = await model.getContextSize();
+
+    expect(ctxSize).toBe(8192);
+    expect(mockRpc).toHaveBeenCalledWith("/v1/models");
+  });
+
+  it("should fall back to n_ctx when max_model_len is 0", async () => {
+    mockRpc.mockResolvedValueOnce({ n_ctx: 4096 });
+    mockRpc.mockResolvedValueOnce({
+      data: [{ max_model_len: 0 }],
+    });
+
+    const model = createModel();
+    const ctxSize = await model.getContextSize();
+
+    expect(ctxSize).toBe(4096);
+  });
+
+  it("should return DEFAULT_CTX when both values are missing/null", async () => {
+    mockRpc.mockResolvedValueOnce({});
+    mockRpc.mockResolvedValueOnce({
+      data: [{ max_model_len: null }],
+    });
+
+    const model = createModel();
+    const ctxSize = await model.getContextSize();
+
+    expect(ctxSize).toBe(128000);
+  });
+});