Sfoglia il codice sorgente

feat(core): add createBookContextTransform for truth file injection

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
fanghanjun 6 mesi fa
parent
commit
4fbf4ee66f

+ 82 - 0
packages/core/src/__tests__/context-transform.test.ts

@@ -0,0 +1,82 @@
+import { describe, it, expect, beforeEach, afterEach } from "vitest";
+import { mkdtemp, mkdir, writeFile, rm } from "node:fs/promises";
+import { join } from "node:path";
+import { tmpdir } from "node:os";
+import { createBookContextTransform } from "../agent/context-transform.js";
+
+describe("createBookContextTransform", () => {
+  let projectRoot: string;
+  const bookId = "test-book";
+
+  beforeEach(async () => {
+    projectRoot = await mkdtemp(join(tmpdir(), "ctx-test-"));
+    const storyDir = join(projectRoot, "books", bookId, "story");
+    await mkdir(storyDir, { recursive: true });
+    await writeFile(join(storyDir, "story_bible.md"), "# Story Bible\nA hero's journey.");
+    await writeFile(join(storyDir, "current_focus.md"), "Focus on chapter 3.");
+  });
+
+  afterEach(async () => {
+    await rm(projectRoot, { recursive: true, force: true });
+  });
+
+  it("returns messages unchanged when bookId is null", async () => {
+    const transform = createBookContextTransform(null, projectRoot);
+    const messages = [
+      { role: "user" as const, content: "hello", timestamp: Date.now() },
+    ];
+    const result = await transform(messages);
+    expect(result).toBe(messages);
+  });
+
+  it("prepends a user message with truth file contents", async () => {
+    const transform = createBookContextTransform(bookId, projectRoot);
+    const original = [
+      { role: "user" as const, content: "写下一章", timestamp: Date.now() },
+    ];
+    const result = await transform(original);
+
+    expect(original).toHaveLength(1);
+    expect(result).toHaveLength(2);
+    const injected = result[0] as { role: string; content: string };
+    expect(injected.role).toBe("user");
+    expect(injected.content).toContain("story_bible.md");
+    expect(injected.content).toContain("A hero's journey.");
+    expect(injected.content).toContain("current_focus.md");
+    expect(injected.content).toContain("Focus on chapter 3.");
+    expect(result[1]).toBe(original[0]);
+  });
+
+  it("sorts truth files in priority order", async () => {
+    const storyDir = join(projectRoot, "books", bookId, "story");
+    await writeFile(join(storyDir, "volume_outline.md"), "# Volume Outline");
+    await writeFile(join(storyDir, "book_rules.md"), "# Book Rules");
+    await writeFile(join(storyDir, "extra_notes.md"), "# Extra");
+
+    const transform = createBookContextTransform(bookId, projectRoot);
+    const result = await transform([
+      { role: "user" as const, content: "test", timestamp: Date.now() },
+    ]);
+    const content = (result[0] as { content: string }).content;
+
+    const bibleIdx = content.indexOf("story_bible.md");
+    const outlineIdx = content.indexOf("volume_outline.md");
+    const rulesIdx = content.indexOf("book_rules.md");
+    const focusIdx = content.indexOf("current_focus.md");
+    const extraIdx = content.indexOf("extra_notes.md");
+
+    expect(bibleIdx).toBeLessThan(outlineIdx);
+    expect(outlineIdx).toBeLessThan(rulesIdx);
+    expect(rulesIdx).toBeLessThan(focusIdx);
+    expect(focusIdx).toBeLessThan(extraIdx);
+  });
+
+  it("returns original messages when story/ directory does not exist", async () => {
+    const transform = createBookContextTransform("nonexistent-book", projectRoot);
+    const original = [
+      { role: "user" as const, content: "test", timestamp: Date.now() },
+    ];
+    const result = await transform(original);
+    expect(result).toBe(original);
+  });
+});

+ 73 - 0
packages/core/src/agent/context-transform.ts

@@ -0,0 +1,73 @@
+import type { AgentMessage } from "@mariozechner/pi-agent-core";
+import type { UserMessage } from "@mariozechner/pi-ai";
+import { readdir, readFile } from "node:fs/promises";
+import { join } from "node:path";
+
+/** Files read in this order; anything else in story/ comes after, sorted alphabetically. */
+const PRIORITY_FILES = [
+  "story_bible.md",
+  "volume_outline.md",
+  "book_rules.md",
+  "current_focus.md",
+];
+
+export function createBookContextTransform(
+  bookId: string | null,
+  projectRoot: string,
+): (messages: AgentMessage[], signal?: AbortSignal) => Promise<AgentMessage[]> {
+  if (bookId === null) {
+    return async (messages) => messages;
+  }
+
+  const storyDir = join(projectRoot, "books", bookId, "story");
+
+  return async (messages) => {
+    const sections = await readTruthFiles(storyDir);
+    if (sections.length === 0) return messages;
+
+    const body =
+      "[以下是当前书籍的真相文件,每次对话时自动从磁盘读取注入。请基于这些内容进行创作和判断。]\n\n" +
+      sections.map((s) => `=== ${s.name} ===\n${s.content}`).join("\n\n");
+
+    const injected: UserMessage = {
+      role: "user",
+      content: body,
+      timestamp: Date.now(),
+    };
+
+    return [injected, ...messages];
+  };
+}
+
+interface TruthFileSection {
+  name: string;
+  content: string;
+}
+
+async function readTruthFiles(storyDir: string): Promise<TruthFileSection[]> {
+  let entries: string[];
+  try {
+    entries = await readdir(storyDir);
+  } catch {
+    return [];
+  }
+
+  const mdFiles = entries.filter((f) => f.endsWith(".md"));
+  if (mdFiles.length === 0) return [];
+
+  const prioritySet = new Set(PRIORITY_FILES);
+  const prioritized = PRIORITY_FILES.filter((f) => mdFiles.includes(f));
+  const rest = mdFiles.filter((f) => !prioritySet.has(f)).sort();
+  const ordered = [...prioritized, ...rest];
+
+  const sections: TruthFileSection[] = [];
+  for (const fileName of ordered) {
+    try {
+      const content = await readFile(join(storyDir, fileName), "utf-8");
+      sections.push({ name: fileName, content });
+    } catch {
+      // skip unreadable files
+    }
+  }
+  return sections;
+}