From a506e97e68f0186abdecb13c07fb4d81aa5eb0f6 Mon Sep 17 00:00:00 2001 From: Chai Date: Mon, 24 Aug 2026 20:35:15 -0400 Subject: [PATCH] fix(ai): retrieve annotations by chapter and position --- packages/core/src/ai/__tests__/tools.test.ts | 70 ++++++++++++++ packages/core/src/ai/system-prompt.ts | 2 +- .../core/src/ai/tools/annotation-tools.ts | 94 ++++++++++++++++++- 3 files changed, 161 insertions(+), 5 deletions(-) diff --git a/packages/core/src/ai/__tests__/tools.test.ts b/packages/core/src/ai/__tests__/tools.test.ts index c879a3196..a41e8e91b 100644 --- a/packages/core/src/ai/__tests__/tools.test.ts +++ b/packages/core/src/ai/__tests__/tools.test.ts @@ -873,6 +873,76 @@ describe("getAnnotations tool", () => { expect(result.highlights[0].text).toBe("Important text"); expect(result.notes).toHaveLength(1); expect(result.notes[0].title).toBe("Note 1"); + expect(result.pagination).toEqual({ + highlights: { total: 1, returned: 1, offset: 0, limit: 20, hasMore: false }, + notes: { total: 1, returned: 1, offset: 0, limit: 20, hasMore: false }, + }); + }); + + it("filters current-chapter annotations before applying the limit", async () => { + vi.mocked(getHighlights).mockResolvedValue([ + ...Array.from({ length: 20 }, (_, index) => ({ + text: `Early ${index}`, + chapterTitle: "Chapter 1", + color: "yellow", + })), + { + text: "Current note", + note: "Visible", + chapterTitle: "Spending Time Apart", + color: "green", + }, + ] as any); + vi.mocked(getNotes).mockResolvedValue([]); + + const tools = getAvailableTools({ bookId: "book-1", isVectorized: true, enabledSkills: [] }); + const tool = findTool(tools, "getAnnotations"); + const result = (await tool.execute({ + type: "highlights", + chapterTitle: " spending time apart ", + })) as any; + + expect(result.highlights).toEqual([ + expect.objectContaining({ text: "Current note", note: "Visible" }), + ]); + expect(result.pagination.highlights).toMatchObject({ + total: 1, + returned: 1, + hasMore: false, + }); + }); + + it("returns later book positions first for recent chapter requests", async () => { + vi.mocked(getHighlights).mockResolvedValue( + Array.from({ length: 25 }, (_, index) => ({ + text: `Position ${index + 1}`, + chapterTitle: `Chapter ${index + 1}`, + color: "yellow", + })) as any, + ); + vi.mocked(getNotes).mockResolvedValue([]); + + const tools = getAvailableTools({ bookId: "book-1", isVectorized: true, enabledSkills: [] }); + const tool = findTool(tools, "getAnnotations"); + const result = (await tool.execute({ + type: "highlights", + order: "reverse_book", + offset: 0, + limit: 3, + })) as any; + + expect(result.highlights.map((item: any) => item.text)).toEqual([ + "Position 25", + "Position 24", + "Position 23", + ]); + expect(result.pagination.highlights).toEqual({ + total: 25, + returned: 3, + offset: 0, + limit: 3, + hasMore: true, + }); }); it("should return only highlights when type is 'highlights'", async () => { diff --git a/packages/core/src/ai/system-prompt.ts b/packages/core/src/ai/system-prompt.ts index 91aaf916a..9fb073b06 100644 --- a/packages/core/src/ai/system-prompt.ts +++ b/packages/core/src/ai/system-prompt.ts @@ -289,7 +289,7 @@ function buildToolsSection( if (hasBookContext) { pushTool( "getAnnotations", - "- **getAnnotations**: Get user's highlights and notes (params: type)", + "- **getAnnotations**: Get user's highlights and notes. Pass chapterTitle for the current chapter; use order=reverse_book for recent/later chapters (params: type, chapterTitle, order, offset, limit)", ); if (isVectorized && canUse("addCitation")) { const citationSourceHint = canUse("ragSearch") diff --git a/packages/core/src/ai/tools/annotation-tools.ts b/packages/core/src/ai/tools/annotation-tools.ts index 1e7d0dfd5..b7fd3b395 100644 --- a/packages/core/src/ai/tools/annotation-tools.ts +++ b/packages/core/src/ai/tools/annotation-tools.ts @@ -5,43 +5,129 @@ import { getChunks, getHighlights, getNotes } from "../../db/database"; import { resolveFallbackCitationSource } from "../fallback-source-resolver"; import type { ToolDefinition } from "./tool-types"; +const DEFAULT_ANNOTATION_LIMIT = 20; +const MAX_ANNOTATION_LIMIT = 50; + +interface AnnotationPageMetadata { + total: number; + returned: number; + offset: number; + limit: number; + hasMore: boolean; +} + +function clampInteger(value: unknown, fallback: number, min: number, max: number): number { + if (typeof value !== "number" || !Number.isFinite(value)) return fallback; + return Math.min(max, Math.max(min, Math.trunc(value))); +} + +function normalizeChapterTitle(value: unknown): string | undefined { + if (typeof value !== "string") return undefined; + const normalized = value.trim().replace(/\s+/g, " ").toLocaleLowerCase(); + return normalized || undefined; +} + +function pageAnnotations( + items: T[], + options: { chapterTitle?: string; reverse: boolean; offset: number; limit: number }, +): { page: T[]; metadata: AnnotationPageMetadata } { + const matching = options.chapterTitle + ? items.filter((item) => normalizeChapterTitle(item.chapterTitle) === options.chapterTitle) + : items; + const ordered = options.reverse ? [...matching].reverse() : matching; + const page = ordered.slice(options.offset, options.offset + options.limit); + + return { + page, + metadata: { + total: ordered.length, + returned: page.length, + offset: options.offset, + limit: options.limit, + hasMore: options.offset + page.length < ordered.length, + }, + }; +} + /** Create get annotations tool for a specific book */ export function createGetAnnotationsTool(bookId: string): ToolDefinition { return { name: "getAnnotations", description: - "Get the user's highlights and notes from the book. Use this to reference what the user has marked as important.", + "Get the user's highlights and notes from the book. For the current chapter, pass its exact chapterTitle. For recent or later chapters, use order 'reverse_book'. Results are paginated.", parameters: { type: { type: "string", description: "'highlights' for highlights only, 'notes' for notes only, 'all' for both", }, + chapterTitle: { + type: "string", + description: + "Optional exact chapter title to filter before pagination. Match is case-insensitive and ignores surrounding/repeated whitespace.", + }, + order: { + type: "string", + description: + "'book' for beginning-to-end book position (default), or 'reverse_book' for later book positions first", + }, + offset: { + type: "number", + description: + "Zero-based annotation offset after chapter filtering and ordering (default 0)", + }, + limit: { + type: "number", + description: + "Maximum annotations of each requested type to return (default 20, maximum 50)", + }, }, execute: async (args) => { const type = (args.type as string) || "all"; + const chapterTitle = normalizeChapterTitle(args.chapterTitle); + const reverse = args.order === "reverse_book"; + const offset = clampInteger(args.offset, 0, 0, Number.MAX_SAFE_INTEGER); + const limit = clampInteger(args.limit, DEFAULT_ANNOTATION_LIMIT, 1, MAX_ANNOTATION_LIMIT); const result: { highlights?: Array<{ text: string; note?: string; chapterTitle?: string; color: string }>; notes?: Array<{ title: string; content: string; chapterTitle?: string }>; - } = {}; + pagination: { + highlights?: AnnotationPageMetadata; + notes?: AnnotationPageMetadata; + }; + } = { pagination: {} }; if (type === "highlights" || type === "all") { const highlights = await getHighlights(bookId); - result.highlights = highlights.slice(0, 20).map((h) => ({ + const { page, metadata } = pageAnnotations(highlights, { + chapterTitle, + reverse, + offset, + limit, + }); + result.highlights = page.map((h) => ({ text: h.text, note: h.note, chapterTitle: h.chapterTitle, color: h.color, })); + result.pagination.highlights = metadata; } if (type === "notes" || type === "all") { const notes = await getNotes(bookId); - result.notes = notes.slice(0, 20).map((n) => ({ + const { page, metadata } = pageAnnotations(notes, { + chapterTitle, + reverse, + offset, + limit, + }); + result.notes = page.map((n) => ({ title: n.title, content: n.content, chapterTitle: n.chapterTitle, })); + result.pagination.notes = metadata; } return result;