Files
arubinochan-bot/lib/harness.ts
T
2026-09-15 22:57:34 +09:00

107 lines
3.3 KiB
TypeScript

import type { LLamaChatPromptOptions, LlamaGrammar } from "node-llama-cpp";
import { BOT_NAME } from "./prompt";
export const harnessNames = ["json", "speaker", "tagged"] as const;
export type HarnessName = (typeof harnessNames)[number];
type HarnessOptions = Pick<
LLamaChatPromptOptions,
"grammar" | "responsePrefix" | "customStopTriggers"
>;
export type OutputHarness = {
name: HarnessName;
promptSuffix: string;
options: HarnessOptions;
parse: (raw: string) => string | null;
formatHistory: (text: string) => string;
};
const stripPrefix = (raw: string, prefix: string) =>
raw.startsWith(prefix) ? raw.slice(prefix.length) : raw;
export function createHarness(
name: HarnessName,
jsonGrammar: LlamaGrammar,
): OutputHarness {
switch (name) {
case "json":
return {
name,
promptSuffix:
'\n出力は {"name":"あるびのちゃん","text":"投稿本文"} という JSON オブジェクトだけにします。',
options: { grammar: jsonGrammar },
parse(raw) {
try {
const value: unknown = JSON.parse(raw.trim());
if (
typeof value === "object" &&
value !== null &&
"text" in value &&
typeof value.text === "string"
) {
return value.text.trim() || null;
}
} catch {
// Reported as a harness failure by the caller.
}
return null;
},
formatHistory: (text) => JSON.stringify({ name: BOT_NAME, text }),
};
case "speaker": {
const prefix = `${BOT_NAME}:\n`;
return {
name,
promptSuffix:
"\nアプリが書き手名を補うので、あなたは投稿本文だけを続けます。末尾に完了報告などを付けません。",
options: {
responsePrefix: prefix,
// This was the separator used by the pre-JSON harness. Retaining it
// also prevents the model from starting a fabricated next speaker.
customStopTriggers: [
"\n----------",
`\n${BOT_NAME}:`,
"\n<post",
"\n</post>",
],
},
parse(raw) {
const [firstTurn = ""] = stripPrefix(raw.trim(), prefix).split(
`\n${BOT_NAME}:`,
);
const body = firstTurn
.replace(/\n?<\/?post[^>]*>?[\s\S]*$/u, "")
.trim();
return body || null;
},
formatHistory: (text) => `${prefix}${text}`,
};
}
case "tagged": {
const prefix = `<post author="${BOT_NAME}">\n`;
return {
name,
promptSuffix:
"\n<post> はアプリが開始します。投稿本文だけを書き、直後に </post> で閉じます。",
options: {
responsePrefix: prefix,
customStopTriggers: ["</post>"],
},
parse(raw) {
const body = stripPrefix(raw.trim(), prefix)
.replace(/<\/post>[\s\S]*$/u, "")
.replace(/(?:<|&lt;)post(?:\s[^>]*?)?(?:>|&gt;)[\s\S]*$/u, "")
.trim();
return body || null;
},
formatHistory: (text) => `${prefix}${text}\n</post>`,
};
}
}
}
export function isHarnessName(value: string): value is HarnessName {
return harnessNames.some((name) => name === value);
}