pi-llamacpp.ts

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

function getLlamaServerConfig() {
  const host = process.env.LLAMA_ARG_HOST || "127.0.0.1";
  const port = process.env.LLAMA_ARG_PORT || "8080";
  const apiPrefix = process.env.LLAMA_ARG_API_PREFIX || "";
  const sslCert = process.env.LLAMA_ARG_SSL_CERT_FILE;
  const apiKey = process.env.LLAMA_API_KEY || "none";

  const protocol = sslCert ? "https" : "http";
  
  let apiUrl = `${protocol}://${host}:${port}`;
  
  if (apiPrefix) {
    const normalizedPrefix = apiPrefix.startsWith("/") 
      ? apiPrefix 
      : "/" + apiPrefix;
    apiUrl = apiUrl + normalizedPrefix.replace(/\/$/, "");
  }

  return { apiUrl, apiKey };
}

function applyReasoningConfig(payload: any, modelId: string, thinkingLevel: string, maxTokens: number) {
  // gpt-oss models support thinking effort natively
  if (modelId.includes("gpt-oss")) {
    const effortMap: Record<string, string> = { off: "low", minimal: "low", low: "low", medium: "medium", high: "high" };
    payload.chat_template_kwargs = { ...payload.chat_template_kwargs, enable_thinking: true, reasoning_effort: effortMap[thinkingLevel] || "medium" };
    return;
  }

  // all other models
  // when thinking is off, disable it in the chat template
  if (thinkingLevel === "off") {
    payload.chat_template_kwargs = { ...payload.chat_template_kwargs, enable_thinking: false };
    return;
  }

  // up to 1024 thinking tokens offers a small but meaningful improvent over direct with many models  
  const minimal = Math.min(maxTokens / 8, 1024);
  // up to 8192 thinking tokens are enough for most thinking traces
  const medium = Math.min(maxTokens / 2, 8192);
  // geometric mean between minimal and medium
  const low = Math.sqrt(minimal * medium);
  // do not interrupt the model while thinking
  const high = maxTokens;

  const thinkingBudgets = { minimal, low, medium, high };
  payload.chat_template_kwargs = { ...payload.chat_template_kwargs, enable_thinking: true };
  payload.thinking_budget_tokens = thinkingBudgets[thinkingLevel];
}

export default async function (pi: ExtensionAPI) {
  const config = getLlamaServerConfig();
  const apiUrl = config.apiUrl;
  const apiKey = config.apiKey;

  const modelsRes = await fetch(`${apiUrl}/models`);
  if (!modelsRes.ok) {
    throw new Error(`Failed to fetch models: ${modelsRes.statusText}`);
  }
  const modelsJson = await modelsRes.json();
  const models = Array.isArray(modelsJson) ? modelsJson : (modelsJson.data || []);

  if (models.length === 0) {
    console.warn(`No models found at ${apiUrl}/models`);
    return;
  }

  let capturedModels = models
    .map((m: any) => ({
      id: m.id,
      name: m.id.replace(/^.*\/|:[^:]*$/g, ""),
      reasoning: false,
      input: ["text"],
      cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
      // Unknown until the model props are fetched.
      contextWindow: -1,
      maxTokens: -1,
    }));

  function registerProvider() {
    pi.registerProvider("llamacpp", {
      baseUrl: `${apiUrl}/v1`,
      apiKey: apiKey,
      api: "openai-completions",
      models: capturedModels,
    });
  }

  async function updateModelMetadata(modelId: string, ctx?: any) {
    try {
      const propsRes = await fetch(`${apiUrl}/props?model=${encodeURIComponent(modelId)}`);
      if (!propsRes.ok) return false;

      const props = await propsRes.json();
      const n_ctx = props.default_generation_settings?.n_ctx;
      const n_predict = props.default_generation_settings?.params.n_predict;
      const reasoning = props.chat_template_caps?.supports_preserve_reasoning === true;

      // Fallback to pi's default contextWindow and maxTokens for custom models (models.json)
      const contextWindow = n_ctx > 0 ? n_ctx : 128000;
      const maxTokens = n_predict > 0 ? n_predict : 16384;

      const input = ["text"];
      if (props.modalities?.vision) {
        input.push("image");
      }

      const modelIdx = capturedModels.findIndex(m => m.id === modelId);
      if (modelIdx !== -1) {
        capturedModels[modelIdx].contextWindow = contextWindow;
        capturedModels[modelIdx].maxTokens = maxTokens;
        capturedModels[modelIdx].input = input;
        capturedModels[modelIdx].reasoning = reasoning;
        return true;
      }
    } catch (error) {
      console.error(`Error fetching props for model ${modelId}:`, error);
    }
    return false;
  }

  // Detect standard server mode: only one model is loaded, no status field
  // Router mode: models have status.value ("loaded" or "unloaded")
  // In standard mode, we must fetch /props for the single loaded model to get full metadata
  // (chat_template_caps, modalities, etc.) that isn't in /models response
  const isStandardServerMode = models.length > 0 && models[0].status === undefined;
  
  await Promise.all(
    models
      .filter((m: any) => isStandardServerMode || m.status?.value === "loaded")
      .map((m: any) => updateModelMetadata(m.id))
  );

  registerProvider();

  pi.on("session_start", async (_event, ctx) => {
    const model = ctx.model;
    if (model && model.provider === "llamacpp") {
      if (await updateModelMetadata(model.id, ctx)) {
        registerProvider();
      }
    }
  });

  pi.on("model_select", async (event, ctx) => {
    if (event.model.provider !== "llamacpp") return;
    if (await updateModelMetadata(event.model.id, ctx)) {
      registerProvider();
    }
  });

  pi.on("before_provider_request", (event, ctx) => {
    if (ctx.model?.provider !== "llamacpp") return;

    const payload = { ...event.payload };
    applyReasoningConfig(payload, ctx.model.id, pi.getThinkingLevel(), ctx.model.maxTokens);
    return payload;
  });
}

@pyfisch for input modalities thanks to @ngxson you can now get them like this:

https://gist.github.com/julien-c/971d17c6ba1db6495c942e4e2a7e77b2#file-llama-cpp-ts-L69-L72

				const isLoaded = model.status?.value === "loaded";
				const modalities = model.architecture?.input_modalities ?? ["text"];
				const input = modalities.filter((m): m is "text" | "image" => m === "text" || m === "image");