From bed59d325ff691fb6331ab14aa08821f09aa1959 Mon Sep 17 00:00:00 2001 From: Amp Date: Fri, 28 Aug 2026 13:20:03 +0000 Subject: [PATCH] feat(zettel): add graph traversal tool Co-authored-by: Aditya Balakrishnan --- apps/zettel/src/index.ts | 10 ++++- apps/zettel/src/tools/notes.test.ts | 62 +++++++++++++++++++++++++++++ apps/zettel/src/tools/notes.ts | 36 ++++++++++++++++- 3 files changed, 106 insertions(+), 2 deletions(-) create mode 100644 apps/zettel/src/tools/notes.test.ts diff --git a/apps/zettel/src/index.ts b/apps/zettel/src/index.ts index 01331f3..0472bfe 100644 --- a/apps/zettel/src/index.ts +++ b/apps/zettel/src/index.ts @@ -9,7 +9,13 @@ import { serve } from "@hono/node-server"; import { serveStatic } from "@hono/node-server/serve-static"; import { WebSocketServer } from "ws"; import { auth } from "./notes/auth.js"; -import { createNoteTool, linkNotesTool, searchNotesTool, getNoteTool } from "./tools/notes.js"; +import { + createNoteTool, + linkNotesTool, + searchNotesTool, + getNoteTool, + traverseGraphTool, +} from "./tools/notes.js"; import { transcribeAudioTool, transcribeAudio } from "./tools/transcribe.js"; import { listNotes, @@ -346,6 +352,7 @@ function getOrCreateUserAgent(userId: string): AgentEventLoop { "Capture each thought as an atomic note via createNote.", "After creating a note, ALWAYS searchNotes for related existing notes and", "propose/draw links with linkNotes for genuine conceptual connections.", + "Use traverseGraph to explore connections between entities in the knowledge graph.", "Be concise.", ].join(" "), tools: { @@ -353,6 +360,7 @@ function getOrCreateUserAgent(userId: string): AgentEventLoop { linkNotes: linkNotesTool, searchNotes: searchNotesTool, getNote: getNoteTool, + traverseGraph: traverseGraphTool, transcribeAudio: transcribeAudioTool, }, autoTick: true, diff --git a/apps/zettel/src/tools/notes.test.ts b/apps/zettel/src/tools/notes.test.ts new file mode 100644 index 0000000..c54d432 --- /dev/null +++ b/apps/zettel/src/tools/notes.test.ts @@ -0,0 +1,62 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; + +const { traverseGraphStore } = vi.hoisted(() => ({ + traverseGraphStore: vi.fn(), +})); + +vi.mock("../notes/store.js", () => ({ + addLink: vi.fn(), + backlinksOf: vi.fn(), + readNote: vi.fn(), + searchNotes: vi.fn(), + traverseGraphStore, + writeNote: vi.fn(), +})); + +import { traverseGraph, traverseGraphSchema } from "./notes.js"; + +describe("traverseGraph tool", () => { + beforeEach(() => { + traverseGraphStore.mockReset(); + }); + + it("returns a successful graph traversal", async () => { + const result = { + entities: ["Alpha", "Beta"], + relations: [ + { source: "Alpha", target: "Beta", relationship: "connects", noteId: "note-1" }, + ], + }; + traverseGraphStore.mockResolvedValue(result); + + await expect(traverseGraph({ entityName: "Alpha", depth: 3, userId: "tenant-a" })).resolves.toEqual({ + success: true, + result, + }); + expect(traverseGraphStore).toHaveBeenCalledWith("tenant-a", "Alpha", 3); + }); + + it("forwards the tenant and uses defaults for omitted userId and depth", async () => { + traverseGraphStore.mockResolvedValue({ entities: [], relations: [] }); + + await traverseGraph({ entityName: "Alpha" }); + + expect(traverseGraphStore).toHaveBeenCalledWith("default", "Alpha", 2); + }); + + it("validates traversal depth bounds", () => { + expect(traverseGraphSchema.safeParse({ entityName: "Alpha", depth: 1 }).success).toBe(true); + expect(traverseGraphSchema.safeParse({ entityName: "Alpha", depth: 5 }).success).toBe(true); + expect(traverseGraphSchema.safeParse({ entityName: "Alpha", depth: 0 }).success).toBe(false); + expect(traverseGraphSchema.safeParse({ entityName: "Alpha", depth: 6 }).success).toBe(false); + }); + + it("returns a failure when the graph store throws", async () => { + traverseGraphStore.mockRejectedValue(new Error("database unavailable")); + + await expect(traverseGraph({ entityName: "Alpha", userId: "tenant-a" })).resolves.toEqual({ + success: false, + error: "database unavailable", + }); + }); +}); diff --git a/apps/zettel/src/tools/notes.ts b/apps/zettel/src/tools/notes.ts index 2380131..92bf3cb 100644 --- a/apps/zettel/src/tools/notes.ts +++ b/apps/zettel/src/tools/notes.ts @@ -1,7 +1,7 @@ /** * Zettelkasten note tools. * - * Four tools matching the ToolDefinition contract from @agentx/core. Each tool + * Tools matching the ToolDefinition contract from @agentx/core. Each tool * exports a zod schema, an async implementation invoked inside a thread-pool * worker, and a ToolDefinition pointing at this module by file path. */ @@ -14,6 +14,7 @@ import { searchNotes as searchNotesStore, addLink, backlinksOf, + traverseGraphStore, } from "../notes/store.js"; // ── createNote ──────────────────────────────────────────────────────────────── @@ -98,6 +99,31 @@ export async function getNote(args: GetNoteInput & { userId?: string }) { } } +// ── traverseGraph ───────────────────────────────────────────────────────────── + +export const traverseGraphSchema = z.object({ + entityName: z.string().min(1).describe("The entity name from which to start traversal."), + depth: z + .number() + .int() + .min(1) + .max(5) + .optional() + .describe("Traversal depth (default 2, maximum 5)."), +}); +export type TraverseGraphInput = z.infer; + +export async function traverseGraph(args: TraverseGraphInput & { userId?: string }) { + const { entityName, depth } = traverseGraphSchema.parse(args); + const userId = args.userId ?? "default"; + try { + const result = await traverseGraphStore(userId, entityName, depth ?? 2); + return { success: true, result }; + } catch (err) { + return { success: false, error: err instanceof Error ? err.message : String(err) }; + } +} + // ── ToolDefinitions ─────────────────────────────────────────────────────────── export const createNoteTool: ToolDefinition = { @@ -133,3 +159,11 @@ export const getNoteTool: ToolDefinition = { modulePath: new URL(import.meta.url).pathname, exportName: "getNote", }; + +export const traverseGraphTool: ToolDefinition = { + name: "traverseGraph", + description: "Traverse the knowledge graph from an entity to discover related entities and concepts.", + inputSchema: traverseGraphSchema, + modulePath: new URL(import.meta.url).pathname, + exportName: "traverseGraph", +};