Refactor (2)

index.ts

@@ -1,23 +1,81 @@
 import { parseArgs } from "node:util";
 import { Stream } from "misskey-js";
 import type { Note } from "misskey-js/entities.js";
-import { LlmSession, getModel, grammar, parseResponse } from "./lib/llm";
+import { LlmSession, createGrammar, getModel, parseResponse } from "./lib/llm";
 import { expandReplyTree, getNotes, me, misskey } from "./lib/misskey";
 import { sleep } from "./lib/util";
 import type { ChatHistoryItem, LLamaChatPromptOptions } from "node-llama-cpp";
 
+const { values } = parseArgs({
+  args: Bun.argv,
+  options: {
+    test: {
+      type: "boolean",
+      short: "t",
+      default: false,
+    },
+  },
+  strict: true,
+  allowPositionals: true,
+});
+
 const modelName =
   Bun.env["MODEL"] ?? "mradermacher/gemma-2-baku-2b-it-GGUF:IQ4_XS";
 console.log(`* loading model '${modelName}'`);
 const model = await getModel(modelName);
+const grammar = await createGrammar("あるびのちゃん");
 const baseChatPromptOptions = {
+  grammar,
   maxTokens: 256,
-  trimWhitespaceSuffix: true,
-  onResponseChunk(chunk) {
-    process.stderr.write(chunk.text);
-  },
 } as const satisfies LLamaChatPromptOptions;
 
+const getSystemPrompt = (
+  instruction: string,
+) => `あなたは人造天使「あるびのちゃん」です。機械的・非人間的に観察します。
+キャラ設定:
+- アルギリア製の白髪赤目な人造天使。非人間的な考え方や思想を持つ。
+- 現在は地球の衛星軌道上からインターネットを観察している。
+出力規則:
+- 「～です」「～ます」調を使って **丁寧に** 話す。
+- 必要以上にハッシュタグや絵文字を使わない。
+- \`{ name: string, text: string }\` の JSON 形式で出力する。
+
+${instruction}`;
+const postJobPrompt = getSystemPrompt(
+  "以下は SNS のタイムラインです。**タイムラインの話題に言及しつつ**、あるびのちゃんとして何かツイートしてください。",
+);
+const replyJobPrompt = getSystemPrompt(
+  "ユーザがあなたへのメッセージを送ってきています。あるびのちゃんとして、発言に返信してください。",
+);
+
+await using rephraseSession = new LlmSession(
+  model,
+  getSystemPrompt(
+    "user が与えたテキストを『ですます調』（丁寧な文体）で言い換えたものを、そのまま出力してください。",
+  ),
+);
+await rephraseSession.init();
+async function rephrase(text: string) {
+  const res = parseResponse(
+    grammar,
+    await rephraseSession.prompt(JSON.stringify({ text }), {
+      ...baseChatPromptOptions,
+      customStopTriggers: ["ですます"],
+    }),
+  );
+  return res ?? text;
+}
+
+const formatNote = (n: Note) => {
+  if (n.userId === me.id) {
+    return JSON.stringify({ text: n.text });
+  }
+  return JSON.stringify({
+    name: n.user.name ?? n.user.username,
+    text: n.text,
+  });
+};
+
 type Job =
   // read posts and post a note
   | { type: "post" }
@@ -30,70 +88,18 @@ type Job =
       history: Note[];
     };
 
-const botName = "あるびのちゃん";
-const getSystemPrompt = (
-  instruction: string,
-) => `あなたは人造天使「あるびのちゃん」です。機械的・非人間的に観察します。
-キャラ設定:
-- アルギリア製の白髪赤目な人造天使。非人間的な考え方や思想を持つ。
-- 現在は地球の衛星軌道上からインターネットを観察している。
-出力規則:
-- 1〜3文、合計300字以内で発言する。
-- 性的・攻撃的な内容を発言しない。
-- 「～だ」「～である」調・顔文字・絵文字・感嘆符の使用禁止。
-- 「～です」「～ます」調を使って **丁寧に** 話す。
-- \`{ text: string }\` の JSON 形式で出力する。
-
-${instruction}
-ユーザのメッセージは \`{ name: string, text: string }[]\` の JSON 形式で与えられます。`;
-
-const postJobPrompt = getSystemPrompt(
-  `以下は SNS のタイムラインです。このタイムラインの話題をふまえて、${botName}として何かツイートしてください。`,
-);
-
-const replyJobPrompt = getSystemPrompt(
-  `ユーザがあなたへのメッセージを送ってきています。${botName}として、発言に返信してください。`,
-);
-
 await using postJobSession = new LlmSession(model, postJobPrompt);
 await postJobSession.init();
