From 6130edf43e565169c32a2faeb60cb6db4ae6f690 Mon Sep 17 00:00:00 2001 From: Amp Date: Fri, 28 Aug 2026 13:28:35 +0000 Subject: [PATCH] feat(zettel): add graph extraction boundary Co-authored-by: Aditya Balakrishnan --- .../zettel/src/notes/graph-extraction.test.ts | 68 +++++++++++++++ apps/zettel/src/notes/graph-extraction.ts | 84 +++++++++++++++++++ 2 files changed, 152 insertions(+) create mode 100644 apps/zettel/src/notes/graph-extraction.test.ts create mode 100644 apps/zettel/src/notes/graph-extraction.ts diff --git a/apps/zettel/src/notes/graph-extraction.test.ts b/apps/zettel/src/notes/graph-extraction.test.ts new file mode 100644 index 0000000..491755b --- /dev/null +++ b/apps/zettel/src/notes/graph-extraction.test.ts @@ -0,0 +1,68 @@ +import { describe, expect, it, vi } from "vitest"; +import { extractAndReplaceNoteGraph, extractNoteGraph } from "./graph-extraction.js"; + +describe("note graph extraction", () => { + it("normalizes valid provider output into NoteGraphInput", async () => { + const provider = { + extractGraph: vi.fn().mockResolvedValue({ + entities: [ + { name: " Ada ", type: " person ", description: " mathematician " }, + { name: " Engine ", type: " machine ", description: "" }, + ], + relations: [{ source: " Ada ", target: " Engine ", relationship: " designed " }], + }), + }; + + await expect(extractNoteGraph(provider, "tenant-1", "note-1", "note text")).resolves.toEqual({ + entities: [ + { name: "Ada", type: "person", description: "mathematician" }, + { name: "Engine", type: "machine", description: "" }, + ], + relations: [{ source: "Ada", target: "Engine", relationship: "designed" }], + }); + }); + + it("rejects malformed provider output", async () => { + const provider = { extractGraph: vi.fn().mockResolvedValue({ entities: "Ada", relations: [] }) }; + + await expect(extractNoteGraph(provider, "tenant-1", "note-1", "text")).rejects.toThrow(); + }); + + it("rejects relations whose normalized endpoints are absent", async () => { + const provider = { + extractGraph: vi.fn().mockResolvedValue({ + entities: [{ name: "Ada", type: "person", description: "" }], + relations: [{ source: "Ada", target: "Engine", relationship: "designed" }], + }), + }; + + await expect(extractNoteGraph(provider, "tenant-1", "note-1", "text")).rejects.toThrow( + "relation endpoint", + ); + }); + + it("forwards tenant, note, and text to the provider and persistence boundary", async () => { + const graph = { entities: [], relations: [] }; + const provider = { extractGraph: vi.fn().mockResolvedValue(graph) }; + const replaceGraph = vi.fn().mockResolvedValue(undefined); + + await extractAndReplaceNoteGraph("tenant-7", "note-9", "full note", provider, replaceGraph); + + expect(provider.extractGraph).toHaveBeenCalledWith({ + userId: "tenant-7", + noteId: "note-9", + text: "full note", + }); + expect(replaceGraph).toHaveBeenCalledWith("tenant-7", "note-9", graph); + }); + + it("does not persist when extraction fails", async () => { + const provider = { extractGraph: vi.fn().mockRejectedValue(new Error("model unavailable")) }; + const replaceGraph = vi.fn(); + + await expect( + extractAndReplaceNoteGraph("tenant-1", "note-1", "text", provider, replaceGraph), + ).rejects.toThrow("model unavailable"); + expect(replaceGraph).not.toHaveBeenCalled(); + }); +}); diff --git a/apps/zettel/src/notes/graph-extraction.ts b/apps/zettel/src/notes/graph-extraction.ts new file mode 100644 index 0000000..95d3c50 --- /dev/null +++ b/apps/zettel/src/notes/graph-extraction.ts @@ -0,0 +1,84 @@ +import { z } from "zod"; +import { replaceNoteGraph, type NoteGraphInput } from "./store.js"; + +const graphSchema = z + .object({ + entities: z.array( + z + .object({ + name: z.string().trim().min(1), + type: z.string().trim().min(1), + description: z.string().trim(), + }) + .strict(), + ), + relations: z.array( + z + .object({ + source: z.string().trim().min(1), + target: z.string().trim().min(1), + relationship: z.string().trim().min(1), + }) + .strict(), + ), + }) + .strict() + .superRefine((graph, context) => { + const entityNames = new Set(); + for (const entity of graph.entities) { + if (entityNames.has(entity.name)) { + context.addIssue({ + code: "custom", + message: `duplicate entity: ${entity.name}`, + path: ["entities"], + }); + } + entityNames.add(entity.name); + } + + graph.relations.forEach((relation, index) => { + if (!entityNames.has(relation.source) || !entityNames.has(relation.target)) { + context.addIssue({ + code: "custom", + message: "relation endpoint must be included in entities", + path: ["relations", index], + }); + } + }); + }); + +export interface GraphExtractionRequest { + userId: string; + noteId: string; + text: string; +} + +/** Model-independent boundary. Implementations may call any structured-output provider. */ +export interface GraphExtractionProvider { + extractGraph(request: GraphExtractionRequest): Promise; +} + +export async function extractNoteGraph( + provider: GraphExtractionProvider, + userId: string, + noteId: string, + text: string, +): Promise { + const output = await provider.extractGraph({ userId, noteId, text }); + return graphSchema.parse(output); +} + +type ReplaceGraph = typeof replaceNoteGraph; + +/** Persist only a fully extracted and validated graph, preserving prior data on extraction failure. */ +export async function extractAndReplaceNoteGraph( + userId: string, + noteId: string, + text: string, + provider: GraphExtractionProvider, + replaceGraph: ReplaceGraph = replaceNoteGraph, +): Promise { + const graph = await extractNoteGraph(provider, userId, noteId, text); + await replaceGraph(userId, noteId, graph); + return graph; +}