diff --git a/packages/app-expo/package.json b/packages/app-expo/package.json index 8117f1158..0198ff4a7 100644 --- a/packages/app-expo/package.json +++ b/packages/app-expo/package.json @@ -7,10 +7,7 @@ "entryPoint": "./index.js", "doctor": { "reactNativeDirectoryCheck": { - "exclude": [ - "react-native-tcp-socket", - "react-native-track-player" - ], + "exclude": ["react-native-tcp-socket", "react-native-track-player"], "listUnknownPackages": false } } @@ -73,6 +70,7 @@ "expo-document-picker": "~14.0.8", "expo-file-system": "~19.0.21", "expo-font": "~14.0.11", + "expo-keep-awake": "~15.0.8", "expo-image-picker": "~17.0.11", "expo-linear-gradient": "~15.0.8", "expo-linking": "~8.0.11", diff --git a/packages/app-expo/src/App.tsx b/packages/app-expo/src/App.tsx index 5e9ecbe6d..ab7e31c49 100644 --- a/packages/app-expo/src/App.tsx +++ b/packages/app-expo/src/App.tsx @@ -31,6 +31,7 @@ import { GestureHandlerRootView } from "react-native-gesture-handler"; import { KeyboardProvider } from "react-native-keyboard-controller"; import { SafeAreaProvider } from "react-native-safe-area-context"; +import { AIStreamingKeepAwake } from "@/components/chat/AIStreamingKeepAwake"; import { AnimatedSplash } from "@/components/splash/AnimatedSplash"; import { rnSessionEventSource } from "@/hooks"; import { setStreamingFetch } from "@readany/core/ai/llm-provider"; @@ -290,6 +291,7 @@ function AppInner() { + diff --git a/packages/app-expo/src/components/chat/AIStreamingKeepAwake.tsx b/packages/app-expo/src/components/chat/AIStreamingKeepAwake.tsx new file mode 100644 index 000000000..87dff045e --- /dev/null +++ b/packages/app-expo/src/components/chat/AIStreamingKeepAwake.tsx @@ -0,0 +1,14 @@ +import { useKeepAwake } from "expo-keep-awake"; + +import { useChatStore } from "@/stores/chat-store"; +import { hasActiveAIStream } from "./ai-streaming-keep-awake-state"; + +function ActiveAIKeepAwakeClaim() { + useKeepAwake("readany-ai-stream"); + return null; +} + +export function AIStreamingKeepAwake() { + const isActive = useChatStore((state) => hasActiveAIStream(state.streamingSessions)); + return isActive ? : null; +} diff --git a/packages/app-expo/src/components/chat/ai-streaming-keep-awake-owner.test.ts b/packages/app-expo/src/components/chat/ai-streaming-keep-awake-owner.test.ts new file mode 100644 index 000000000..43deb8685 --- /dev/null +++ b/packages/app-expo/src/components/chat/ai-streaming-keep-awake-owner.test.ts @@ -0,0 +1,26 @@ +import { readFileSync } from "node:fs"; +import { resolve } from "node:path"; +import { describe, expect, it } from "vitest"; + +const root = resolve(import.meta.dirname, "../../../../.."); +const read = (path: string) => readFileSync(resolve(root, path), "utf8"); + +describe("AI streaming keep-awake ownership", () => { + it("owns the native wake claim from the shared streaming-session state", () => { + const component = read("packages/app-expo/src/components/chat/AIStreamingKeepAwake.tsx"); + + expect(component).toContain('import { useKeepAwake } from "expo-keep-awake"'); + expect(component).toContain('import { useChatStore } from "@/stores/chat-store"'); + expect(component).toContain("hasActiveAIStream(state.streamingSessions)"); + expect(component).toContain('useKeepAwake("readany-ai-stream")'); + }); + + it("mounts one owner at app scope beside other global hosts", () => { + const app = read("packages/app-expo/src/App.tsx"); + + expect(app).toContain("import { AIStreamingKeepAwake }"); + expect(app).toMatch( + /[\s\S]*[\s\S]*/, + ); + }); +}); diff --git a/packages/app-expo/src/components/chat/ai-streaming-keep-awake-state.test.ts b/packages/app-expo/src/components/chat/ai-streaming-keep-awake-state.test.ts new file mode 100644 index 000000000..f32af6240 --- /dev/null +++ b/packages/app-expo/src/components/chat/ai-streaming-keep-awake-state.test.ts @@ -0,0 +1,18 @@ +import { describe, expect, it } from "vitest"; +import { hasActiveAIStream } from "./ai-streaming-keep-awake-state"; + +describe("hasActiveAIStream", () => { + it("is false with no sessions or only completed sessions", () => { + expect(hasActiveAIStream({})).toBe(false); + expect(hasActiveAIStream({ done: { isStreaming: false } })).toBe(false); + }); + + it("stays true until the final concurrent stream completes", () => { + expect( + hasActiveAIStream({ + book: { isStreaming: false }, + general: { isStreaming: true }, + }), + ).toBe(true); + }); +}); diff --git a/packages/app-expo/src/components/chat/ai-streaming-keep-awake-state.ts b/packages/app-expo/src/components/chat/ai-streaming-keep-awake-state.ts new file mode 100644 index 000000000..31b0abf92 --- /dev/null +++ b/packages/app-expo/src/components/chat/ai-streaming-keep-awake-state.ts @@ -0,0 +1,3 @@ +export function hasActiveAIStream(sessions: Record): boolean { + return Object.values(sessions).some((session) => session.isStreaming); +} diff --git a/packages/core/src/ai/__tests__/reading-agent-tools.test.ts b/packages/core/src/ai/__tests__/reading-agent-tools.test.ts index d96d8ac1c..eda78fbd5 100644 --- a/packages/core/src/ai/__tests__/reading-agent-tools.test.ts +++ b/packages/core/src/ai/__tests__/reading-agent-tools.test.ts @@ -1,6 +1,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import { z } from "zod"; import type { AIConfig } from "../../types"; +import type { ReadingContext } from "../../types/chat"; import { isOutputLimitTermination, streamReadingAgent } from "../agents/reading-agent"; import type { ToolDefinition } from "../tools"; import { getAvailableTools } from "../tools"; @@ -17,7 +18,9 @@ vi.mock("../llm-provider", () => ({ })), })); -const getReadingContextSnapshotMock = vi.hoisted(() => vi.fn(() => null)); +const getReadingContextSnapshotMock = vi.hoisted(() => + vi.fn<() => ReadingContext | null>(() => null), +); vi.mock("../reading-context-service", () => ({ getReadingContextSnapshot: getReadingContextSnapshotMock, @@ -576,6 +579,226 @@ describe("streamReadingAgent tool registration", () => { expect(fourth.attemptedQueries).toEqual(["张三疯那一章讲了什么", "张三疯", "张三疯"]); }); + it.each([ + ["Review my notes for this chapter", "Spending Time Apart"], + ["点评我对本章的笔记", "Spending Time Apart"], + ["点评我对这一章的笔记", "Spending Time Apart"], + ["Please review my notes from chapter 2", "chapter 2"], + ["look at my notes", undefined], + ])("prefetches annotations before the model answers: %s", async (prompt, chapterTitle) => { + getReadingContextSnapshotMock.mockReturnValue({ + bookId: "book-1", + bookTitle: "Test Book", + currentChapter: { + index: 1, + title: "Spending Time Apart", + href: "chapter-2.xhtml", + }, + currentPosition: { cfi: "epubcfi(/6/4!/4/2)", percentage: 0.2 }, + surroundingText: "", + recentHighlights: [], + operationType: "reading", + timestamp: Date.now(), + }); + const annotationResult = { + highlights: [ + { + text: "Time apart can strengthen desire.", + note: "Distance creates room to want.", + chapterTitle: "Spending Time Apart", + color: "yellow", + }, + ], + notes: [], + pagination: { + highlights: { total: 1, returned: 1, offset: 0, limit: 50, hasMore: false }, + notes: { total: 0, returned: 0, offset: 0, limit: 50, hasMore: false }, + }, + }; + const getAnnotations = vi.fn(async () => annotationResult); + const annotationTool: ToolDefinition = { + name: "getAnnotations", + description: "Get user annotations", + parameters: { + type: { type: "string", description: "annotation type" }, + chapterTitle: { type: "string", description: "chapter title" }, + order: { type: "string", description: "sort order" }, + offset: { type: "number", description: "offset" }, + limit: { type: "number", description: "limit" }, + }, + execute: getAnnotations, + }; + let capturedAgentInput: + | { + messages: Array<{ + _getType: () => string; + content: unknown; + tool_calls?: Array<{ id?: string }>; + tool_call_id?: string; + }>; + } + | undefined; + createReactAgentMock.mockReturnValue({ + streamEvents: vi.fn((input) => { + capturedAgentInput = input; + return { + [Symbol.asyncIterator]: async function* () { + // no-op stream + }, + }; + }), + }); + + const events = []; + for await (const event of streamReadingAgent( + { + aiConfig: makeAIConfig(), + book: null, + bookId: "book-1", + semanticContext: null, + enabledSkills: [], + isVectorized: true, + getAvailableTools: () => [annotationTool], + }, + prompt, + )) { + events.push(event); + } + + const expectedArgs = { + type: "all", + order: "book", + offset: 0, + limit: 50, + ...(chapterTitle ? { chapterTitle } : {}), + }; + expect(getAnnotations).toHaveBeenCalledOnce(); + expect(getAnnotations).toHaveBeenCalledWith(expectedArgs); + expect(events.slice(0, 2)).toEqual([ + { type: "tool_call", name: "getAnnotations", args: expectedArgs }, + { type: "tool_result", name: "getAnnotations", result: annotationResult }, + ]); + + expect(capturedAgentInput).toBeDefined(); + const messages = capturedAgentInput?.messages ?? []; + const humanMessage = messages[messages.length - 3]; + const aiMessage = messages[messages.length - 2]; + const toolMessage = messages[messages.length - 1]; + expect(humanMessage?._getType()).toBe("human"); + expect(humanMessage?.content).toBe(prompt); + expect(aiMessage?._getType()).toBe("ai"); + expect(aiMessage?.tool_calls).toEqual([ + expect.objectContaining({ name: "getAnnotations", args: expectedArgs, type: "tool_call" }), + ]); + expect(toolMessage?._getType()).toBe("tool"); + expect(toolMessage?.tool_call_id).toBe(aiMessage?.tool_calls?.[0]?.id); + expect(toolMessage?.content).toBe(JSON.stringify(annotationResult)); + }); + + it("does not run annotation preflight for a normal chapter summary", async () => { + const getAnnotations = vi.fn(async () => ({ highlights: [], notes: [], pagination: {} })); + const tools: ToolDefinition[] = [ + { + name: "getAnnotations", + description: "Get user annotations", + parameters: {}, + execute: getAnnotations, + }, + { + name: "resolveChapterReference", + description: "Resolve chapter references", + parameters: {}, + execute: vi.fn(async () => ({ matched: true })), + }, + ]; + createReactAgentMock.mockReturnValue({ + streamEvents: vi.fn(() => ({ + [Symbol.asyncIterator]: async function* () { + // no-op stream + }, + })), + }); + + for await (const event of streamReadingAgent( + { + aiConfig: makeAIConfig(), + book: null, + bookId: "book-1", + semanticContext: null, + enabledSkills: [], + isVectorized: true, + getAvailableTools: () => tools, + }, + "summarize chapter 2", + )) { + void event; + } + + expect(getAnnotations).not.toHaveBeenCalled(); + }); + + it("keeps indexed content retrieval available when comparing annotations with the book", async () => { + const tools: ToolDefinition[] = [ + { + name: "getAnnotations", + description: "Get user annotations", + parameters: {}, + execute: vi.fn(async () => ({ highlights: [], notes: [], pagination: {} })), + }, + { + name: "ragSearch", + description: "Search book content", + parameters: {}, + execute: vi.fn(async () => ({ results: [] })), + }, + { + name: "ragContext", + description: "Get book context", + parameters: {}, + execute: vi.fn(async () => ({ chunks: [] })), + }, + { + name: "addCitation", + description: "Register citations", + parameters: {}, + execute: vi.fn(async () => ({ type: "citation" })), + }, + ]; + let capturedTools: Array<{ name: string }> = []; + createReactAgentMock.mockImplementation((config) => { + capturedTools = config.tools; + return { + streamEvents: vi.fn(() => ({ + [Symbol.asyncIterator]: async function* () { + // no-op stream + }, + })), + }; + }); + + for await (const event of streamReadingAgent( + { + aiConfig: makeAIConfig(), + book: null, + bookId: "book-1", + semanticContext: null, + enabledSkills: [], + isVectorized: true, + getAvailableTools: () => tools, + }, + "Compare my notes with the chapter text", + )) { + void event; + } + + expect(capturedTools.map((tool) => tool.name)).toEqual([ + "getAnnotations", + "ragSearch", + "ragContext", + "addCitation", + ]); + }); + it("keeps RAG fallback available for current-page questions on indexed books", async () => { let capturedTools: any[] = []; createReactAgentMock.mockImplementation((config) => { @@ -724,6 +947,144 @@ describe("streamReadingAgent tool registration", () => { expect(toolNames).not.toContain("summarize"); }); + it("keeps the citation budget available after the non-citation budget is exhausted", async () => { + const search = vi.fn(async (args) => ({ query: args.query })); + const addCitation = vi.fn(async (args) => ({ + type: "citation", + citationIndex: args.citationIndex, + })); + const tools: ToolDefinition[] = [ + { + name: "ragSearch", + description: "Search book content", + parameters: { query: { type: "string", description: "query", required: true } }, + execute: search, + }, + { + name: "addCitation", + description: "Register a citation", + parameters: { + citationIndex: { type: "number", description: "citation index", required: true }, + }, + execute: addCitation, + }, + ]; + let capturedTools: Array<{ + name: string; + func: (input: Record) => Promise; + }> = []; + createReactAgentMock.mockImplementation((config) => { + capturedTools = config.tools; + return { + streamEvents: vi.fn(() => ({ + [Symbol.asyncIterator]: async function* () { + // no-op stream + }, + })), + }; + }); + + for await (const event of streamReadingAgent( + { + aiConfig: makeAIConfig(), + book: null, + bookId: "book-1", + semanticContext: null, + enabledSkills: [], + isVectorized: true, + getAvailableTools: () => tools, + }, + "Analyze themes across this book", + )) { + void event; + } + + const wrappedSearch = capturedTools.find((tool) => tool.name === "ragSearch"); + const wrappedCitation = capturedTools.find((tool) => tool.name === "addCitation"); + for (let index = 0; index < 12; index += 1) { + const result = JSON.parse(await wrappedSearch.func({ query: `query-${index}` })); + expect(result.stopToolCalls).toBeUndefined(); + } + const blockedSearch = JSON.parse(await wrappedSearch.func({ query: "query-12" })); + expect(blockedSearch.stopToolCalls).toBe(true); + + for (let index = 1; index <= 16; index += 1) { + const result = JSON.parse(await wrappedCitation.func({ citationIndex: index })); + expect(result).toEqual({ type: "citation", citationIndex: index }); + } + expect(addCitation).toHaveBeenCalledTimes(16); + + const blockedCitation = JSON.parse(await wrappedCitation.func({ citationIndex: 17 })); + expect(blockedCitation).toMatchObject({ + citationLimitReached: true, + stopCitationCalls: true, + }); + expect(blockedCitation.instruction).toContain("Finish the answer now"); + }); + + it("does not let citation calls consume the non-citation budget", async () => { + const search = vi.fn(async (args) => ({ query: args.query })); + const addCitation = vi.fn(async (args) => ({ + type: "citation", + citationIndex: args.citationIndex, + })); + const tools: ToolDefinition[] = [ + { + name: "ragSearch", + description: "Search book content", + parameters: { query: { type: "string", description: "query", required: true } }, + execute: search, + }, + { + name: "addCitation", + description: "Register a citation", + parameters: { + citationIndex: { type: "number", description: "citation index", required: true }, + }, + execute: addCitation, + }, + ]; + let capturedTools: Array<{ + name: string; + func: (input: Record) => Promise; + }> = []; + createReactAgentMock.mockImplementation((config) => { + capturedTools = config.tools; + return { + streamEvents: vi.fn(() => ({ + [Symbol.asyncIterator]: async function* () { + // no-op stream + }, + })), + }; + }); + + for await (const event of streamReadingAgent( + { + aiConfig: makeAIConfig(), + book: null, + bookId: "book-1", + semanticContext: null, + enabledSkills: [], + isVectorized: true, + getAvailableTools: () => tools, + }, + "Analyze themes across this book", + )) { + void event; + } + + const wrappedSearch = capturedTools.find((tool) => tool.name === "ragSearch"); + const wrappedCitation = capturedTools.find((tool) => tool.name === "addCitation"); + for (let index = 1; index <= 16; index += 1) { + await wrappedCitation.func({ citationIndex: index }); + } + + const result = JSON.parse(await wrappedSearch.func({ query: "still available" })); + expect(result).toEqual({ query: "still available" }); + expect(search).toHaveBeenCalledOnce(); + }); + it("reuses duplicate search requests within the same turn", async () => { const searchCalls: string[] = []; const searchTool: ToolDefinition = { diff --git a/packages/core/src/ai/__tests__/system-prompt.test.ts b/packages/core/src/ai/__tests__/system-prompt.test.ts index 41f26b3fd..54f5ae127 100644 --- a/packages/core/src/ai/__tests__/system-prompt.test.ts +++ b/packages/core/src/ai/__tests__/system-prompt.test.ts @@ -133,4 +133,23 @@ describe("buildSystemPrompt citations", () => { expect(prompt).not.toContain("addCitation"); expect(prompt).not.toContain("mindmap"); }); + + it("treats prefetched annotations as available user data without citation registration", () => { + const prompt = buildSystemPrompt({ + book: makeBook(), + semanticContext: null, + enabledSkills: [], + isVectorized: true, + userLanguage: "en", + questionCategory: "annotation_request", + routeHint: "The user's annotations have already been fetched for this turn.", + allowedToolNames: ["getAnnotations"], + }); + + expect(prompt).toContain("Treat the matching getAnnotations tool result as current user data"); + expect(prompt).toContain("never claim annotation access is unavailable"); + expect(prompt).toContain("do not call addCitation or invent [N] markers"); + expect(prompt).toContain("pagination.hasMore"); + expect(prompt).not.toContain("- **addCitation**"); + }); }); diff --git a/packages/core/src/ai/agents/reading-agent.ts b/packages/core/src/ai/agents/reading-agent.ts index 37bf5f654..f80346bff 100644 --- a/packages/core/src/ai/agents/reading-agent.ts +++ b/packages/core/src/ai/agents/reading-agent.ts @@ -1,4 +1,4 @@ -import { AIMessage, HumanMessage } from "@langchain/core/messages"; +import { AIMessage, HumanMessage, ToolMessage } from "@langchain/core/messages"; import type { BaseMessage } from "@langchain/core/messages"; import { createReactAgent } from "@langchain/langgraph/prebuilt"; import i18n from "i18next"; @@ -23,12 +23,15 @@ import type { ToolDefinition, ToolParameter } from "../tools/tool-types"; const CHAPTER_REFERENCE_RE = /(?:第\s*)?[零〇一二两三四五六七八九十百千万\d]{1,8}\s*(?:章|卷|节|回|讲|篇|话)|这一章|这一节|chapter\s*\d+/iu; +const EXPLICIT_CHAPTER_REFERENCE_RE = + /(?:第\s*[零〇一二两三四五六七八九十百千万\d]{1,8}|\d{1,8})\s*(?:章|卷|节|回|讲|篇|话)|chapter\s*\d+/iu; const CHAPTER_REFERENCE_EXECUTION_LIMIT = 3; const CHAPTER_TOOL_EXECUTION_LIMIT = 8; const DEFAULT_RECURSION_LIMIT = 24; const CHAPTER_TASK_RECURSION_LIMIT = 24; const DEFAULT_TOOL_TIMEOUT_MS = 45_000; -const TOOL_EXECUTION_LIMIT = 12; +const NON_CITATION_TOOL_EXECUTION_LIMIT = 12; +const CITATION_TOOL_EXECUTION_LIMIT = 16; const REPEATED_TOOL_CALL_LIMIT = 2; const TOOL_TIMEOUT_MS_BY_NAME: Record = { getSelection: 5_000, @@ -112,6 +115,11 @@ const CURRENT_PAGE_CONTEXT_RE = /(?:这里|這裡|当前页|當前頁|这一页|這一頁|这页|這頁|当前位置|當前位置|目前看到|我看到这里|我看到這裡)/u; const CURRENT_CHAPTER_CONTEXT_RE = /(?:这一章|這一章|这章|這章|当前章节|當前章節|当前章|當前章|現在這章|现在这章|本章)/u; +const CURRENT_CHAPTER_CONTEXT_EN_RE = /\b(?:this|current)\s+chapter\b/iu; +const ANNOTATION_REQUEST_RE = + /(?:\bmy\s+(?:notes?|highlights?|annotations?)\b|\b(?:look\s+at|review|analy[sz]e|summari[sz]e|discuss|read|check)\s+(?:my\s+)?(?:notes?|highlights?|annotations?)\b|(?:我(?:的|对.{0,24})?)?(?:笔记|筆記|高亮|标注|標註|划线|劃線))/iu; +const ANNOTATION_BOOK_COMPARISON_RE = + /(?:\b(?:compare|contrast|connect|relate)\b.{0,80}\b(?:book|text|chapter|passage|author)\b|\b(?:book|text|chapter|passage|author)\b.{0,80}\b(?:compare|contrast|connect|relate)\b|(?:对照|對照|比较|比較|结合|結合).{0,80}(?:原文|正文|章节|章節|书中|書中))/iu; const IMMEDIATE_CONTEXT_RE = /(?:什么意思|什麼意思|看不懂|沒看懂|没看懂|解释一下|解釋一下|怎么理解|怎麼理解)/u; const BOOK_CONTENT_RE = @@ -146,6 +154,7 @@ const CATEGORY_TOOL_ORDER: Record = { "updateBookMetadata", "manageBookGroups", ], + annotation_request: ["getAnnotations"], current_selection: [ "getSelection", "getSurroundingContext", @@ -218,6 +227,7 @@ const CATEGORY_TOOL_ORDER: Record = { type ReadingQuestionCategory = | "general_chat" | "library_request" + | "annotation_request" | "current_selection" | "current_page_context" | "current_chapter_context" @@ -270,7 +280,8 @@ function detectQuestionCategory(options: { if (LIBRARY_REQUEST_RE.test(text) || !options.hasBookContext) return "library_request"; const hasExplicitCurrentSelectionCue = CURRENT_SELECTION_RE.test(text); const hasExplicitCurrentPageCue = CURRENT_PAGE_CONTEXT_RE.test(text); - const hasExplicitCurrentChapterCue = CURRENT_CHAPTER_CONTEXT_RE.test(text); + const hasExplicitCurrentChapterCue = + CURRENT_CHAPTER_CONTEXT_RE.test(text) || CURRENT_CHAPTER_CONTEXT_EN_RE.test(text); const asksForImmediateExplanation = IMMEDIATE_CONTEXT_RE.test(text); if (options.selectionActive && hasExplicitCurrentSelectionCue) { @@ -279,6 +290,7 @@ function detectQuestionCategory(options: { if (hasExplicitCurrentPageCue || (asksForImmediateExplanation && hasExplicitCurrentPageCue)) { return "current_page_context"; } + if (ANNOTATION_REQUEST_RE.test(text)) return "annotation_request"; if (CHAPTER_REFERENCE_RE.test(text)) return "specific_chapter_request"; if (hasExplicitCurrentChapterCue) return "current_chapter_context"; if (BOOK_CONTENT_RE.test(text)) return "book_wide_search"; @@ -288,12 +300,21 @@ function detectQuestionCategory(options: { function getFocusedToolNames( category: ReadingQuestionCategory, isVectorized: boolean, + annotationNeedsBookContent: boolean, ): Set | null { switch (category) { case "general_chat": return new Set(); case "library_request": return GENERAL_TOOL_NAMES; + case "annotation_request": + return new Set( + annotationNeedsBookContent + ? isVectorized + ? ["getAnnotations", "ragSearch", "ragContext", "addCitation"] + : ["getAnnotations", "fallbackSearch", "fallbackChapterContext", "addCitation"] + : ["getAnnotations"], + ); case "current_selection": return new Set( isVectorized @@ -397,8 +418,13 @@ function filterToolsForQuestion(options: { tools: ToolDefinition[]; category: ReadingQuestionCategory; isVectorized: boolean; + userInput: string; }): ToolDefinition[] { - const focusedNames = getFocusedToolNames(options.category, options.isVectorized); + const focusedNames = getFocusedToolNames( + options.category, + options.isVectorized, + ANNOTATION_BOOK_COMPARISON_RE.test(options.userInput), + ); if (focusedNames === null) { return sortToolsForCategory( options.tools.filter((tool) => !GENERAL_TOOL_NAMES.has(tool.name)), @@ -428,6 +454,8 @@ function buildRouteHint( return isVectorized ? "This question is about the chapter the user is currently reading. Get the current chapter first, then prefer indexed chapter/content retrieval." : "This question is about the chapter the user is currently reading. Get the current chapter first, then use fallback chapter content."; + case "annotation_request": + return "The user's annotations have already been fetched for this turn. Use that getAnnotations result as current user data and paginate only if its metadata says more results are available."; case "specific_chapter_request": return isVectorized ? "This question targets a specific chapter reference. Resolve the chapter reference first; if resolution is weak or the user asks for content, use ragSearch/ragToc/ragContext instead of guessing." @@ -454,6 +482,8 @@ function getRecursionLimitForCategory(category: ReadingQuestionCategory): number return CHAPTER_TASK_RECURSION_LIMIT; case "library_request": return 20; + case "annotation_request": + return 20; case "book_wide_search": return DEFAULT_RECURSION_LIMIT; default: @@ -461,6 +491,35 @@ function getRecursionLimitForCategory(category: ReadingQuestionCategory): number } } +function extractAnnotationChapterQuery(userInput: string): string | undefined { + const match = userInput.normalize("NFKC").match(EXPLICIT_CHAPTER_REFERENCE_RE)?.[0]?.trim(); + return match || undefined; +} + +function buildAnnotationPreflightArgs(options: { + userInput: string; + currentChapterTitle?: string; +}): Record { + const explicitChapter = extractAnnotationChapterQuery(options.userInput); + const currentChapter = + CURRENT_CHAPTER_CONTEXT_RE.test(options.userInput) || + CURRENT_CHAPTER_CONTEXT_EN_RE.test(options.userInput) + ? options.currentChapterTitle?.trim() + : undefined; + + return { + type: "all", + order: "book", + offset: 0, + limit: 50, + ...(explicitChapter + ? { chapterTitle: explicitChapter } + : currentChapter + ? { chapterTitle: currentChapter } + : {}), + }; +} + function simplifyChapterLookupQuery(query: string, fallback: string): string { const source = (query || fallback).normalize("NFKC"); const sanitizedSource = source.replace(/(?:这|那|哪)\s*一\s*(?:章|卷|节|回|讲|篇|话)/gu, " "); @@ -549,6 +608,17 @@ function buildRepeatedToolCallResult( }; } +function buildCitationLimitResult(): Record { + return { + type: "notice", + citationLimitReached: true, + stopCitationCalls: true, + reason: "Citation registration limit reached for this turn.", + instruction: + "Stop registering citations. Finish the answer now and use plain chapter/source references for any remaining sources.", + }; +} + // --- Stream Event Types --- export type AgentStreamEvent = @@ -842,7 +912,8 @@ export async function* streamReadingAgent( const searchResultCache = new Map(); const toolExecutionCounts = new Map(); const lastToolResults = new Map(); - let totalToolExecutions = 0; + let totalNonCitationToolExecutions = 0; + let totalCitationToolExecutions = 0; const pendingToolCallNames: string[] = []; const isChapterTask = questionCategory === "specific_chapter_request" || CHAPTER_REFERENCE_RE.test(userInput); @@ -877,6 +948,7 @@ export async function* streamReadingAgent( }), category: questionCategory, isVectorized, + userInput, }); console.log( "[ReadingAgent] tools", @@ -930,6 +1002,45 @@ export async function* streamReadingAgent( new HumanMessage(userInput), ]; + if (questionCategory === "annotation_request") { + const annotationTool = tools.find((tool) => tool.name === "getAnnotations"); + if (annotationTool) { + const preflightArgs = buildAnnotationPreflightArgs({ + userInput, + currentChapterTitle: + readingContextSnapshot?.currentChapter.title || semanticContext?.currentChapter, + }); + yield { type: "tool_call", name: annotationTool.name, args: preflightArgs }; + const preflightResult = await executeTool( + annotationTool, + preflightArgs, + getToolTimeoutMs(annotationTool, toolTimeoutMs), + ); + toolResultCache.set(buildToolCacheKey(annotationTool.name, preflightArgs), preflightResult); + + const toolCallId = `annotation-preflight-${Date.now()}`; + inputMessages.push( + new AIMessage({ + content: "", + tool_calls: [ + { + id: toolCallId, + name: annotationTool.name, + args: preflightArgs, + type: "tool_call", + }, + ], + }), + new ToolMessage({ + content: JSON.stringify(preflightResult), + tool_call_id: toolCallId, + name: annotationTool.name, + }), + ); + yield { type: "tool_result", name: annotationTool.name, result: preflightResult }; + } + } + // If no tools available, stream directly without agent graph if (tools.length === 0) { const { SystemMessage } = await import("@langchain/core/messages"); @@ -1060,14 +1171,6 @@ export async function* streamReadingAgent( buildRepeatedToolCallResult(lastToolResults.get(progressKey), tool.name, "duplicate"), ); } - if (totalToolExecutions >= TOOL_EXECUTION_LIMIT) { - return JSON.stringify( - buildRepeatedToolCallResult(lastToolResults.get(progressKey), tool.name, "limit"), - ); - } - toolExecutionCounts.set(progressKey, previousExecutions + 1); - totalToolExecutions += 1; - const skipExactCache = tool.name === "addCitation" || tool.name === "resolveChapterReference"; const exactCacheKey = skipExactCache @@ -1075,6 +1178,7 @@ export async function* streamReadingAgent( : buildToolCacheKey(tool.name, toolInput); if (exactCacheKey && toolResultCache.has(exactCacheKey)) { const cachedResult = toolResultCache.get(exactCacheKey); + toolExecutionCounts.set(progressKey, previousExecutions + 1); lastToolResults.set(progressKey, cachedResult); return JSON.stringify(cachedResult); } @@ -1085,10 +1189,26 @@ export async function* streamReadingAgent( : undefined; if (searchCacheKey && searchResultCache.has(searchCacheKey)) { const cachedResult = searchResultCache.get(searchCacheKey); + toolExecutionCounts.set(progressKey, previousExecutions + 1); lastToolResults.set(progressKey, cachedResult); return JSON.stringify(cachedResult); } + if (tool.name === "addCitation") { + if (totalCitationToolExecutions >= CITATION_TOOL_EXECUTION_LIMIT) { + return JSON.stringify(buildCitationLimitResult()); + } + totalCitationToolExecutions += 1; + } else { + if (totalNonCitationToolExecutions >= NON_CITATION_TOOL_EXECUTION_LIMIT) { + return JSON.stringify( + buildRepeatedToolCallResult(lastToolResults.get(progressKey), tool.name, "limit"), + ); + } + totalNonCitationToolExecutions += 1; + } + toolExecutionCounts.set(progressKey, previousExecutions + 1); + const result = await executeTool(tool, toolInput, getToolTimeoutMs(tool, toolTimeoutMs)); if (exactCacheKey) { toolResultCache.set(exactCacheKey, result); diff --git a/packages/core/src/ai/system-prompt.ts b/packages/core/src/ai/system-prompt.ts index 91aaf916a..833f420f2 100644 --- a/packages/core/src/ai/system-prompt.ts +++ b/packages/core/src/ai/system-prompt.ts @@ -13,6 +13,7 @@ import { getBookProgressPercent } from "../utils/book-progress"; type ReadingQuestionCategory = | "general_chat" | "library_request" + | "annotation_request" | "current_selection" | "current_page_context" | "current_chapter_context" @@ -42,6 +43,7 @@ export function buildSystemPrompt(ctx: PromptContext): string { buildMemorySection(ctx.memorySummary), buildSemanticSection(ctx.semanticContext), buildRouteSection(ctx.questionCategory, ctx.selectionActive, ctx.routeHint), + buildAnnotationRequestSection(ctx.questionCategory), buildTurnAvailableToolsSection(ctx.allowedToolNames), buildToolsSection( ctx.enabledSkills, @@ -63,6 +65,20 @@ export function buildSystemPrompt(ctx: PromptContext): string { return sections.filter(Boolean).join("\n\n---\n\n"); } +function buildAnnotationRequestSection(category?: ReadingQuestionCategory): string { + if (category !== "annotation_request") return ""; + + return [ + "## Annotation Request Contract", + "- The application fetched the user's annotations before model generation.", + "- Treat the matching getAnnotations tool result as current user data and answer from it directly.", + "- When that result is present, never claim annotation access is unavailable or ask the user to paste annotations that were already returned.", + "- User-authored notes and highlights are not book-content citations: do not call addCitation or invent [N] markers for them.", + "- If pagination.hasMore is true and more annotations are needed, call getAnnotations with the next offset. Otherwise, stop retrieving and answer.", + "- If separately retrieved book text is used to compare against the annotations, follow the normal citation rules for that book text only.", + ].join("\n"); +} + function buildMemorySection(memorySummary?: string): string { if (!memorySummary?.trim()) return ""; return ["## Conversation Memory", memorySummary.trim()].join("\n"); diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 560ef437a..f4c35e4d3 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -343,6 +343,9 @@ importers: expo-image-picker: specifier: ~17.0.11 version: 17.0.11(expo@54.0.33(@babel/core@7.29.0)(graphql@16.8.1)(react-native-webview@13.15.0(react-native@0.81.5(@babel/core@7.29.0)(@types/react@19.1.17)(react@19.1.0))(react@19.1.0))(react-native@0.81.5(@babel/core@7.29.0)(@types/react@19.1.17)(react@19.1.0))(react@19.1.0)) + expo-keep-awake: + specifier: ~15.0.8 + version: 15.0.8(expo@54.0.33(@babel/core@7.29.0)(graphql@16.8.1)(react-native-webview@13.15.0(react-native@0.81.5(@babel/core@7.29.0)(@types/react@19.1.17)(react@19.1.0))(react@19.1.0))(react-native@0.81.5(@babel/core@7.29.0)(@types/react@19.1.17)(react@19.1.0))(react@19.1.0))(react@19.1.0) expo-linear-gradient: specifier: ~15.0.8 version: 15.0.8(expo@54.0.33(@babel/core@7.29.0)(graphql@16.8.1)(react-native-webview@13.15.0(react-native@0.81.5(@babel/core@7.29.0)(@types/react@19.1.17)(react@19.1.0))(react@19.1.0))(react-native@0.81.5(@babel/core@7.29.0)(@types/react@19.1.17)(react@19.1.0))(react@19.1.0))(react-native@0.81.5(@babel/core@7.29.0)(@types/react@19.1.17)(react@19.1.0))(react@19.1.0)