From 32ecabb4f1252ff47032106a4f92d85bcbe7c41d Mon Sep 17 00:00:00 2001
From: Amp <amp@ampcode.com>
Date: Fri, 28 Aug 2026 13:34:38 +0000
Subject: [PATCH] feat(zettel): add graph model adapter

Amp-Thread-ID: https://ampcode.com/threads/T-01a047e7-d402-72a8-ad22-eb5d3d3df6b1
Co-authored-by: Aditya Balakrishnan <adibalak1925@gmail.com>
---
 .../src/notes/graph-model-provider.test.ts    | 44 ++++++++++++++++++
 apps/zettel/src/notes/graph-model-provider.ts | 45 +++++++++++++++++++
 2 files changed, 89 insertions(+)
 create mode 100644 apps/zettel/src/notes/graph-model-provider.test.ts
 create mode 100644 apps/zettel/src/notes/graph-model-provider.ts

diff --git a/apps/zettel/src/notes/graph-model-provider.test.ts b/apps/zettel/src/notes/graph-model-provider.test.ts
new file mode 100644
index 0000000..d9865f8
--- /dev/null
+++ b/apps/zettel/src/notes/graph-model-provider.test.ts
@@ -0,0 +1,44 @@
+import { beforeEach, describe, expect, it, vi } from "vitest";
+
+const { generateObject } = vi.hoisted(() => ({ generateObject: vi.fn() }));
+
+vi.mock("ai", () => ({ generateObject }));
+
+import { createGraphModelProvider } from "./graph-model-provider.js";
+
+describe("graph model provider", () => {
+  beforeEach(() => {
+    generateObject.mockReset();
+  });
+
+  it("requests structured graph output from the configured model", async () => {
+    const model = { modelId: "test-model" } as never;
+    const object = {
+      entities: [{ name: "Ada", type: "person", description: "mathematician" }],
+      relations: [],
+    };
+    generateObject.mockResolvedValue({ object });
+
+    const provider = createGraphModelProvider(model);
+
+    await expect(
+      provider.extractGraph({ userId: "tenant-secret", noteId: "note-1", text: "Ada wrote notes." }),
+    ).resolves.toEqual(object);
+    expect(generateObject).toHaveBeenCalledOnce();
+    const request = generateObject.mock.calls[0][0];
+    expect(request.model).toBe(model);
+    expect(request.schema.safeParse(object).success).toBe(true);
+    expect(request.prompt).toContain("Ada wrote notes.");
+    expect(request.prompt).not.toContain("tenant-secret");
+    expect(request.prompt).not.toContain("note-1");
+  });
+
+  it("lets model errors reach the extraction boundary", async () => {
+    generateObject.mockRejectedValue(new Error("provider unavailable"));
+    const provider = createGraphModelProvider({} as never);
+
+    await expect(
+      provider.extractGraph({ userId: "tenant", noteId: "note", text: "text" }),
+    ).rejects.toThrow("provider unavailable");
+  });
+});
diff --git a/apps/zettel/src/notes/graph-model-provider.ts b/apps/zettel/src/notes/graph-model-provider.ts
new file mode 100644
index 0000000..b5b7d0f
--- /dev/null
+++ b/apps/zettel/src/notes/graph-model-provider.ts
@@ -0,0 +1,45 @@
+import { generateObject, type LanguageModel } from "ai";
+import { groq } from "@ai-sdk/groq";
+import { z } from "zod";
+import type { GraphExtractionProvider } from "./graph-extraction.js";
+
+const graphOutputSchema = z.object({
+  entities: z.array(
+    z.object({
+      name: z.string(),
+      type: z.string(),
+      description: z.string(),
+    }),
+  ),
+  relations: z.array(
+    z.object({
+      source: z.string(),
+      target: z.string(),
+      relationship: z.string(),
+    }),
+  ),
+});
+
+/** Create a structured-output provider without coupling graph persistence to a model vendor. */
+export function createGraphModelProvider(model: LanguageModel): GraphExtractionProvider {
+  return {
+    async extractGraph({ text }) {
+      const { object } = await generateObject({
+        model,
+        schema: graphOutputSchema,
+        prompt: [
+          "Extract only entities and explicit relationships stated in this note.",
+          "Every relation source and target must exactly match an entity name.",
+          "Do not infer facts that are not present.",
+          "",
+          text,
+        ].join("\n"),
+      });
+      return object;
+    },
+  };
+}
+
+export function createDefaultGraphModelProvider(): GraphExtractionProvider {
+  return createGraphModelProvider(groq(process.env.GRAPH_EXTRACTION_MODEL ?? "llama-3.3-70b-versatile"));
+}