b4ec727210a33cfac0482c8488381c4be100c0b6

Author
Julien Chaumond <julien@huggingface.co>
Committer
Julien Chaumond <julien@huggingface.co>
Date

Message

Initial commit: add llama.ts pi extension

Diff

This diff is truncated to protect this page.

  1diff --git a/llama.ts b/llama.ts
  2new file mode 100644
  3index 0000000000000000000000000000000000000000..810cf748099a07deda55979da499586c90e249f7
  4--- /dev/null
  5+++ b/llama.ts
  6@@ -0,0 +1,216 @@
  7+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
  8+import { Type } from "typebox";
  9+import { Compile } from "typebox/compile";
 10+
 11+const DEFAULT_BASE_URL = "http://localhost:8080/v1";
 12+const PROPS_TIMEOUT_MS = 120_000;
 13+
 14+const ModelsResponseSchema = Type.Object({
 15+	data: Type.Optional(
 16+		Type.Array(
 17+			Type.Object({
 18+				id: Type.String(),
 19+				status: Type.Optional(
 20+					Type.Object({
 21+						value: Type.Optional(
 22+							Type.Union([
 23+								Type.Literal("unloaded"),
 24+								Type.Literal("loading"),
 25+								Type.Literal("loaded"),
 26+								Type.Literal("sleeping"),
 27+								Type.Literal("unknown"),
 28+							]),
 29+						),
 30+					}),
 31+				),
 32+				architecture: Type.Optional(
 33+					Type.Object({
 34+						input_modalities: Type.Optional(Type.Array(Type.String())),
 35+					}),
 36+				),
 37+				meta: Type.Optional(
 38+					Type.Object({
 39+						n_ctx: Type.Optional(Type.Number()),
 40+						n_params: Type.Optional(Type.Number()),
 41+					}),
 42+				),
 43+			}),
 44+		),
 45+	),
 46+});
 47+
 48+const validateModelsResponse = Compile(ModelsResponseSchema);
 49+
 50+const PropsResponseSchema = Type.Object({
 51+	default_generation_settings: Type.Optional(
 52+		Type.Object({
 53+			n_ctx: Type.Optional(Type.Number()),
 54+		}),
 55+	),
 56+});
 57+
 58+const validatePropsResponse = Compile(PropsResponseSchema);
 59+
 60+type LlamaModel = NonNullable<Parameters<ExtensionAPI["registerProvider"]>[1]["models"]>[number];
 61+type ExtensionCtx = Parameters<Parameters<ExtensionAPI["on"]>[1]>[1];
 62+
 63+export default async function (pi: ExtensionAPI) {
 64+	let currentModels: LlamaModel[] = [];
 65+
 66+	pi.registerCommand("llama-version", {
 67+		description: "Print llama-server --version output",
 68+		handler: async (_args, ctx) => {
 69+			const result = await pi.exec("llama-server", ["--version"]);
 70+			const output = `${result.stderr ?? ""}\n${result.stdout ?? ""}`;
 71+			const versionLine = output
 72+				.split("\n")
 73+				.map((l) => l.trim())
 74+				.find((l) => /^version:\s/i.test(l));
 75+			ctx.ui.notify(
 76+				versionLine ?? `llama-server exited with code ${result.code}`,
 77+				versionLine ? "info" : "error",
 78+			);
 79+		},
 80+	});
 81+
 82+	const baseUrl = (process?.env?.LLAMA_BASE_URL ?? DEFAULT_BASE_URL).replace(/\/+$/, "");
 83+
 84+	async function refreshProvider(): Promise<void> {
 85+		try {
 86+			const response = await fetch(`${baseUrl}/models`);
 87+			if (!response.ok) {
 88+				console.warn(`[llama-cpp] ${baseUrl}/models returned ${response.status}`);
 89+				return;
 90+			}
 91+
 92+			const payload: unknown = await response.json();
 93+			if (!validateModelsResponse.Check(payload)) {
 94+				const errors = [...validateModelsResponse.Errors(payload)]
 95+					.map((e) => `${"path" in e ? e.path : ""} ${e.message}`)
 96+					.join("; ");
 97+				console.warn(`[llama-cpp] invalid /models response: ${errors}`);
 98+				return;
 99+			}
100+
101+			const previousById = new Map(currentModels.map((m) => [m.id, m]));
102+
103+			currentModels = (payload.data ?? []).map((model) => {
104+				const isLoaded = model.status?.value === "loaded";
105+				const modalities = model.architecture?.input_modalities ?? ["text"];