llama-cpp-provider/llama-cpp-provider.ts

/**
 * Llama.cpp Provider Extension
 *
 * Dynamically registers models from a llama.cpp llama-server instance
 * by fetching /v1/models, loading them, and extracting context window sizes
 *
 * Usage:
 *   pi -e ./llama-cpp-provider.ts
 *
 * Or add to ~/.pi/agent/extensions/ for auto-discovery
 */

import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import type { Model } from "@mariozechner/pi-ai";

// Configuration - can be overridden via environment variable
const LLAMA_SERVER_URL =
  process.env.LLAMA_SERVER_URL || "http://example.com:8080/v1";

interface LlamaModel {
  id: string;
  aliases: string[];
  tags: string[];
  object: string;
  owned_by: string;
  created: number;
  status: {
    value: string;
    args: string[];
    preset: string;
  };
}

interface LlamaModelsResponse {
  data: LlamaModel[];
  object: string;
}

/**
 * Parse context window size from llama.cpp preset string
 *
 * Example preset:
 * ```
 * [Qwen3.5-35B]
 * chat-template-kwargs = {"enable_thinking":false}
 * jinja = 1
 * min-p = 0.0
 * temperature = 0.6
 * top-k = 20
 * top-p = 0.95
 * ctx-size = 64000
 * model = ./models/unsloth/unsloth_Qwen3.5-35B-A3B-GGUF_Qwen3.5-35B-A3B-UD-Q4_K_XL.gguf
 * parallel = 1
 * ```
 */
function parseContextWindow(preset: string): number {
  const match = preset.match(/ctx-size\s*=\s*(\d+)/);
  return match ? parseInt(match[1], 10) : 32000; // Default to 32k if not found
}

/**
 * Fetch models from llama.cpp server
 */
async function fetchLlamaModels(): Promise<LlamaModel[]> {
  try {
    const response = await fetch(`${LLAMA_SERVER_URL}/models`);
    if (!response.ok) {
      throw new Error(`Failed to fetch models: ${response.status} ${response.statusText}`);
    }
    const data = (await response.json()) as LlamaModelsResponse;
    return data.data || [];
  } catch (error) {
    console.error("[llama-cpp-provider] Error fetching models:", error);
    return [];
  }
}

/**
 * Load a model on the llama.cpp server
 * POST /models/load with { "model": "model_id" }
 */
async function loadLlamaModel(modelId: string): Promise<boolean> {
  try {
    const response = await fetch(`${LLAMA_SERVER_URL}/models/load`, {
      method: "POST",
      headers: {
        "Content-Type": "application/json",
      },
      body: JSON.stringify({ model: modelId }),
    });

    if (!response.ok) {
      throw new Error(`Failed to load model ${modelId}: ${response.status} ${response.statusText}`);
    }

    const result = await response.json();
    return result.success || false;
  } catch (error) {
    console.error(`[llama-cpp-provider] Error loading model ${modelId}:`, error);
    return false;
  }
}

/**
 * Convert llama.cpp model to pi-ai Model configuration
 */
function llamaModelToPiModel(llamaModel: LlamaModel): Model<"openai-completions"> {
  const contextWindow = parseContextWindow(llamaModel.status.preset);

  return {
    id: llamaModel.id,
    name: llamaModel.id,
    api: "openai-completions",
    provider: "llama-cpp",
    baseUrl: LLAMA_SERVER_URL.replace("/v1", ""),
    reasoning: false, // llama.cpp doesn't support reasoning in the pi-ai sense
    input: ["text"], // Check if model supports images based on ID or preset
    cost: {
      input: 0,
      output: 0,
      cacheRead: 0,
      cacheWrite: 0,
    },
    contextWindow,
    maxTokens: Math.floor(contextWindow / 2), // Conservative estimate
  };
}

/**
 * Extension entry point
 */
export default function (pi: ExtensionAPI) {
  let registered = false;
  let models: Model<"openai-completions">[] = [];

  /**
   * Fetch and register models from llama.cpp server
   */
  async function registerModels(ctx?: import("@mariozechner/pi-coding-agent").ExtensionContext) {
    const llamaModels = await fetchLlamaModels();

    if (llamaModels.length === 0) {
      ctx?.ui.notify(
        "No models found from llama.cpp server. Check URL and server status.",
        "warning"
      );
      return;
    }

    models = llamaModels.map(llamaModelToPiModel);

    pi.registerProvider("llama-cpp", {
      baseUrl: LLAMA_SERVER_URL.replace("/v1", ""),
      apiKey: "llama", // llama.cpp doesn't require auth, any value works
      api: "openai-completions",
      models,
    });

    ctx?.ui.notify(
      `Registered ${models.length} models from llama.cpp server`,
      "info"
    );

    registered = true;
  }

  /**
   * Reload models (useful if models are added/removed from server)
   */
  async function reloadModels(ctx?: import("@mariozechner/pi-coding-agent").ExtensionContext) {
    if (!registered) {
      await registerModels(ctx);
      return;
    }

    // Unregister and re-register to update models
    pi.unregisterProvider("llama-cpp");
    await registerModels(ctx);
  }

  // Register models on load (no ctx available at extension load time)
  registerModels().then(() => {
    // Set runtime API key override after models are registered
    // This prevents the "No models available" warning in the TUI
    // Note: ctx is not available here, so we'll set it via session_start or commands
  });

  // Register a command to reload models
  pi.registerCommand("llama-reload", {
    description: "Reload models from llama.cpp server",
    handler: async (_args: string, ctx) => {
      ctx.ui.setWorkingMessage("Fetching and loading models from llama.cpp server...");
      await reloadModels(ctx);
      ctx.ui.setWorkingMessage();
    },
  });

  // Also register a command to list models
  pi.registerCommand("llama-list", {
    description: "List available models from llama.cpp server",
    handler: async (_args: string, ctx) => {
      const llamaModels = await fetchLlamaModels();

      if (llamaModels.length === 0) {
        ctx.ui.notify("No models found from llama.cpp server", "warning");
        return;
      }

      const lines = llamaModels.map((m) => {
        const ctxSize = parseContextWindow(m.status.preset);
        const status = m.status.value;
        return `  ${m.id} - ${ctxSize} tokens - ${status}`;
      });

      ctx.ui.notify(
        `Available models (${llamaModels.length}):\n${lines.join("\n")}`,
        "info"
      );
    },
  });
}