122 lines
4.4 KiB
TypeScript
122 lines
4.4 KiB
TypeScript
import { spawn as nodeSpawn, type ChildProcessWithoutNullStreams } from "node:child_process";
|
|
import type { AppConfig } from "../config.js";
|
|
import { RpcClient } from "../rpc/rpc-client.js";
|
|
import { buildPiChildEnv } from "./provider-credentials.js";
|
|
import {
|
|
loadPiEnabledModels,
|
|
type PiEnabledModelsResult,
|
|
} from "./enabled-models.js";
|
|
import {
|
|
readConfiguredPiAgentFile,
|
|
validateDeclarativePiConfig,
|
|
} from "./managed-config.js";
|
|
import type { RuntimeModelCatalog } from "../models/runtime-model-catalog.js";
|
|
import { secretValue } from "../config/secret-bundle.js";
|
|
|
|
export interface PiModel {
|
|
provider: string;
|
|
id: string;
|
|
name: string;
|
|
reasoning: boolean;
|
|
}
|
|
|
|
export type ListModelsFn = () => Promise<PiModel[]>;
|
|
|
|
interface Opts {
|
|
spawnFn?: (
|
|
command: string,
|
|
args: string[],
|
|
options: { cwd: string; env: NodeJS.ProcessEnv },
|
|
) => ChildProcessWithoutNullStreams;
|
|
ttlMs?: number;
|
|
nowMs?: () => number;
|
|
loadEnabledModels?: () => PiEnabledModelsResult;
|
|
readModelsStore?: () => string | undefined;
|
|
modelCatalog?: RuntimeModelCatalog;
|
|
warn?: (message: string) => void;
|
|
}
|
|
|
|
/**
|
|
* Returns a function that lists Pi's available models (those with auth
|
|
* configured) via an ephemeral `pi --mode rpc` process. Result is cached for
|
|
* `ttlMs`. The returned function rejects on timeout/error; callers degrade.
|
|
*/
|
|
export function createPiModelLister(cfg: AppConfig, opts: Opts = {}): ListModelsFn {
|
|
const ttlMs = opts.ttlMs ?? 60_000;
|
|
const now = opts.nowMs ?? (() => Date.now());
|
|
const spawnFn = opts.spawnFn ?? nodeSpawn;
|
|
|
|
let cache: { at: number; models: PiModel[] } | null = null;
|
|
|
|
return async function listModels(): Promise<PiModel[]> {
|
|
const managedModels = opts.readModelsStore
|
|
? opts.readModelsStore()
|
|
: readConfiguredPiAgentFile("models.json", true);
|
|
if (managedModels !== undefined) validateDeclarativePiConfig(managedModels);
|
|
|
|
if (cache && now() - cache.at < ttlMs) return cache.models;
|
|
|
|
const enabled = (opts.loadEnabledModels
|
|
?? (() => loadPiEnabledModels({ harnessDir: cfg.harnessDir })))();
|
|
for (const warning of enabled.warnings) opts.warn?.(warning);
|
|
if (enabled.ids.length === 0) {
|
|
cache = { at: now(), models: [] };
|
|
return [];
|
|
}
|
|
|
|
const env = buildPiChildEnv({});
|
|
// Pi's availability enumeration also needs the catalog-owned credentials for built-in
|
|
// providers. It must keep working after their obsolete Pi auth entries are removed.
|
|
for (const model of opts.modelCatalog?.sessionModels() ?? []) {
|
|
const name = model.authentication.mode === "secret_env" ? model.authentication.apiKeyEnv : undefined;
|
|
if (!name) continue;
|
|
const value = secretValue(cfg, name);
|
|
if (value) env[name] = value;
|
|
}
|
|
delete env.THT_DATA_ROOT;
|
|
if (cfg.dataRoot !== undefined) env.THT_DATA_ROOT = cfg.dataRoot;
|
|
const child = spawnFn(cfg.piBin, ["--mode", "rpc"], { cwd: cfg.harnessDir, env });
|
|
child.stderr.resume();
|
|
const rpc = new RpcClient(child);
|
|
try {
|
|
const resp = await Promise.race([
|
|
rpc.request({ type: "get_available_models" }),
|
|
new Promise<never>((_, rej) => setTimeout(() => rej(new Error("pi model list timeout")), 8000)),
|
|
]);
|
|
const raw = (resp?.data?.models ?? []) as Array<Record<string, unknown>>;
|
|
const available: PiModel[] = raw.map((m) => ({
|
|
provider: String(m.provider ?? ""),
|
|
id: String(m.id ?? ""),
|
|
name: String(m.name ?? m.id ?? ""),
|
|
reasoning: Boolean(m.reasoning),
|
|
}));
|
|
const byCompositeId = new Map(
|
|
available.map((model) => [`${model.provider}/${model.id}`, model]),
|
|
);
|
|
const catalogByPiId = new Map(
|
|
(opts.modelCatalog?.sessionModels() ?? []).map((model) => [
|
|
`${model.provider}/${model.upstreamModel}`,
|
|
model,
|
|
]),
|
|
);
|
|
const models = enabled.ids.flatMap((id) => {
|
|
const model = byCompositeId.get(id);
|
|
if (!model) return [];
|
|
const catalogModel = catalogByPiId.get(id);
|
|
return [catalogModel ? {
|
|
...model,
|
|
provider: catalogModel.provider,
|
|
id: catalogModel.model,
|
|
name: catalogModel.label,
|
|
reasoning: catalogModel.session?.reasoning ?? model.reasoning,
|
|
} : model];
|
|
});
|
|
if (models.length === 0) opts.warn?.("No Pi-enabled models are currently available");
|
|
cache = { at: now(), models };
|
|
return models;
|
|
} finally {
|
|
child.kill();
|
|
}
|
|
};
|
|
}
|