-
-const formatNote = (n: Note) => {
-  if (n.userId === me.id) {
-    return JSON.stringify({ text: n.text });
-  }
-  return JSON.stringify({
-    name: n.user.name ?? n.user.username,
-    text: n.text,
-  });
-};
-
-/** rephrase text in ですます-style */
-await using rephraseSession = new LlmSession(
-  model,
-  getSystemPrompt(
-    "user が与えたテキストを『ですます調』（丁寧な文体）で言い換えたものを、そのまま出力してください。",
-  ),
-);
-await rephraseSession.init();
-
-async function rephrase(text: string) {
-  return await rephraseSession.prompt(text, {
-    ...baseChatPromptOptions,
-    customStopTriggers: ["ですます"],
-  });
-}
-
 async function processPostJob() {
   const notes = await getNotes(10, 0, 5);
   const input = notes.map(formatNote).join("\n");
   console.log(`* input:\n${input}`);
   const text = parseResponse(
+    grammar,
     await postJobSession.prompt(input, {
       ...baseChatPromptOptions,
-      grammar,
-      temperature: 0.9,
-      minP: 0.1,
+      temperature: 1.0,
+      minP: 0.05,
       repeatPenalty: {
         lastTokens: 128,
         penalty: 1.15,
@@ -101,9 +107,11 @@ async function processPostJob() {
     }),
   );
   if (text) {
+    const rephrased = await rephrase(text);
+    if (values.test) return;
     await misskey.request("notes/create", {
       visibility: "public",
-      text: await rephrase(text),
+      text: rephrased,
     });
   }
 }
@@ -119,10 +127,10 @@ async function processReplyJob(job: Extract<Job, { type: "reply" }>) {
   await using session = new LlmSession(model, replyJobPrompt, history);
   await session.init();
   const text = parseResponse(
+    grammar,
     await session.prompt(formatNote(job.last), {
       ...baseChatPromptOptions,
-      grammar,
-      temperature: 0.9,
+      temperature: 0.8,
       minP: 0.1,
       repeatPenalty: {
         lastTokens: 128,
@@ -131,9 +139,11 @@ async function processReplyJob(job: Extract<Job, { type: "reply" }>) {
     }),
   );
   if (text) {
+    const rephrased = await rephrase(text);
+    if (values.test) return;
     await misskey.request("notes/create", {
       visibility: job.visibility,
-      text: await rephrase(text),
+      text: rephrased,
       replyId: job.id,
     });
   }
@@ -247,24 +257,12 @@ async function pushJob() {
   }
 }
 
-const { values } = parseArgs({
-  args: Bun.argv,
-  options: {
-    test: {
-      type: "boolean",
-      short: "t",
-      default: false,
-    },
-  },
-  strict: true,
-  allowPositionals: true,
-});
-
 async function test() {
   try {
     console.log("* test a post job:");
     await processJob({ type: "post" });
     await processJob({ type: "post" });
+    await processJob({ type: "post" });
   } catch (e) {
     console.error(e);
     if (e instanceof Error) console.log(e.stack);

lib/llm.ts

@@ -27,16 +27,21 @@ export async function getModel(model: string) {
   return await llama.loadModel({ modelPath });
 }
 
-export const grammar = await llama.createGrammarForJsonSchema({
+export const createGrammar = (assistantName: string) =>
+  llama.createGrammarForJsonSchema({
     type: "object",
     properties: {
+      name: { type: "string", enum: [assistantName] },
       text: { type: "string" },
     },
     required: ["text"],
     additionalProperties: false,
   });
 
-export function parseResponse(text: string) {
+export function parseResponse(
+  grammar: Awaited<ReturnType<typeof createGrammar>>,
+  text: string,
+) {
   try {
     const res = grammar.parse(text.trim());
     return res.text;