Fix Pi package updates and merge feynman-model
This commit is contained in:
Advait Paliwal
309
extensions/research-tools/feynman-model.ts
Normal file
309
extensions/research-tools/feynman-model.ts
Normal file
@@ -0,0 +1,309 @@
|
||||
import { type Dirent, existsSync, readdirSync, readFileSync, writeFileSync } from "node:fs";
|
||||
import { homedir } from "node:os";
|
||||
import { basename, join, resolve } from "node:path";
|
||||
|
||||
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
|
||||
|
||||
const FRONTMATTER_PATTERN = /^---\n([\s\S]*?)\n---\n?([\s\S]*)$/;
|
||||
const INHERIT_MAIN = "__inherit_main__";
|
||||
|
||||
type FrontmatterDocument = {
|
||||
lines: string[];
|
||||
body: string;
|
||||
eol: string;
|
||||
trailingNewline: boolean;
|
||||
};
|
||||
|
||||
type SubagentModelConfig = {
|
||||
agent: string;
|
||||
model?: string;
|
||||
filePath: string;
|
||||
};
|
||||
|
||||
type SelectOption<T> = {
|
||||
label: string;
|
||||
value: T;
|
||||
};
|
||||
|
||||
type CommandContext = Parameters<Parameters<ExtensionAPI["registerCommand"]>[1]["handler"]>[1];
|
||||
|
||||
type TargetChoice =
|
||||
| { type: "main" }
|
||||
| { type: "subagent"; agent: string; model?: string };
|
||||
|
||||
function expandHomePath(value: string): string {
|
||||
if (value === "~") return homedir();
|
||||
if (value.startsWith("~/")) return resolve(homedir(), value.slice(2));
|
||||
return value;
|
||||
}
|
||||
|
||||
function resolveFeynmanAgentDir(): string {
|
||||
const configured = process.env.PI_CODING_AGENT_DIR ?? process.env.FEYNMAN_CODING_AGENT_DIR;
|
||||
if (configured?.trim()) {
|
||||
return resolve(expandHomePath(configured.trim()));
|
||||
}
|
||||
return resolve(homedir(), ".feynman", "agent");
|
||||
}
|
||||
|
||||
function formatModelSpec(model: { provider: string; id: string }): string {
|
||||
return `${model.provider}/${model.id}`;
|
||||
}
|
||||
|
||||
function detectEol(text: string): string {
|
||||
return text.includes("\r\n") ? "\r\n" : "\n";
|
||||
}
|
||||
|
||||
function normalizeLineEndings(text: string): string {
|
||||
return text.replace(/\r\n/g, "\n");
|
||||
}
|
||||
|
||||
function parseFrontmatterDocument(text: string): FrontmatterDocument | null {
|
||||
const normalized = normalizeLineEndings(text);
|
||||
const match = normalized.match(FRONTMATTER_PATTERN);
|
||||
if (!match) return null;
|
||||
|
||||
return {
|
||||
lines: match[1].split("\n"),
|
||||
body: match[2] ?? "",
|
||||
eol: detectEol(text),
|
||||
trailingNewline: normalized.endsWith("\n"),
|
||||
};
|
||||
}
|
||||
|
||||
function serializeFrontmatterDocument(document: FrontmatterDocument): string {
|
||||
const normalized = `---\n${document.lines.join("\n")}\n---\n${document.body}`;
|
||||
const withTrailingNewline =
|
||||
document.trailingNewline && !normalized.endsWith("\n") ? `${normalized}\n` : normalized;
|
||||
return document.eol === "\n" ? withTrailingNewline : withTrailingNewline.replace(/\n/g, "\r\n");
|
||||
}
|
||||
|
||||
function parseFrontmatterKey(line: string): string | undefined {
|
||||
const match = line.match(/^\s*([A-Za-z0-9_-]+)\s*:/);
|
||||
return match?.[1]?.toLowerCase();
|
||||
}
|
||||
|
||||
function getFrontmatterValue(lines: string[], key: string): string | undefined {
|
||||
const normalizedKey = key.toLowerCase();
|
||||
for (const line of lines) {
|
||||
const parsedKey = parseFrontmatterKey(line);
|
||||
if (parsedKey !== normalizedKey) continue;
|
||||
const separatorIndex = line.indexOf(":");
|
||||
if (separatorIndex === -1) return undefined;
|
||||
const value = line.slice(separatorIndex + 1).trim();
|
||||
return value.length > 0 ? value : undefined;
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
function upsertFrontmatterValue(lines: string[], key: string, value: string): string[] {
|
||||
const normalizedKey = key.toLowerCase();
|
||||
const nextLines = [...lines];
|
||||
const existingIndex = nextLines.findIndex((line) => parseFrontmatterKey(line) === normalizedKey);
|
||||
const serialized = `${key}: ${value}`;
|
||||
|
||||
if (existingIndex !== -1) {
|
||||
nextLines[existingIndex] = serialized;
|
||||
return nextLines;
|
||||
}
|
||||
|
||||
const descriptionIndex = nextLines.findIndex((line) => parseFrontmatterKey(line) === "description");
|
||||
const nameIndex = nextLines.findIndex((line) => parseFrontmatterKey(line) === "name");
|
||||
const insertIndex = descriptionIndex !== -1 ? descriptionIndex + 1 : nameIndex !== -1 ? nameIndex + 1 : nextLines.length;
|
||||
nextLines.splice(insertIndex, 0, serialized);
|
||||
return nextLines;
|
||||
}
|
||||
|
||||
function removeFrontmatterKey(lines: string[], key: string): string[] {
|
||||
const normalizedKey = key.toLowerCase();
|
||||
return lines.filter((line) => parseFrontmatterKey(line) !== normalizedKey);
|
||||
}
|
||||
|
||||
function normalizeAgentName(name: string): string {
|
||||
return name.trim().toLowerCase();
|
||||
}
|
||||
|
||||
function getAgentsDir(agentDir: string): string {
|
||||
return join(agentDir, "agents");
|
||||
}
|
||||
|
||||
function listAgentFiles(agentsDir: string): string[] {
|
||||
if (!existsSync(agentsDir)) return [];
|
||||
|
||||
return readdirSync(agentsDir, { withFileTypes: true })
|
||||
.filter((entry: Dirent) => (entry.isFile() || entry.isSymbolicLink()) && entry.name.endsWith(".md"))
|
||||
.filter((entry) => !entry.name.endsWith(".chain.md"))
|
||||
.map((entry) => join(agentsDir, entry.name));
|
||||
}
|
||||
|
||||
function readAgentConfig(filePath: string): SubagentModelConfig {
|
||||
const content = readFileSync(filePath, "utf8");
|
||||
const parsed = parseFrontmatterDocument(content);
|
||||
const fallbackName = basename(filePath, ".md");
|
||||
if (!parsed) return { agent: fallbackName, filePath };
|
||||
|
||||
return {
|
||||
agent: getFrontmatterValue(parsed.lines, "name") ?? fallbackName,
|
||||
model: getFrontmatterValue(parsed.lines, "model"),
|
||||
filePath,
|
||||
};
|
||||
}
|
||||
|
||||
function listSubagentModelConfigs(agentDir: string): SubagentModelConfig[] {
|
||||
return listAgentFiles(getAgentsDir(agentDir))
|
||||
.map((filePath) => readAgentConfig(filePath))
|
||||
.sort((left, right) => left.agent.localeCompare(right.agent));
|
||||
}
|
||||
|
||||
function findAgentConfig(configs: SubagentModelConfig[], agentName: string): SubagentModelConfig | undefined {
|
||||
const normalized = normalizeAgentName(agentName);
|
||||
return (
|
||||
configs.find((config) => normalizeAgentName(config.agent) === normalized) ??
|
||||
configs.find((config) => normalizeAgentName(basename(config.filePath, ".md")) === normalized)
|
||||
);
|
||||
}
|
||||
|
||||
function getAgentConfigOrThrow(agentDir: string, agentName: string): SubagentModelConfig {
|
||||
const configs = listSubagentModelConfigs(agentDir);
|
||||
const target = findAgentConfig(configs, agentName);
|
||||
if (target) return target;
|
||||
|
||||
if (configs.length === 0) {
|
||||
throw new Error(`No subagent definitions found in ${getAgentsDir(agentDir)}.`);
|
||||
}
|
||||
|
||||
const availableAgents = configs.map((config) => config.agent).join(", ");
|
||||
throw new Error(`Unknown subagent: ${agentName}. Available agents: ${availableAgents}`);
|
||||
}
|
||||
|
||||
function setSubagentModel(agentDir: string, agentName: string, modelSpec: string): void {
|
||||
const normalizedModelSpec = modelSpec.trim();
|
||||
if (!normalizedModelSpec) throw new Error("Model spec cannot be empty.");
|
||||
|
||||
const target = getAgentConfigOrThrow(agentDir, agentName);
|
||||
const content = readFileSync(target.filePath, "utf8");
|
||||
const parsed = parseFrontmatterDocument(content);
|
||||
|
||||
if (!parsed) {
|
||||
const eol = detectEol(content);
|
||||
const injected = `---${eol}name: ${target.agent}${eol}model: ${normalizedModelSpec}${eol}---${eol}${content}`;
|
||||
writeFileSync(target.filePath, injected, "utf8");
|
||||
return;
|
||||
}
|
||||
|
||||
const nextLines = upsertFrontmatterValue(parsed.lines, "model", normalizedModelSpec);
|
||||
if (nextLines.join("\n") !== parsed.lines.join("\n")) {
|
||||
writeFileSync(target.filePath, serializeFrontmatterDocument({ ...parsed, lines: nextLines }), "utf8");
|
||||
}
|
||||
}
|
||||
|
||||
function unsetSubagentModel(agentDir: string, agentName: string): void {
|
||||
const target = getAgentConfigOrThrow(agentDir, agentName);
|
||||
const content = readFileSync(target.filePath, "utf8");
|
||||
const parsed = parseFrontmatterDocument(content);
|
||||
if (!parsed) return;
|
||||
|
||||
const nextLines = removeFrontmatterKey(parsed.lines, "model");
|
||||
if (nextLines.join("\n") !== parsed.lines.join("\n")) {
|
||||
writeFileSync(target.filePath, serializeFrontmatterDocument({ ...parsed, lines: nextLines }), "utf8");
|
||||
}
|
||||
}
|
||||
|
||||
async function selectOption<T>(
|
||||
ctx: CommandContext,
|
||||
title: string,
|
||||
options: SelectOption<T>[],
|
||||
): Promise<T | undefined> {
|
||||
const selected = await ctx.ui.select(
|
||||
title,
|
||||
options.map((option) => option.label),
|
||||
);
|
||||
if (!selected) return undefined;
|
||||
return options.find((option) => option.label === selected)?.value;
|
||||
}
|
||||
|
||||
export function registerFeynmanModelCommand(pi: ExtensionAPI): void {
|
||||
pi.registerCommand("feynman-model", {
|
||||
description: "Open Feynman model menu (main + per-subagent overrides).",
|
||||
handler: async (_args, ctx) => {
|
||||
if (!ctx.hasUI) {
|
||||
ctx.ui.notify("feynman-model requires interactive mode.", "error");
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
ctx.modelRegistry.refresh();
|
||||
const availableModels = [...ctx.modelRegistry.getAvailable()].sort((left, right) =>
|
||||
formatModelSpec(left).localeCompare(formatModelSpec(right)),
|
||||
);
|
||||
if (availableModels.length === 0) {
|
||||
ctx.ui.notify("No models available.", "error");
|
||||
return;
|
||||
}
|
||||
|
||||
const agentDir = resolveFeynmanAgentDir();
|
||||
const subagentConfigs = listSubagentModelConfigs(agentDir);
|
||||
const currentMain = ctx.model ? formatModelSpec(ctx.model) : "(none)";
|
||||
|
||||
const targetOptions: SelectOption<TargetChoice>[] = [
|
||||
{ label: `main (default): ${currentMain}`, value: { type: "main" } },
|
||||
...subagentConfigs.map((config) => ({
|
||||
label: `${config.agent}: ${config.model ?? "default"}`,
|
||||
value: { type: "subagent" as const, agent: config.agent, model: config.model },
|
||||
})),
|
||||
];
|
||||
|
||||
const target = await selectOption(ctx, "Choose target", targetOptions);
|
||||
if (!target) return;
|
||||
|
||||
if (target.type === "main") {
|
||||
const selectedModel = await selectOption(
|
||||
ctx,
|
||||
"Select main model",
|
||||
availableModels.map((model) => {
|
||||
const spec = formatModelSpec(model);
|
||||
const suffix = spec === currentMain ? " (current)" : "";
|
||||
return { label: `${spec}${suffix}`, value: model };
|
||||
}),
|
||||
);
|
||||
if (!selectedModel) return;
|
||||
|
||||
const success = await pi.setModel(selectedModel);
|
||||
if (!success) {
|
||||
ctx.ui.notify(`No API key found for ${selectedModel.provider}.`, "error");
|
||||
return;
|
||||
}
|
||||
ctx.ui.notify(`Main model set to ${formatModelSpec(selectedModel)}.`, "info");
|
||||
return;
|
||||
}
|
||||
|
||||
const selectedSubagentModel = await selectOption(
|
||||
ctx,
|
||||
`Select model for ${target.agent}`,
|
||||
[
|
||||
{
|
||||
label: target.model ? "(inherit main default)" : "(inherit main default) (current)",
|
||||
value: INHERIT_MAIN,
|
||||
},
|
||||
...availableModels.map((model) => {
|
||||
const spec = formatModelSpec(model);
|
||||
const suffix = spec === target.model ? " (current)" : "";
|
||||
return { label: `${spec}${suffix}`, value: spec };
|
||||
}),
|
||||
],
|
||||
);
|
||||
if (!selectedSubagentModel) return;
|
||||
|
||||
if (selectedSubagentModel === INHERIT_MAIN) {
|
||||
unsetSubagentModel(agentDir, target.agent);
|
||||
ctx.ui.notify(`${target.agent} now inherits the main model.`, "info");
|
||||
return;
|
||||
}
|
||||
|
||||
setSubagentModel(agentDir, target.agent, selectedSubagentModel);
|
||||
ctx.ui.notify(`${target.agent} model set to ${selectedSubagentModel}.`, "info");
|
||||
} catch (error) {
|
||||
ctx.ui.notify(error instanceof Error ? error.message : String(error), "error");
|
||||
}
|
||||
},
|
||||
});
|
||||
}
|
||||
Reference in New Issue
Block a user