119 lines
4.0 KiB
TypeScript
119 lines
4.0 KiB
TypeScript
import { readFile } from "node:fs/promises";
|
|
import { homedir } from "node:os";
|
|
import { join } from "node:path";
|
|
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
|
|
|
|
type OmlxSettings = {
|
|
server?: { host?: string; port?: number };
|
|
auth?: { api_key?: string };
|
|
sampling?: { max_context_window?: number; max_tokens?: number };
|
|
};
|
|
|
|
type ListedModel = {
|
|
id: string;
|
|
max_model_len?: number | null;
|
|
};
|
|
|
|
type ModelStatus = {
|
|
id: string;
|
|
model_type?: string | null;
|
|
thinking_default?: boolean | null;
|
|
preserve_thinking_default?: boolean | null;
|
|
max_context_window?: number | null;
|
|
max_tokens?: number | null;
|
|
};
|
|
|
|
const ZERO_COST = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 };
|
|
|
|
function maySupportReasoning(id: string): boolean {
|
|
return /(?:thinking|reasoning|deepseek[-_. ]?r1|\b[or][13]\b|qwen3(?:[._-][56])?)/i.test(id);
|
|
}
|
|
|
|
export default async function (pi: ExtensionAPI) {
|
|
const settingsPath = join(homedir(), ".omlx", "settings.json");
|
|
|
|
try {
|
|
const settings = JSON.parse(await readFile(settingsPath, "utf8")) as OmlxSettings;
|
|
const configuredHost = settings.server?.host ?? "127.0.0.1";
|
|
const host = configuredHost === "0.0.0.0" || configuredHost === "::" ? "127.0.0.1" : configuredHost;
|
|
const port = settings.server?.port ?? 8127;
|
|
const apiKey = settings.auth?.api_key || "omlx";
|
|
const baseUrl = `http://${host}:${port}/v1`;
|
|
const headers = { Authorization: `Bearer ${apiKey}` };
|
|
const fetchOptions = { headers, signal: AbortSignal.timeout(5000) };
|
|
|
|
const modelsResponse = await fetch(`${baseUrl}/models`, fetchOptions);
|
|
if (!modelsResponse.ok) {
|
|
throw new Error(`${modelsResponse.status} ${modelsResponse.statusText}`);
|
|
}
|
|
|
|
const listed = (await modelsResponse.json()) as { data?: ListedModel[] };
|
|
|
|
let statuses = new Map<string, ModelStatus>();
|
|
try {
|
|
const statusResponse = await fetch(`${baseUrl}/models/status`, fetchOptions);
|
|
if (statusResponse.ok) {
|
|
const payload = (await statusResponse.json()) as { models?: ModelStatus[] };
|
|
statuses = new Map((payload.models ?? []).map((model) => [model.id, model]));
|
|
}
|
|
} catch {
|
|
// /v1/models is sufficient; status only enriches capabilities and limits.
|
|
}
|
|
|
|
const models = (listed.data ?? [])
|
|
.filter((model) => {
|
|
if (model.id === "MarkItDown") return false;
|
|
const type = statuses.get(model.id)?.model_type;
|
|
return type == null || type === "llm" || type === "vlm";
|
|
})
|
|
.map((model) => {
|
|
const status = statuses.get(model.id);
|
|
const reasoning = status?.thinking_default != null || maySupportReasoning(model.id);
|
|
const compat: Record<string, unknown> = {
|
|
supportsDeveloperRole: false,
|
|
supportsReasoningEffort: false,
|
|
};
|
|
|
|
if (reasoning) {
|
|
compat.thinkingFormat = "qwen-chat-template";
|
|
compat.chatTemplateKwargs = {
|
|
enable_thinking: { $var: "thinking.enabled" },
|
|
preserve_thinking: status?.preserve_thinking_default === true,
|
|
};
|
|
}
|
|
|
|
return {
|
|
id: model.id,
|
|
name: `${model.id} (oMLX Local)`,
|
|
reasoning,
|
|
input: status?.model_type === "vlm" ? (["text", "image"] as const) : (["text"] as const),
|
|
cost: ZERO_COST,
|
|
contextWindow:
|
|
status?.max_context_window ??
|
|
model.max_model_len ??
|
|
settings.sampling?.max_context_window ??
|
|
128000,
|
|
maxTokens: status?.max_tokens ?? settings.sampling?.max_tokens ?? 16384,
|
|
compat,
|
|
};
|
|
});
|
|
|
|
if (models.length === 0) {
|
|
console.warn("oMLX model discovery returned no chat models");
|
|
return;
|
|
}
|
|
|
|
pi.registerProvider("omlx", {
|
|
name: "oMLX (Local)",
|
|
baseUrl,
|
|
api: "openai-completions",
|
|
apiKey,
|
|
authHeader: true,
|
|
models,
|
|
});
|
|
} catch (error) {
|
|
const message = error instanceof Error ? error.message : String(error);
|
|
console.warn(`Could not discover oMLX models from ~/.omlx: ${message}`);
|
|
}
|
|
}
|