From 660204fed6544786848ac22752918b5ecc1d634b Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 7 Jul 2026 13:35:59 +0800 Subject: [PATCH 01/14] Present page-mode sources as page assets --- .../[sourceId]/page-assets/route.test.ts | 271 +++++++++++++++++ .../sources/[sourceId]/page-assets/route.ts | 32 ++ src/components/chat-message-list.test.ts | 14 +- src/components/chat-message-list.tsx | 34 +-- src/components/chunks-panel.test.ts | 110 +++++++ src/components/chunks-panel.tsx | 286 +++++++++++++++++- src/components/source-row.test.ts | 21 ++ src/components/source-row.tsx | 3 + src/components/sources-panel.tsx | 1 + .../workspace-citation-focus.test.ts | 43 +++ src/components/workspace-citation-focus.ts | 32 ++ src/components/workspace-citation-state.ts | 37 +++ .../workspace-selected-chunks.test.ts | 22 ++ src/components/workspace-selected-chunks.ts | 4 +- src/components/workspace-shell-layout.tsx | 16 + src/components/workspace-shell.tsx | 1 + src/domains/chat/chat-citation-persistence.ts | 1 + src/domains/chat/citations.ts | 3 + src/domains/chat/page-citation-assets.ts | 8 +- src/domains/chat/types.ts | 1 + src/domains/chat/view.ts | 4 + src/domains/sources/counts.test.ts | 61 ++++ src/domains/sources/counts.ts | 138 ++++++++- src/domains/sources/page-assets.ts | 104 +++++++ src/domains/sources/route-page-assets.ts | 161 ++++++++++ src/domains/sources/route-service.ts | 5 + src/domains/sources/route-types.ts | 32 +- src/domains/sources/types.ts | 17 ++ src/domains/sources/view.test.ts | 12 + src/domains/sources/view.ts | 9 +- src/domains/workspace/client.ts | 36 ++- 31 files changed, 1448 insertions(+), 71 deletions(-) create mode 100644 src/app/api/sources/[sourceId]/page-assets/route.test.ts create mode 100644 src/app/api/sources/[sourceId]/page-assets/route.ts create mode 100644 src/domains/sources/page-assets.ts create mode 100644 src/domains/sources/route-page-assets.ts diff --git a/src/app/api/sources/[sourceId]/page-assets/route.test.ts b/src/app/api/sources/[sourceId]/page-assets/route.test.ts new file mode 100644 index 0000000..4bf5c5e --- /dev/null +++ b/src/app/api/sources/[sourceId]/page-assets/route.test.ts @@ -0,0 +1,271 @@ +import { NextRequest } from "next/server" +import { beforeEach, describe, expect, it, vi } from "vitest" + +const mocks = vi.hoisted(() => ({ + deleteBlob: vi.fn(), + ensureApiKeyForWorkspace: vi.fn(), + ensureWorkspace: vi.fn(), + findSourceInWorkspace: vi.fn(), + getCurrentUser: vi.fn(), + makeKnowhereClient: vi.fn(), + makeKnowhereClientWithParsedStorage: vi.fn(), + readChunks: vi.fn(), + requireUser: vi.fn(), +})) + +vi.mock("next/headers", () => ({ + headers: vi.fn(async () => new Headers({ cookie: "session=abc" })), +})) + +vi.mock("@/integrations/dashboard/api-key-service", () => ({ + ensureApiKeyForWorkspace: mocks.ensureApiKeyForWorkspace, +})) + +vi.mock("@/integrations/knowhere-demo", () => ({ + knowhereDemoApi: { + fetchCatalog: vi.fn(), + fetchChunkPage: vi.fn(), + }, +})) + +vi.mock("@/infrastructure/auth", () => ({ + getCurrentUser: mocks.getCurrentUser, + requireUser: mocks.requireUser, +})) + +vi.mock("@/integrations/knowhere", () => ({ + makeKnowhereClient: mocks.makeKnowhereClient, + makeKnowhereClientWithParsedStorage: + mocks.makeKnowhereClientWithParsedStorage, +})) + +vi.mock("@vercel/blob", () => ({ + del: mocks.deleteBlob, +})) + +vi.mock("@/domains/sources/service", () => ({ + sourceService: { + findInWorkspace: mocks.findSourceInWorkspace, + localizeRemoteDocument: vi.fn(), + }, +})) + +vi.mock("@/domains/workspace/service", () => ({ + workspaceService: { + ensureWorkspace: mocks.ensureWorkspace, + }, +})) + +import { GET } from "./route" + +describe("GET /api/sources/[sourceId]/page-assets", () => { + beforeEach(() => { + vi.clearAllMocks() + mocks.makeKnowhereClientWithParsedStorage.mockReturnValue({ + client: { documents: { listChunks: vi.fn() } }, + knowledge: { readChunks: mocks.readChunks }, + }) + }) + + it("returns durable page assets for a ready workspace source", async () => { + mocks.getCurrentUser.mockResolvedValue({ + id: "user_1", + email: null, + name: null, + }) + mocks.ensureWorkspace.mockResolvedValue({ + id: "workspace_1", + userId: "user_1", + namespace: "notebook-workspace_1", + createdAt: new Date("2026-05-10T00:00:00.000Z"), + }) + mocks.findSourceInWorkspace.mockResolvedValue( + makeReadySource({ + id: "00000000-0000-0000-0000-000000000002", + knowhereJobId: "job_1", + knowhereDocumentId: "doc_1", + }), + ) + mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") + mocks.readChunks.mockResolvedValue({ + chunks: [ + { + chunkId: "page_1", + chunkType: "page", + assetUrl: "https://assets.example/fallback.png", + metadata: { + pageAssets: [ + { + pageNum: 1, + assetUrl: "https://assets.example/page-000001.png", + contentType: "image/png", + width: 1200, + height: 1600, + }, + ], + }, + }, + ], + page: 1, + pageSize: 1, + totalChunks: 3, + totalPages: 3, + }) + + const response = await GET( + new NextRequest( + "http://localhost:3001/api/sources/00000000-0000-0000-0000-000000000002/page-assets?page=1&pageSize=1", + ), + { + params: Promise.resolve({ + sourceId: "00000000-0000-0000-0000-000000000002", + }), + }, + ) + + await expect(response.json()).resolves.toEqual({ + pages: [ + { + pageNumber: 1, + assetUrl: "https://assets.example/page-000001.png", + contentType: "image/png", + width: 1200, + height: 1600, + }, + ], + pagination: { + page: 1, + pageSize: 1, + total: 3, + totalPages: 3, + }, + }) + expect(response.status).toBe(200) + expect(mocks.readChunks).toHaveBeenCalledWith({ + documentId: "doc_1", + revisionKey: "job_1", + chunkType: "page", + page: 1, + pageSize: 1, + assetUrlPolicy: "durable", + }) + }) + + it("rejects non-ready workspace sources", async () => { + mocks.getCurrentUser.mockResolvedValue({ + id: "user_1", + email: null, + name: null, + }) + mocks.ensureWorkspace.mockResolvedValue({ + id: "workspace_1", + userId: "user_1", + namespace: "notebook-workspace_1", + createdAt: new Date("2026-05-10T00:00:00.000Z"), + }) + mocks.findSourceInWorkspace.mockResolvedValue( + makeReadySource({ + id: "00000000-0000-0000-0000-000000000002", + status: "parsing", + knowhereDocumentId: "doc_1", + }), + ) + + const response = await GET( + new NextRequest( + "http://localhost:3001/api/sources/00000000-0000-0000-0000-000000000002/page-assets?page=1&pageSize=1", + ), + { + params: Promise.resolve({ + sourceId: "00000000-0000-0000-0000-000000000002", + }), + }, + ) + + await expect(response.json()).resolves.toEqual({ + message: "Source is not ready.", + }) + expect(response.status).toBe(409) + expect(mocks.readChunks).not.toHaveBeenCalled() + }) + + it("returns an empty page list when page chunks have no usable assets", async () => { + mocks.getCurrentUser.mockResolvedValue({ + id: "user_1", + email: null, + name: null, + }) + mocks.ensureWorkspace.mockResolvedValue({ + id: "workspace_1", + userId: "user_1", + namespace: "notebook-workspace_1", + createdAt: new Date("2026-05-10T00:00:00.000Z"), + }) + mocks.findSourceInWorkspace.mockResolvedValue( + makeReadySource({ + id: "00000000-0000-0000-0000-000000000002", + knowhereDocumentId: "doc_1", + }), + ) + mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") + mocks.readChunks.mockResolvedValue({ + chunks: [ + { + chunkId: "page_1", + chunkType: "page", + metadata: { pageAssets: [] }, + }, + ], + page: 1, + pageSize: 1, + totalChunks: 1, + totalPages: 1, + }) + + const response = await GET( + new NextRequest( + "http://localhost:3001/api/sources/00000000-0000-0000-0000-000000000002/page-assets?page=1&pageSize=1", + ), + { + params: Promise.resolve({ + sourceId: "00000000-0000-0000-0000-000000000002", + }), + }, + ) + + await expect(response.json()).resolves.toEqual({ + pages: [], + pagination: { + page: 1, + pageSize: 1, + total: 1, + totalPages: 1, + }, + }) + expect(response.status).toBe(200) + }) +}) + +function makeReadySource(overrides: Record) { + return { + id: "00000000-0000-0000-0000-000000000002", + workspaceId: "workspace_1", + title: "notes.pdf", + mimeType: "application/pdf", + sizeBytes: 1024, + status: "ready", + failureReason: null, + failureStage: null, + knowhereJobId: "job_1", + knowhereDocumentId: "doc_1", + stagedBlobPathname: null, + stagedBlobUrl: null, + originalBlobPathname: null, + originalBlobUrl: null, + demoKey: null, + createdAt: new Date("2026-05-10T00:00:00.000Z"), + updatedAt: new Date("2026-05-10T00:00:00.000Z"), + deletedAt: null, + ...overrides, + } +} diff --git a/src/app/api/sources/[sourceId]/page-assets/route.ts b/src/app/api/sources/[sourceId]/page-assets/route.ts new file mode 100644 index 0000000..b7606f7 --- /dev/null +++ b/src/app/api/sources/[sourceId]/page-assets/route.ts @@ -0,0 +1,32 @@ +import type { NextRequest, NextResponse } from "next/server" + +import { getChunkPageParams } from "@/domains/chunks" +import { createSourceRouteService } from "@/domains/sources/route-service" +import { withApiErrorResponse } from "@/lib/api-error-response" +import { nextRouteContext } from "@/lib/next-route-context" +import { nextRouteResponse } from "@/lib/next-route-response" + +type RouteContext = { + params: Promise<{ + sourceId: string + }> +} + +const sourceRouteService = createSourceRouteService() + +export async function GET( + request: NextRequest, + context: RouteContext, +): Promise { + return withApiErrorResponse("sources:page-assets", async () => { + const { sourceId } = await context.params + const routeContext = await nextRouteContext.read() + const result = await sourceRouteService.loadSourcePageAssets({ + cookieHeader: routeContext.cookieHeader, + pageParams: getChunkPageParams(request.nextUrl.searchParams), + sourceId, + }) + + return nextRouteResponse.toNextResponse(result) + }) +} diff --git a/src/components/chat-message-list.test.ts b/src/components/chat-message-list.test.ts index f94478e..9444128 100644 --- a/src/components/chat-message-list.test.ts +++ b/src/components/chat-message-list.test.ts @@ -146,6 +146,7 @@ describe("ChatMessageList", () => { score: 0.9, pageCitationAssetUrl: "https://blob.example/pages/page-000004.png", + pageCitationPageNumber: 4, source: { documentId: "doc_1", sourceFileName: "report.pdf", @@ -162,18 +163,17 @@ describe("ChatMessageList", () => { const citationButton = screen.getByRole("button", { name: "Open source report.pdf", }); - const pageImageLink = screen.getByRole("link", { - name: "Open page image for report.pdf", - }); - - expect(pageImageLink.getAttribute("href")).toBe( - "https://blob.example/pages/page-000004.png", - ); + expect( + screen.queryByRole("link", { + name: "Open page image for report.pdf", + }), + ).toBeNull(); await user.click(citationButton); expect(onCitationClick).toHaveBeenCalledWith( expect.objectContaining({ pageCitationAssetUrl: "https://blob.example/pages/page-000004.png", + pageCitationPageNumber: 4, }), "assistant_1:0", ); diff --git a/src/components/chat-message-list.tsx b/src/components/chat-message-list.tsx index 531c71f..38c5059 100644 --- a/src/components/chat-message-list.tsx +++ b/src/components/chat-message-list.tsx @@ -2,7 +2,7 @@ import { type CSSProperties, type ReactElement } from "react"; import { type VirtualItem } from "@tanstack/react-virtual"; -import { ExternalLink, ImageIcon, MessageCircle } from "lucide-react"; +import { ImageIcon, MessageCircle } from "lucide-react"; import ReactMarkdown, { defaultUrlTransform, type Components, @@ -633,10 +633,6 @@ function CitationChip({ citationId: string, ) => void; }): ReactElement { - const pageCitationAssetUrl = getTrimmedCitationField( - citation.pageCitationAssetUrl, - ); - return ( @@ -646,11 +642,7 @@ function CitationChip({ disabled={!onCitationClick || isPending} onClick={() => onCitationClick?.(citation, citationId)} aria-busy={isPending} - className={`inline-flex h-8 max-w-[250px] cursor-pointer items-center border border-primary/20 bg-primary/10 px-3 text-left font-mono text-xs font-semibold leading-none text-primary shadow-[0_1px_0_rgba(15,23,42,0.06)] transition-[background-color,border-color,color,box-shadow,transform] hover:border-primary/35 hover:bg-primary/15 hover:text-primary hover:shadow-[0_0_0_2px_rgba(37,99,235,0.12)] active:translate-y-px active:bg-primary/20 focus:outline-none focus:ring-4 focus:ring-ring/15 focus:ring-offset-2 focus:ring-offset-background disabled:cursor-wait disabled:opacity-75 disabled:hover:border-primary/20 disabled:hover:bg-primary/10 disabled:hover:text-primary disabled:hover:shadow-[0_1px_0_rgba(15,23,42,0.06)] dark:border-transparent dark:bg-[#5c606b] dark:text-[#cfd3dc] dark:shadow-none dark:hover:border-[#8f96a8] dark:hover:bg-[#4f535e] dark:hover:text-white dark:hover:shadow-[0_0_0_2px_rgba(143,150,168,0.22)] dark:active:bg-[#454955] dark:disabled:hover:border-transparent dark:disabled:hover:bg-[#5c606b] dark:disabled:hover:text-[#cfd3dc] dark:disabled:hover:shadow-none ${ - pageCitationAssetUrl - ? "rounded-l-md rounded-r-none" - : "rounded-md" - }`} + className="inline-flex h-8 max-w-[250px] cursor-pointer items-center rounded-md border border-primary/20 bg-primary/10 px-3 text-left font-mono text-xs font-semibold leading-none text-primary shadow-[0_1px_0_rgba(15,23,42,0.06)] transition-[background-color,border-color,color,box-shadow,transform] hover:border-primary/35 hover:bg-primary/15 hover:text-primary hover:shadow-[0_0_0_2px_rgba(37,99,235,0.12)] active:translate-y-px active:bg-primary/20 focus:outline-none focus:ring-4 focus:ring-ring/15 focus:ring-offset-2 focus:ring-offset-background disabled:cursor-wait disabled:opacity-75 disabled:hover:border-primary/20 disabled:hover:bg-primary/10 disabled:hover:text-primary disabled:hover:shadow-[0_1px_0_rgba(15,23,42,0.06)] dark:border-transparent dark:bg-[#5c606b] dark:text-[#cfd3dc] dark:shadow-none dark:hover:border-[#8f96a8] dark:hover:bg-[#4f535e] dark:hover:text-white dark:hover:shadow-[0_0_0_2px_rgba(143,150,168,0.22)] dark:active:bg-[#454955] dark:disabled:hover:border-transparent dark:disabled:hover:bg-[#5c606b] dark:disabled:hover:text-[#cfd3dc] dark:disabled:hover:shadow-none" aria-label={`Open source ${label}`} > {label} @@ -664,28 +656,6 @@ function CitationChip({ {tooltipLabel} - {pageCitationAssetUrl ? ( - - - - - - - - Open page image - - - ) : null} ); } diff --git a/src/components/chunks-panel.test.ts b/src/components/chunks-panel.test.ts index f99a68b..9e5bb4b 100644 --- a/src/components/chunks-panel.test.ts +++ b/src/components/chunks-panel.test.ts @@ -14,6 +14,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { ChunksPanel } from "./chunks-panel"; import { sourceOriginalPreviewRequest } from "./source-original-preview-request"; +const fetchPageAssetPageMock = vi.hoisted(() => vi.fn()); const C = ChunksPanel as React.FC>; const virtualizerScrollResetDelayMs = 150; @@ -29,14 +30,31 @@ vi.mock("react-pdf", () => ({ Page: () => React.createElement("div", { "data-testid": "pdf-page" }), })); +vi.mock("@/domains/workspace/client", () => ({ + workspaceClient: { + fetchPageAssetPage: fetchPageAssetPageMock, + }, +})); + describe("ChunksPanel", () => { beforeEach(() => { shouldFlushVirtualizerTimers = false; + fetchPageAssetPageMock.mockReset(); + fetchPageAssetPageMock.mockResolvedValue({ + pages: [], + pagination: { + page: 1, + pageSize: 50, + total: 0, + totalPages: 1, + }, + }); globalThis.ResizeObserver = class ResizeObserver { observe() {} unobserve() {} disconnect() {} }; + Element.prototype.scrollIntoView = vi.fn(); }); afterEach(async () => { @@ -71,6 +89,98 @@ describe("ChunksPanel", () => { expect(screen.getByText(/Showing all parsed chunks from/)).toBeTruthy(); }); + it("renders page assets instead of parsed chunk controls for page-mode documents", async () => { + fetchPageAssetPageMock.mockResolvedValue({ + pages: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-000004.png", + contentType: "image/png", + width: 1200, + height: 1600, + }, + ], + pagination: { + page: 1, + pageSize: 50, + total: 4, + totalPages: 1, + }, + }); + + render( + React.createElement(C, { + chunks: [], + selectedSource: "report.pdf", + selectedSourceView: { + id: "source_1", + title: "report.pdf", + mimeType: "application/pdf", + status: "ready", + documentPresentation: { kind: "page-assets", pageCount: 4 }, + }, + }), + ); + + await waitFor(() => + expect( + screen.getByTestId("page-asset-document-viewer"), + ).toBeTruthy(), + ); + expect(screen.getByRole("heading", { name: "Original File" })).toBeTruthy(); + expect(screen.getByText("Page 4")).toBeTruthy(); + expect(screen.queryByRole("button", { name: "Parsed" })).toBeNull(); + expect(screen.queryByRole("button", { name: "Tree" })).toBeNull(); + expect(fetchPageAssetPageMock).toHaveBeenCalledWith("source_1", 1); + }); + + it("loads missing page buckets after focusing a later page asset", async () => { + const user = userEvent.setup(); + fetchPageAssetPageMock.mockImplementation( + async (_sourceId: string, page: number) => ({ + pages: [ + { + pageNumber: page === 3 ? 101 : page, + assetUrl: `https://assets.example/page-${page}.png`, + contentType: "image/png", + }, + ], + pagination: { + page, + pageSize: 50, + total: 120, + totalPages: 3, + }, + }), + ); + + render( + React.createElement(C, { + chunks: [], + selectedSource: "report.pdf", + selectedSourceView: { + id: "source_1", + title: "report.pdf", + mimeType: "application/pdf", + status: "ready", + documentPresentation: { kind: "page-assets", pageCount: 120 }, + }, + focusedPageNumber: 101, + focusedPageRequestId: 1, + }), + ); + + await waitFor(() => + expect(fetchPageAssetPageMock).toHaveBeenCalledWith("source_1", 3), + ); + + await user.click(screen.getByRole("button", { name: "Load more pages" })); + + await waitFor(() => + expect(fetchPageAssetPageMock).toHaveBeenCalledWith("source_1", 2), + ); + }); + it("defaults parsed chunks into a section tree view", () => { render( React.createElement(C, { diff --git a/src/components/chunks-panel.tsx b/src/components/chunks-panel.tsx index 78cd170..5561c4e 100644 --- a/src/components/chunks-panel.tsx +++ b/src/components/chunks-panel.tsx @@ -41,17 +41,25 @@ import { chunksPanelState } from "@/components/chunks-panel-state"; import { MAX_UPLOAD_MB } from "@/domains/sources/validation"; import { useSourceOriginalPreviewWarmup } from "@/components/source-original-preview-warmup"; import { sourceOriginalPreviewModel } from "@/components/source-original-preview-model"; +import { workspaceClient } from "@/domains/workspace/client"; import type { ParsedChunkView } from "@/domains/chunks/types"; -import type { SourceOriginalFileView, SourceView } from "@/domains/sources/types"; +import type { + SourceOriginalFileView, + SourcePageAssetView, + SourceView, +} from "@/domains/sources/types"; import type { AnalyticsContext } from "@/lib/posthog"; import { cn } from "@/lib/utils"; export type ChunksPanelProps = { chunks: ParsedChunkView[]; selectedSource?: string | null; + selectedSourceView?: SourceView | null; selectedSourceFile?: SourceOriginalFileView | null; focusedChunkId?: string | null; focusedChunkRequestId?: number; + focusedPageNumber?: number | null; + focusedPageRequestId?: number; citationListViewRequestId?: number; isLoading?: boolean; isLoadingMore?: boolean; @@ -77,9 +85,12 @@ type ChunkDisplayModeState = { export function ChunksPanel({ chunks = [], selectedSource = null, + selectedSourceView = null, selectedSourceFile = null, focusedChunkId = null, focusedChunkRequestId = 0, + focusedPageNumber = null, + focusedPageRequestId = 0, citationListViewRequestId = 0, isLoading = false, isLoadingMore = false, @@ -93,6 +104,8 @@ export function ChunksPanel({ analyticsContext, sourceCountSnapshot = 0, }: Partial = {}) { + const isPageAssetSource = + selectedSourceView?.documentPresentation?.kind === "page-assets"; const originalPreviewCacheKey = selectedSourceFile?.url ?? null; const isOriginalPreviewAvailable = sourceOriginalPreviewModel.canPreviewOriginalFile( @@ -240,14 +253,29 @@ export function ChunksPanel({ citationListViewRequestId ? "list" : chunkDisplayModeState.mode; - const headerTitle = focusedChunkId ? "Referenced Chunks" : "Parsed Chunks"; + const headerTitle = isPageAssetSource || visibleView === "original" + ? "Original File" + : focusedChunkId + ? "Referenced Chunks" + : "Parsed Chunks"; const shouldMountOriginalPreview = visibleView === "original" || (originalPreviewCacheKey !== null && mountedOriginalPreviewKey === originalPreviewCacheKey); const isTreeModeVisible = - visibleView === "parsed" && chunkDisplayMode === "tree"; - const headerSubtitle = visibleView === "original" ? ( + !isPageAssetSource && visibleView === "parsed" && chunkDisplayMode === "tree"; + const headerSubtitle = isPageAssetSource ? ( + selectedSource ? ( + <> + Showing page images for{" "} + + {selectedSource} + + + ) : ( + "Select a source to preview its page images." + ) + ) : visibleView === "original" ? ( selectedSource ? ( <> Showing the original file for{" "} @@ -299,14 +327,14 @@ export function ChunksPanel({

- {visibleView === "original" ? "Original File" : headerTitle} + {headerTitle}

{headerSubtitle}

- {visibleView === "parsed" && chunks.length > 0 ? ( + {!isPageAssetSource && visibleView === "parsed" && chunks.length > 0 ? (
) : null} - {hasOriginalView ? ( + {!isPageAssetSource && hasOriginalView ? (
+
+ ) : null} + {pageState.isLoading ? ( +
+ Loading page images... +
+ ) : null} +
+ ); +} + +function PageAssetImage({ + isFocused, + page, + refCallback, +}: { + readonly isFocused: boolean; + readonly page: SourcePageAssetView; + readonly refCallback: (element: HTMLDivElement | null) => void; +}): ReactNode { + const aspectRatio = + page.width && page.height ? `${page.width} / ${page.height}` : undefined; + + return ( +
+
+ Page {page.pageNumber} + {page.contentType} +
+
+ {`Page +
+
+ ); +} + +function getNextPageAssetPageIndex( + loadedPageIndexes: ReadonlySet, + totalPages: number, +): number | null { + for (let pageIndex = 1; pageIndex <= totalPages; pageIndex += 1) { + if (!loadedPageIndexes.has(pageIndex)) return pageIndex; + } + + return null; +} + +function EmptyPageAssets(): ReactNode { + return ( +
+ +

+ No page images available. +

+
+ ); +} + function truncateTreeLabel(value: string): string { const maxLength = 42; if (value.length <= maxLength) return value; diff --git a/src/components/source-row.test.ts b/src/components/source-row.test.ts index 427e09a..97c4bf9 100644 --- a/src/components/source-row.test.ts +++ b/src/components/source-row.test.ts @@ -87,6 +87,27 @@ describe("SourceRow", () => { ); }); + it("labels page-asset documents with page counts", () => { + render( + React.createElement(SourceRow, { + isArchiving: false, + isSelected: false, + onSelect: vi.fn(), + source: { + id: "source_1", + mimeType: "application/pdf", + title: "scan.pdf", + status: "ready", + chunkCount: 4, + documentPresentation: { kind: "page-assets", pageCount: 4 }, + }, + }), + ); + + expect(screen.getByText("Processed · 4 pages")).toBeTruthy(); + expect(screen.queryByText("Processed · 4 chunks")).toBeNull(); + }); + it("shows source archive loading locally", () => { render( React.createElement(SourceRow, { diff --git a/src/components/source-row.tsx b/src/components/source-row.tsx index 9fb87d3..473453f 100644 --- a/src/components/source-row.tsx +++ b/src/components/source-row.tsx @@ -197,6 +197,9 @@ function getReadySourceLabel(source: SourceView): string { function getReadySourceStatusText(source: SourceView): string { if (source.kind === "remote") return getReadySourceLabel(source); + if (source.documentPresentation?.kind === "page-assets") { + return `${getReadySourceLabel(source)} · ${source.documentPresentation.pageCount} pages`; + } if (typeof source.chunkCount !== "number") return getReadySourceLabel(source); return `${getReadySourceLabel(source)} · ${source.chunkCount} chunks`; } diff --git a/src/components/sources-panel.tsx b/src/components/sources-panel.tsx index 70a0285..d1950ef 100644 --- a/src/components/sources-panel.tsx +++ b/src/components/sources-panel.tsx @@ -351,6 +351,7 @@ function getSelectedSourcePage( } function getChunkTreeHref(source: SourceView): string | undefined { + if (source.documentPresentation?.kind === "page-assets") return undefined; return source.documentId ? `/inspect/${encodeURIComponent(source.documentId)}/chunks` : undefined; diff --git a/src/components/workspace-citation-focus.test.ts b/src/components/workspace-citation-focus.test.ts index baa5a1a..bfc4654 100644 --- a/src/components/workspace-citation-focus.test.ts +++ b/src/components/workspace-citation-focus.test.ts @@ -169,6 +169,49 @@ describe("useWorkspaceCitationFocus", () => { expect(result.current.pendingCitationId).toBeNull(); }); + it("focuses page-asset citations without fetching full chunks", async () => { + const fetchChunks = vi.fn(async () => [prefetchedChunk]); + const selectSource = vi.fn(); + const pageAssetSource: SourceView = { + ...readySource, + documentPresentation: { kind: "page-assets", pageCount: 8 }, + }; + const pageCitation: ChatCitationView = { + chunkType: "page", + score: 0.9, + pageCitationAssetUrl: "https://assets.example/page-000004.png", + pageCitationPageNumber: 4, + source: { + documentId: "document_1", + sourceFileName: "Contract.pdf", + sectionPath: "Page 4", + }, + }; + + const { result } = renderHook(() => + useWorkspaceCitationFocus({ + fetchChunks, + onSelectSource: selectSource, + selectedSourceId: null, + sources: [pageAssetSource], + }), + { wrapper: createSWRWrapper }, + ); + + await act(async () => { + await result.current.handleCitationClick(pageCitation, "message_1:0"); + }); + + expect(fetchChunks).not.toHaveBeenCalled(); + expect(selectSource).toHaveBeenCalledWith("source_1"); + expect(result.current.focusedChunk.chunkId).toBeNull(); + expect(result.current.focusedPage).toEqual({ + pageNumber: 4, + requestId: 1, + }); + expect(result.current.citationListViewRequestId).toBe(0); + }); + it("reuses cached chunks for a different source without refetching", async () => { const fetchChunks = vi.fn(async () => [prefetchedChunk]); const selectSource = vi.fn(); diff --git a/src/components/workspace-citation-focus.ts b/src/components/workspace-citation-focus.ts index 357cd4e..c233671 100644 --- a/src/components/workspace-citation-focus.ts +++ b/src/components/workspace-citation-focus.ts @@ -13,6 +13,11 @@ type FocusedChunkState = { readonly requestId: number } +type FocusedPageState = { + readonly pageNumber: number | null + readonly requestId: number +} + type PrefetchedChunksBySourceId = Readonly> type PrefetchedChunksUpdater = ( current: PrefetchedChunksBySourceId, @@ -30,6 +35,7 @@ type WorkspaceCitationFocusInput = { type WorkspaceCitationFocus = { readonly citationListViewRequestId: number readonly focusedChunk: FocusedChunkState + readonly focusedPage: FocusedPageState readonly handleCitationClick: ( citation: ChatCitationView, citationId: string, @@ -61,6 +67,10 @@ export function useWorkspaceCitationFocus({ chunkId: null, requestId: 0, }) + const [focusedPage, setFocusedPage] = useState({ + pageNumber: null, + requestId: 0, + }) const [pendingCitationId, setPendingCitationId] = useState( null, ) @@ -102,6 +112,12 @@ export function useWorkspaceCitationFocus({ }, [], ) + const requestPageFocus = useCallback((pageNumber: number | null): void => { + setFocusedPage((current) => ({ + pageNumber, + requestId: current.requestId + 1, + })) + }, []) const updatePrefetchedChunksBySourceId = useCallback( (updater: PrefetchedChunksUpdater): void => { @@ -122,10 +138,12 @@ export function useWorkspaceCitationFocus({ ) } requestChunkFocus(null) + requestPageFocus(null) }, [ onSelectSource, requestChunkFocus, + requestPageFocus, selectedSourceId, updatePrefetchedChunksBySourceId, ], @@ -192,6 +210,18 @@ export function useWorkspaceCitationFocus({ citation, ) if (!source) return + + if ( + workspaceCitationState.isPageAssetCitationTarget(source, citation) + ) { + if (selectedSourceId !== source.id) onSelectSource(source.id) + requestChunkFocus(null) + requestPageFocus( + workspaceCitationState.getCitationPageNumber(citation), + ) + return + } + setCitationListViewRequestId((current) => current + 1) const loadedChunkId = workspaceCitationState.getLoadedCitationChunkId({ @@ -261,6 +291,7 @@ export function useWorkspaceCitationFocus({ loadAllChunksForSource, onSelectSource, requestChunkFocus, + requestPageFocus, selectedChunks, selectedSourceId, sources, @@ -271,6 +302,7 @@ export function useWorkspaceCitationFocus({ return { citationListViewRequestId, focusedChunk, + focusedPage, handleCitationClick, handleLoadAllChunks, handleLoadMoreChunks, diff --git a/src/components/workspace-citation-state.ts b/src/components/workspace-citation-state.ts index d8331af..cf33dad 100644 --- a/src/components/workspace-citation-state.ts +++ b/src/components/workspace-citation-state.ts @@ -27,6 +27,11 @@ type WorkspaceCitationStateModule = { readonly hasExactCitationTargetHint: ( citation: ChatCitationView, ) => boolean + readonly isPageAssetCitationTarget: ( + source: SourceView, + citation: ChatCitationView, + ) => boolean + readonly getCitationPageNumber: (citation: ChatCitationView) => number | null readonly upsertPrefetchedChunks: ( current: PrefetchedChunksBySourceId, sourceId: string, @@ -79,6 +84,36 @@ function hasExactCitationTargetHint(citation: ChatCitationView): boolean { return true } +function isPageAssetCitationTarget( + source: SourceView, + citation: ChatCitationView, +): boolean { + return ( + source.documentPresentation?.kind === "page-assets" && + (getCitationPageNumber(citation) !== null || + typeof citation.pageCitationAssetUrl === "string") + ) +} + +function getCitationPageNumber(citation: ChatCitationView): number | null { + if ( + typeof citation.pageCitationPageNumber === "number" && + Number.isSafeInteger(citation.pageCitationPageNumber) && + citation.pageCitationPageNumber > 0 + ) { + return citation.pageCitationPageNumber + } + + const sectionPath = citation.source.sectionPath + if (typeof sectionPath !== "string") return null + + const match = /\bpage\s+(\d+)\b/i.exec(sectionPath) + if (!match) return null + + const pageNumber = Number.parseInt(match[1]!, 10) + return Number.isSafeInteger(pageNumber) && pageNumber > 0 ? pageNumber : null +} + function upsertPrefetchedChunks( current: PrefetchedChunksBySourceId, sourceId: string, @@ -120,7 +155,9 @@ function removePrefetchedChunks( export const workspaceCitationState: WorkspaceCitationStateModule = { findCitationSource, getLoadedCitationChunkId, + getCitationPageNumber, hasExactCitationTargetHint, + isPageAssetCitationTarget, upsertPrefetchedChunks, removePrefetchedChunks, } diff --git a/src/components/workspace-selected-chunks.test.ts b/src/components/workspace-selected-chunks.test.ts index 45db25a..b8675d4 100644 --- a/src/components/workspace-selected-chunks.test.ts +++ b/src/components/workspace-selected-chunks.test.ts @@ -161,6 +161,28 @@ describe("useWorkspaceSelectedChunks", () => { expect(result.current.selectedChunks).toEqual([]); }); + it("does not fetch chunk pages for page-asset sources", () => { + const pageAssetSource: SourceView = { + ...readySource, + documentPresentation: { kind: "page-assets", pageCount: 4 }, + }; + + const { result } = renderHook( + () => + useWorkspaceSelectedChunks({ + selectedSourceId: "source_1", + sources: [pageAssetSource], + prefetchedChunksBySourceId: {}, + }), + { wrapper: createSWRWrapper }, + ); + + expect(result.current.selectedSource?.id).toBe("source_1"); + expect(result.current.selectedChunks).toEqual([]); + expect(result.current.isSelectedChunksLoading).toBe(false); + expect(fetchChunkPageMock).not.toHaveBeenCalled(); + }); + it("requests a source refresh after loading an unlocalized remote source", async () => { const onRemoteSourceChunksLoaded = vi.fn(); const remoteSource: SourceView = { diff --git a/src/components/workspace-selected-chunks.ts b/src/components/workspace-selected-chunks.ts index f00ac68..0b920de 100644 --- a/src/components/workspace-selected-chunks.ts +++ b/src/components/workspace-selected-chunks.ts @@ -42,7 +42,9 @@ export function useWorkspaceSelectedChunks({ ? prefetchedChunksBySourceId[selectedSourceId] : undefined const selectedChunkSourceId = - selectedSource && selectedSource.status === "ready" + selectedSource && + selectedSource.status === "ready" && + selectedSource.documentPresentation?.kind !== "page-assets" ? selectedSource.id : null const { diff --git a/src/components/workspace-shell-layout.tsx b/src/components/workspace-shell-layout.tsx index 3b9e997..b9ccc6c 100644 --- a/src/components/workspace-shell-layout.tsx +++ b/src/components/workspace-shell-layout.tsx @@ -40,6 +40,11 @@ type FocusedChunkState = { readonly requestId: number } +type FocusedPageState = { + readonly pageNumber: number | null + readonly requestId: number +} + type WorkspaceShellUser = { readonly id: string readonly name: string | null @@ -66,6 +71,7 @@ export type WorkspaceShellLayoutProps = { readonly dashboardUrl?: string readonly desktopPanelWidths: Readonly readonly focusedChunk: FocusedChunkState + readonly focusedPage?: FocusedPageState readonly hasMessages: boolean readonly hasMoreSelectedChunks: boolean readonly contentView: ContentView @@ -146,6 +152,10 @@ export function WorkspaceShellLayout( props.desktopPanelWidths.chat <= workspaceShellState.desktopSidePanelCompactThreshold const isSourcesPanelNarrow = props.desktopPanelWidths.sources < 220 + const selectedSource = props.sources.find( + (source) => source.id === props.selectedSourceId, + ) + const focusedPage = props.focusedPage ?? { pageNumber: null, requestId: 0 } const handleDesktopLayoutRef = useCallback( (element: HTMLDivElement | null): void => { onDesktopLayoutElementChange(element) @@ -265,10 +275,13 @@ export function WorkspaceShellLayout( = {}): Source { describe("countChunksBySourceId", () => { it("counts ready source chunks from the document total", async () => { const listChunks = vi.fn(async () => ({ pagination: { total: 12 } })) + const readChunks = vi.fn(async () => ({ + chunks: [], + totalChunks: 0, + })) const mockClient = { documents: { listChunks }, + knowledge: { readChunks }, } as unknown as Knowhere const { countChunksBySourceId } = await import("./counts") @@ -54,6 +59,14 @@ describe("countChunksBySourceId", () => { ) expect(listChunks).toHaveBeenCalledTimes(1) + expect(readChunks).toHaveBeenCalledWith({ + documentId: "doc_ready", + revisionKey: "job_1", + chunkType: "page", + page: 1, + pageSize: 1, + assetUrlPolicy: "durable", + }) expect(listChunks).toHaveBeenCalledWith("doc_ready", { page: 1, pageSize: 1, @@ -65,8 +78,10 @@ describe("countChunksBySourceId", () => { const listChunks = vi.fn(async () => { throw new Error("temporary outage") }) + const readChunks = vi.fn(async () => ({ chunks: [], totalChunks: 0 })) const mockClient = { documents: { listChunks }, + knowledge: { readChunks }, } as unknown as Knowhere const { countChunksBySourceId } = await import("./counts") @@ -86,8 +101,10 @@ describe("countChunksBySourceId", () => { const listChunks = vi.fn().mockResolvedValue({ pagination: { total: 70 }, }) + const readChunks = vi.fn(async () => ({ chunks: [], totalChunks: 0 })) const mockClient = { documents: { listChunks }, + knowledge: { readChunks }, } as unknown as Knowhere const { countChunksBySourceId } = await import("./counts") @@ -106,6 +123,50 @@ describe("countChunksBySourceId", () => { ) expect(listChunks).not.toHaveBeenCalled() + expect(readChunks).not.toHaveBeenCalled() expect(counts.size).toBe(0) }) }) + +describe("sourceViewOptionsBySourceId", () => { + it("detects page-asset documents from SDK page chunks", async () => { + const listChunks = vi.fn(async () => ({ pagination: { total: 12 } })) + const readChunks = vi.fn(async () => ({ + chunks: [ + { + chunkId: "page_1", + chunkType: "page", + metadata: { + pageAssets: [ + { + pageNum: 1, + artifactRef: "pages/page-000001.png", + assetUrl: "https://assets.example/page-000001.png", + }, + ], + }, + }, + ], + totalChunks: 4, + })) + const mockClient = { + documents: { listChunks }, + knowledge: { readChunks }, + } as unknown as Knowhere + + const { sourceViewOptionsBySourceId } = await import("./counts") + + const options = await Effect.runPromise( + sourceViewOptionsBySourceId( + [makeSource({ id: "ready", knowhereDocumentId: "doc_ready" })], + mockClient, + ), + ) + + expect(options.get("ready")).toEqual({ + chunkCount: 4, + documentPresentation: { kind: "page-assets", pageCount: 4 }, + }) + expect(listChunks).not.toHaveBeenCalled() + }) +}) diff --git a/src/domains/sources/counts.ts b/src/domains/sources/counts.ts index e6e0c9b..110b3ae 100644 --- a/src/domains/sources/counts.ts +++ b/src/domains/sources/counts.ts @@ -2,8 +2,15 @@ import "server-only" import { Effect } from "effect" import type Knowhere from "@ontos-ai/knowhere-sdk" +import type { KnowledgeReadChunk } from "@ontos-ai/knowhere-sdk" import type { Source } from "@/infrastructure/db/schema" +import type { SourceDocumentPresentation } from "./types" + +type PageAssetDocumentPresentation = Extract< + SourceDocumentPresentation, + { readonly kind: "page-assets" } +> type CountChunksClient = { readonly documents: { @@ -14,9 +21,27 @@ type CountChunksClient = { readonly pagination?: { readonly total?: number } }> } + readonly knowledge: { + readChunks(params: { + readonly documentId: string + readonly revisionKey?: string + readonly chunkType: "page" + readonly page: number + readonly pageSize: number + readonly assetUrlPolicy: "durable" + }): Promise<{ + readonly chunks: readonly KnowledgeReadChunk[] + readonly totalChunks?: number + }> + } } -export const countChunksBySourceId = ( +export type SourceViewOptions = { + readonly chunkCount?: number + readonly documentPresentation?: SourceDocumentPresentation +} + +export const sourceViewOptionsBySourceId = ( sources: readonly Source[], client: Knowhere, ) => @@ -28,15 +53,19 @@ export const countChunksBySourceId = ( source.status === "ready" && source.knowhereDocumentId, ) - if (readySources.length === 0) return new Map() + if (readySources.length === 0) return new Map() const entries = yield* Effect.all( readySources.map((source) => Effect.gen(function* () { - const total = yield* Effect.tryPromise(() => - loadSourceChunkCount(countClient, source.knowhereDocumentId!), - ).pipe(Effect.catchAll(() => Effect.succeed(undefined))) - return [source.id, total] as const + const options = yield* Effect.tryPromise(() => + loadSourceViewOptions(countClient, source), + ).pipe( + Effect.catchAll(() => + Effect.sync((): SourceViewOptions | undefined => undefined), + ), + ) + return [source.id, options] as const }), ), { concurrency: "unbounded" }, @@ -44,26 +73,81 @@ export const countChunksBySourceId = ( return new Map( entries.filter( - (entry): entry is readonly [string, number] => - typeof entry[1] === "number", + (entry): entry is readonly [string, SourceViewOptions] => + entry[1] !== undefined, ), ) }) -export const sourceViewOptionsBySourceId = ( +export const countChunksBySourceId = ( sources: readonly Source[], client: Knowhere, ) => Effect.gen(function* () { - const counts = yield* countChunksBySourceId(sources, client) - return new Map( - sources.map((source) => [ - source.id, - { chunkCount: counts.get(source.id) }, - ]), - ) + const sourceOptions = yield* sourceViewOptionsBySourceId(sources, client) + const countEntries: [string, number][] = [] + for (const [sourceId, options] of sourceOptions.entries()) { + if (typeof options.chunkCount === "number") { + countEntries.push([sourceId, options.chunkCount]) + } + } + return new Map(countEntries) }) +async function loadSourceViewOptions( + client: CountChunksClient, + source: Source, +): Promise { + const documentId = source.knowhereDocumentId + if (!documentId) return undefined + + const pagePresentation = await loadPageAssetPresentation( + client, + documentId, + source.knowhereJobId, + ) + if (pagePresentation) { + return { + chunkCount: pagePresentation.pageCount, + documentPresentation: pagePresentation, + } + } + + const chunkCount = await loadSourceChunkCount(client, documentId) + return typeof chunkCount === "number" ? { chunkCount } : undefined +} + +async function loadPageAssetPresentation( + client: CountChunksClient, + documentId: string, + revisionKey: string | null, +): Promise { + try { + const response = await client.knowledge.readChunks({ + documentId, + ...(revisionKey ? { revisionKey } : {}), + chunkType: "page", + page: 1, + pageSize: 1, + assetUrlPolicy: "durable", + }) + const firstChunk = response.chunks[0] + if (!firstChunk || firstChunk.chunkType !== "page") return undefined + if (!hasUsablePageAssets(firstChunk.metadata.pageAssets)) return undefined + + const pageCount = + typeof response.totalChunks === "number" && + Number.isFinite(response.totalChunks) && + response.totalChunks > 0 + ? response.totalChunks + : 1 + + return { kind: "page-assets", pageCount } + } catch { + return undefined + } +} + async function loadSourceChunkCount( client: CountChunksClient, documentId: string, @@ -75,3 +159,25 @@ async function loadSourceChunkCount( const total = response.pagination?.total return typeof total === "number" && Number.isFinite(total) ? total : undefined } + +function hasUsablePageAssets(value: unknown): boolean { + if (!Array.isArray(value)) return false + + return value.some((item) => { + if (!isRecord(item)) return false + const pageNum = item.pageNum + const artifactRef = item.artifactRef + const assetUrl = item.assetUrl + return ( + typeof pageNum === "number" && + Number.isSafeInteger(pageNum) && + pageNum > 0 && + ((typeof artifactRef === "string" && artifactRef.trim().length > 0) || + (typeof assetUrl === "string" && assetUrl.trim().length > 0)) + ) + }) +} + +function isRecord(value: unknown): value is Readonly> { + return typeof value === "object" && value !== null +} diff --git a/src/domains/sources/page-assets.ts b/src/domains/sources/page-assets.ts new file mode 100644 index 0000000..55c108b --- /dev/null +++ b/src/domains/sources/page-assets.ts @@ -0,0 +1,104 @@ +import "server-only" + +import type { Knowledge } from "@ontos-ai/knowhere-sdk" + +import type { ChunkPageParams } from "@/domains/chunks" +import type { SourcePageAssetView } from "./route-types" + +type ReadableSource = { + readonly documentId: string + readonly revisionKey?: string | null +} + +export type SourcePageAssetsPage = { + readonly pages: readonly SourcePageAssetView[] + readonly pagination: { + readonly page: number + readonly pageSize: number + readonly total: number + readonly totalPages: number + } +} + +export async function readSourcePageAssets(input: { + readonly knowledge: Knowledge + readonly source: ReadableSource + readonly params: ChunkPageParams +}): Promise { + const response = await input.knowledge.readChunks({ + documentId: input.source.documentId, + ...(input.source.revisionKey ? { revisionKey: input.source.revisionKey } : {}), + chunkType: "page", + page: input.params.page, + pageSize: input.params.pageSize, + assetUrlPolicy: "durable", + }) + const pages = response.chunks.flatMap((chunk): SourcePageAssetView[] => + readPageAssetViews(chunk.metadata.pageAssets, chunk.assetUrl), + ) + + return { + pages, + pagination: { + page: response.page ?? input.params.page, + pageSize: response.pageSize ?? input.params.pageSize, + total: response.totalChunks ?? pages.length, + totalPages: + response.totalPages ?? + Math.max(1, Math.ceil(pages.length / input.params.pageSize)), + }, + } +} + +function readPageAssetViews( + value: unknown, + fallbackAssetUrl?: string, +): SourcePageAssetView[] { + if (!Array.isArray(value)) return [] + + return value.flatMap((item): SourcePageAssetView[] => { + if (!isRecord(item)) return [] + const pageNumber = getPositiveInteger(item.pageNum) + const assetUrl = getTrimmedString(item.assetUrl) ?? fallbackAssetUrl + const contentType = getTrimmedString(item.contentType) + if (!pageNumber || !assetUrl || !contentType) return [] + + return [ + { + pageNumber, + assetUrl, + contentType, + ...(getPositiveNumber(item.width) !== undefined + ? { width: getPositiveNumber(item.width) } + : {}), + ...(getPositiveNumber(item.height) !== undefined + ? { height: getPositiveNumber(item.height) } + : {}), + }, + ] + }) +} + +function getTrimmedString(value: unknown): string | undefined { + if (typeof value !== "string") return undefined + const trimmed = value.trim() + return trimmed.length > 0 ? trimmed : undefined +} + +function getPositiveInteger(value: unknown): number | undefined { + return typeof value === "number" && + Number.isSafeInteger(value) && + value > 0 + ? value + : undefined +} + +function getPositiveNumber(value: unknown): number | undefined { + return typeof value === "number" && Number.isFinite(value) && value > 0 + ? value + : undefined +} + +function isRecord(value: unknown): value is Readonly> { + return typeof value === "object" && value !== null +} diff --git a/src/domains/sources/route-page-assets.ts b/src/domains/sources/route-page-assets.ts new file mode 100644 index 0000000..e356702 --- /dev/null +++ b/src/domains/sources/route-page-assets.ts @@ -0,0 +1,161 @@ +import { Effect } from "effect" + +import { readSourcePageAssets } from "./page-assets" +import { routeResult } from "@/lib/route-result" +import { + decodeRemoteSourceId, + findRemoteLibraryDocumentBySourceId, +} from "./remote-library" +import { + getClientForWorkspace, + getKnowledgeForSource, +} from "./route-dependencies" +import { sourceRowRepository } from "./source-row-repository" +import type { + JsonRouteResult, + LoadSourcePageAssetsInput, + SourcePageAssetsBody, + SourceRouteServiceDependencies, +} from "./route-types" + +type RoutePageAssetsDependencies = Pick< + SourceRouteServiceDependencies, + | "ensureApiKeyForWorkspace" + | "ensureWorkspace" + | "getCurrentUser" + | "makeKnowhereClient" + | "sourceService" +> + +type RoutePageAssets = { + readonly loadSourcePageAssets: ( + input: LoadSourcePageAssetsInput, + ) => Promise> +} + +function createRoutePageAssets( + deps: RoutePageAssetsDependencies, +): RoutePageAssets { + return { + loadSourcePageAssets: (input: LoadSourcePageAssetsInput) => + Effect.runPromise(loadSourcePageAssetsEffect(input, deps)), + } +} + +const loadSourcePageAssetsEffect = ( + input: LoadSourcePageAssetsInput, + deps: RoutePageAssetsDependencies, +) => + Effect.gen(function* () { + if (!sourceRowRepository.isWorkspaceSourceId(input.sourceId)) { + const remoteResult = yield* loadRemotePageAssetsEffect(input, deps) + return remoteResult ?? sourceNotFound() + } + + const user = yield* Effect.tryPromise(() => deps.getCurrentUser()) + if (!user) return sourceNotFound() + + const workspace = yield* Effect.tryPromise(() => + deps.ensureWorkspace(user.id), + ) + const source = yield* Effect.tryPromise(() => + deps.sourceService.findInWorkspace(workspace.id, input.sourceId), + ) + if (!source) return sourceNotFound() + if (source.status !== "ready" || !source.knowhereDocumentId) { + return sourceNotReady() + } + + const documentId = source.knowhereDocumentId + const apiKey = yield* Effect.tryPromise(() => + deps.ensureApiKeyForWorkspace(workspace.id, input.cookieHeader), + ) + const knowledge = getKnowledgeForSource({ + apiKey, + workspaceId: workspace.id, + sourceId: source.id, + documentId, + revisionKey: source.knowhereJobId, + }) + const pageAssets = yield* Effect.tryPromise(() => + readSourcePageAssets({ + knowledge, + source: { + documentId, + revisionKey: source.knowhereJobId, + }, + params: input.pageParams, + }), + ) + + return routeResult.ok(pageAssets) + }) + +const loadRemotePageAssetsEffect = ( + input: LoadSourcePageAssetsInput, + deps: RoutePageAssetsDependencies, +) => + Effect.gen(function* () { + if (!decodeRemoteSourceId(input.sourceId)) return null + + const user = yield* Effect.tryPromise(() => deps.getCurrentUser()) + if (!user) return null + + const workspace = yield* Effect.tryPromise(() => + deps.ensureWorkspace(user.id), + ) + const apiKey = yield* Effect.tryPromise(() => + deps.ensureApiKeyForWorkspace(workspace.id, input.cookieHeader), + ) + const client = yield* Effect.tryPromise(() => + getClientForWorkspace(workspace.id, input.cookieHeader, deps), + ) + const remoteDocument = yield* findRemoteLibraryDocumentBySourceId({ + sourceId: input.sourceId, + workspace, + client, + localSources: [], + }) + if (!remoteDocument) return null + + const source = yield* Effect.tryPromise(() => + deps.sourceService.localizeRemoteDocument(workspace.id, { + documentId: remoteDocument.documentId, + namespace: remoteDocument.namespace, + status: remoteDocument.status, + title: remoteDocument.title, + mimeType: remoteDocument.mimeType, + sizeBytes: remoteDocument.sizeBytes, + revisionKey: remoteDocument.revisionKey ?? null, + }), + ) + const documentId = source.knowhereDocumentId ?? remoteDocument.documentId + const revisionKey = + source.knowhereJobId ?? remoteDocument.revisionKey ?? null + const knowledge = getKnowledgeForSource({ + apiKey, + workspaceId: workspace.id, + sourceId: source.id, + documentId, + revisionKey, + }) + const pageAssets = yield* Effect.tryPromise(() => + readSourcePageAssets({ + knowledge, + source: { documentId, revisionKey }, + params: input.pageParams, + }), + ) + + return routeResult.ok(pageAssets) + }) + +function sourceNotFound(): JsonRouteResult<{ readonly message: string }> { + return routeResult.error(404, "Source not found.") +} + +function sourceNotReady(): JsonRouteResult<{ readonly message: string }> { + return routeResult.error(409, "Source is not ready.") +} + +export { createRoutePageAssets } diff --git a/src/domains/sources/route-service.ts b/src/domains/sources/route-service.ts index 9a19bef..7e3b85f 100644 --- a/src/domains/sources/route-service.ts +++ b/src/domains/sources/route-service.ts @@ -4,12 +4,14 @@ import { createRouteArchive } from "./route-archive" import { createRouteChunks } from "./route-chunks" import { createSourceRouteDependencies } from "./route-dependencies" import { createRouteListing } from "./route-listing" +import { createRoutePageAssets } from "./route-page-assets" import { createRouteRetry } from "./route-retry" import { createRouteUpload } from "./route-upload" import type { ArchiveSourceInput, ListSourcesInput, LoadSourceChunksInput, + LoadSourcePageAssetsInput, RetrySourceInput, SourceRouteService, SourceRouteServiceOverrides, @@ -25,6 +27,7 @@ export function createSourceRouteService( const archive = createRouteArchive(deps) const retry = createRouteRetry(deps) const chunks = createRouteChunks(deps) + const pageAssets = createRoutePageAssets(deps) return { listSources: (input: ListSourcesInput) => listing.listSources(input), @@ -33,5 +36,7 @@ export function createSourceRouteService( retrySource: (input: RetrySourceInput) => retry.retrySource(input), loadSourceChunks: (input: LoadSourceChunksInput) => chunks.loadSourceChunks(input), + loadSourcePageAssets: (input: LoadSourcePageAssetsInput) => + pageAssets.loadSourcePageAssets(input), } } diff --git a/src/domains/sources/route-types.ts b/src/domains/sources/route-types.ts index 96bd80c..48401a0 100644 --- a/src/domains/sources/route-types.ts +++ b/src/domains/sources/route-types.ts @@ -4,7 +4,11 @@ import type { ChunkPageParams, } from "@/domains/chunks" import type { ParsedChunkView } from "@/domains/chunks/types" -import type { SourceStatus, SourceView } from "@/domains/sources/types" +import type { + SourcePageAssetView, + SourceStatus, + SourceView, +} from "@/domains/sources/types" import type { AuthUser } from "@/infrastructure/auth" import type { Source, Workspace } from "@/infrastructure/db/schema" import type { @@ -106,6 +110,20 @@ type SourceChunksBody = readonly message: string } +type SourcePageAssetsBody = + | { + readonly pages: readonly SourcePageAssetView[] + readonly pagination: { + readonly page: number + readonly pageSize: number + readonly total: number + readonly totalPages: number + } + } + | { + readonly message: string + } + type ListSourcesInput = { readonly cookieHeader: string } @@ -133,6 +151,12 @@ type LoadSourceChunksInput = { readonly pageParams: ChunkPageParams } +type LoadSourcePageAssetsInput = { + readonly cookieHeader: string + readonly sourceId: string + readonly pageParams: ChunkPageParams +} + type SourceRouteService = { readonly listSources: ( input: ListSourcesInput, @@ -149,6 +173,9 @@ type SourceRouteService = { readonly loadSourceChunks: ( input: LoadSourceChunksInput, ) => Promise> + readonly loadSourcePageAssets: ( + input: LoadSourcePageAssetsInput, + ) => Promise> } type SourceWorkflowService = { @@ -251,9 +278,12 @@ export type { ListSourcesBody, ListSourcesInput, LoadSourceChunksInput, + LoadSourcePageAssetsInput, RetrySourceBody, RetrySourceInput, SourceChunksBody, + SourcePageAssetsBody, + SourcePageAssetView, SourceRouteDemoApi, SourceRouteKnowhereClient, SourceRouteService, diff --git a/src/domains/sources/types.ts b/src/domains/sources/types.ts index 8479f87..e0ff9be 100644 --- a/src/domains/sources/types.ts +++ b/src/domains/sources/types.ts @@ -16,6 +16,18 @@ export type SourceOfficialLibraryView = { readonly sourceUrl: string } +export type SourceDocumentPresentation = + | { readonly kind: "parsed-chunks" } + | { readonly kind: "page-assets"; readonly pageCount: number } + +export type SourcePageAssetView = { + readonly pageNumber: number + readonly assetUrl: string + readonly contentType: string + readonly width?: number + readonly height?: number +} + export type OfficialLibrarySourceView = { readonly librarySourceId: string readonly categoryId: string @@ -50,6 +62,11 @@ export type SourceView = { readonly officialLibrary?: SourceOfficialLibraryView /** Count from the Notebook parsed snapshot manifest when available. */ readonly chunkCount?: number + /** + * Preferred source content presentation. Missing values are treated as + * parsed chunks so older views and cached responses remain compatible. + */ + readonly documentPresentation?: SourceDocumentPresentation /** User opt-out for this query session. Drives excludeDocumentIds. */ readonly excludedFromQuery?: boolean } diff --git a/src/domains/sources/view.test.ts b/src/domains/sources/view.test.ts index cc5b223..0069f6a 100644 --- a/src/domains/sources/view.test.ts +++ b/src/domains/sources/view.test.ts @@ -121,6 +121,18 @@ describe("toSourceView", () => { }); }); + it("includes document presentation options when provided", () => { + expect( + toSourceView(makeSource(), { + chunkCount: 4, + documentPresentation: { kind: "page-assets", pageCount: 4 }, + }), + ).toMatchObject({ + chunkCount: 4, + documentPresentation: { kind: "page-assets", pageCount: 4 }, + }); + }); + it("does not expose legacy demo original proxy routes", () => { const view = toSourceView( makeSource({ diff --git a/src/domains/sources/view.ts b/src/domains/sources/view.ts index 4096a81..5c6fe50 100644 --- a/src/domains/sources/view.ts +++ b/src/domains/sources/view.ts @@ -3,6 +3,7 @@ import { Schema } from "effect"; import type { Source } from "@/infrastructure/db/schema"; import type { SourceView } from "@/domains/sources/types"; import { sourceFailureMessage } from "./failure-message"; +import type { SourceDocumentPresentation } from "./types"; const SourceStatus = Schema.Literal( "uploading", @@ -13,7 +14,10 @@ const SourceStatus = Schema.Literal( export function toSourceView( source: Source, - options: { chunkCount?: number } = {}, + options: { + readonly chunkCount?: number + readonly documentPresentation?: SourceDocumentPresentation + } = {}, ): SourceView { const originalFile = getSourceOriginalFile(source) const status = toSourceStatus(source.status) @@ -35,6 +39,9 @@ export function toSourceView( ...(options.chunkCount !== undefined ? { chunkCount: options.chunkCount } : {}), + ...(options.documentPresentation !== undefined + ? { documentPresentation: options.documentPresentation } + : {}), }; } diff --git a/src/domains/workspace/client.ts b/src/domains/workspace/client.ts index a650b6f..efd8e32 100644 --- a/src/domains/workspace/client.ts +++ b/src/domains/workspace/client.ts @@ -4,7 +4,10 @@ import type { ChatThreadView, } from "@/domains/chat/types" import type { ParsedChunkView } from "@/domains/chunks/types" -import type { SourceView } from "@/domains/sources/types" +import type { + SourcePageAssetView, + SourceView, +} from "@/domains/sources/types" import { workspaceRouteClient } from "./route-client" const workspaceClientKeys = { @@ -34,6 +37,17 @@ type SourceChunksResponse = { } } +type SourcePageAssetsResponse = { + pages?: SourcePageAssetView[] + message?: string + pagination?: { + page: number + pageSize: number + total: number + totalPages: number + } +} + type ChatThreadResponse = { thread?: ChatThreadView messages?: ChatMessageView[] @@ -91,6 +105,7 @@ export const workspaceClient = { keys: workspaceClientKeys, fetchChunks, fetchChunkPage, + fetchPageAssetPage, fetchSources, fetchChatThreads, fetchChatThread, @@ -136,6 +151,25 @@ async function fetchChunkPage( } } +async function fetchPageAssetPage( + sourceId: string, + page: number, +): Promise { + const searchParams = new URLSearchParams({ + page: String(page), + pageSize: String(workspaceClientConfig.sourceChunkPageSize), + }) + const body = await workspaceRouteClient.getJson( + `/api/sources/${encodeURIComponent(sourceId)}/page-assets?${searchParams.toString()}`, + ) + + return { + pages: Array.isArray(body.pages) ? body.pages : [], + ...(typeof body.message === "string" ? { message: body.message } : {}), + pagination: body.pagination, + } +} + async function fetchSources(): Promise { const body = await workspaceRouteClient.getJson( workspaceClientKeys.sources, From d955a8104398fc7af07dbf9f26f685cbf9d54efb Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 7 Jul 2026 15:00:17 +0800 Subject: [PATCH 02/14] Tolerate missing remote display reads --- .../sources/[sourceId]/chunks/route.test.ts | 50 +++++++++++++ .../[sourceId]/page-assets/route.test.ts | 49 +++++++++++++ src/components/chunks-panel.test.ts | 38 ++++++++++ src/components/chunks-panel.tsx | 25 +++++++ .../workspace-selected-chunks.test.ts | 33 +++++++++ src/components/workspace-selected-chunks.ts | 7 +- .../sources/display-read-unavailable.ts | 47 ++++++++++++ src/domains/sources/route-chunks.ts | 72 ++++++++++++++++--- src/domains/sources/route-page-assets.ts | 57 +++++++++++++-- src/domains/sources/route-types.ts | 22 ++++++ src/domains/workspace/client.test.ts | 60 ++++++++++++++++ src/domains/workspace/client.ts | 8 ++- 12 files changed, 451 insertions(+), 17 deletions(-) create mode 100644 src/domains/sources/display-read-unavailable.ts diff --git a/src/app/api/sources/[sourceId]/chunks/route.test.ts b/src/app/api/sources/[sourceId]/chunks/route.test.ts index 85fb1bd..d13677d 100644 --- a/src/app/api/sources/[sourceId]/chunks/route.test.ts +++ b/src/app/api/sources/[sourceId]/chunks/route.test.ts @@ -548,6 +548,56 @@ describe("GET /api/sources/[sourceId]/chunks", () => { expect(mocks.readChunks).not.toHaveBeenCalled() }) + it("marks workspace chunks unavailable when the parsed document is missing remotely", async () => { + const notFoundError = new Error("Document not found") + notFoundError.name = "NotFoundError" + mocks.getCurrentUser.mockResolvedValue({ + id: "user_1", + email: null, + name: null, + }) + mocks.ensureWorkspace.mockResolvedValue({ + id: "workspace_1", + userId: "user_1", + namespace: "notebook-workspace_1", + createdAt: new Date("2026-05-10T00:00:00.000Z"), + }) + mocks.findSourceInWorkspace.mockResolvedValue( + makeReadySource({ + id: "00000000-0000-0000-0000-000000000002", + knowhereJobId: "job_1", + knowhereDocumentId: "doc_missing", + }), + ) + mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") + mocks.readChunks.mockRejectedValue(notFoundError) + + const response = await GET( + new NextRequest( + "http://localhost:3001/api/sources/00000000-0000-0000-0000-000000000002/chunks?page=1&pageSize=1", + ), + { + params: Promise.resolve({ + sourceId: "00000000-0000-0000-0000-000000000002", + }), + }, + ) + + await expect(response.json()).resolves.toEqual({ + chunks: [], + pagination: { + page: 1, + pageSize: 1, + total: 0, + totalPages: 0, + }, + message: + "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", + isUnavailable: true, + }) + expect(response.status).toBe(200) + }) + it("materializes a remote source id and reads chunks through the SDK", async () => { const knowhereClient = { documents: { diff --git a/src/app/api/sources/[sourceId]/page-assets/route.test.ts b/src/app/api/sources/[sourceId]/page-assets/route.test.ts index 4bf5c5e..073a20e 100644 --- a/src/app/api/sources/[sourceId]/page-assets/route.test.ts +++ b/src/app/api/sources/[sourceId]/page-assets/route.test.ts @@ -244,6 +244,55 @@ describe("GET /api/sources/[sourceId]/page-assets", () => { }) expect(response.status).toBe(200) }) + + it("marks page assets unavailable when the parsed document is missing remotely", async () => { + const notFoundError = new Error("Document not found") + notFoundError.name = "NotFoundError" + mocks.getCurrentUser.mockResolvedValue({ + id: "user_1", + email: null, + name: null, + }) + mocks.ensureWorkspace.mockResolvedValue({ + id: "workspace_1", + userId: "user_1", + namespace: "notebook-workspace_1", + createdAt: new Date("2026-05-10T00:00:00.000Z"), + }) + mocks.findSourceInWorkspace.mockResolvedValue( + makeReadySource({ + id: "00000000-0000-0000-0000-000000000002", + knowhereDocumentId: "doc_missing", + }), + ) + mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") + mocks.readChunks.mockRejectedValue(notFoundError) + + const response = await GET( + new NextRequest( + "http://localhost:3001/api/sources/00000000-0000-0000-0000-000000000002/page-assets?page=1&pageSize=1", + ), + { + params: Promise.resolve({ + sourceId: "00000000-0000-0000-0000-000000000002", + }), + }, + ) + + await expect(response.json()).resolves.toEqual({ + pages: [], + pagination: { + page: 1, + pageSize: 1, + total: 0, + totalPages: 0, + }, + message: + "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", + isUnavailable: true, + }) + expect(response.status).toBe(200) + }) }) function makeReadySource(overrides: Record) { diff --git a/src/components/chunks-panel.test.ts b/src/components/chunks-panel.test.ts index 9e5bb4b..238ca0c 100644 --- a/src/components/chunks-panel.test.ts +++ b/src/components/chunks-panel.test.ts @@ -181,6 +181,44 @@ describe("ChunksPanel", () => { ); }); + it("shows page asset unavailable messages", async () => { + fetchPageAssetPageMock.mockResolvedValue({ + pages: [], + isUnavailable: true, + message: + "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", + pagination: { + page: 1, + pageSize: 50, + total: 0, + totalPages: 0, + }, + }); + + render( + React.createElement(C, { + chunks: [], + selectedSource: "report.pdf", + selectedSourceView: { + id: "source_1", + title: "report.pdf", + mimeType: "application/pdf", + status: "ready", + documentPresentation: { kind: "page-assets", pageCount: 4 }, + }, + }), + ); + + await waitFor(() => + expect( + screen.getByText( + "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", + ), + ).toBeTruthy(), + ); + expect(screen.queryByText("No page images available.")).toBeNull(); + }); + it("defaults parsed chunks into a section tree view", () => { render( React.createElement(C, { diff --git a/src/components/chunks-panel.tsx b/src/components/chunks-panel.tsx index 5561c4e..70f6448 100644 --- a/src/components/chunks-panel.tsx +++ b/src/components/chunks-panel.tsx @@ -397,6 +397,8 @@ export function ChunksPanel({ /> ) : isLoading ? ( + ) : chunks.length === 0 && processingMessage ? ( + ) : chunks.length === 0 ? ( selectedSource ? ( @@ -1036,6 +1038,7 @@ type PageAssetPageState = { readonly pages: readonly SourcePageAssetView[]; readonly loadedPageIndexes: ReadonlySet; readonly isLoading: boolean; + readonly message: string | null; readonly totalPages: number; }; @@ -1054,6 +1057,7 @@ function PageAssetDocumentViewer({ pages: [], loadedPageIndexes: new Set(), isLoading: false, + message: null, totalPages: Math.max( 1, Math.ceil( @@ -1092,6 +1096,7 @@ function PageAssetDocumentViewer({ ), loadedPageIndexes, isLoading: false, + message: response.message ?? null, totalPages: response.pagination?.totalPages ?? current.totalPages, }; @@ -1115,6 +1120,7 @@ function PageAssetDocumentViewer({ pages: [], loadedPageIndexes: new Set(), isLoading: false, + message: null, totalPages: Math.max( 1, Math.ceil( @@ -1158,6 +1164,10 @@ function PageAssetDocumentViewer({ return ; } + if (pageState.pages.length === 0 && pageState.message) { + return ; + } + if (pageState.pages.length === 0) { return ; } @@ -1264,6 +1274,21 @@ function EmptyPageAssets(): ReactNode { ); } +function UnavailableSourceMessage({ + message, +}: { + readonly message: string; +}): ReactNode { + return ( +
+ +

+ {message} +

+
+ ); +} + function truncateTreeLabel(value: string): string { const maxLength = 42; if (value.length <= maxLength) return value; diff --git a/src/components/workspace-selected-chunks.test.ts b/src/components/workspace-selected-chunks.test.ts index b8675d4..b57a788 100644 --- a/src/components/workspace-selected-chunks.test.ts +++ b/src/components/workspace-selected-chunks.test.ts @@ -161,6 +161,39 @@ describe("useWorkspaceSelectedChunks", () => { expect(result.current.selectedChunks).toEqual([]); }); + it("surfaces unavailable chunk messages without a loading state", async () => { + fetchChunkPageMock.mockResolvedValue({ + chunks: [], + isUnavailable: true, + message: + "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", + pagination: { + page: 1, + pageSize: 50, + total: 0, + totalPages: 0, + }, + }); + + const { result } = renderHook( + () => + useWorkspaceSelectedChunks({ + selectedSourceId: "source_1", + sources: [readySource], + prefetchedChunksBySourceId: {}, + }), + { wrapper: createSWRWrapper }, + ); + + await waitFor(() => + expect(result.current.selectedChunksMessage).toBe( + "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", + ), + ); + expect(result.current.isSelectedChunksLoading).toBe(false); + expect(result.current.selectedChunks).toEqual([]); + }); + it("does not fetch chunk pages for page-asset sources", () => { const pageAssetSource: SourceView = { ...readySource, diff --git a/src/components/workspace-selected-chunks.ts b/src/components/workspace-selected-chunks.ts index 0b920de..f1a1283 100644 --- a/src/components/workspace-selected-chunks.ts +++ b/src/components/workspace-selected-chunks.ts @@ -68,6 +68,8 @@ export function useWorkspaceSelectedChunks({ }, ) const selectedChunksMessage = getSelectedChunksMessage(selectedChunkPages) + const hasProcessingSelectedChunkPage = + hasProcessingChunkPage(selectedChunkPages) const pagedSelectedChunks = useMemo( () => resolveChunkConnectionTargets( @@ -99,7 +101,7 @@ export function useWorkspaceSelectedChunks({ typeof selectedChunkPages[selectedChunkPageCount - 1] === "undefined", ) const isSelectedChunksLoading = - Boolean(selectedChunksMessage) || + hasProcessingSelectedChunkPage || (selectedChunkSourceId !== null && !prefetchedSelectedChunks && !selectedChunkPages && @@ -150,7 +152,8 @@ function getSelectedChunksMessage( ): string | null { const page = pages?.find( (candidate) => - candidate.isProcessing && typeof candidate.message === "string", + (candidate.isProcessing || candidate.isUnavailable) && + typeof candidate.message === "string", ) return page?.message ?? null } diff --git a/src/domains/sources/display-read-unavailable.ts b/src/domains/sources/display-read-unavailable.ts new file mode 100644 index 0000000..302cd85 --- /dev/null +++ b/src/domains/sources/display-read-unavailable.ts @@ -0,0 +1,47 @@ +const sourceUnavailableMessage: string = + "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere." + +function isDisplayReadUnavailableError(error: unknown): boolean { + const details = readErrorDetails(error) + if ( + details.name === "NotFoundError" && + details.message.toLowerCase().includes("document not found") + ) { + return true + } + + const causeDetails = readErrorDetails(details.cause) + return ( + causeDetails.name === "NotFoundError" && + causeDetails.message.toLowerCase().includes("document not found") + ) +} + +function readErrorDetails(error: unknown): { + readonly name: string + readonly message: string + readonly cause: unknown +} { + if (!isRecord(error)) { + return { + name: "", + message: String(error), + cause: undefined, + } + } + + return { + name: typeof error.name === "string" ? error.name : "", + message: typeof error.message === "string" ? error.message : "", + cause: error.cause, + } +} + +function isRecord(value: unknown): value is Readonly> { + return typeof value === "object" && value !== null +} + +export const displayReadUnavailable = { + isError: isDisplayReadUnavailableError, + message: sourceUnavailableMessage, +} as const diff --git a/src/domains/sources/route-chunks.ts b/src/domains/sources/route-chunks.ts index 7d97920..8d9d315 100644 --- a/src/domains/sources/route-chunks.ts +++ b/src/domains/sources/route-chunks.ts @@ -6,6 +6,7 @@ import { resolveChunkConnectionTargets } from "@/domains/chunks" import type { DemoChunkPage } from "@/integrations/knowhere-demo" import { logger } from "@/lib/logger" import { routeResult } from "@/lib/route-result" +import { displayReadUnavailable } from "./display-read-unavailable" import { decodeRemoteSourceId, findRemoteLibraryDocumentBySourceId, @@ -112,20 +113,26 @@ const loadSourceChunksEffect = ( } if (input.shouldLoadAll) { - const chunks = yield* Effect.tryPromise(() => + return yield* Effect.tryPromise(() => readAllSourceChunks({ knowledge, source: readableSource }), + ).pipe( + Effect.map((chunks) => + routeResult.ok({ chunks: resolveChunkConnectionTargets(chunks) }), + ), + Effect.catchAll((error) => recoverUnavailableChunks(input, error)), ) - return routeResult.ok({ chunks: resolveChunkConnectionTargets(chunks) }) } - const chunkPage = yield* Effect.tryPromise(() => + return yield* Effect.tryPromise(() => readSourceChunkPage({ knowledge, source: readableSource, params: input.pageParams, }), + ).pipe( + Effect.map((chunkPage) => routeResult.ok(chunkPage)), + Effect.catchAll((error) => recoverUnavailableChunks(input, error)), ) - return routeResult.ok(chunkPage) }) const loadRemoteChunkPageEffect = ( @@ -182,20 +189,26 @@ const loadRemoteChunkPageEffect = ( } if (input.shouldLoadAll) { - const chunks = yield* Effect.tryPromise(() => + return yield* Effect.tryPromise(() => readAllSourceChunks({ knowledge, source: readableSource }), + ).pipe( + Effect.map((chunks) => + routeResult.ok({ chunks: resolveChunkConnectionTargets(chunks) }), + ), + Effect.catchAll((error) => recoverUnavailableChunks(input, error)), ) - return routeResult.ok({ chunks: resolveChunkConnectionTargets(chunks) }) } - const chunkPage = yield* Effect.tryPromise(() => + return yield* Effect.tryPromise(() => readSourceChunkPage({ knowledge, source: readableSource, params: input.pageParams, }), + ).pipe( + Effect.map((chunkPage) => routeResult.ok(chunkPage)), + Effect.catchAll((error) => recoverUnavailableChunks(input, error)), ) - return routeResult.ok(chunkPage) }) const loadDemoChunkPageEffect = ( @@ -337,4 +350,47 @@ function sourceSnapshotProcessing( ) } +function sourceChunksUnavailable( + input: LoadSourceChunksInput, +): JsonRouteResult<{ + readonly chunks: [] + readonly pagination?: { + readonly page: number + readonly pageSize: number + readonly total: 0 + readonly totalPages: 0 + } + readonly message: string + readonly isUnavailable: true +}> { + if (input.shouldLoadAll) { + return routeResult.ok({ + chunks: [], + message: displayReadUnavailable.message, + isUnavailable: true, + }) + } + + return routeResult.ok({ + chunks: [], + pagination: { + page: input.pageParams.page, + pageSize: input.pageParams.pageSize, + total: 0, + totalPages: 0, + }, + message: displayReadUnavailable.message, + isUnavailable: true, + }) +} + +function recoverUnavailableChunks( + input: LoadSourceChunksInput, + error: unknown, +): Effect.Effect, unknown> { + return displayReadUnavailable.isError(error) + ? Effect.succeed(sourceChunksUnavailable(input)) + : Effect.fail(error) +} + export { createRouteChunks } diff --git a/src/domains/sources/route-page-assets.ts b/src/domains/sources/route-page-assets.ts index e356702..3077a83 100644 --- a/src/domains/sources/route-page-assets.ts +++ b/src/domains/sources/route-page-assets.ts @@ -2,6 +2,7 @@ import { Effect } from "effect" import { readSourcePageAssets } from "./page-assets" import { routeResult } from "@/lib/route-result" +import { displayReadUnavailable } from "./display-read-unavailable" import { decodeRemoteSourceId, findRemoteLibraryDocumentBySourceId, @@ -77,7 +78,7 @@ const loadSourcePageAssetsEffect = ( documentId, revisionKey: source.knowhereJobId, }) - const pageAssets = yield* Effect.tryPromise(() => + return yield* Effect.tryPromise(() => readSourcePageAssets({ knowledge, source: { @@ -86,9 +87,12 @@ const loadSourcePageAssetsEffect = ( }, params: input.pageParams, }), + ).pipe( + Effect.map((pageAssets) => routeResult.ok(pageAssets)), + Effect.catchAll((error) => + recoverUnavailablePageAssets(input, error), + ), ) - - return routeResult.ok(pageAssets) }) const loadRemotePageAssetsEffect = ( @@ -139,15 +143,18 @@ const loadRemotePageAssetsEffect = ( documentId, revisionKey, }) - const pageAssets = yield* Effect.tryPromise(() => + return yield* Effect.tryPromise(() => readSourcePageAssets({ knowledge, source: { documentId, revisionKey }, params: input.pageParams, }), + ).pipe( + Effect.map((pageAssets) => routeResult.ok(pageAssets)), + Effect.catchAll((error) => + recoverUnavailablePageAssets(input, error), + ), ) - - return routeResult.ok(pageAssets) }) function sourceNotFound(): JsonRouteResult<{ readonly message: string }> { @@ -158,4 +165,42 @@ function sourceNotReady(): JsonRouteResult<{ readonly message: string }> { return routeResult.error(409, "Source is not ready.") } +function sourcePageAssetsUnavailable( + input: LoadSourcePageAssetsInput, +): { + readonly pages: [] + readonly pagination: { + readonly page: number + readonly pageSize: number + readonly total: 0 + readonly totalPages: 0 + } + readonly message: string + readonly isUnavailable: true +} { + return { + pages: [], + pagination: { + page: input.pageParams.page, + pageSize: input.pageParams.pageSize, + total: 0, + totalPages: 0, + }, + message: displayReadUnavailable.message, + isUnavailable: true, + } +} + +function recoverUnavailablePageAssets( + input: LoadSourcePageAssetsInput, + error: unknown, +): Effect.Effect< + JsonRouteResult>, + unknown +> { + return displayReadUnavailable.isError(error) + ? Effect.succeed(routeResult.ok(sourcePageAssetsUnavailable(input))) + : Effect.fail(error) +} + export { createRoutePageAssets } diff --git a/src/domains/sources/route-types.ts b/src/domains/sources/route-types.ts index 48401a0..bd8ea61 100644 --- a/src/domains/sources/route-types.ts +++ b/src/domains/sources/route-types.ts @@ -106,6 +106,17 @@ type SourceChunksBody = readonly chunks: readonly ParsedChunkView[] } | ChunkPage + | { + readonly chunks: readonly [] + readonly pagination?: { + readonly page: number + readonly pageSize: number + readonly total: 0 + readonly totalPages: 0 + } + readonly message: string + readonly isUnavailable: true + } | { readonly message: string } @@ -120,6 +131,17 @@ type SourcePageAssetsBody = readonly totalPages: number } } + | { + readonly pages: readonly [] + readonly pagination: { + readonly page: number + readonly pageSize: number + readonly total: 0 + readonly totalPages: 0 + } + readonly message: string + readonly isUnavailable: true + } | { readonly message: string } diff --git a/src/domains/workspace/client.test.ts b/src/domains/workspace/client.test.ts index bc19bba..0c9df2b 100644 --- a/src/domains/workspace/client.test.ts +++ b/src/domains/workspace/client.test.ts @@ -90,6 +90,66 @@ describe("workspaceClient", () => { }) }) + it("preserves source chunk unavailable messages without processing state", async () => { + mockRouteClient.getJson.mockResolvedValue({ + chunks: [], + message: + "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", + isUnavailable: true, + pagination: { + page: 1, + pageSize: 50, + total: 0, + totalPages: 0, + }, + }) + + const page = await workspaceClient.fetchChunkPage("source_1", 1) + + expect(page).toEqual({ + chunks: [], + isUnavailable: true, + message: + "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", + pagination: { + page: 1, + pageSize: 50, + total: 0, + totalPages: 0, + }, + }) + }) + + it("preserves page asset unavailable messages", async () => { + mockRouteClient.getJson.mockResolvedValue({ + pages: [], + message: + "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", + isUnavailable: true, + pagination: { + page: 1, + pageSize: 50, + total: 0, + totalPages: 0, + }, + }) + + const page = await workspaceClient.fetchPageAssetPage("source_1", 1) + + expect(page).toEqual({ + pages: [], + isUnavailable: true, + message: + "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", + pagination: { + page: 1, + pageSize: 50, + total: 0, + totalPages: 0, + }, + }) + }) + it("throws materialization route errors instead of treating them as empty sources", async () => { mockRouteClient.postJsonWithStatus.mockResolvedValue({ status: 502, diff --git a/src/domains/workspace/client.ts b/src/domains/workspace/client.ts index efd8e32..5490d33 100644 --- a/src/domains/workspace/client.ts +++ b/src/domains/workspace/client.ts @@ -28,6 +28,7 @@ const workspaceClientConfig = { type SourceChunksResponse = { chunks?: ParsedChunkView[] isProcessing?: boolean + isUnavailable?: boolean message?: string pagination?: { page: number @@ -39,6 +40,7 @@ type SourceChunksResponse = { type SourcePageAssetsResponse = { pages?: SourcePageAssetView[] + isUnavailable?: boolean message?: string pagination?: { page: number @@ -146,7 +148,10 @@ async function fetchChunkPage( return { chunks: Array.isArray(body.chunks) ? body.chunks : [], ...(typeof body.message === "string" ? { message: body.message } : {}), - ...(typeof body.message === "string" ? { isProcessing: true } : {}), + ...(body.isUnavailable === true ? { isUnavailable: true } : {}), + ...(typeof body.message === "string" && body.isUnavailable !== true + ? { isProcessing: true } + : {}), pagination: body.pagination, } } @@ -165,6 +170,7 @@ async function fetchPageAssetPage( return { pages: Array.isArray(body.pages) ? body.pages : [], + ...(body.isUnavailable === true ? { isUnavailable: true } : {}), ...(typeof body.message === "string" ? { message: body.message } : {}), pagination: body.pagination, } From b5e0c2d30413b1bff8073bb2926c70f456836242 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 7 Jul 2026 17:40:49 +0800 Subject: [PATCH 03/14] Lazy detect page asset documents --- .../workspace-selected-chunks.test.ts | 59 +++++++ src/components/workspace-selected-chunks.ts | 78 ++++++++- src/components/workspace-shell-layout.tsx | 7 +- src/components/workspace-shell.tsx | 1 + src/domains/sources/counts.test.ts | 163 ++++++++++++++++++ src/domains/sources/counts.ts | 80 +++++---- src/domains/sources/page-assets.ts | 21 ++- src/domains/sources/route-dependencies.ts | 3 +- src/domains/sources/route-listing.ts | 3 + src/domains/sources/route-service.test.ts | 12 ++ src/domains/sources/route-types.ts | 6 +- src/domains/workspace/client-cache.ts | 4 + src/domains/workspace/initial-state.test.ts | 27 ++- src/domains/workspace/initial-state.ts | 4 + 14 files changed, 423 insertions(+), 45 deletions(-) diff --git a/src/components/workspace-selected-chunks.test.ts b/src/components/workspace-selected-chunks.test.ts index b57a788..be575fb 100644 --- a/src/components/workspace-selected-chunks.test.ts +++ b/src/components/workspace-selected-chunks.test.ts @@ -10,10 +10,12 @@ import type { ParsedChunkView } from "@/domains/chunks/types"; import type { SourceView } from "@/domains/sources/types"; const fetchChunkPageMock = vi.hoisted(() => vi.fn()); +const fetchPageAssetPageMock = vi.hoisted(() => vi.fn()); vi.mock("@/domains/workspace/client", () => ({ workspaceClient: { fetchChunkPage: fetchChunkPageMock, + fetchPageAssetPage: fetchPageAssetPageMock, }, })); @@ -28,6 +30,7 @@ const readySource: SourceView = { describe("useWorkspaceSelectedChunks", () => { beforeEach(() => { fetchChunkPageMock.mockReset(); + fetchPageAssetPageMock.mockReset(); fetchChunkPageMock.mockResolvedValue({ chunks: [], pagination: { @@ -37,6 +40,15 @@ describe("useWorkspaceSelectedChunks", () => { totalPages: 1, }, }); + fetchPageAssetPageMock.mockResolvedValue({ + pages: [], + pagination: { + page: 1, + pageSize: 50, + total: 0, + totalPages: 0, + }, + }); }); it("returns prefetched chunks while checking the visible page for media", async () => { @@ -59,6 +71,9 @@ describe("useWorkspaceSelectedChunks", () => { { wrapper: createSWRWrapper }, ); + await waitFor(() => + expect(fetchPageAssetPageMock).toHaveBeenCalledWith("source_1", 1), + ); await waitFor(() => expect(fetchChunkPageMock).toHaveBeenCalledWith("source_1", 1), ); @@ -213,6 +228,50 @@ describe("useWorkspaceSelectedChunks", () => { expect(result.current.selectedSource?.id).toBe("source_1"); expect(result.current.selectedChunks).toEqual([]); expect(result.current.isSelectedChunksLoading).toBe(false); + expect(fetchPageAssetPageMock).not.toHaveBeenCalled(); + expect(fetchChunkPageMock).not.toHaveBeenCalled(); + }); + + it("detects page assets for selected sources without presentation metadata", async () => { + fetchPageAssetPageMock.mockResolvedValue({ + pages: [ + { + pageNumber: 1, + assetUrl: "https://blob.example/page-1.png", + contentType: "image/png", + }, + { + pageNumber: 20, + assetUrl: "https://blob.example/page-20.png", + contentType: "image/png", + }, + ], + pagination: { + page: 1, + pageSize: 50, + total: 1, + totalPages: 1, + }, + }); + + const { result } = renderHook( + () => + useWorkspaceSelectedChunks({ + selectedSourceId: "source_1", + sources: [readySource], + prefetchedChunksBySourceId: {}, + }), + { wrapper: createSWRWrapper }, + ); + + await waitFor(() => + expect(result.current.selectedSource?.documentPresentation).toEqual({ + kind: "page-assets", + pageCount: 20, + }), + ); + expect(result.current.selectedChunks).toEqual([]); + expect(fetchPageAssetPageMock).toHaveBeenCalledWith("source_1", 1); expect(fetchChunkPageMock).not.toHaveBeenCalled(); }); diff --git a/src/components/workspace-selected-chunks.ts b/src/components/workspace-selected-chunks.ts index f1a1283..acbcd50 100644 --- a/src/components/workspace-selected-chunks.ts +++ b/src/components/workspace-selected-chunks.ts @@ -1,6 +1,6 @@ "use client" -import { useEffect, useMemo, useRef } from "react" +import { useEffect, useMemo, useRef, useState } from "react" import useSWRInfinite from "swr/infinite" import { workspaceClient } from "@/domains/workspace/client" @@ -8,6 +8,7 @@ import { workspaceClientCache, type SourceChunksKey, type SourceChunksResponse, + type SourcePageAssetsResponse, } from "@/domains/workspace/client-cache" import { resolveChunkConnectionTargets } from "@/domains/chunks" import type { ParsedChunkView } from "@/domains/chunks/types" @@ -30,21 +31,51 @@ type WorkspaceSelectedChunks = { readonly selectedSource: SourceView | undefined } +type PageAssetProbeState = + | { readonly status: "page-assets"; readonly pageCount: number } + | { readonly status: "parsed-chunks" } + export function useWorkspaceSelectedChunks({ selectedSourceId, sources, prefetchedChunksBySourceId, onRemoteSourceChunksLoaded, }: WorkspaceSelectedChunksInput): WorkspaceSelectedChunks { - const selectedSource = sources.find((source) => source.id === selectedSourceId) + const rawSelectedSource = sources.find( + (source) => source.id === selectedSourceId, + ) const remoteSourceRefreshRequestedIdsRef = useRef>(new Set()) + const requestedPageAssetProbeIdsRef = useRef>(new Set()) + const [pageAssetProbeBySourceId, setPageAssetProbeBySourceId] = useState< + Readonly> + >({}) + const pageAssetProbeState = rawSelectedSource + ? pageAssetProbeBySourceId[rawSelectedSource.id] + : undefined + const selectedSource = + rawSelectedSource && pageAssetProbeState?.status === "page-assets" + ? { + ...rawSelectedSource, + chunkCount: pageAssetProbeState.pageCount, + documentPresentation: { + kind: "page-assets" as const, + pageCount: pageAssetProbeState.pageCount, + }, + } + : rawSelectedSource const prefetchedSelectedChunks = selectedSourceId ? prefetchedChunksBySourceId[selectedSourceId] : undefined + const shouldProbePageAssets = + rawSelectedSource !== undefined && + rawSelectedSource.status === "ready" && + rawSelectedSource.documentPresentation === undefined && + pageAssetProbeState === undefined const selectedChunkSourceId = selectedSource && selectedSource.status === "ready" && - selectedSource.documentPresentation?.kind !== "page-assets" + selectedSource.documentPresentation?.kind !== "page-assets" && + !shouldProbePageAssets ? selectedSource.id : null const { @@ -101,12 +132,40 @@ export function useWorkspaceSelectedChunks({ typeof selectedChunkPages[selectedChunkPageCount - 1] === "undefined", ) const isSelectedChunksLoading = + shouldProbePageAssets || hasProcessingSelectedChunkPage || (selectedChunkSourceId !== null && !prefetchedSelectedChunks && !selectedChunkPages && isChunksLoading) + useEffect(() => { + const source = rawSelectedSource + if (!source || !shouldProbePageAssets) return + if (requestedPageAssetProbeIdsRef.current.has(source.id)) return + + requestedPageAssetProbeIdsRef.current.add(source.id) + void workspaceClient + .fetchPageAssetPage(source.id, 1) + .then((response) => { + setPageAssetProbeBySourceId((current) => ({ + ...current, + [source.id]: getPageAssetProbeState(response), + })) + + if (source.kind === "remote" && (response.pages?.length ?? 0) > 0) { + onRemoteSourceChunksLoaded?.(source.id) + } + }) + .catch(() => { + requestedPageAssetProbeIdsRef.current.delete(source.id) + setPageAssetProbeBySourceId((current) => ({ + ...current, + [source.id]: { status: "parsed-chunks" }, + })) + }) + }, [onRemoteSourceChunksLoaded, rawSelectedSource, shouldProbePageAssets]) + useEffect(() => { const sourceId = selectedSource?.id if (!sourceId || selectedSource.kind !== "remote") return @@ -133,6 +192,19 @@ export function useWorkspaceSelectedChunks({ } } +function getPageAssetProbeState( + response: SourcePageAssetsResponse, +): PageAssetProbeState { + const pages = response.pages ?? [] + if (pages.length === 0) return { status: "parsed-chunks" } + + const maxPageNumber = Math.max( + ...pages.map((page) => page.pageNumber), + ) + const pageCount = Math.max(response.pagination?.total ?? 0, maxPageNumber) + return { status: "page-assets", pageCount } +} + function fetchChunksByKey([ , sourceId, diff --git a/src/components/workspace-shell-layout.tsx b/src/components/workspace-shell-layout.tsx index b9ccc6c..0a71d90 100644 --- a/src/components/workspace-shell-layout.tsx +++ b/src/components/workspace-shell-layout.tsx @@ -90,6 +90,7 @@ export type WorkspaceShellLayoutProps = { readonly selectedSourceFile: SourceOriginalFileView | null readonly selectedSourceId: string | null readonly selectedSourceTitle: string | null + readonly selectedSourceView?: SourceView readonly sourceTitlesByDocumentId: Readonly> readonly sources: readonly SourceView[] readonly officialLibrarySources: readonly OfficialLibrarySourceView[] @@ -152,9 +153,9 @@ export function WorkspaceShellLayout( props.desktopPanelWidths.chat <= workspaceShellState.desktopSidePanelCompactThreshold const isSourcesPanelNarrow = props.desktopPanelWidths.sources < 220 - const selectedSource = props.sources.find( - (source) => source.id === props.selectedSourceId, - ) + const selectedSource = + props.selectedSourceView ?? + props.sources.find((source) => source.id === props.selectedSourceId) const focusedPage = props.focusedPage ?? { pageNumber: null, requestId: 0 } const handleDesktopLayoutRef = useCallback( (element: HTMLDivElement | null): void => { diff --git a/src/components/workspace-shell.tsx b/src/components/workspace-shell.tsx index 6fc842a..bd45e2f 100644 --- a/src/components/workspace-shell.tsx +++ b/src/components/workspace-shell.tsx @@ -242,6 +242,7 @@ function WorkspaceShellContent({ selectedSourceFile={citationFocus.selectedSource?.originalFile ?? null} selectedSourceId={sourceWorkflow.selectedSourceId} selectedSourceTitle={selectedSourceTitle} + selectedSourceView={citationFocus.selectedSource} sourceTitlesByDocumentId={sourceWorkflow.sourceTitlesByDocumentId} sources={sourceWorkflow.sources} officialLibrarySources={officialLibrarySources ?? []} diff --git a/src/domains/sources/counts.test.ts b/src/domains/sources/counts.test.ts index 652f404..72d132f 100644 --- a/src/domains/sources/counts.test.ts +++ b/src/domains/sources/counts.test.ts @@ -2,6 +2,7 @@ import { describe, expect, it, vi } from "vitest" import { Effect } from "effect" import type Knowhere from "@ontos-ai/knowhere-sdk" +import type { Knowledge } from "@ontos-ai/knowhere-sdk" import type { Source } from "@/infrastructure/db/schema" @@ -129,6 +130,45 @@ describe("countChunksBySourceId", () => { }) describe("sourceViewOptionsBySourceId", () => { + it("detects page count from many page assets in a single SDK page chunk", async () => { + const listChunks = vi.fn(async () => ({ pagination: { total: 12 } })) + const readChunks = vi.fn(async () => ({ + chunks: [ + { + chunkId: "page_bundle", + chunkType: "page", + metadata: { + pageAssets: Array.from({ length: 20 }, (_, index) => ({ + pageNum: index + 1, + artifactRef: `pages/page-${String(index + 1).padStart(6, "0")}.png`, + assetUrl: `https://assets.example/page-${index + 1}.png`, + })), + }, + }, + ], + totalChunks: 1, + })) + const mockClient = { + documents: { listChunks }, + knowledge: { readChunks }, + } as unknown as Knowhere + + const { sourceViewOptionsBySourceId } = await import("./counts") + + const options = await Effect.runPromise( + sourceViewOptionsBySourceId( + [makeSource({ id: "ready", knowhereDocumentId: "doc_ready" })], + mockClient, + ), + ) + + expect(options.get("ready")).toEqual({ + chunkCount: 20, + documentPresentation: { kind: "page-assets", pageCount: 20 }, + }) + expect(listChunks).not.toHaveBeenCalled() + }) + it("detects page-asset documents from SDK page chunks", async () => { const listChunks = vi.fn(async () => ({ pagination: { total: 12 } })) const readChunks = vi.fn(async () => ({ @@ -169,4 +209,127 @@ describe("sourceViewOptionsBySourceId", () => { }) expect(listChunks).not.toHaveBeenCalled() }) + + it("uses a source-specific knowledge reader for presentation detection", async () => { + const defaultReadChunks = vi.fn(async () => ({ chunks: [], totalChunks: 0 })) + const parsedStorageReadChunks = vi.fn(async () => ({ + chunks: [ + { + chunkId: "page_1", + chunkType: "page", + metadata: { + pageAssets: [ + { + pageNum: 1, + artifactRef: "pages/page-000001.png", + assetUrl: "https://assets.example/page-000001.png", + }, + ], + }, + }, + ], + totalChunks: 1, + })) + const listChunks = vi.fn(async () => ({ pagination: { total: 12 } })) + const mockClient = { + documents: { listChunks }, + knowledge: { readChunks: defaultReadChunks }, + } as unknown as Knowhere + + const { sourceViewOptionsBySourceId } = await import("./counts") + + const options = await Effect.runPromise( + sourceViewOptionsBySourceId( + [makeSource({ id: "ready", knowhereDocumentId: "doc_ready" })], + mockClient, + { + getKnowledgeForSource: () => + ({ readChunks: parsedStorageReadChunks }) as unknown as Knowledge, + }, + ), + ) + + expect(options.get("ready")).toEqual({ + chunkCount: 1, + documentPresentation: { kind: "page-assets", pageCount: 1 }, + }) + expect(parsedStorageReadChunks).toHaveBeenCalledWith({ + documentId: "doc_ready", + revisionKey: "job_1", + chunkType: "page", + page: 1, + pageSize: 1, + assetUrlPolicy: "durable", + }) + expect(defaultReadChunks).not.toHaveBeenCalled() + }) + + it("falls back to chunk counts when page presentation detection fails", async () => { + const listChunks = vi.fn(async () => ({ pagination: { total: 12 } })) + const readChunks = vi.fn(async () => { + throw new Error("parsed storage and remote unavailable") + }) + const mockClient = { + documents: { listChunks }, + knowledge: { readChunks }, + } as unknown as Knowhere + + const { sourceViewOptionsBySourceId } = await import("./counts") + + const options = await Effect.runPromise( + sourceViewOptionsBySourceId( + [makeSource({ id: "ready", knowhereDocumentId: "doc_ready" })], + mockClient, + ), + ) + + expect(options.get("ready")).toEqual({ chunkCount: 12 }) + expect(listChunks).toHaveBeenCalledWith("doc_ready", { + page: 1, + pageSize: 1, + }) + }) + + it("skips page presentation reads when detection is disabled", async () => { + const listChunks = vi.fn(async () => ({ pagination: { total: 12 } })) + const readChunks = vi.fn(async () => ({ + chunks: [ + { + chunkId: "page_1", + chunkType: "page", + metadata: { + pageAssets: [ + { + pageNum: 1, + artifactRef: "pages/page-000001.png", + assetUrl: "https://assets.example/page-000001.png", + }, + ], + }, + }, + ], + totalChunks: 1, + })) + const mockClient = { + documents: { listChunks }, + knowledge: { readChunks }, + } as unknown as Knowhere + + const { sourceViewOptionsBySourceId } = await import("./counts") + + const options = await Effect.runPromise( + sourceViewOptionsBySourceId( + [makeSource({ id: "ready", knowhereDocumentId: "doc_ready" })], + mockClient, + { documentPresentationDetection: "disabled" }, + ), + ) + + expect(readChunks).not.toHaveBeenCalled() + expect(listChunks).toHaveBeenCalledWith("doc_ready", { + page: 1, + pageSize: 1, + }) + expect(options.get("ready")).toEqual({ chunkCount: 12 }) + }) }) diff --git a/src/domains/sources/counts.ts b/src/domains/sources/counts.ts index 110b3ae..d161adf 100644 --- a/src/domains/sources/counts.ts +++ b/src/domains/sources/counts.ts @@ -2,7 +2,7 @@ import "server-only" import { Effect } from "effect" import type Knowhere from "@ontos-ai/knowhere-sdk" -import type { KnowledgeReadChunk } from "@ontos-ai/knowhere-sdk" +import type { Knowledge, KnowledgeReadChunk } from "@ontos-ai/knowhere-sdk" import type { Source } from "@/infrastructure/db/schema" import type { SourceDocumentPresentation } from "./types" @@ -36,6 +36,11 @@ type CountChunksClient = { } } +export type SourceViewOptionsLoadOptions = { + readonly documentPresentationDetection?: "enabled" | "disabled" + readonly getKnowledgeForSource?: (source: Source) => Knowledge +} + export type SourceViewOptions = { readonly chunkCount?: number readonly documentPresentation?: SourceDocumentPresentation @@ -44,6 +49,7 @@ export type SourceViewOptions = { export const sourceViewOptionsBySourceId = ( sources: readonly Source[], client: Knowhere, + options: SourceViewOptionsLoadOptions = {}, ) => Effect.gen(function* () { const countClient = client as unknown as CountChunksClient @@ -58,14 +64,14 @@ export const sourceViewOptionsBySourceId = ( const entries = yield* Effect.all( readySources.map((source) => Effect.gen(function* () { - const options = yield* Effect.tryPromise(() => - loadSourceViewOptions(countClient, source), + const loadedOptions = yield* Effect.tryPromise(() => + loadSourceViewOptions(countClient, source, options), ).pipe( Effect.catchAll(() => Effect.sync((): SourceViewOptions | undefined => undefined), ), ) - return [source.id, options] as const + return [source.id, loadedOptions] as const }), ), { concurrency: "unbounded" }, @@ -97,19 +103,22 @@ export const countChunksBySourceId = ( async function loadSourceViewOptions( client: CountChunksClient, source: Source, + options: SourceViewOptionsLoadOptions, ): Promise { const documentId = source.knowhereDocumentId if (!documentId) return undefined - const pagePresentation = await loadPageAssetPresentation( - client, - documentId, - source.knowhereJobId, - ) - if (pagePresentation) { - return { - chunkCount: pagePresentation.pageCount, - documentPresentation: pagePresentation, + if (options.documentPresentationDetection !== "disabled") { + const pagePresentation = await loadPageAssetPresentation( + client, + source, + options, + ) + if (pagePresentation) { + return { + chunkCount: pagePresentation.pageCount, + documentPresentation: pagePresentation, + } } } @@ -119,13 +128,17 @@ async function loadSourceViewOptions( async function loadPageAssetPresentation( client: CountChunksClient, - documentId: string, - revisionKey: string | null, + source: Source, + options: SourceViewOptionsLoadOptions, ): Promise { + const documentId = source.knowhereDocumentId + if (!documentId) return undefined + try { - const response = await client.knowledge.readChunks({ + const knowledge = options.getKnowledgeForSource?.(source) ?? client.knowledge + const response = await knowledge.readChunks({ documentId, - ...(revisionKey ? { revisionKey } : {}), + ...(source.knowhereJobId ? { revisionKey: source.knowhereJobId } : {}), chunkType: "page", page: 1, pageSize: 1, @@ -133,14 +146,13 @@ async function loadPageAssetPresentation( }) const firstChunk = response.chunks[0] if (!firstChunk || firstChunk.chunkType !== "page") return undefined - if (!hasUsablePageAssets(firstChunk.metadata.pageAssets)) return undefined + const maxPageAssetNumber = getMaxUsablePageAssetNumber( + firstChunk.metadata.pageAssets, + ) + if (!maxPageAssetNumber) return undefined - const pageCount = - typeof response.totalChunks === "number" && - Number.isFinite(response.totalChunks) && - response.totalChunks > 0 - ? response.totalChunks - : 1 + const totalChunks = getPositiveFiniteNumber(response.totalChunks) ?? 0 + const pageCount = Math.max(totalChunks, maxPageAssetNumber) return { kind: "page-assets", pageCount } } catch { @@ -160,22 +172,30 @@ async function loadSourceChunkCount( return typeof total === "number" && Number.isFinite(total) ? total : undefined } -function hasUsablePageAssets(value: unknown): boolean { - if (!Array.isArray(value)) return false +function getMaxUsablePageAssetNumber(value: unknown): number | undefined { + if (!Array.isArray(value)) return undefined - return value.some((item) => { - if (!isRecord(item)) return false + const pageNumbers = value.flatMap((item): number[] => { + if (!isRecord(item)) return [] const pageNum = item.pageNum const artifactRef = item.artifactRef const assetUrl = item.assetUrl - return ( + const isUsable = typeof pageNum === "number" && Number.isSafeInteger(pageNum) && pageNum > 0 && ((typeof artifactRef === "string" && artifactRef.trim().length > 0) || (typeof assetUrl === "string" && assetUrl.trim().length > 0)) - ) + return isUsable ? [pageNum] : [] }) + if (pageNumbers.length === 0) return undefined + return Math.max(...pageNumbers) +} + +function getPositiveFiniteNumber(value: unknown): number | undefined { + return typeof value === "number" && Number.isFinite(value) && value > 0 + ? value + : undefined } function isRecord(value: unknown): value is Readonly> { diff --git a/src/domains/sources/page-assets.ts b/src/domains/sources/page-assets.ts index 55c108b..cfb2765 100644 --- a/src/domains/sources/page-assets.ts +++ b/src/domains/sources/page-assets.ts @@ -36,16 +36,22 @@ export async function readSourcePageAssets(input: { const pages = response.chunks.flatMap((chunk): SourcePageAssetView[] => readPageAssetViews(chunk.metadata.pageAssets, chunk.assetUrl), ) + const maxPageNumber = getMaxPageNumber(pages) + const total = Math.max(response.totalChunks ?? 0, maxPageNumber ?? 0) + const resolvedTotal = total > 0 ? total : pages.length + const computedTotalPages = Math.max( + 1, + Math.ceil(resolvedTotal / input.params.pageSize), + ) + const totalPages = Math.max(response.totalPages ?? 0, computedTotalPages) return { pages, pagination: { page: response.page ?? input.params.page, pageSize: response.pageSize ?? input.params.pageSize, - total: response.totalChunks ?? pages.length, - totalPages: - response.totalPages ?? - Math.max(1, Math.ceil(pages.length / input.params.pageSize)), + total: resolvedTotal, + totalPages, }, } } @@ -99,6 +105,13 @@ function getPositiveNumber(value: unknown): number | undefined { : undefined } +function getMaxPageNumber( + pages: readonly SourcePageAssetView[], +): number | undefined { + if (pages.length === 0) return undefined + return Math.max(...pages.map((page) => page.pageNumber)) +} + function isRecord(value: unknown): value is Readonly> { return typeof value === "object" && value !== null } diff --git a/src/domains/sources/route-dependencies.ts b/src/domains/sources/route-dependencies.ts index a39635f..a51a9b7 100644 --- a/src/domains/sources/route-dependencies.ts +++ b/src/domains/sources/route-dependencies.ts @@ -28,10 +28,11 @@ const defaultDependencies: SourceRouteServiceDependencies = { ensureApiKeyForWorkspace, ensureWorkspace: workspaceService.ensureWorkspace, getCurrentUser, - getSourceViewOptionsBySourceId: (sources, client) => + getSourceViewOptionsBySourceId: (sources, client, options) => getDefaultSourceViewOptionsBySourceId( sources, client as ReturnType, + options, ), makeKnowhereClient: (apiKey: string) => makeDefaultKnowhereClient(apiKey) as SourceRouteKnowhereClient, diff --git a/src/domains/sources/route-listing.ts b/src/domains/sources/route-listing.ts index 35c0c72..08c17b2 100644 --- a/src/domains/sources/route-listing.ts +++ b/src/domains/sources/route-listing.ts @@ -113,6 +113,9 @@ const listSourcesEffect = ( const sourceOptions = yield* deps.getSourceViewOptionsBySourceId( sourcesNeedingChunkCount, client, + { + documentPresentationDetection: "disabled", + }, ) const hiddenDemoSourceIds = new Set( yield* Effect.tryPromise(() => diff --git a/src/domains/sources/route-service.test.ts b/src/domains/sources/route-service.test.ts index a056427..cb9c7da 100644 --- a/src/domains/sources/route-service.test.ts +++ b/src/domains/sources/route-service.test.ts @@ -472,6 +472,9 @@ describe("source route service", () => { expect(getSourceViewOptionsBySourceId).toHaveBeenCalledWith( [source], knowhereClient, + expect.objectContaining({ + documentPresentationDetection: "disabled", + }), ); expect(result).toEqual({ status: 200, @@ -546,6 +549,9 @@ describe("source route service", () => { expect(getSourceViewOptionsBySourceId).toHaveBeenCalledWith( [], knowhereClient, + expect.objectContaining({ + documentPresentationDetection: "disabled", + }), ); expect(result).toEqual({ status: 200, @@ -628,6 +634,9 @@ describe("source route service", () => { expect(getSourceViewOptionsBySourceId).toHaveBeenCalledWith( [], knowhereClient, + expect.objectContaining({ + documentPresentationDetection: "disabled", + }), ); expect(result).toEqual({ status: 200, @@ -700,6 +709,9 @@ describe("source route service", () => { expect(getSourceViewOptionsBySourceId).toHaveBeenCalledWith( [], knowhereClient, + expect.objectContaining({ + documentPresentationDetection: "disabled", + }), ); expect(knowhereClient.documents.listChunks).not.toHaveBeenCalled(); expect(result).toEqual({ diff --git a/src/domains/sources/route-types.ts b/src/domains/sources/route-types.ts index bd8ea61..8ce806d 100644 --- a/src/domains/sources/route-types.ts +++ b/src/domains/sources/route-types.ts @@ -17,7 +17,10 @@ import type { } from "@/integrations/knowhere-demo" import type { RouteResult } from "@/lib/route-result" import type { SourceBlobUploadInput } from "./blob-upload" -import type { sourceViewOptionsBySourceId } from "./counts" +import type { + sourceViewOptionsBySourceId, + SourceViewOptionsLoadOptions, +} from "./counts" import type { UploadKnowhereClient } from "./upload" type SourceRouteKnowhereClient = UploadKnowhereClient & @@ -275,6 +278,7 @@ type SourceRouteServiceDependencies = { readonly getSourceViewOptionsBySourceId: ( sources: readonly Source[], client: SourceRouteKnowhereClient, + options?: SourceViewOptionsLoadOptions, ) => ReturnType readonly makeKnowhereClient: (apiKey: string) => SourceRouteKnowhereClient readonly listSourcesForWorkspace: (workspaceId: string) => Promise diff --git a/src/domains/workspace/client-cache.ts b/src/domains/workspace/client-cache.ts index af5baae..beaa41d 100644 --- a/src/domains/workspace/client-cache.ts +++ b/src/domains/workspace/client-cache.ts @@ -11,6 +11,9 @@ import type { SourceView } from "@/domains/sources/types" type SourceChunksResponse = Awaited< ReturnType > +type SourcePageAssetsResponse = Awaited< + ReturnType +> type ChatThreadDetailResponse = Awaited< ReturnType > @@ -114,4 +117,5 @@ export type { ChatThreadKey, SourceChunksKey, SourceChunksResponse, + SourcePageAssetsResponse, } diff --git a/src/domains/workspace/initial-state.test.ts b/src/domains/workspace/initial-state.test.ts index 23bd50f..9fc1450 100644 --- a/src/domains/workspace/initial-state.test.ts +++ b/src/domains/workspace/initial-state.test.ts @@ -172,6 +172,9 @@ describe("loadWorkspaceShellInitialState", () => { expect(sourceViewOptionsBySourceId).toHaveBeenCalledWith( [source], expect.any(Object), + expect.objectContaining({ + documentPresentationDetection: "disabled", + }), ) expect(state.sources).toEqual([ { @@ -203,7 +206,13 @@ describe("loadWorkspaceShellInitialState", () => { const state = await loadWorkspaceShellInitialState(deps) - expect(sourceViewOptionsBySourceId).toHaveBeenCalledWith([], expect.any(Object)) + expect(sourceViewOptionsBySourceId).toHaveBeenCalledWith( + [], + expect.any(Object), + expect.objectContaining({ + documentPresentationDetection: "disabled", + }), + ) expect(state.sources).toEqual([ expect.objectContaining({ id: "source_demo", @@ -231,7 +240,13 @@ describe("loadWorkspaceShellInitialState", () => { const state = await loadWorkspaceShellInitialState(deps) - expect(sourceViewOptionsBySourceId).toHaveBeenCalledWith([], expect.any(Object)) + expect(sourceViewOptionsBySourceId).toHaveBeenCalledWith( + [], + expect.any(Object), + expect.objectContaining({ + documentPresentationDetection: "disabled", + }), + ) expect(state.sources).toEqual([ expect.objectContaining({ id: "demo-tsla-q4-2025", @@ -259,7 +274,13 @@ describe("loadWorkspaceShellInitialState", () => { const state = await loadWorkspaceShellInitialState(deps) - expect(sourceViewOptionsBySourceId).toHaveBeenCalledWith([], expect.any(Object)) + expect(sourceViewOptionsBySourceId).toHaveBeenCalledWith( + [], + expect.any(Object), + expect.objectContaining({ + documentPresentationDetection: "disabled", + }), + ) expect(state.sources).toEqual([ expect.objectContaining({ id: "demo-tsla-q4-2025", diff --git a/src/domains/workspace/initial-state.ts b/src/domains/workspace/initial-state.ts index 851a34b..f959ad9 100644 --- a/src/domains/workspace/initial-state.ts +++ b/src/domains/workspace/initial-state.ts @@ -149,6 +149,7 @@ type WorkspaceShellInitialStateDependencies = { readonly sourceViewOptionsBySourceId: ( sources: readonly Source[], client: WorkspaceShellInitialStateClient, + options?: Parameters[2], ) => ReturnType } @@ -356,6 +357,9 @@ export const loadWorkspaceShellInitialStateEffect = ( deps.sourceViewOptionsBySourceId( sourcesNeedingChunkCount, client, + { + documentPresentationDetection: "disabled", + }, ), ) From f6198d93543d1962ab2a0b698325da589176c11f Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 7 Jul 2026 18:17:14 +0800 Subject: [PATCH 04/14] Make page asset viewing non-blocking --- .../sources/[sourceId]/chunks/route.test.ts | 125 ++++++++++++- .../[sourceId]/page-assets/route.test.ts | 153 +++++++++++++++- src/components/chunks-panel.test.ts | 42 +++++ src/components/chunks-panel.tsx | 41 ++++- src/domains/chunks/index.ts | 1 + src/domains/chunks/read.test.ts | 109 +++++++++++- src/domains/chunks/read.ts | 167 ++++++++++++++++-- src/domains/sources/page-assets.ts | 111 +++++++++++- src/domains/sources/route-chunks.ts | 24 ++- src/domains/sources/route-dependencies.ts | 15 +- src/domains/sources/route-page-assets.ts | 12 +- 11 files changed, 743 insertions(+), 57 deletions(-) diff --git a/src/app/api/sources/[sourceId]/chunks/route.test.ts b/src/app/api/sources/[sourceId]/chunks/route.test.ts index d13677d..f410285 100644 --- a/src/app/api/sources/[sourceId]/chunks/route.test.ts +++ b/src/app/api/sources/[sourceId]/chunks/route.test.ts @@ -8,6 +8,7 @@ const mocks = vi.hoisted(() => ({ fetchDemoChunkPage: vi.fn(), findSourceInWorkspace: vi.fn(), getCurrentUser: vi.fn(), + listChunks: vi.fn(), localizeRemoteDocument: vi.fn(), makeKnowhereClient: vi.fn(), makeKnowhereClientWithParsedStorage: vi.fn(), @@ -64,7 +65,7 @@ describe("GET /api/sources/[sourceId]/chunks", () => { beforeEach(() => { vi.clearAllMocks() mocks.makeKnowhereClientWithParsedStorage.mockReturnValue({ - client: { documents: { listChunks: vi.fn() } }, + client: { documents: { listChunks: mocks.listChunks } }, knowledge: { readChunks: mocks.readChunks }, }) }) @@ -421,7 +422,7 @@ describe("GET /api/sources/[sourceId]/chunks", () => { } }) - it("loads authenticated workspace chunks through the SDK durable read", async () => { + it("loads authenticated workspace chunks from parsed storage without durable hardening", async () => { mocks.getCurrentUser.mockResolvedValue({ id: "user_1", email: null, @@ -442,7 +443,10 @@ describe("GET /api/sources/[sourceId]/chunks", () => { ) mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") mocks.readChunks.mockResolvedValue({ - document: { localDocumentId: "doc_1" }, + document: { + localDocumentId: "doc_1", + resultDirectoryPath: "parsed-storage:doc_1", + }, chunks: [ { position: 1, @@ -498,7 +502,114 @@ describe("GET /api/sources/[sourceId]/chunks", () => { revisionKey: "job_1", page: 1, pageSize: 1, - assetUrlPolicy: "durable", + }) + expect(mocks.listChunks).not.toHaveBeenCalled() + }) + + it("falls back to raw Knowhere asset URLs for workspace chunk display when storage misses", async () => { + mocks.getCurrentUser.mockResolvedValue({ + id: "user_1", + email: null, + name: null, + }) + mocks.ensureWorkspace.mockResolvedValue({ + id: "workspace_1", + userId: "user_1", + namespace: "notebook-workspace_1", + createdAt: new Date("2026-05-10T00:00:00.000Z"), + }) + mocks.findSourceInWorkspace.mockResolvedValue( + makeReadySource({ + id: "00000000-0000-0000-0000-000000000002", + knowhereJobId: "job_1", + knowhereDocumentId: "doc_1", + }), + ) + mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") + mocks.readChunks.mockResolvedValue({ + document: { + localDocumentId: "doc_1", + resultDirectoryPath: "remote:doc_1", + }, + chunks: [ + { + position: 1, + chunkId: "parser_1", + chunkType: "image", + content: "Workspace chunk", + readableContent: "Workspace chunk", + sectionPath: "Summary", + sourceChunkPath: "Summary", + filePath: "images/chart.png", + metadata: {}, + }, + ], + page: 1, + pageSize: 1, + totalChunks: 1, + totalPages: 1, + }) + mocks.listChunks.mockResolvedValue({ + chunks: [ + { + id: "document_chunk_1", + chunkId: "parser_1", + chunkType: "image", + content: "Workspace chunk", + sectionId: null, + sectionPath: "Summary", + sourceChunkPath: "Summary", + filePath: "images/chart.png", + sortOrder: 0, + metadata: {}, + assetUrl: "https://knowhere.example/assets/chart.png", + }, + ], + pagination: { + page: 1, + pageSize: 1, + total: 1, + totalPages: 1, + }, + }) + + const response = await GET( + new NextRequest( + "http://localhost:3001/api/sources/00000000-0000-0000-0000-000000000002/chunks?page=1&pageSize=1", + ), + { + params: Promise.resolve({ + sourceId: "00000000-0000-0000-0000-000000000002", + }), + }, + ) + + await expect(response.json()).resolves.toMatchObject({ + chunks: [ + { + parserChunkId: "parser_1", + documentId: "doc_1", + assetUrl: "https://knowhere.example/assets/chart.png", + sourceTitle: "notes.pdf", + }, + ], + pagination: { + page: 1, + pageSize: 1, + total: 1, + }, + }) + expect(response.status).toBe(200) + expect(mocks.readChunks).toHaveBeenCalledWith({ + documentId: "doc_1", + revisionKey: "job_1", + page: 1, + pageSize: 1, + }) + expect(mocks.listChunks).toHaveBeenCalledWith("doc_1", { + page: 1, + pageSize: 1, + includeAssetUrls: true, }) }) @@ -641,7 +752,10 @@ describe("GET /api/sources/[sourceId]/chunks", () => { }), ) mocks.readChunks.mockResolvedValue({ - document: { localDocumentId: "doc_remote" }, + document: { + localDocumentId: "doc_remote", + resultDirectoryPath: "parsed-storage:doc_remote", + }, chunks: [ { position: 1, @@ -710,7 +824,6 @@ describe("GET /api/sources/[sourceId]/chunks", () => { revisionKey: "job_result_1", page: 1, pageSize: 1, - assetUrlPolicy: "durable", }) }) }) diff --git a/src/app/api/sources/[sourceId]/page-assets/route.test.ts b/src/app/api/sources/[sourceId]/page-assets/route.test.ts index 073a20e..76ef16a 100644 --- a/src/app/api/sources/[sourceId]/page-assets/route.test.ts +++ b/src/app/api/sources/[sourceId]/page-assets/route.test.ts @@ -7,6 +7,7 @@ const mocks = vi.hoisted(() => ({ ensureWorkspace: vi.fn(), findSourceInWorkspace: vi.fn(), getCurrentUser: vi.fn(), + listChunks: vi.fn(), makeKnowhereClient: vi.fn(), makeKnowhereClientWithParsedStorage: vi.fn(), readChunks: vi.fn(), @@ -62,12 +63,12 @@ describe("GET /api/sources/[sourceId]/page-assets", () => { beforeEach(() => { vi.clearAllMocks() mocks.makeKnowhereClientWithParsedStorage.mockReturnValue({ - client: { documents: { listChunks: vi.fn() } }, + client: { documents: { listChunks: mocks.listChunks } }, knowledge: { readChunks: mocks.readChunks }, }) }) - it("returns durable page assets for a ready workspace source", async () => { + it("returns stored page assets for a ready workspace source without durable hardening", async () => { mocks.getCurrentUser.mockResolvedValue({ id: "user_1", email: null, @@ -88,6 +89,10 @@ describe("GET /api/sources/[sourceId]/page-assets", () => { ) mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") mocks.readChunks.mockResolvedValue({ + document: { + localDocumentId: "doc_1", + resultDirectoryPath: "parsed-storage:doc_1", + }, chunks: [ { chunkId: "page_1", @@ -147,8 +152,8 @@ describe("GET /api/sources/[sourceId]/page-assets", () => { chunkType: "page", page: 1, pageSize: 1, - assetUrlPolicy: "durable", }) + expect(mocks.listChunks).not.toHaveBeenCalled() }) it("rejects non-ready workspace sources", async () => { @@ -209,6 +214,10 @@ describe("GET /api/sources/[sourceId]/page-assets", () => { ) mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") mocks.readChunks.mockResolvedValue({ + document: { + localDocumentId: "doc_1", + resultDirectoryPath: "remote:doc_1", + }, chunks: [ { chunkId: "page_1", @@ -221,6 +230,15 @@ describe("GET /api/sources/[sourceId]/page-assets", () => { totalChunks: 1, totalPages: 1, }) + mocks.listChunks.mockResolvedValue({ + chunks: [], + pagination: { + page: 1, + pageSize: 1, + total: 1, + totalPages: 1, + }, + }) const response = await GET( new NextRequest( @@ -243,6 +261,135 @@ describe("GET /api/sources/[sourceId]/page-assets", () => { }, }) expect(response.status).toBe(200) + expect(mocks.listChunks).toHaveBeenCalledWith("doc_1", { + page: 1, + pageSize: 1, + chunkType: "page", + includeAssetUrls: true, + }) + }) + + it("returns Knowhere page asset URLs when the storage probe falls through to remote chunks", async () => { + mocks.getCurrentUser.mockResolvedValue({ + id: "user_1", + email: null, + name: null, + }) + mocks.ensureWorkspace.mockResolvedValue({ + id: "workspace_1", + userId: "user_1", + namespace: "notebook-workspace_1", + createdAt: new Date("2026-05-10T00:00:00.000Z"), + }) + mocks.findSourceInWorkspace.mockResolvedValue( + makeReadySource({ + id: "00000000-0000-0000-0000-000000000002", + knowhereDocumentId: "doc_1", + }), + ) + mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") + mocks.readChunks.mockResolvedValue({ + document: { + localDocumentId: "doc_1", + resultDirectoryPath: "remote:doc_1", + }, + chunks: [ + { + chunkId: "page_1", + chunkType: "page", + metadata: { + pageAssets: [ + { + pageNum: 1, + contentType: "image/png", + width: 1200, + height: 1600, + }, + ], + }, + }, + ], + page: 1, + pageSize: 1, + totalChunks: 1, + totalPages: 1, + }) + mocks.listChunks.mockResolvedValue({ + chunks: [ + { + id: "document_page_1", + chunkId: "page_1", + chunkType: "page", + content: "Page 1", + sectionId: null, + sectionPath: "pages/1", + sourceChunkPath: "pages/1", + filePath: "pages/page-000001.png", + sortOrder: 0, + metadata: { + pageAssets: [ + { + pageNum: 1, + assetUrl: "https://knowhere.example/page-000001.png", + contentType: "image/png", + width: 1200, + height: 1600, + }, + ], + }, + assetUrl: "https://knowhere.example/page-000001.png", + }, + ], + pagination: { + page: 1, + pageSize: 1, + total: 1, + totalPages: 1, + }, + }) + + const response = await GET( + new NextRequest( + "http://localhost:3001/api/sources/00000000-0000-0000-0000-000000000002/page-assets?page=1&pageSize=1", + ), + { + params: Promise.resolve({ + sourceId: "00000000-0000-0000-0000-000000000002", + }), + }, + ) + + await expect(response.json()).resolves.toEqual({ + pages: [ + { + pageNumber: 1, + assetUrl: "https://knowhere.example/page-000001.png", + contentType: "image/png", + width: 1200, + height: 1600, + }, + ], + pagination: { + page: 1, + pageSize: 1, + total: 1, + totalPages: 1, + }, + }) + expect(response.status).toBe(200) + expect(mocks.readChunks).toHaveBeenCalledWith({ + documentId: "doc_1", + revisionKey: "job_1", + chunkType: "page", + page: 1, + pageSize: 1, + }) + expect(mocks.listChunks).toHaveBeenCalledWith("doc_1", { + page: 1, + pageSize: 1, + chunkType: "page", + includeAssetUrls: true, + }) }) it("marks page assets unavailable when the parsed document is missing remotely", async () => { diff --git a/src/components/chunks-panel.test.ts b/src/components/chunks-panel.test.ts index 238ca0c..8e06155 100644 --- a/src/components/chunks-panel.test.ts +++ b/src/components/chunks-panel.test.ts @@ -134,6 +134,48 @@ describe("ChunksPanel", () => { expect(fetchPageAssetPageMock).toHaveBeenCalledWith("source_1", 1); }); + it("shows a page-level placeholder when a page asset image fails to load", async () => { + fetchPageAssetPageMock.mockResolvedValue({ + pages: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-000004.png", + contentType: "image/png", + width: 1200, + height: 1600, + }, + ], + pagination: { + page: 1, + pageSize: 50, + total: 4, + totalPages: 1, + }, + }); + + render( + React.createElement(C, { + chunks: [], + selectedSource: "report.pdf", + selectedSourceView: { + id: "source_1", + title: "report.pdf", + mimeType: "application/pdf", + status: "ready", + documentPresentation: { kind: "page-assets", pageCount: 4 }, + }, + }), + ); + + const pageImage = await screen.findByRole("img", { name: "Page 4" }); + fireEvent.error(pageImage); + + expect( + screen.getByTestId("page-asset-image-unavailable-4"), + ).toBeTruthy(); + expect(screen.getByText("Page image unavailable.")).toBeTruthy(); + }); + it("loads missing page buckets after focusing a later page asset", async () => { const user = userEvent.setup(); fetchPageAssetPageMock.mockImplementation( diff --git a/src/components/chunks-panel.tsx b/src/components/chunks-panel.tsx index 70f6448..fc785a2 100644 --- a/src/components/chunks-panel.tsx +++ b/src/components/chunks-panel.tsx @@ -19,6 +19,7 @@ import { } from "d3-hierarchy"; import { FilePlus2, + ImageOff, Layers, RotateCcw, UploadCloud, @@ -1221,6 +1222,8 @@ function PageAssetImage({ readonly page: SourcePageAssetView; readonly refCallback: (element: HTMLDivElement | null) => void; }): ReactNode { + const [failedAssetUrl, setFailedAssetUrl] = useState(null); + const hasImageError = failedAssetUrl === page.assetUrl; const aspectRatio = page.width && page.height ? `${page.width} / ${page.height}` : undefined; @@ -1241,17 +1244,43 @@ function PageAssetImage({ className="overflow-hidden rounded-md bg-muted/30" style={aspectRatio ? { aspectRatio } : undefined} > - {`Page + {hasImageError ? ( + + ) : ( + // eslint-disable-next-line @next/next/no-img-element -- Page assets can be short-lived Knowhere URLs outside Next image optimization. + {`Page setFailedAssetUrl(page.assetUrl)} + /> + )} ); } +function PageAssetImageUnavailable({ + pageNumber, +}: { + readonly pageNumber: number; +}): ReactNode { + return ( +
+
+ +
+

+ Page image unavailable. +

+
+ ); +} + function getNextPageAssetPageIndex( loadedPageIndexes: ReadonlySet, totalPages: number, diff --git a/src/domains/chunks/index.ts b/src/domains/chunks/index.ts index eca1830..e78293e 100644 --- a/src/domains/chunks/index.ts +++ b/src/domains/chunks/index.ts @@ -17,6 +17,7 @@ export type ChunkKnowhereClient = { params: { page: number pageSize: number + chunkType?: "text" | "image" | "table" | "page" includeAssetUrls: boolean }, ): Promise<{ diff --git a/src/domains/chunks/read.test.ts b/src/domains/chunks/read.test.ts index e3c1000..9ba2203 100644 --- a/src/domains/chunks/read.test.ts +++ b/src/domains/chunks/read.test.ts @@ -1,5 +1,5 @@ import { describe, expect, it, vi } from "vitest" -import type { Knowledge } from "@ontos-ai/knowhere-sdk" +import type { DocumentChunk, Knowledge } from "@ontos-ai/knowhere-sdk" import { readAllSourceChunks, readSourceChunkPage } from "./read" @@ -20,9 +20,13 @@ function makeReadChunk(overrides: Record = {}) { } describe("readSourceChunkPage", () => { - it("reads a page with durable asset URLs and maps chunks to the view model", async () => { + it("uses parsed-storage chunks without durable asset hardening", async () => { + const listChunks = vi.fn() const readChunks = vi.fn(async () => ({ - document: { localDocumentId: "doc_1" }, + document: { + localDocumentId: "doc_1", + resultDirectoryPath: "parsed-storage:doc_1", + }, chunks: [ makeReadChunk({ chunkType: "image", @@ -38,6 +42,7 @@ describe("readSourceChunkPage", () => { const knowledge = { readChunks } as unknown as Knowledge const result = await readSourceChunkPage({ + client: { documents: { listChunks } }, knowledge, source: { documentId: "doc_1", title: "notes.pdf", revisionKey: "rev_1" }, params: { page: 2, pageSize: 50 }, @@ -48,8 +53,8 @@ describe("readSourceChunkPage", () => { revisionKey: "rev_1", page: 2, pageSize: 50, - assetUrlPolicy: "durable", }) + expect(listChunks).not.toHaveBeenCalled() expect(result.pagination).toEqual({ page: 2, pageSize: 50, @@ -64,9 +69,68 @@ describe("readSourceChunkPage", () => { }) }) + it("falls back to Knowhere asset URLs when the parsed-storage probe reads remote chunks", async () => { + const readChunks = vi.fn(async () => ({ + document: { + localDocumentId: "doc_1", + resultDirectoryPath: "remote:doc_1", + }, + chunks: [makeReadChunk({ chunkType: "image", filePath: "images/a.png" })], + page: 1, + pageSize: 50, + totalChunks: 1, + totalPages: 1, + })) + const listChunks = vi.fn(async () => ({ + chunks: [ + makeRemoteDocumentChunk({ + id: "document_chunk_1", + chunkId: "parser_1", + chunkType: "image", + content: "Body", + sectionId: null, + sectionPath: "Summary", + sourceChunkPath: "Summary", + filePath: "images/a.png", + sortOrder: 0, + metadata: {}, + assetUrl: "https://knowhere.example/assets/a.png", + }), + ], + pagination: { page: 1, pageSize: 50, total: 1, totalPages: 1 }, + })) + const knowledge = { readChunks } as unknown as Knowledge + + const result = await readSourceChunkPage({ + client: { documents: { listChunks } }, + knowledge, + source: { documentId: "doc_1", title: "notes.pdf", revisionKey: "rev_1" }, + params: { page: 1, pageSize: 50 }, + }) + + expect(readChunks).toHaveBeenCalledWith({ + documentId: "doc_1", + revisionKey: "rev_1", + page: 1, + pageSize: 50, + }) + expect(listChunks).toHaveBeenCalledWith("doc_1", { + page: 1, + pageSize: 50, + includeAssetUrls: true, + }) + expect(result.chunks[0]?.assetUrl).toBe( + "https://knowhere.example/assets/a.png", + ) + }) + it("omits revisionKey when the source has none", async () => { + const listChunks = vi.fn() const readChunks = vi.fn(async () => ({ - document: { localDocumentId: "doc_1" }, + document: { + localDocumentId: "doc_1", + resultDirectoryPath: "parsed-storage:doc_1", + }, chunks: [], page: 1, pageSize: 50, @@ -76,6 +140,7 @@ describe("readSourceChunkPage", () => { const knowledge = { readChunks } as unknown as Knowledge await readSourceChunkPage({ + client: { documents: { listChunks } }, knowledge, source: { documentId: "doc_1", title: "notes.pdf", revisionKey: null }, params: { page: 1, pageSize: 50 }, @@ -85,17 +150,39 @@ describe("readSourceChunkPage", () => { documentId: "doc_1", page: 1, pageSize: 50, - assetUrlPolicy: "durable", }) }) }) +function makeRemoteDocumentChunk( + overrides: Partial = {}, +): DocumentChunk { + return { + id: "document_chunk_1", + chunkId: "parser_1", + chunkType: "text", + content: "Body", + sectionId: null, + sectionPath: "Summary", + sourceChunkPath: "Summary", + filePath: null, + sortOrder: 0, + metadata: {}, + assetUrl: null, + ...overrides, + } +} + describe("readAllSourceChunks", () => { it("pages the SDK to exhaustion", async () => { + const listChunks = vi.fn() const readChunks = vi .fn() .mockResolvedValueOnce({ - document: { localDocumentId: "doc_1" }, + document: { + localDocumentId: "doc_1", + resultDirectoryPath: "parsed-storage:doc_1", + }, chunks: [makeReadChunk({ chunkId: "c1" })], page: 1, pageSize: 200, @@ -103,7 +190,10 @@ describe("readAllSourceChunks", () => { totalPages: 2, }) .mockResolvedValueOnce({ - document: { localDocumentId: "doc_1" }, + document: { + localDocumentId: "doc_1", + resultDirectoryPath: "parsed-storage:doc_1", + }, chunks: [makeReadChunk({ chunkId: "c2" })], page: 2, pageSize: 200, @@ -113,6 +203,7 @@ describe("readAllSourceChunks", () => { const knowledge = { readChunks } as unknown as Knowledge const chunks = await readAllSourceChunks({ + client: { documents: { listChunks } }, knowledge, source: { documentId: "doc_1", title: "notes.pdf", revisionKey: "rev_1" }, }) @@ -123,8 +214,8 @@ describe("readAllSourceChunks", () => { revisionKey: "rev_1", page: 1, pageSize: 200, - assetUrlPolicy: "durable", }) + expect(listChunks).not.toHaveBeenCalled() expect(chunks.map((chunk) => chunk.parserChunkId)).toEqual(["c1", "c2"]) }) }) diff --git a/src/domains/chunks/read.ts b/src/domains/chunks/read.ts index 1bf4790..ae474cf 100644 --- a/src/domains/chunks/read.ts +++ b/src/domains/chunks/read.ts @@ -1,25 +1,61 @@ import "server-only" -import type { Knowledge } from "@ontos-ai/knowhere-sdk" +import type { + DocumentChunk, + Knowledge, + KnowledgeReadChunk, + KnowledgeReadResponse, +} from "@ontos-ai/knowhere-sdk" -import { toParsedChunkViewFromReadChunk, type ChunkPage, type ChunkPageParams } from "@/domains/chunks" +import { + toParsedChunkView, + toParsedChunkViewFromReadChunk, + type ChunkPage, + type ChunkPageParams, +} from "@/domains/chunks" import type { ParsedChunkView } from "@/domains/chunks/types" const loadAllPageSize = 200 +type DisplayReadChunkType = "text" | "image" | "table" | "page" + type ReadableSource = { readonly documentId: string readonly title: string readonly revisionKey?: string | null } +type DisplayReadClient = { + readonly documents: { + listChunks( + documentId: string, + params: { + readonly page: number + readonly pageSize: number + readonly chunkType?: DisplayReadChunkType + readonly includeAssetUrls: true + }, + ): Promise<{ + readonly chunks: readonly DocumentChunk[] + readonly pagination?: { + readonly page?: number + readonly pageSize?: number + readonly total?: number + readonly totalPages?: number + } + }> + } +} + /** * Read a single display page of parsed chunks through the SDK. The SDK serves - * from configured Blob storage when fresh and falls back to Knowhere remote - * otherwise. Display reads request durable asset URLs so media chunks remain - * usable whether the page came from Blob or the remote Knowhere fallback. + * from configured Blob storage when fresh and falls back to Knowhere remote. + * If the probe falls through to remote without asset URLs, fetch the same page + * from Knowhere with short-lived asset URLs instead of hardening assets into + * Notebook storage. */ export async function readSourceChunkPage(input: { + readonly client: DisplayReadClient readonly knowledge: Knowledge readonly source: ReadableSource readonly params: ChunkPageParams @@ -29,22 +65,59 @@ export async function readSourceChunkPage(input: { ...(input.source.revisionKey ? { revisionKey: input.source.revisionKey } : {}), page: input.params.page, pageSize: input.params.pageSize, - assetUrlPolicy: "durable", }) + if (shouldUseKnowledgeChunkResponse(response)) { + return toChunkPageFromKnowledgeResponse(response, input.source, input.params) + } + + const remoteResponse = await input.client.documents.listChunks( + input.source.documentId, + { + page: input.params.page, + pageSize: input.params.pageSize, + includeAssetUrls: true, + }, + ) + const chunks = remoteResponse.chunks.map((chunk) => + toParsedChunkView( + chunk, + input.source.title, + input.source.documentId, + ), + ) + + return { + chunks, + pagination: { + page: remoteResponse.pagination?.page ?? input.params.page, + pageSize: remoteResponse.pagination?.pageSize ?? input.params.pageSize, + total: remoteResponse.pagination?.total ?? chunks.length, + totalPages: + remoteResponse.pagination?.totalPages ?? + Math.max(1, Math.ceil(chunks.length / input.params.pageSize)), + }, + } +} + +function toChunkPageFromKnowledgeResponse( + response: KnowledgeReadResponse, + source: ReadableSource, + params: ChunkPageParams, +): ChunkPage { const chunks = response.chunks.map((chunk) => - toParsedChunkViewFromReadChunk(chunk, input.source.title, input.source.documentId), + toParsedChunkViewFromReadChunk(chunk, source.title, source.documentId), ) return { chunks, pagination: { - page: response.page ?? input.params.page, - pageSize: response.pageSize ?? input.params.pageSize, + page: response.page ?? params.page, + pageSize: response.pageSize ?? params.pageSize, total: response.totalChunks ?? chunks.length, totalPages: response.totalPages ?? - Math.max(1, Math.ceil(chunks.length / input.params.pageSize)), + Math.max(1, Math.ceil(chunks.length / params.pageSize)), }, } } @@ -54,6 +127,7 @@ export async function readSourceChunkPage(input: { * the tree view and load-all display mode. */ export async function readAllSourceChunks(input: { + readonly client: DisplayReadClient readonly knowledge: Knowledge readonly source: ReadableSource }): Promise { @@ -69,20 +143,85 @@ export async function readAllSourceChunks(input: { : {}), page, pageSize: loadAllPageSize, - assetUrlPolicy: "durable", }) - for (const chunk of response.chunks) { + if (shouldUseKnowledgeChunkResponse(response)) { + for (const chunk of response.chunks) { + chunks.push( + toParsedChunkViewFromReadChunk( + chunk, + input.source.title, + input.source.documentId, + ), + ) + } + totalPages = Math.max(1, response.totalPages ?? 1) + page += 1 + continue + } + + const remoteResponse = await input.client.documents.listChunks( + input.source.documentId, + { + page, + pageSize: loadAllPageSize, + includeAssetUrls: true, + }, + ) + for (const chunk of remoteResponse.chunks) { chunks.push( - toParsedChunkViewFromReadChunk( + toParsedChunkView( chunk, input.source.title, input.source.documentId, ), ) } - totalPages = Math.max(1, response.totalPages ?? 1) + totalPages = Math.max(1, remoteResponse.pagination?.totalPages ?? 1) page += 1 } while (page <= totalPages) return chunks } + +function shouldUseKnowledgeChunkResponse( + response: KnowledgeReadResponse, +): boolean { + return isParsedStorageReadResponse(response) || hasUsableReadChunkAssetUrl(response.chunks) +} + +function isParsedStorageReadResponse(response: KnowledgeReadResponse): boolean { + const resultDirectoryPath = response.document.resultDirectoryPath + return ( + typeof resultDirectoryPath === "string" && + resultDirectoryPath.startsWith("parsed-storage:") + ) +} + +function hasUsableReadChunkAssetUrl( + chunks: readonly KnowledgeReadChunk[], +): boolean { + return chunks.some( + (chunk) => + hasNonEmptyString(chunk.assetUrl) || + hasMetadataPageAssetUrl(chunk.metadata), + ) +} + +function hasMetadataPageAssetUrl( + metadata: Readonly>, +): boolean { + const value = metadata.pageAssets + if (!Array.isArray(value)) return false + + return value.some( + (item) => isRecord(item) && hasNonEmptyString(item.assetUrl), + ) +} + +function hasNonEmptyString(value: unknown): value is string { + return typeof value === "string" && value.trim().length > 0 +} + +function isRecord(value: unknown): value is Readonly> { + return typeof value === "object" && value !== null +} diff --git a/src/domains/sources/page-assets.ts b/src/domains/sources/page-assets.ts index cfb2765..c23a702 100644 --- a/src/domains/sources/page-assets.ts +++ b/src/domains/sources/page-assets.ts @@ -1,6 +1,10 @@ import "server-only" -import type { Knowledge } from "@ontos-ai/knowhere-sdk" +import type { + DocumentChunk, + Knowledge, + KnowledgeReadResponse, +} from "@ontos-ai/knowhere-sdk" import type { ChunkPageParams } from "@/domains/chunks" import type { SourcePageAssetView } from "./route-types" @@ -20,7 +24,30 @@ export type SourcePageAssetsPage = { } } +type PageAssetReadClient = { + readonly documents: { + listChunks( + documentId: string, + params: { + readonly page: number + readonly pageSize: number + readonly chunkType: "page" + readonly includeAssetUrls: true + }, + ): Promise<{ + readonly chunks: readonly DocumentChunk[] + readonly pagination?: { + readonly page?: number + readonly pageSize?: number + readonly total?: number + readonly totalPages?: number + } + }> + } +} + export async function readSourcePageAssets(input: { + readonly client: PageAssetReadClient readonly knowledge: Knowledge readonly source: ReadableSource readonly params: ChunkPageParams @@ -31,8 +58,36 @@ export async function readSourcePageAssets(input: { chunkType: "page", page: input.params.page, pageSize: input.params.pageSize, - assetUrlPolicy: "durable", }) + + const knowledgePage = toSourcePageAssetsPageFromKnowledgeResponse( + response, + input.params, + ) + if ( + isParsedStorageReadResponse(response) && + knowledgePage.pages.length > 0 + ) { + return knowledgePage + } + if (knowledgePage.pages.length > 0) return knowledgePage + + const remoteResponse = await input.client.documents.listChunks( + input.source.documentId, + { + page: input.params.page, + pageSize: input.params.pageSize, + chunkType: "page", + includeAssetUrls: true, + }, + ) + return toSourcePageAssetsPageFromRemoteResponse(remoteResponse, input.params) +} + +function toSourcePageAssetsPageFromKnowledgeResponse( + response: KnowledgeReadResponse, + params: ChunkPageParams, +): SourcePageAssetsPage { const pages = response.chunks.flatMap((chunk): SourcePageAssetView[] => readPageAssetViews(chunk.metadata.pageAssets, chunk.assetUrl), ) @@ -41,21 +96,67 @@ export async function readSourcePageAssets(input: { const resolvedTotal = total > 0 ? total : pages.length const computedTotalPages = Math.max( 1, - Math.ceil(resolvedTotal / input.params.pageSize), + Math.ceil(resolvedTotal / params.pageSize), ) const totalPages = Math.max(response.totalPages ?? 0, computedTotalPages) return { pages, pagination: { - page: response.page ?? input.params.page, - pageSize: response.pageSize ?? input.params.pageSize, + page: response.page ?? params.page, + pageSize: response.pageSize ?? params.pageSize, total: resolvedTotal, totalPages, }, } } +function toSourcePageAssetsPageFromRemoteResponse( + response: { + readonly chunks: readonly DocumentChunk[] + readonly pagination?: { + readonly page?: number + readonly pageSize?: number + readonly total?: number + readonly totalPages?: number + } + }, + params: ChunkPageParams, +): SourcePageAssetsPage { + const pages = response.chunks.flatMap((chunk): SourcePageAssetView[] => + readPageAssetViews(chunk.metadata.pageAssets, chunk.assetUrl ?? undefined), + ) + const maxPageNumber = getMaxPageNumber(pages) + const total = Math.max(response.pagination?.total ?? 0, maxPageNumber ?? 0) + const resolvedTotal = total > 0 ? total : pages.length + const computedTotalPages = Math.max( + 1, + Math.ceil(resolvedTotal / params.pageSize), + ) + const totalPages = Math.max( + response.pagination?.totalPages ?? 0, + computedTotalPages, + ) + + return { + pages, + pagination: { + page: response.pagination?.page ?? params.page, + pageSize: response.pagination?.pageSize ?? params.pageSize, + total: resolvedTotal, + totalPages, + }, + } +} + +function isParsedStorageReadResponse(response: KnowledgeReadResponse): boolean { + const resultDirectoryPath = response.document.resultDirectoryPath + return ( + typeof resultDirectoryPath === "string" && + resultDirectoryPath.startsWith("parsed-storage:") + ) +} + function readPageAssetViews( value: unknown, fallbackAssetUrl?: string, diff --git a/src/domains/sources/route-chunks.ts b/src/domains/sources/route-chunks.ts index 8d9d315..c636cb6 100644 --- a/src/domains/sources/route-chunks.ts +++ b/src/domains/sources/route-chunks.ts @@ -13,7 +13,7 @@ import { } from "./remote-library" import { getClientForWorkspace, - getKnowledgeForSource, + getKnowledgeResourcesForSource, } from "./route-dependencies" import { sourceRowRepository } from "./source-row-repository" import type { @@ -99,7 +99,7 @@ const loadSourceChunksEffect = ( const apiKey = yield* Effect.tryPromise(() => deps.ensureApiKeyForWorkspace(workspace.id, input.cookieHeader), ) - const knowledge = getKnowledgeForSource({ + const readResources = getKnowledgeResourcesForSource({ apiKey, workspaceId: workspace.id, sourceId: source.id, @@ -114,7 +114,11 @@ const loadSourceChunksEffect = ( if (input.shouldLoadAll) { return yield* Effect.tryPromise(() => - readAllSourceChunks({ knowledge, source: readableSource }), + readAllSourceChunks({ + client: readResources.client, + knowledge: readResources.knowledge, + source: readableSource, + }), ).pipe( Effect.map((chunks) => routeResult.ok({ chunks: resolveChunkConnectionTargets(chunks) }), @@ -125,7 +129,8 @@ const loadSourceChunksEffect = ( return yield* Effect.tryPromise(() => readSourceChunkPage({ - knowledge, + client: readResources.client, + knowledge: readResources.knowledge, source: readableSource, params: input.pageParams, }), @@ -175,7 +180,7 @@ const loadRemoteChunkPageEffect = ( ) const documentId = source.knowhereDocumentId ?? remoteDocument.documentId - const knowledge = getKnowledgeForSource({ + const readResources = getKnowledgeResourcesForSource({ apiKey, workspaceId: workspace.id, sourceId: source.id, @@ -190,7 +195,11 @@ const loadRemoteChunkPageEffect = ( if (input.shouldLoadAll) { return yield* Effect.tryPromise(() => - readAllSourceChunks({ knowledge, source: readableSource }), + readAllSourceChunks({ + client: readResources.client, + knowledge: readResources.knowledge, + source: readableSource, + }), ).pipe( Effect.map((chunks) => routeResult.ok({ chunks: resolveChunkConnectionTargets(chunks) }), @@ -201,7 +210,8 @@ const loadRemoteChunkPageEffect = ( return yield* Effect.tryPromise(() => readSourceChunkPage({ - knowledge, + client: readResources.client, + knowledge: readResources.knowledge, source: readableSource, params: input.pageParams, }), diff --git a/src/domains/sources/route-dependencies.ts b/src/domains/sources/route-dependencies.ts index a51a9b7..b0cc887 100644 --- a/src/domains/sources/route-dependencies.ts +++ b/src/domains/sources/route-dependencies.ts @@ -99,6 +99,16 @@ function getKnowledgeForSource(input: { readonly documentId: string readonly revisionKey?: string | null }): Knowledge { + return getKnowledgeResourcesForSource(input).knowledge +} + +function getKnowledgeResourcesForSource(input: { + readonly apiKey: string + readonly workspaceId: string + readonly sourceId: string + readonly documentId: string + readonly revisionKey?: string | null +}): { readonly client: SourceRouteKnowhereClient; readonly knowledge: Knowledge } { const scheduler = createParsedDocumentSyncScheduler({ workspaceId: input.workspaceId, sourceId: input.sourceId, @@ -106,15 +116,16 @@ function getKnowledgeForSource(input: { apiKey: input.apiKey, revisionKey: input.revisionKey ?? undefined, }) - const { knowledge } = makeKnowhereClientWithParsedStorage(input.apiKey, { + const resources = makeKnowhereClientWithParsedStorage(input.apiKey, { workspaceId: input.workspaceId, scheduler, }) - return knowledge + return { client: resources.client, knowledge: resources.knowledge } } export { createSourceRouteDependencies, getClientForWorkspace, getKnowledgeForSource, + getKnowledgeResourcesForSource, } diff --git a/src/domains/sources/route-page-assets.ts b/src/domains/sources/route-page-assets.ts index 3077a83..a2316fd 100644 --- a/src/domains/sources/route-page-assets.ts +++ b/src/domains/sources/route-page-assets.ts @@ -9,7 +9,7 @@ import { } from "./remote-library" import { getClientForWorkspace, - getKnowledgeForSource, + getKnowledgeResourcesForSource, } from "./route-dependencies" import { sourceRowRepository } from "./source-row-repository" import type { @@ -71,7 +71,7 @@ const loadSourcePageAssetsEffect = ( const apiKey = yield* Effect.tryPromise(() => deps.ensureApiKeyForWorkspace(workspace.id, input.cookieHeader), ) - const knowledge = getKnowledgeForSource({ + const readResources = getKnowledgeResourcesForSource({ apiKey, workspaceId: workspace.id, sourceId: source.id, @@ -80,7 +80,8 @@ const loadSourcePageAssetsEffect = ( }) return yield* Effect.tryPromise(() => readSourcePageAssets({ - knowledge, + client: readResources.client, + knowledge: readResources.knowledge, source: { documentId, revisionKey: source.knowhereJobId, @@ -136,7 +137,7 @@ const loadRemotePageAssetsEffect = ( const documentId = source.knowhereDocumentId ?? remoteDocument.documentId const revisionKey = source.knowhereJobId ?? remoteDocument.revisionKey ?? null - const knowledge = getKnowledgeForSource({ + const readResources = getKnowledgeResourcesForSource({ apiKey, workspaceId: workspace.id, sourceId: source.id, @@ -145,7 +146,8 @@ const loadRemotePageAssetsEffect = ( }) return yield* Effect.tryPromise(() => readSourcePageAssets({ - knowledge, + client: readResources.client, + knowledge: readResources.knowledge, source: { documentId, revisionKey }, params: input.pageParams, }), From 2d825942e48de7cdded607837e43d6cdb1f783d6 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 7 Jul 2026 20:08:06 +0800 Subject: [PATCH 05/14] Adapt parsed document storage layout --- .../[sourceId]/page-assets/route.test.ts | 87 +++++++++ .../chat/media-asset-hardening.test.ts | 8 +- src/domains/chat/route-service.test.ts | 8 +- src/domains/chunks/read.test.ts | 32 ++++ .../parsed-document-blob-storage.test.ts | 73 ++++++- .../sources/parsed-document-blob-storage.ts | 179 ++++++++++++++---- 6 files changed, 335 insertions(+), 52 deletions(-) diff --git a/src/app/api/sources/[sourceId]/page-assets/route.test.ts b/src/app/api/sources/[sourceId]/page-assets/route.test.ts index 76ef16a..4f00bd3 100644 --- a/src/app/api/sources/[sourceId]/page-assets/route.test.ts +++ b/src/app/api/sources/[sourceId]/page-assets/route.test.ts @@ -156,6 +156,93 @@ describe("GET /api/sources/[sourceId]/page-assets", () => { expect(mocks.listChunks).not.toHaveBeenCalled() }) + it("uses SDK page asset URLs directly when the SDK remote fallback is usable", async () => { + mocks.getCurrentUser.mockResolvedValue({ + id: "user_1", + email: null, + name: null, + }) + mocks.ensureWorkspace.mockResolvedValue({ + id: "workspace_1", + userId: "user_1", + namespace: "notebook-workspace_1", + createdAt: new Date("2026-05-10T00:00:00.000Z"), + }) + mocks.findSourceInWorkspace.mockResolvedValue( + makeReadySource({ + id: "00000000-0000-0000-0000-000000000002", + knowhereJobId: "job_1", + knowhereDocumentId: "doc_1", + }), + ) + mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") + mocks.readChunks.mockResolvedValue({ + document: { + localDocumentId: "doc_1", + resultDirectoryPath: "remote:doc_1", + }, + chunks: [ + { + chunkId: "page_1", + chunkType: "page", + metadata: { + pageAssets: [ + { + pageNum: 1, + assetUrl: "https://sdk.example/page-000001.png", + contentType: "image/png", + width: 1200, + height: 1600, + }, + ], + }, + }, + ], + page: 1, + pageSize: 1, + totalChunks: 1, + totalPages: 1, + }) + + const response = await GET( + new NextRequest( + "http://localhost:3001/api/sources/00000000-0000-0000-0000-000000000002/page-assets?page=1&pageSize=1", + ), + { + params: Promise.resolve({ + sourceId: "00000000-0000-0000-0000-000000000002", + }), + }, + ) + + await expect(response.json()).resolves.toEqual({ + pages: [ + { + pageNumber: 1, + assetUrl: "https://sdk.example/page-000001.png", + contentType: "image/png", + width: 1200, + height: 1600, + }, + ], + pagination: { + page: 1, + pageSize: 1, + total: 1, + totalPages: 1, + }, + }) + expect(response.status).toBe(200) + expect(mocks.readChunks).toHaveBeenCalledWith({ + documentId: "doc_1", + revisionKey: "job_1", + chunkType: "page", + page: 1, + pageSize: 1, + }) + expect(mocks.listChunks).not.toHaveBeenCalled() + }) + it("rejects non-ready workspace sources", async () => { mocks.getCurrentUser.mockResolvedValue({ id: "user_1", diff --git a/src/domains/chat/media-asset-hardening.test.ts b/src/domains/chat/media-asset-hardening.test.ts index 7c106c9..2afa6d9 100644 --- a/src/domains/chat/media-asset-hardening.test.ts +++ b/src/domains/chat/media-asset-hardening.test.ts @@ -24,7 +24,7 @@ afterEach(() => { describe("hardenChatMediaAssetUrls", () => { it("keeps an already Notebook-owned asset URL without calling the hardener", async () => { const ownedUrl = - "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_1/rev_1/assets/images/a.png" + "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_1/rev_1/images/a.png" const hardenChatAssetUrl = vi.fn(async () => null) const result = await hardenChatMediaAssetUrls({ @@ -52,7 +52,7 @@ describe("hardenChatMediaAssetUrls", () => { const rawAssetUrl = "https://knowhere-storage.example/results/job_1/images/id-front.jpg?AWSAccessKeyId=test" const durableUrl = - "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_identity/rev_1/assets/images/id-front.jpg" + "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_identity/rev_1/images/id-front.jpg" const hardenChatAssetUrl = vi.fn().mockResolvedValue(durableUrl) const result = await hardenChatMediaAssetUrls({ @@ -116,7 +116,7 @@ describe("hardenChatMediaAssetUrls", () => { const rawAssetUrl = "https://knowhere-storage.example/results/job_1/images/front.jpg?AWSAccessKeyId=test" const durableUrl = - "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_identity/rev_1/assets/images/front.jpg" + "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_identity/rev_1/images/front.jpg" const hardenChatAssetUrl = vi.fn().mockResolvedValue(durableUrl) const result = await hardenChatMediaAssetUrls({ @@ -164,7 +164,7 @@ describe("isNotebookOwnedAssetUrl", () => { ).toBe(true) expect( isNotebookOwnedAssetUrl( - "https://cdn.example/workspaces/w/parsed-documents/d/r/assets/a.png", + "https://cdn.example/workspaces/w/parsed-documents/d/r/images/a.png", ), ).toBe(true) expect( diff --git a/src/domains/chat/route-service.test.ts b/src/domains/chat/route-service.test.ts index c771f12..d66ae68 100644 --- a/src/domains/chat/route-service.test.ts +++ b/src/domains/chat/route-service.test.ts @@ -183,7 +183,7 @@ describe("chat route services", () => { }) const rawUrl = "https://knowhere-storage.example/results/job_1/pages/page-1.png" const durableUrl = - "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_legacy/job_1/assets/pages/page-1.png" + "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_legacy/job_1/pages/page-1.png" mocks.parsedStorageWriteAsset.mockResolvedValue({ url: durableUrl }) mocks.getAuthenticatedWithClient.mockResolvedValue({ user: { id: "user_1" }, @@ -249,7 +249,7 @@ describe("chat route services", () => { const rawUrl = "https://knowhere-storage.example/results/job_1/images/id-front.png?AWSAccessKeyId=test" const durableUrl = - "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_identity/job_1/assets/images/id-front.png" + "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_identity/job_1/images/id-front.png" mocks.parsedStorageWriteAsset.mockResolvedValue({ url: durableUrl }) mocks.generateText.mockResolvedValue({ text: `The card number is visible. ${durableUrl} ${rawUrl}`, @@ -364,7 +364,7 @@ describe("chat route services", () => { knowhereJobId: "job_1", }) const durableUrl = - "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_contract/job_1/assets/page_citation_assets/page-8.png" + "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_contract/job_1/page_citation_assets/page-8.png" mocks.parsedStorageGetAssetUrl.mockResolvedValue(durableUrl) mocks.generateText.mockResolvedValue({ text: "The page states 5000 yuan per occurrence.", @@ -467,7 +467,7 @@ describe("chat route services", () => { const rawUrl = "https://knowhere-storage.example/results/job_remote/page_citation_assets/page-8.png?AWSAccessKeyId=test" const durableUrl = - "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_remote/job_remote/assets/page_citation_assets/page-8.png" + "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_remote/job_remote/page_citation_assets/page-8.png" mocks.parsedStorageWriteAsset.mockResolvedValue({ url: durableUrl }) mocks.generateText.mockResolvedValue({ text: "The page states 5000 yuan per occurrence.", diff --git a/src/domains/chunks/read.test.ts b/src/domains/chunks/read.test.ts index 9ba2203..f63c897 100644 --- a/src/domains/chunks/read.test.ts +++ b/src/domains/chunks/read.test.ts @@ -124,6 +124,38 @@ describe("readSourceChunkPage", () => { ) }) + it("uses SDK remote chunks directly when the SDK returns usable asset URLs", async () => { + const listChunks = vi.fn() + const readChunks = vi.fn(async () => ({ + document: { + localDocumentId: "doc_1", + resultDirectoryPath: "remote:doc_1", + }, + chunks: [ + makeReadChunk({ + chunkType: "image", + filePath: "images/a.png", + assetUrl: "https://sdk.example/assets/a.png", + }), + ], + page: 1, + pageSize: 50, + totalChunks: 1, + totalPages: 1, + })) + const knowledge = { readChunks } as unknown as Knowledge + + const result = await readSourceChunkPage({ + client: { documents: { listChunks } }, + knowledge, + source: { documentId: "doc_1", title: "notes.pdf", revisionKey: "rev_1" }, + params: { page: 1, pageSize: 50 }, + }) + + expect(listChunks).not.toHaveBeenCalled() + expect(result.chunks[0]?.assetUrl).toBe("https://sdk.example/assets/a.png") + }) + it("omits revisionKey when the source has none", async () => { const listChunks = vi.fn() const readChunks = vi.fn(async () => ({ diff --git a/src/domains/sources/parsed-document-blob-storage.test.ts b/src/domains/sources/parsed-document-blob-storage.test.ts index 8a3c427..4c3ee16 100644 --- a/src/domains/sources/parsed-document-blob-storage.test.ts +++ b/src/domains/sources/parsed-document-blob-storage.test.ts @@ -34,6 +34,7 @@ function createFakeBlobStore(): { return { statusCode: 200, stream: new Response(body).body as ReadableStream, + contentType: object.contentType, } }, put: async (pathname, body, options) => { @@ -93,7 +94,7 @@ const chunkPage: KnowhereParsedSnapshotChunkPage = { } describe("BlobParsedDocumentStorage", () => { - it("round-trips a manifest keyed by workspace/document/revision", async () => { + it("round-trips a manifest at the result-relative manifest path", async () => { const { store, objects } = createFakeBlobStore() const storage = new BlobParsedDocumentStorage({ workspaceId: "ws_1", @@ -103,20 +104,23 @@ describe("BlobParsedDocumentStorage", () => { await storage.writeManifest({ documentId, revisionKey, manifest }) expect([...objects.keys()]).toContain( - "workspaces/ws_1/parsed-documents/doc_123/job_result_1/manifest/current.json", + "workspaces/ws_1/parsed-documents/doc_123/job_result_1/manifest.json", ) const read = await storage.readManifest({ documentId, revisionKey }) expect(read).toEqual(manifest) }) - it("round-trips a chunk page", async () => { - const { store } = createFakeBlobStore() + it("round-trips a chunk page at its result-relative chunk path", async () => { + const { store, objects } = createFakeBlobStore() const storage = new BlobParsedDocumentStorage({ workspaceId: "ws_1", blobStore: store, }) await storage.writeChunkPage({ documentId, revisionKey, page: chunkPage }) + expect([...objects.keys()]).toContain( + "workspaces/ws_1/parsed-documents/doc_123/job_result_1/chunks/page-1.json", + ) const read = await storage.readChunkPage({ documentId, revisionKey, page: 1 }) expect(read).toEqual(chunkPage) }) @@ -141,7 +145,7 @@ describe("BlobParsedDocumentStorage", () => { expect(read).toEqual(progress) }) - it("writes an asset and resolves a durable URL", async () => { + it("writes an asset without adding a Notebook logical assets prefix", async () => { const { store } = createFakeBlobStore() const storage = new BlobParsedDocumentStorage({ workspaceId: "ws_1", @@ -156,7 +160,7 @@ describe("BlobParsedDocumentStorage", () => { contentType: "image/png", }) expect(written.url).toContain( - "workspaces/ws_1/parsed-documents/doc_123/job_result_1/assets/images/fig-1.png", + "workspaces/ws_1/parsed-documents/doc_123/job_result_1/images/fig-1.png", ) const url = await storage.getAssetUrl({ @@ -167,6 +171,56 @@ describe("BlobParsedDocumentStorage", () => { expect(url).toBe(written.url) }) + it("round-trips arbitrary result-relative objects", async () => { + const { store, objects } = createFakeBlobStore() + const storage = new BlobParsedDocumentStorage({ + workspaceId: "ws_1", + blobStore: store, + }) + + await storage.writeObject({ + documentId, + revisionKey, + path: "page_citation_assets/page-1.png", + body: new Uint8Array([4, 5, 6]), + contentType: "image/png", + }) + + expect([...objects.keys()]).toContain( + "workspaces/ws_1/parsed-documents/doc_123/job_result_1/page_citation_assets/page-1.png", + ) + const read = await storage.readObject({ + documentId, + revisionKey, + path: "page_citation_assets/page-1.png", + }) + expect(read?.body).toEqual(Buffer.from([4, 5, 6])) + expect(read?.contentType).toBe("image/png") + }) + + it("resolves legacy asset-prefixed URLs without writing new assets there", async () => { + const { store, objects } = createFakeBlobStore() + const storage = new BlobParsedDocumentStorage({ + workspaceId: "ws_1", + blobStore: store, + }) + const legacyKey = + "workspaces/ws_1/parsed-documents/doc_123/job_result_1/assets/images/fig-1.png" + objects.set(legacyKey, { + body: Buffer.from([1]), + contentType: "image/png", + }) + + const url = await storage.getAssetUrl({ + documentId, + revisionKey, + sourcePath: "images/fig-1.png", + }) + expect(url).toBe( + "https://fake.public.blob.vercel-storage.com/workspaces/ws_1/parsed-documents/doc_123/job_result_1/assets/images/fig-1.png", + ) + }) + it("returns null for a missing manifest, chunk page, progress, and asset", async () => { const { store } = createFakeBlobStore() const storage = new BlobParsedDocumentStorage({ @@ -219,5 +273,12 @@ describe("BlobParsedDocumentStorage", () => { sourcePath: "../../etc/passwd", }), ).rejects.toThrow(/Invalid parsed storage path/) + await expect( + storage.readObject({ + documentId, + revisionKey, + path: "images/../escape.png", + }), + ).rejects.toThrow(/Invalid parsed storage path/) }) }) diff --git a/src/domains/sources/parsed-document-blob-storage.ts b/src/domains/sources/parsed-document-blob-storage.ts index 23a93a0..8f19837 100644 --- a/src/domains/sources/parsed-document-blob-storage.ts +++ b/src/domains/sources/parsed-document-blob-storage.ts @@ -26,14 +26,17 @@ import type { */ const parsedDocumentsDirectoryName = "parsed-documents" -const manifestStoragePath = "manifest/current.json" +const manifestStoragePath = "manifest.json" +const legacyManifestStoragePath = "manifest/current.json" const syncProgressStoragePath = "sync-progress.json" const jsonContentType = "application/json; charset=utf-8" +const binaryContentType = "application/octet-stream" type BlobGetResult = | { readonly statusCode: 200 readonly stream: ReadableStream + readonly contentType?: string } | { readonly statusCode: 304 @@ -77,8 +80,34 @@ export type BlobParsedDocumentStorageInput = { readonly blobStore?: ParsedDocumentBlobStore } +export type ParsedDocumentStorageObject = ParsedDocumentStorageDocument & { + readonly path: string +} + +export type ParsedDocumentStorageWritableObject = + ParsedDocumentStorageObject & { + readonly body: string | Buffer | Uint8Array + readonly contentType?: string + } + +export type ParsedDocumentStorageReadableObject = { + readonly body: Buffer + readonly contentType?: string +} + const vercelBlobStore: ParsedDocumentBlobStore = { - get: (pathname, options) => get(pathname, options), + get: async (pathname, options) => { + const blob = await get(pathname, options) + if (!blob) return null + if (blob.statusCode === 304) { + return { statusCode: 304, stream: null } + } + return { + statusCode: 200, + stream: blob.stream, + contentType: blob.blob.contentType, + } + }, put: async (pathname, body, options) => { const blob = await put(pathname, body, { access: options.access, @@ -113,8 +142,21 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { async readManifest( params: ParsedDocumentStorageManifestParams, ): Promise { + const current = await this.readJson( + this.getObjectKey({ + documentId: params.documentId, + revisionKey: params.revisionKey, + path: manifestStoragePath, + }), + ) + if (current) return current + return this.readJson( - this.getManifestKey(params.documentId, params.revisionKey), + this.getObjectKey({ + documentId: params.documentId, + revisionKey: params.revisionKey, + path: legacyManifestStoragePath, + }), ) } @@ -124,7 +166,11 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { readonly manifest: KnowhereParsedSnapshotManifest }): Promise { await this.writeJson( - this.getManifestKey(params.documentId, params.revisionKey), + this.getObjectKey({ + documentId: params.documentId, + revisionKey: params.revisionKey, + path: manifestStoragePath, + }), params.manifest, ) } @@ -134,7 +180,11 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { ): Promise { // chunkType filtering happens SDK-side after read; storage returns the full page. return this.readJson( - this.getChunkPageKey(params.documentId, params.revisionKey, params.page), + this.getObjectKey({ + documentId: params.documentId, + revisionKey: params.revisionKey, + path: getChunkPageStoragePath(params.page), + }), ) } @@ -144,11 +194,11 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { readonly page: KnowhereParsedSnapshotChunkPage }): Promise { await this.writeJson( - this.getChunkPageKey( - params.documentId, - params.revisionKey, - params.page.page, - ), + this.getObjectKey({ + documentId: params.documentId, + revisionKey: params.revisionKey, + path: getChunkPageStoragePath(params.page.page), + }), params.page, ) } @@ -157,7 +207,11 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { params: ParsedDocumentStorageDocument & ParsedDocumentStorageAsset, ): Promise<{ readonly sourcePath: string; readonly url?: string }> { const blob = await this.blobStore.put( - this.getAssetKey(params.documentId, params.revisionKey, params.sourcePath), + this.getObjectKey({ + documentId: params.documentId, + revisionKey: params.revisionKey, + path: params.sourcePath, + }), Buffer.from(params.body), { access: "public", @@ -173,26 +227,80 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { params: ParsedDocumentStorageAssetParams, ): Promise { const result = await this.blobStore.head( - this.getAssetKey(params.documentId, params.revisionKey, params.sourcePath), + this.getObjectKey({ + documentId: params.documentId, + revisionKey: params.revisionKey, + path: params.sourcePath, + }), ) - return result?.url ?? null + if (result) return result.url + + const legacyResult = await this.blobStore.head( + this.getObjectKey({ + documentId: params.documentId, + revisionKey: params.revisionKey, + path: `assets/${params.sourcePath}`, + }), + ) + return legacyResult?.url ?? null } async readSyncProgress( params: ParsedDocumentStorageDocument, ): Promise { return this.readJson( - this.getSyncProgressKey(params.documentId, params.revisionKey), + this.getObjectKey({ + documentId: params.documentId, + revisionKey: params.revisionKey, + path: syncProgressStoragePath, + }), ) } async writeSyncProgress(params: ParsedDocumentSyncProgress): Promise { await this.writeJson( - this.getSyncProgressKey(params.documentId, params.revisionKey), + this.getObjectKey({ + documentId: params.documentId, + revisionKey: params.revisionKey, + path: syncProgressStoragePath, + }), params, ) } + async readObject( + params: ParsedDocumentStorageObject, + ): Promise { + return this.readBlobObject(this.getObjectKey(params)) + } + + async writeObject( + params: ParsedDocumentStorageWritableObject, + ): Promise<{ readonly path: string; readonly url?: string }> { + const blob = await this.blobStore.put( + this.getObjectKey(params), + Buffer.from(params.body), + { + access: "public", + allowOverwrite: true, + contentType: params.contentType ?? binaryContentType, + multipart: true, + }, + ) + return { path: params.path, url: blob.url } + } + + async getObjectUrl( + params: ParsedDocumentStorageObject, + ): Promise { + const result = await this.blobStore.head(this.getObjectKey(params)) + return result?.url ?? null + } + + async deleteObject(params: ParsedDocumentStorageObject): Promise { + await this.blobStore.del(this.getObjectKey(params)) + } + private getRevisionPrefix(documentId: string, revisionKey: string): string { return [ "workspaces", @@ -203,28 +311,8 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { ].join("/") } - private getManifestKey(documentId: string, revisionKey: string): string { - return `${this.getRevisionPrefix(documentId, revisionKey)}/${manifestStoragePath}` - } - - private getChunkPageKey( - documentId: string, - revisionKey: string, - page: number, - ): string { - return `${this.getRevisionPrefix(documentId, revisionKey)}/chunks/page-${page}.json` - } - - private getSyncProgressKey(documentId: string, revisionKey: string): string { - return `${this.getRevisionPrefix(documentId, revisionKey)}/${syncProgressStoragePath}` - } - - private getAssetKey( - documentId: string, - revisionKey: string, - sourcePath: string, - ): string { - return `${this.getRevisionPrefix(documentId, revisionKey)}/assets/${normalizeRelativeStoragePath(sourcePath)}` + private getObjectKey(params: ParsedDocumentStorageObject): string { + return `${this.getRevisionPrefix(params.documentId, params.revisionKey)}/${normalizeRelativeStoragePath(params.path)}` } private async readJson(key: string): Promise { @@ -246,10 +334,21 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { } private async readBlobText(key: string): Promise { + const object = await this.readBlobObject(key) + return object ? object.body.toString("utf8") : null + } + + private async readBlobObject( + key: string, + ): Promise { try { const result = await this.blobStore.get(key, { access: "public" }) if (!result || result.statusCode !== 200) return null - return await new Response(result.stream).text() + const body = Buffer.from(await new Response(result.stream).arrayBuffer()) + return { + body, + ...(result.contentType ? { contentType: result.contentType } : {}), + } } catch (error) { if (error instanceof BlobNotFoundError) return null throw error @@ -257,6 +356,10 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { } } +function getChunkPageStoragePath(page: number): string { + return `chunks/page-${page}.json` +} + /** * Reject path segments that could traverse outside the intended prefix. Mirrors * the SDK `DiskParsedDocumentStorage` guard so blob keys stay well-formed. From fde80272615ef99f5ff5800cfd31c5080a7ca106 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 7 Jul 2026 23:24:40 +0800 Subject: [PATCH 06/14] Align Notebook with SDK result layout storage --- src/domains/chat/index.test.ts | 10 +- src/domains/chat/index.ts | 5 +- src/domains/chat/service.test.ts | 4 +- .../parsed-document-blob-storage.test.ts | 128 +++++++++----- .../sources/parsed-document-blob-storage.ts | 165 +++++++----------- src/integrations/knowhere.ts | 2 - 6 files changed, 154 insertions(+), 160 deletions(-) diff --git a/src/domains/chat/index.test.ts b/src/domains/chat/index.test.ts index 8c239a9..9a34a4b 100644 --- a/src/domains/chat/index.test.ts +++ b/src/domains/chat/index.test.ts @@ -78,7 +78,7 @@ describe("answerQuestionWithRetrieval", () => { namespace: "notebook-workspace", query: "What does the document say?", topK: 8, - useAgentic: true, + useAgentic: false, dataType: 1, excludeDocumentIds: ["doc_excluded", "doc_remote"], }); @@ -579,7 +579,7 @@ describe("answerQuestionWithRetrieval", () => { namespace: "notebook-workspace", query: "SpaceX rocket photos", topK: 8, - useAgentic: true, + useAgentic: false, dataType: 3, }); expect(answer.answer).toBe("Use this launch photo."); @@ -1573,7 +1573,7 @@ describe("answerQuestionWithRetrieval", () => { namespace: "notebook-workspace", query: "公民身份证明 图片", topK: 8, - useAgentic: true, + useAgentic: false, dataType: 3, }); const imageCitations = answer.citations.filter( @@ -1664,7 +1664,7 @@ describe("answerQuestionWithRetrieval", () => { namespace: "notebook-workspace", query: "Tesla Q4 2025 Update energy generation and storage deployments", topK: 8, - useAgentic: true, + useAgentic: false, dataType: 1, }); expect(generateAnswer).toHaveBeenCalledWith({ @@ -1720,7 +1720,7 @@ describe("answerQuestionWithRetrieval", () => { namespace: "notebook-workspace", query: "Tesla energy storage deployments", topK: 8, - useAgentic: true, + useAgentic: false, dataType: 1, }); expect(JSON.stringify(queryInput)).not.toContain( diff --git a/src/domains/chat/index.ts b/src/domains/chat/index.ts index fa07579..4f5dee6 100644 --- a/src/domains/chat/index.ts +++ b/src/domains/chat/index.ts @@ -42,6 +42,9 @@ import { enrichRetrievalResultsWithPageCitationAssetUrls } from "./page-citation import type { HardenableRetrievalResult } from "./media-asset-hardening" const DEFAULT_TOP_K = 8 +const NOTEBOOK_USE_AGENTIC_RETRIEVAL: NonNullable< + RetrievalQueryParams["useAgentic"] +> = false const MAX_AGENTIC_TOP_K = 12 const MAX_AGENTIC_MERGED_RESULT_COUNT = 24 const MAX_AGENTIC_MERGED_REFERENCED_CHUNK_COUNT = 24 @@ -795,7 +798,7 @@ function buildRetrievalQueryParams(input: { namespace: input.namespace, query, topK: normalizeTopK(input.input.topK), - useAgentic: true, + useAgentic: NOTEBOOK_USE_AGENTIC_RETRIEVAL, dataType, ...(input.input.signalPaths && input.input.signalPaths.length > 0 ? { signalPaths: input.input.signalPaths } diff --git a/src/domains/chat/service.test.ts b/src/domains/chat/service.test.ts index 7c0ce92..12df745 100644 --- a/src/domains/chat/service.test.ts +++ b/src/domains/chat/service.test.ts @@ -57,7 +57,7 @@ describe("handleChatTurn", () => { namespace: "notebook-namespace", query: "What does the document say?", topK: 8, - useAgentic: true, + useAgentic: false, dataType: 1, excludeDocumentIds: ["doc_excluded"], }); @@ -221,7 +221,7 @@ describe("handleChatTurn", () => { namespace: "notebook-namespace", query: "Tesla Q4 2025 Update energy generation and storage deployments", topK: 8, - useAgentic: true, + useAgentic: false, dataType: 1, }); }); diff --git a/src/domains/sources/parsed-document-blob-storage.test.ts b/src/domains/sources/parsed-document-blob-storage.test.ts index 4c3ee16..b5fb005 100644 --- a/src/domains/sources/parsed-document-blob-storage.test.ts +++ b/src/domains/sources/parsed-document-blob-storage.test.ts @@ -1,9 +1,5 @@ import { describe, expect, it } from "vitest" -import type { - KnowhereParsedSnapshotChunkPage, - KnowhereParsedSnapshotManifest, - ParsedDocumentSyncProgress, -} from "@ontos-ai/knowhere-sdk" +import type { ParsedDocumentSyncProgress } from "@ontos-ai/knowhere-sdk" import { BlobParsedDocumentStorage, @@ -34,6 +30,7 @@ function createFakeBlobStore(): { return { statusCode: 200, stream: new Response(body).body as ReadableStream, + url: toUrl(pathname), contentType: object.contentType, } }, @@ -56,42 +53,30 @@ function createFakeBlobStore(): { const documentId = "doc_123" const revisionKey = "job_result_1" -const manifest: KnowhereParsedSnapshotManifest = { - version: 1, - kind: "knowhere-parsed-result-snapshot", +const manifest = { + version: "2.0", jobId: "job_1", - revisionKey, - documentId, sourceFileName: "example.pdf", - totalChunks: 2, - chunkPageSize: 200, - chunkPages: [{ page: 1, pageSize: 200, chunkCount: 2, key: "chunks/page-1.json" }], - assetUrlsByFilePath: {}, - createdAt: "2026-07-04T00:00:00.000Z", -} + statistics: { + totalChunks: 2, + textChunks: 1, + imageChunks: 1, + tableChunks: 0, + pageChunks: 0, + }, +} satisfies Record -const chunkPage: KnowhereParsedSnapshotChunkPage = { - version: 1, - jobId: "job_1", - revisionKey, - documentId, - sourceFileName: "example.pdf", - page: 1, - pageSize: 200, - total: 2, - totalPages: 1, +const chunks = { chunks: [ { - id: "c1", - chunkId: "c1", - chunkType: "text", + chunk_id: "c1", + type: "text", content: "hello", - sourceChunkPath: "example.pdf", - sortOrder: 0, + path: "example.pdf", metadata: {}, }, ], -} +} satisfies Record describe("BlobParsedDocumentStorage", () => { it("round-trips a manifest at the result-relative manifest path", async () => { @@ -101,28 +86,52 @@ describe("BlobParsedDocumentStorage", () => { blobStore: store, }) - await storage.writeManifest({ documentId, revisionKey, manifest }) + await storage.writeObject({ + documentId, + revisionKey, + path: "manifest.json", + body: Buffer.from(JSON.stringify(manifest), "utf8"), + contentType: "application/json; charset=utf-8", + }) expect([...objects.keys()]).toContain( "workspaces/ws_1/parsed-documents/doc_123/job_result_1/manifest.json", ) - const read = await storage.readManifest({ documentId, revisionKey }) - expect(read).toEqual(manifest) + const read = await storage.readObject({ + documentId, + revisionKey, + path: "manifest.json", + }) + expect(read ? Buffer.from(read.body).toString("utf8") : null).toBe( + JSON.stringify(manifest), + ) }) - it("round-trips a chunk page at its result-relative chunk path", async () => { + it("round-trips chunks at the result-relative chunks path", async () => { const { store, objects } = createFakeBlobStore() const storage = new BlobParsedDocumentStorage({ workspaceId: "ws_1", blobStore: store, }) - await storage.writeChunkPage({ documentId, revisionKey, page: chunkPage }) + await storage.writeObject({ + documentId, + revisionKey, + path: "chunks.json", + body: Buffer.from(JSON.stringify(chunks), "utf8"), + contentType: "application/json; charset=utf-8", + }) expect([...objects.keys()]).toContain( - "workspaces/ws_1/parsed-documents/doc_123/job_result_1/chunks/page-1.json", + "workspaces/ws_1/parsed-documents/doc_123/job_result_1/chunks.json", + ) + const read = await storage.readObject({ + documentId, + revisionKey, + path: "chunks.json", + }) + expect(read ? Buffer.from(read.body).toString("utf8") : null).toBe( + JSON.stringify(chunks), ) - const read = await storage.readChunkPage({ documentId, revisionKey, page: 1 }) - expect(read).toEqual(chunkPage) }) it("round-trips sync progress", async () => { @@ -135,7 +144,6 @@ describe("BlobParsedDocumentStorage", () => { documentId, revisionKey, nextChunkPage: 3, - nextAssetIndex: 0, status: "running", updatedAt: "2026-07-04T00:00:00.000Z", } @@ -185,10 +193,20 @@ describe("BlobParsedDocumentStorage", () => { body: new Uint8Array([4, 5, 6]), contentType: "image/png", }) + await storage.writeObject({ + documentId, + revisionKey, + path: "images/fig-1.png", + body: new Uint8Array([7, 8, 9]), + contentType: "image/png", + }) expect([...objects.keys()]).toContain( "workspaces/ws_1/parsed-documents/doc_123/job_result_1/page_citation_assets/page-1.png", ) + expect([...objects.keys()]).toContain( + "workspaces/ws_1/parsed-documents/doc_123/job_result_1/images/fig-1.png", + ) const read = await storage.readObject({ documentId, revisionKey, @@ -221,16 +239,15 @@ describe("BlobParsedDocumentStorage", () => { ) }) - it("returns null for a missing manifest, chunk page, progress, and asset", async () => { + it("returns null for a missing object, progress, and asset", async () => { const { store } = createFakeBlobStore() const storage = new BlobParsedDocumentStorage({ workspaceId: "ws_1", blobStore: store, }) - expect(await storage.readManifest({ documentId, revisionKey })).toBeNull() expect( - await storage.readChunkPage({ documentId, revisionKey, page: 9 }), + await storage.readObject({ documentId, revisionKey, path: "manifest.json" }), ).toBeNull() expect(await storage.readSyncProgress({ documentId, revisionKey })).toBeNull() expect( @@ -245,10 +262,17 @@ describe("BlobParsedDocumentStorage", () => { blobStore: store, }) - await storage.writeManifest({ documentId, revisionKey, manifest }) - const staleRead = await storage.readManifest({ + await storage.writeObject({ + documentId, + revisionKey, + path: "manifest.json", + body: Buffer.from(JSON.stringify(manifest), "utf8"), + contentType: "application/json; charset=utf-8", + }) + const staleRead = await storage.readObject({ documentId, revisionKey: "job_result_2", + path: "manifest.json", }) expect(staleRead).toBeNull() }) @@ -261,10 +285,18 @@ describe("BlobParsedDocumentStorage", () => { }) await expect( - storage.readManifest({ documentId: "../escape", revisionKey }), + storage.readObject({ + documentId: "../escape", + revisionKey, + path: "manifest.json", + }), ).rejects.toThrow(/Invalid parsed storage segment/) await expect( - storage.readManifest({ documentId, revisionKey: "a/b" }), + storage.readObject({ + documentId, + revisionKey: "a/b", + path: "manifest.json", + }), ).rejects.toThrow(/Invalid parsed storage segment/) await expect( storage.getAssetUrl({ diff --git a/src/domains/sources/parsed-document-blob-storage.ts b/src/domains/sources/parsed-document-blob-storage.ts index 8f19837..5a05e5c 100644 --- a/src/domains/sources/parsed-document-blob-storage.ts +++ b/src/domains/sources/parsed-document-blob-storage.ts @@ -2,15 +2,15 @@ import "server-only" import { del, get, head, put, BlobNotFoundError } from "@vercel/blob" import type { - KnowhereParsedSnapshotChunkPage, - KnowhereParsedSnapshotManifest, + ParsedDocumentObject, + ParsedDocumentObjectHead, + ParsedDocumentObjectParams, + ParsedDocumentRevisionParams, ParsedDocumentStorage, - ParsedDocumentStorageAsset, - ParsedDocumentStorageAssetParams, - ParsedDocumentStorageChunkPageParams, ParsedDocumentStorageDocument, - ParsedDocumentStorageManifestParams, ParsedDocumentSyncProgress, + ParsedDocumentWriteObjectParams, + ParsedDocumentWriteObjectResult, } from "@ontos-ai/knowhere-sdk" /** @@ -26,8 +26,6 @@ import type { */ const parsedDocumentsDirectoryName = "parsed-documents" -const manifestStoragePath = "manifest.json" -const legacyManifestStoragePath = "manifest/current.json" const syncProgressStoragePath = "sync-progress.json" const jsonContentType = "application/json; charset=utf-8" const binaryContentType = "application/octet-stream" @@ -36,6 +34,7 @@ type BlobGetResult = | { readonly statusCode: 200 readonly stream: ReadableStream + readonly url: string readonly contentType?: string } | { @@ -80,19 +79,19 @@ export type BlobParsedDocumentStorageInput = { readonly blobStore?: ParsedDocumentBlobStore } -export type ParsedDocumentStorageObject = ParsedDocumentStorageDocument & { +type ParsedDocumentStorageObject = ParsedDocumentStorageDocument & { readonly path: string } -export type ParsedDocumentStorageWritableObject = - ParsedDocumentStorageObject & { - readonly body: string | Buffer | Uint8Array - readonly contentType?: string - } +type LegacyParsedDocumentStorageAsset = { + readonly sourcePath: string + readonly body: string | Buffer | Uint8Array + readonly contentType: string + readonly metadata?: Readonly> +} -export type ParsedDocumentStorageReadableObject = { - readonly body: Buffer - readonly contentType?: string +type LegacyParsedDocumentStorageAssetParams = ParsedDocumentStorageDocument & { + readonly sourcePath: string } const vercelBlobStore: ParsedDocumentBlobStore = { @@ -105,6 +104,7 @@ const vercelBlobStore: ParsedDocumentBlobStore = { return { statusCode: 200, stream: blob.stream, + url: blob.blob.url, contentType: blob.blob.contentType, } }, @@ -139,72 +139,8 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { this.blobStore = input.blobStore ?? vercelBlobStore } - async readManifest( - params: ParsedDocumentStorageManifestParams, - ): Promise { - const current = await this.readJson( - this.getObjectKey({ - documentId: params.documentId, - revisionKey: params.revisionKey, - path: manifestStoragePath, - }), - ) - if (current) return current - - return this.readJson( - this.getObjectKey({ - documentId: params.documentId, - revisionKey: params.revisionKey, - path: legacyManifestStoragePath, - }), - ) - } - - async writeManifest(params: { - readonly documentId: string - readonly revisionKey: string - readonly manifest: KnowhereParsedSnapshotManifest - }): Promise { - await this.writeJson( - this.getObjectKey({ - documentId: params.documentId, - revisionKey: params.revisionKey, - path: manifestStoragePath, - }), - params.manifest, - ) - } - - async readChunkPage( - params: ParsedDocumentStorageChunkPageParams, - ): Promise { - // chunkType filtering happens SDK-side after read; storage returns the full page. - return this.readJson( - this.getObjectKey({ - documentId: params.documentId, - revisionKey: params.revisionKey, - path: getChunkPageStoragePath(params.page), - }), - ) - } - - async writeChunkPage(params: { - readonly documentId: string - readonly revisionKey: string - readonly page: KnowhereParsedSnapshotChunkPage - }): Promise { - await this.writeJson( - this.getObjectKey({ - documentId: params.documentId, - revisionKey: params.revisionKey, - path: getChunkPageStoragePath(params.page.page), - }), - params.page, - ) - } - async writeAsset( - params: ParsedDocumentStorageDocument & ParsedDocumentStorageAsset, + params: ParsedDocumentStorageDocument & LegacyParsedDocumentStorageAsset, ): Promise<{ readonly sourcePath: string; readonly url?: string }> { const blob = await this.blobStore.put( this.getObjectKey({ @@ -224,7 +160,7 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { } async getAssetUrl( - params: ParsedDocumentStorageAssetParams, + params: LegacyParsedDocumentStorageAssetParams, ): Promise { const result = await this.blobStore.head( this.getObjectKey({ @@ -246,7 +182,7 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { } async readSyncProgress( - params: ParsedDocumentStorageDocument, + params: ParsedDocumentRevisionParams, ): Promise { return this.readJson( this.getObjectKey({ @@ -269,14 +205,25 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { } async readObject( - params: ParsedDocumentStorageObject, - ): Promise { - return this.readBlobObject(this.getObjectKey(params)) + params: ParsedDocumentObjectParams, + ): Promise { + const object = await this.readBlobObject(this.getObjectKey(params)) + if (!object) return null + + return { + documentId: params.documentId, + revisionKey: params.revisionKey, + path: params.path, + body: object.body, + ...(object.contentType ? { contentType: object.contentType } : {}), + contentLength: object.body.byteLength, + url: object.url, + } } async writeObject( - params: ParsedDocumentStorageWritableObject, - ): Promise<{ readonly path: string; readonly url?: string }> { + params: ParsedDocumentWriteObjectParams, + ): Promise { const blob = await this.blobStore.put( this.getObjectKey(params), Buffer.from(params.body), @@ -287,20 +234,35 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { multipart: true, }, ) - return { path: params.path, url: blob.url } + return { + documentId: params.documentId, + revisionKey: params.revisionKey, + path: params.path, + url: blob.url, + } + } + + async headObject( + params: ParsedDocumentObjectParams, + ): Promise { + const result = await this.blobStore.head(this.getObjectKey(params)) + if (!result) return null + + return { + documentId: params.documentId, + revisionKey: params.revisionKey, + path: params.path, + url: result.url, + } } async getObjectUrl( - params: ParsedDocumentStorageObject, + params: ParsedDocumentObjectParams, ): Promise { const result = await this.blobStore.head(this.getObjectKey(params)) return result?.url ?? null } - async deleteObject(params: ParsedDocumentStorageObject): Promise { - await this.blobStore.del(this.getObjectKey(params)) - } - private getRevisionPrefix(documentId: string, revisionKey: string): string { return [ "workspaces", @@ -338,9 +300,11 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { return object ? object.body.toString("utf8") : null } - private async readBlobObject( - key: string, - ): Promise { + private async readBlobObject(key: string): Promise<{ + readonly body: Buffer + readonly contentType?: string + readonly url?: string + } | null> { try { const result = await this.blobStore.get(key, { access: "public" }) if (!result || result.statusCode !== 200) return null @@ -348,6 +312,7 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { return { body, ...(result.contentType ? { contentType: result.contentType } : {}), + url: result.url, } } catch (error) { if (error instanceof BlobNotFoundError) return null @@ -356,10 +321,6 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { } } -function getChunkPageStoragePath(page: number): string { - return `chunks/page-${page}.json` -} - /** * Reject path segments that could traverse outside the intended prefix. Mirrors * the SDK `DiskParsedDocumentStorage` guard so blob keys stay well-formed. diff --git a/src/integrations/knowhere.ts b/src/integrations/knowhere.ts index c0d074d..e55fc0b 100644 --- a/src/integrations/knowhere.ts +++ b/src/integrations/knowhere.ts @@ -14,10 +14,8 @@ import { BlobParsedDocumentStorage } from "@/domains/sources/parsed-document-blo * SDK returns `completed:false` and the caller re-enqueues to continue. */ const defaultParsedStorageLimits: ParsedDocumentStorageLimits = { - chunkPageSize: 200, remotePageSize: 100, maxPagesPerSync: 10, - maxAssetsPerSync: 20, syncDeadlineMs: 8000, grepMaxPages: 50, grepDeadlineMs: 8000, From 7c2bba4df77544f5b1cbd2593200fbe610ab4915 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 7 Jul 2026 23:27:53 +0800 Subject: [PATCH 07/14] Keep home page route exports valid --- src/app/home-content.tsx | 28 ++++++++++++++++++++++++++++ src/app/page.test.ts | 2 +- src/app/page.tsx | 29 +---------------------------- 3 files changed, 30 insertions(+), 29 deletions(-) create mode 100644 src/app/home-content.tsx diff --git a/src/app/home-content.tsx b/src/app/home-content.tsx new file mode 100644 index 0000000..c2f1a64 --- /dev/null +++ b/src/app/home-content.tsx @@ -0,0 +1,28 @@ +import { connection } from "next/server" +import { WorkspaceShell } from "@/components/workspace-shell" +import { loadWorkspaceShellInitialState } from "@/domains/workspace/initial-state" +import { effectOperation } from "@/lib/effect-operation" +import { summarizeUnknownError } from "@/lib/format-log-value" +import { logger } from "@/lib/logger" + +export async function HomeContent() { + await connection() + const initialState = await loadWorkspaceInitialState() + return +} + +async function loadWorkspaceInitialState(): ReturnType< + typeof loadWorkspaceShellInitialState +> { + try { + return await loadWorkspaceShellInitialState() + } catch (error) { + logger.error("workspace: initial state failed", { + error: summarizeUnknownError(error), + }) + throw effectOperation.createBoundaryError( + "Workspace initial state failed", + error, + ) + } +} diff --git a/src/app/page.test.ts b/src/app/page.test.ts index f7e3525..3e9f811 100644 --- a/src/app/page.test.ts +++ b/src/app/page.test.ts @@ -20,7 +20,7 @@ vi.mock("@/domains/workspace/initial-state", () => ({ loadWorkspaceShellInitialState: mocks.loadWorkspaceShellInitialState, })) -import { HomeContent } from "./page" +import { HomeContent } from "./home-content" import { makeWorkspaceInitialStateFailureFixture } from "@/test/workspace-initial-state-failure-fixture" describe("Home", () => { diff --git a/src/app/page.tsx b/src/app/page.tsx index 2ebc597..0a3c967 100644 --- a/src/app/page.tsx +++ b/src/app/page.tsx @@ -1,10 +1,5 @@ import { Suspense } from "react" -import { WorkspaceShell } from "@/components/workspace-shell" -import { loadWorkspaceShellInitialState } from "@/domains/workspace/initial-state" -import { effectOperation } from "@/lib/effect-operation" -import { summarizeUnknownError } from "@/lib/format-log-value" -import { logger } from "@/lib/logger" -import { connection } from "next/server" +import { HomeContent } from "./home-content" export default function Home() { return ( @@ -13,25 +8,3 @@ export default function Home() { ) } - -export async function HomeContent() { - await connection() - const initialState = await loadWorkspaceInitialState() - return -} - -async function loadWorkspaceInitialState(): ReturnType< - typeof loadWorkspaceShellInitialState -> { - try { - return await loadWorkspaceShellInitialState() - } catch (error) { - logger.error("workspace: initial state failed", { - error: summarizeUnknownError(error), - }) - throw effectOperation.createBoundaryError( - "Workspace initial state failed", - error, - ) - } -} From 203f72cbad08237e4e41c814e0b95fcd612f6d6a Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 7 Jul 2026 23:48:45 +0800 Subject: [PATCH 08/14] Install published Knowhere SDK 2.1.2 --- package.json | 2 +- pnpm-lock.yaml | 10 +++++----- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/package.json b/package.json index 94d1c3a..cf48c4a 100644 --- a/package.json +++ b/package.json @@ -28,7 +28,7 @@ "@effect/platform": "^0.96.1", "@napi-rs/canvas": "^1.0.2", "@neondatabase/serverless": "^1.1.0", - "@ontos-ai/knowhere-sdk": "^2.1.1", + "@ontos-ai/knowhere-sdk": "^2.1.2", "@radix-ui/react-alert-dialog": "^1.1.15", "@radix-ui/react-checkbox": "^1.3.3", "@radix-ui/react-dialog": "^1.1.15", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 48aabe5..12dc39f 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -24,8 +24,8 @@ importers: specifier: ^1.1.0 version: 1.1.0 '@ontos-ai/knowhere-sdk': - specifier: ^2.1.1 - version: 2.1.1 + specifier: ^2.1.2 + version: 2.1.2 '@radix-ui/react-alert-dialog': specifier: ^1.1.15 version: 1.1.15(@types/react-dom@19.2.3(@types/react@19.2.14))(@types/react@19.2.14)(react-dom@19.2.4(react@19.2.4))(react@19.2.4) @@ -1603,8 +1603,8 @@ packages: resolution: {integrity: sha512-nn5ozdjYQpUCZlWGuxcJY/KpxkWQs4DcbMCmKojjyrYDEAGy4Ce19NN4v5MduafTwJlbKc99UA8YhSVqq9yPZA==} engines: {node: '>=12.4.0'} - '@ontos-ai/knowhere-sdk@2.1.1': - resolution: {integrity: sha512-K33ylB/QjVYlLgow+/Hw+uNhQbWWrMTIqZ3IAlx5buHSI2EHPVP9+kDsfslwpfffYSn85j9Ab0/E9qTG0QotTw==} + '@ontos-ai/knowhere-sdk@2.1.2': + resolution: {integrity: sha512-m1wbQZNesExcg2yoEnmqFRlML/tprv11y/v7tlTLXEKTqKOY5zpG67idxXtbmw2qjOjQkR02HhgUwo/SdrFWPA==} engines: {node: '>=22.13.0', npm: '>=10.0.0', pnpm: '>=9.0.0'} '@open-draft/deferred-promise@2.2.0': @@ -6731,7 +6731,7 @@ snapshots: '@nolyfill/is-core-module@1.0.39': {} - '@ontos-ai/knowhere-sdk@2.1.1': + '@ontos-ai/knowhere-sdk@2.1.2': dependencies: axios: 1.18.1 jszip: 3.10.1 From 3fe820627c47f61807f31fef192311186c841082 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Wed, 8 Jul 2026 01:03:55 +0800 Subject: [PATCH 09/14] fix(sources): render page assets from chunks --- .../[sourceId]/page-assets/route.test.ts | 554 ------------------ .../sources/[sourceId]/page-assets/route.ts | 32 - src/components/chunks-panel.test.ts | 215 +++---- src/components/chunks-panel.tsx | 332 +---------- src/components/parsed-chunk-card.test.ts | 44 +- src/components/parsed-chunk-card.tsx | 112 +++- .../workspace-citation-focus.test.ts | 28 +- src/components/workspace-citation-focus.ts | 12 - .../workspace-citation-state.test.ts | 54 ++ src/components/workspace-citation-state.ts | 65 +- .../workspace-selected-chunks.test.ts | 89 +-- src/components/workspace-selected-chunks.ts | 123 ++-- src/domains/chunks/index.test.ts | 36 ++ src/domains/chunks/normalization.ts | 51 ++ src/domains/chunks/types.ts | 3 + src/domains/sources/page-assets.ts | 218 ------- src/domains/sources/route-page-assets.ts | 208 ------- src/domains/sources/route-service.ts | 5 - src/domains/sources/route-types.ts | 38 -- src/domains/workspace/client-cache.ts | 4 - src/domains/workspace/client.test.ts | 30 - src/domains/workspace/client.ts | 34 -- 22 files changed, 564 insertions(+), 1723 deletions(-) delete mode 100644 src/app/api/sources/[sourceId]/page-assets/route.test.ts delete mode 100644 src/app/api/sources/[sourceId]/page-assets/route.ts delete mode 100644 src/domains/sources/page-assets.ts delete mode 100644 src/domains/sources/route-page-assets.ts diff --git a/src/app/api/sources/[sourceId]/page-assets/route.test.ts b/src/app/api/sources/[sourceId]/page-assets/route.test.ts deleted file mode 100644 index 4f00bd3..0000000 --- a/src/app/api/sources/[sourceId]/page-assets/route.test.ts +++ /dev/null @@ -1,554 +0,0 @@ -import { NextRequest } from "next/server" -import { beforeEach, describe, expect, it, vi } from "vitest" - -const mocks = vi.hoisted(() => ({ - deleteBlob: vi.fn(), - ensureApiKeyForWorkspace: vi.fn(), - ensureWorkspace: vi.fn(), - findSourceInWorkspace: vi.fn(), - getCurrentUser: vi.fn(), - listChunks: vi.fn(), - makeKnowhereClient: vi.fn(), - makeKnowhereClientWithParsedStorage: vi.fn(), - readChunks: vi.fn(), - requireUser: vi.fn(), -})) - -vi.mock("next/headers", () => ({ - headers: vi.fn(async () => new Headers({ cookie: "session=abc" })), -})) - -vi.mock("@/integrations/dashboard/api-key-service", () => ({ - ensureApiKeyForWorkspace: mocks.ensureApiKeyForWorkspace, -})) - -vi.mock("@/integrations/knowhere-demo", () => ({ - knowhereDemoApi: { - fetchCatalog: vi.fn(), - fetchChunkPage: vi.fn(), - }, -})) - -vi.mock("@/infrastructure/auth", () => ({ - getCurrentUser: mocks.getCurrentUser, - requireUser: mocks.requireUser, -})) - -vi.mock("@/integrations/knowhere", () => ({ - makeKnowhereClient: mocks.makeKnowhereClient, - makeKnowhereClientWithParsedStorage: - mocks.makeKnowhereClientWithParsedStorage, -})) - -vi.mock("@vercel/blob", () => ({ - del: mocks.deleteBlob, -})) - -vi.mock("@/domains/sources/service", () => ({ - sourceService: { - findInWorkspace: mocks.findSourceInWorkspace, - localizeRemoteDocument: vi.fn(), - }, -})) - -vi.mock("@/domains/workspace/service", () => ({ - workspaceService: { - ensureWorkspace: mocks.ensureWorkspace, - }, -})) - -import { GET } from "./route" - -describe("GET /api/sources/[sourceId]/page-assets", () => { - beforeEach(() => { - vi.clearAllMocks() - mocks.makeKnowhereClientWithParsedStorage.mockReturnValue({ - client: { documents: { listChunks: mocks.listChunks } }, - knowledge: { readChunks: mocks.readChunks }, - }) - }) - - it("returns stored page assets for a ready workspace source without durable hardening", async () => { - mocks.getCurrentUser.mockResolvedValue({ - id: "user_1", - email: null, - name: null, - }) - mocks.ensureWorkspace.mockResolvedValue({ - id: "workspace_1", - userId: "user_1", - namespace: "notebook-workspace_1", - createdAt: new Date("2026-05-10T00:00:00.000Z"), - }) - mocks.findSourceInWorkspace.mockResolvedValue( - makeReadySource({ - id: "00000000-0000-0000-0000-000000000002", - knowhereJobId: "job_1", - knowhereDocumentId: "doc_1", - }), - ) - mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") - mocks.readChunks.mockResolvedValue({ - document: { - localDocumentId: "doc_1", - resultDirectoryPath: "parsed-storage:doc_1", - }, - chunks: [ - { - chunkId: "page_1", - chunkType: "page", - assetUrl: "https://assets.example/fallback.png", - metadata: { - pageAssets: [ - { - pageNum: 1, - assetUrl: "https://assets.example/page-000001.png", - contentType: "image/png", - width: 1200, - height: 1600, - }, - ], - }, - }, - ], - page: 1, - pageSize: 1, - totalChunks: 3, - totalPages: 3, - }) - - const response = await GET( - new NextRequest( - "http://localhost:3001/api/sources/00000000-0000-0000-0000-000000000002/page-assets?page=1&pageSize=1", - ), - { - params: Promise.resolve({ - sourceId: "00000000-0000-0000-0000-000000000002", - }), - }, - ) - - await expect(response.json()).resolves.toEqual({ - pages: [ - { - pageNumber: 1, - assetUrl: "https://assets.example/page-000001.png", - contentType: "image/png", - width: 1200, - height: 1600, - }, - ], - pagination: { - page: 1, - pageSize: 1, - total: 3, - totalPages: 3, - }, - }) - expect(response.status).toBe(200) - expect(mocks.readChunks).toHaveBeenCalledWith({ - documentId: "doc_1", - revisionKey: "job_1", - chunkType: "page", - page: 1, - pageSize: 1, - }) - expect(mocks.listChunks).not.toHaveBeenCalled() - }) - - it("uses SDK page asset URLs directly when the SDK remote fallback is usable", async () => { - mocks.getCurrentUser.mockResolvedValue({ - id: "user_1", - email: null, - name: null, - }) - mocks.ensureWorkspace.mockResolvedValue({ - id: "workspace_1", - userId: "user_1", - namespace: "notebook-workspace_1", - createdAt: new Date("2026-05-10T00:00:00.000Z"), - }) - mocks.findSourceInWorkspace.mockResolvedValue( - makeReadySource({ - id: "00000000-0000-0000-0000-000000000002", - knowhereJobId: "job_1", - knowhereDocumentId: "doc_1", - }), - ) - mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") - mocks.readChunks.mockResolvedValue({ - document: { - localDocumentId: "doc_1", - resultDirectoryPath: "remote:doc_1", - }, - chunks: [ - { - chunkId: "page_1", - chunkType: "page", - metadata: { - pageAssets: [ - { - pageNum: 1, - assetUrl: "https://sdk.example/page-000001.png", - contentType: "image/png", - width: 1200, - height: 1600, - }, - ], - }, - }, - ], - page: 1, - pageSize: 1, - totalChunks: 1, - totalPages: 1, - }) - - const response = await GET( - new NextRequest( - "http://localhost:3001/api/sources/00000000-0000-0000-0000-000000000002/page-assets?page=1&pageSize=1", - ), - { - params: Promise.resolve({ - sourceId: "00000000-0000-0000-0000-000000000002", - }), - }, - ) - - await expect(response.json()).resolves.toEqual({ - pages: [ - { - pageNumber: 1, - assetUrl: "https://sdk.example/page-000001.png", - contentType: "image/png", - width: 1200, - height: 1600, - }, - ], - pagination: { - page: 1, - pageSize: 1, - total: 1, - totalPages: 1, - }, - }) - expect(response.status).toBe(200) - expect(mocks.readChunks).toHaveBeenCalledWith({ - documentId: "doc_1", - revisionKey: "job_1", - chunkType: "page", - page: 1, - pageSize: 1, - }) - expect(mocks.listChunks).not.toHaveBeenCalled() - }) - - it("rejects non-ready workspace sources", async () => { - mocks.getCurrentUser.mockResolvedValue({ - id: "user_1", - email: null, - name: null, - }) - mocks.ensureWorkspace.mockResolvedValue({ - id: "workspace_1", - userId: "user_1", - namespace: "notebook-workspace_1", - createdAt: new Date("2026-05-10T00:00:00.000Z"), - }) - mocks.findSourceInWorkspace.mockResolvedValue( - makeReadySource({ - id: "00000000-0000-0000-0000-000000000002", - status: "parsing", - knowhereDocumentId: "doc_1", - }), - ) - - const response = await GET( - new NextRequest( - "http://localhost:3001/api/sources/00000000-0000-0000-0000-000000000002/page-assets?page=1&pageSize=1", - ), - { - params: Promise.resolve({ - sourceId: "00000000-0000-0000-0000-000000000002", - }), - }, - ) - - await expect(response.json()).resolves.toEqual({ - message: "Source is not ready.", - }) - expect(response.status).toBe(409) - expect(mocks.readChunks).not.toHaveBeenCalled() - }) - - it("returns an empty page list when page chunks have no usable assets", async () => { - mocks.getCurrentUser.mockResolvedValue({ - id: "user_1", - email: null, - name: null, - }) - mocks.ensureWorkspace.mockResolvedValue({ - id: "workspace_1", - userId: "user_1", - namespace: "notebook-workspace_1", - createdAt: new Date("2026-05-10T00:00:00.000Z"), - }) - mocks.findSourceInWorkspace.mockResolvedValue( - makeReadySource({ - id: "00000000-0000-0000-0000-000000000002", - knowhereDocumentId: "doc_1", - }), - ) - mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") - mocks.readChunks.mockResolvedValue({ - document: { - localDocumentId: "doc_1", - resultDirectoryPath: "remote:doc_1", - }, - chunks: [ - { - chunkId: "page_1", - chunkType: "page", - metadata: { pageAssets: [] }, - }, - ], - page: 1, - pageSize: 1, - totalChunks: 1, - totalPages: 1, - }) - mocks.listChunks.mockResolvedValue({ - chunks: [], - pagination: { - page: 1, - pageSize: 1, - total: 1, - totalPages: 1, - }, - }) - - const response = await GET( - new NextRequest( - "http://localhost:3001/api/sources/00000000-0000-0000-0000-000000000002/page-assets?page=1&pageSize=1", - ), - { - params: Promise.resolve({ - sourceId: "00000000-0000-0000-0000-000000000002", - }), - }, - ) - - await expect(response.json()).resolves.toEqual({ - pages: [], - pagination: { - page: 1, - pageSize: 1, - total: 1, - totalPages: 1, - }, - }) - expect(response.status).toBe(200) - expect(mocks.listChunks).toHaveBeenCalledWith("doc_1", { - page: 1, - pageSize: 1, - chunkType: "page", - includeAssetUrls: true, - }) - }) - - it("returns Knowhere page asset URLs when the storage probe falls through to remote chunks", async () => { - mocks.getCurrentUser.mockResolvedValue({ - id: "user_1", - email: null, - name: null, - }) - mocks.ensureWorkspace.mockResolvedValue({ - id: "workspace_1", - userId: "user_1", - namespace: "notebook-workspace_1", - createdAt: new Date("2026-05-10T00:00:00.000Z"), - }) - mocks.findSourceInWorkspace.mockResolvedValue( - makeReadySource({ - id: "00000000-0000-0000-0000-000000000002", - knowhereDocumentId: "doc_1", - }), - ) - mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") - mocks.readChunks.mockResolvedValue({ - document: { - localDocumentId: "doc_1", - resultDirectoryPath: "remote:doc_1", - }, - chunks: [ - { - chunkId: "page_1", - chunkType: "page", - metadata: { - pageAssets: [ - { - pageNum: 1, - contentType: "image/png", - width: 1200, - height: 1600, - }, - ], - }, - }, - ], - page: 1, - pageSize: 1, - totalChunks: 1, - totalPages: 1, - }) - mocks.listChunks.mockResolvedValue({ - chunks: [ - { - id: "document_page_1", - chunkId: "page_1", - chunkType: "page", - content: "Page 1", - sectionId: null, - sectionPath: "pages/1", - sourceChunkPath: "pages/1", - filePath: "pages/page-000001.png", - sortOrder: 0, - metadata: { - pageAssets: [ - { - pageNum: 1, - assetUrl: "https://knowhere.example/page-000001.png", - contentType: "image/png", - width: 1200, - height: 1600, - }, - ], - }, - assetUrl: "https://knowhere.example/page-000001.png", - }, - ], - pagination: { - page: 1, - pageSize: 1, - total: 1, - totalPages: 1, - }, - }) - - const response = await GET( - new NextRequest( - "http://localhost:3001/api/sources/00000000-0000-0000-0000-000000000002/page-assets?page=1&pageSize=1", - ), - { - params: Promise.resolve({ - sourceId: "00000000-0000-0000-0000-000000000002", - }), - }, - ) - - await expect(response.json()).resolves.toEqual({ - pages: [ - { - pageNumber: 1, - assetUrl: "https://knowhere.example/page-000001.png", - contentType: "image/png", - width: 1200, - height: 1600, - }, - ], - pagination: { - page: 1, - pageSize: 1, - total: 1, - totalPages: 1, - }, - }) - expect(response.status).toBe(200) - expect(mocks.readChunks).toHaveBeenCalledWith({ - documentId: "doc_1", - revisionKey: "job_1", - chunkType: "page", - page: 1, - pageSize: 1, - }) - expect(mocks.listChunks).toHaveBeenCalledWith("doc_1", { - page: 1, - pageSize: 1, - chunkType: "page", - includeAssetUrls: true, - }) - }) - - it("marks page assets unavailable when the parsed document is missing remotely", async () => { - const notFoundError = new Error("Document not found") - notFoundError.name = "NotFoundError" - mocks.getCurrentUser.mockResolvedValue({ - id: "user_1", - email: null, - name: null, - }) - mocks.ensureWorkspace.mockResolvedValue({ - id: "workspace_1", - userId: "user_1", - namespace: "notebook-workspace_1", - createdAt: new Date("2026-05-10T00:00:00.000Z"), - }) - mocks.findSourceInWorkspace.mockResolvedValue( - makeReadySource({ - id: "00000000-0000-0000-0000-000000000002", - knowhereDocumentId: "doc_missing", - }), - ) - mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") - mocks.readChunks.mockRejectedValue(notFoundError) - - const response = await GET( - new NextRequest( - "http://localhost:3001/api/sources/00000000-0000-0000-0000-000000000002/page-assets?page=1&pageSize=1", - ), - { - params: Promise.resolve({ - sourceId: "00000000-0000-0000-0000-000000000002", - }), - }, - ) - - await expect(response.json()).resolves.toEqual({ - pages: [], - pagination: { - page: 1, - pageSize: 1, - total: 0, - totalPages: 0, - }, - message: - "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", - isUnavailable: true, - }) - expect(response.status).toBe(200) - }) -}) - -function makeReadySource(overrides: Record) { - return { - id: "00000000-0000-0000-0000-000000000002", - workspaceId: "workspace_1", - title: "notes.pdf", - mimeType: "application/pdf", - sizeBytes: 1024, - status: "ready", - failureReason: null, - failureStage: null, - knowhereJobId: "job_1", - knowhereDocumentId: "doc_1", - stagedBlobPathname: null, - stagedBlobUrl: null, - originalBlobPathname: null, - originalBlobUrl: null, - demoKey: null, - createdAt: new Date("2026-05-10T00:00:00.000Z"), - updatedAt: new Date("2026-05-10T00:00:00.000Z"), - deletedAt: null, - ...overrides, - } -} diff --git a/src/app/api/sources/[sourceId]/page-assets/route.ts b/src/app/api/sources/[sourceId]/page-assets/route.ts deleted file mode 100644 index b7606f7..0000000 --- a/src/app/api/sources/[sourceId]/page-assets/route.ts +++ /dev/null @@ -1,32 +0,0 @@ -import type { NextRequest, NextResponse } from "next/server" - -import { getChunkPageParams } from "@/domains/chunks" -import { createSourceRouteService } from "@/domains/sources/route-service" -import { withApiErrorResponse } from "@/lib/api-error-response" -import { nextRouteContext } from "@/lib/next-route-context" -import { nextRouteResponse } from "@/lib/next-route-response" - -type RouteContext = { - params: Promise<{ - sourceId: string - }> -} - -const sourceRouteService = createSourceRouteService() - -export async function GET( - request: NextRequest, - context: RouteContext, -): Promise { - return withApiErrorResponse("sources:page-assets", async () => { - const { sourceId } = await context.params - const routeContext = await nextRouteContext.read() - const result = await sourceRouteService.loadSourcePageAssets({ - cookieHeader: routeContext.cookieHeader, - pageParams: getChunkPageParams(request.nextUrl.searchParams), - sourceId, - }) - - return nextRouteResponse.toNextResponse(result) - }) -} diff --git a/src/components/chunks-panel.test.ts b/src/components/chunks-panel.test.ts index 8e06155..78a35f7 100644 --- a/src/components/chunks-panel.test.ts +++ b/src/components/chunks-panel.test.ts @@ -14,7 +14,6 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { ChunksPanel } from "./chunks-panel"; import { sourceOriginalPreviewRequest } from "./source-original-preview-request"; -const fetchPageAssetPageMock = vi.hoisted(() => vi.fn()); const C = ChunksPanel as React.FC>; const virtualizerScrollResetDelayMs = 150; @@ -30,25 +29,9 @@ vi.mock("react-pdf", () => ({ Page: () => React.createElement("div", { "data-testid": "pdf-page" }), })); -vi.mock("@/domains/workspace/client", () => ({ - workspaceClient: { - fetchPageAssetPage: fetchPageAssetPageMock, - }, -})); - describe("ChunksPanel", () => { beforeEach(() => { shouldFlushVirtualizerTimers = false; - fetchPageAssetPageMock.mockReset(); - fetchPageAssetPageMock.mockResolvedValue({ - pages: [], - pagination: { - page: 1, - pageSize: 50, - total: 0, - totalPages: 1, - }, - }); globalThis.ResizeObserver = class ResizeObserver { observe() {} unobserve() {} @@ -89,28 +72,35 @@ describe("ChunksPanel", () => { expect(screen.getByText(/Showing all parsed chunks from/)).toBeTruthy(); }); - it("renders page assets instead of parsed chunk controls for page-mode documents", async () => { - fetchPageAssetPageMock.mockResolvedValue({ - pages: [ - { - pageNumber: 4, - assetUrl: "https://assets.example/page-000004.png", - contentType: "image/png", - width: 1200, - height: 1600, - }, - ], - pagination: { - page: 1, - pageSize: 50, - total: 4, - totalPages: 1, - }, - }); + it("renders page chunk assets inside the normal chunk list", async () => { + mockVisibleVirtualViewport(); render( React.createElement(C, { - chunks: [], + chunks: [ + { + chunkId: "page_4", + type: "page", + content: "Page 4 summary", + sourceTitle: "report.pdf", + pageNums: [4], + pageAssets: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-000004.png", + contentType: "image/png", + width: 1200, + height: 1600, + }, + ], + }, + { + chunkId: "table_1", + type: "table", + content: "
Budget
", + sourceTitle: "report.pdf", + }, + ], selectedSource: "report.pdf", selectedSourceView: { id: "source_1", @@ -122,124 +112,74 @@ describe("ChunksPanel", () => { }), ); - await waitFor(() => - expect( - screen.getByTestId("page-asset-document-viewer"), - ).toBeTruthy(), - ); - expect(screen.getByRole("heading", { name: "Original File" })).toBeTruthy(); - expect(screen.getByText("Page 4")).toBeTruthy(); - expect(screen.queryByRole("button", { name: "Parsed" })).toBeNull(); + expect(screen.getByRole("heading", { name: "Parsed Chunks" })).toBeTruthy(); + expect( + await screen.findByRole("img", { name: "Page 4" }), + ).toBeTruthy(); + expect(screen.getByText("image/png")).toBeTruthy(); + expect(screen.getByText("Budget")).toBeTruthy(); + expect( + screen.queryByRole("button", { name: "Original" }), + ).toBeNull(); expect(screen.queryByRole("button", { name: "Tree" })).toBeNull(); - expect(fetchPageAssetPageMock).toHaveBeenCalledWith("source_1", 1); + expect( + screen.queryByRole("button", { name: /original file/i }), + ).toBeNull(); }); - it("shows a page-level placeholder when a page asset image fails to load", async () => { - fetchPageAssetPageMock.mockResolvedValue({ - pages: [ - { - pageNumber: 4, - assetUrl: "https://assets.example/page-000004.png", - contentType: "image/png", - width: 1200, - height: 1600, - }, - ], - pagination: { - page: 1, - pageSize: 50, - total: 4, - totalPages: 1, - }, - }); + it("renders page chunks normally when no page assets exist", async () => { + mockVisibleVirtualViewport(); render( React.createElement(C, { - chunks: [], - selectedSource: "report.pdf", - selectedSourceView: { - id: "source_1", - title: "report.pdf", - mimeType: "application/pdf", - status: "ready", - documentPresentation: { kind: "page-assets", pageCount: 4 }, - }, - }), - ); - - const pageImage = await screen.findByRole("img", { name: "Page 4" }); - fireEvent.error(pageImage); - - expect( - screen.getByTestId("page-asset-image-unavailable-4"), - ).toBeTruthy(); - expect(screen.getByText("Page image unavailable.")).toBeTruthy(); - }); - - it("loads missing page buckets after focusing a later page asset", async () => { - const user = userEvent.setup(); - fetchPageAssetPageMock.mockImplementation( - async (_sourceId: string, page: number) => ({ - pages: [ + chunks: [ { - pageNumber: page === 3 ? 101 : page, - assetUrl: `https://assets.example/page-${page}.png`, - contentType: "image/png", + chunkId: "page_4", + type: "page", + content: "Page 4 summary", + readableContent: "Page 4 summary", + sourceTitle: "report.pdf", + pageNums: [4], }, ], - pagination: { - page, - pageSize: 50, - total: 120, - totalPages: 3, - }, - }), - ); - - render( - React.createElement(C, { - chunks: [], selectedSource: "report.pdf", selectedSourceView: { id: "source_1", title: "report.pdf", mimeType: "application/pdf", status: "ready", - documentPresentation: { kind: "page-assets", pageCount: 120 }, + documentPresentation: { kind: "page-assets", pageCount: 4 }, }, - focusedPageNumber: 101, - focusedPageRequestId: 1, }), ); - await waitFor(() => - expect(fetchPageAssetPageMock).toHaveBeenCalledWith("source_1", 3), - ); - - await user.click(screen.getByRole("button", { name: "Load more pages" })); - - await waitFor(() => - expect(fetchPageAssetPageMock).toHaveBeenCalledWith("source_1", 2), - ); + expect(await screen.findByText("Page 4 summary")).toBeTruthy(); + expect(screen.queryByRole("img", { name: "Page 4" })).toBeNull(); }); - it("shows page asset unavailable messages", async () => { - fetchPageAssetPageMock.mockResolvedValue({ - pages: [], - isUnavailable: true, - message: - "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", - pagination: { - page: 1, - pageSize: 50, - total: 0, - totalPages: 0, - }, - }); + it("shows a page-level placeholder when a page asset image fails to load", async () => { + mockVisibleVirtualViewport(); render( React.createElement(C, { - chunks: [], + chunks: [ + { + chunkId: "page_4", + type: "page", + content: "Page 4 summary", + sourceTitle: "report.pdf", + pageNums: [4], + pageAssets: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-000004.png", + contentType: "image/png", + width: 1200, + height: 1600, + }, + ], + }, + ], selectedSource: "report.pdf", selectedSourceView: { id: "source_1", @@ -251,14 +191,13 @@ describe("ChunksPanel", () => { }), ); - await waitFor(() => - expect( - screen.getByText( - "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", - ), - ).toBeTruthy(), - ); - expect(screen.queryByText("No page images available.")).toBeNull(); + const pageImage = await screen.findByRole("img", { name: "Page 4" }); + fireEvent.error(pageImage); + + expect( + screen.getByTestId("page-asset-image-unavailable-4"), + ).toBeTruthy(); + expect(screen.getByText("Page image unavailable.")).toBeTruthy(); }); it("defaults parsed chunks into a section tree view", () => { diff --git a/src/components/chunks-panel.tsx b/src/components/chunks-panel.tsx index fc785a2..71037c9 100644 --- a/src/components/chunks-panel.tsx +++ b/src/components/chunks-panel.tsx @@ -19,7 +19,6 @@ import { } from "d3-hierarchy"; import { FilePlus2, - ImageOff, Layers, RotateCcw, UploadCloud, @@ -42,11 +41,9 @@ import { chunksPanelState } from "@/components/chunks-panel-state"; import { MAX_UPLOAD_MB } from "@/domains/sources/validation"; import { useSourceOriginalPreviewWarmup } from "@/components/source-original-preview-warmup"; import { sourceOriginalPreviewModel } from "@/components/source-original-preview-model"; -import { workspaceClient } from "@/domains/workspace/client"; import type { ParsedChunkView } from "@/domains/chunks/types"; import type { SourceOriginalFileView, - SourcePageAssetView, SourceView, } from "@/domains/sources/types"; import type { AnalyticsContext } from "@/lib/posthog"; @@ -90,8 +87,6 @@ export function ChunksPanel({ selectedSourceFile = null, focusedChunkId = null, focusedChunkRequestId = 0, - focusedPageNumber = null, - focusedPageRequestId = 0, citationListViewRequestId = 0, isLoading = false, isLoadingMore = false, @@ -105,10 +100,16 @@ export function ChunksPanel({ analyticsContext, sourceCountSnapshot = 0, }: Partial = {}) { + const hasPageAssetChunks = chunks.some( + (chunk) => (chunk.pageAssets?.length ?? 0) > 0, + ); const isPageAssetSource = - selectedSourceView?.documentPresentation?.kind === "page-assets"; + selectedSourceView?.documentPresentation?.kind === "page-assets" || + hasPageAssetChunks; + const effectiveVisibleView = isPageAssetSource ? "parsed" : undefined; const originalPreviewCacheKey = selectedSourceFile?.url ?? null; const isOriginalPreviewAvailable = + !isPageAssetSource && sourceOriginalPreviewModel.canPreviewOriginalFile( selectedSource, selectedSourceFile, @@ -153,6 +154,7 @@ export function ChunksPanel({ isLoadingMore, onLoadMore, }); + const activeVisibleView = effectiveVisibleView ?? visibleView; useSourceOriginalPreviewWarmup({ sourceTitle: selectedSource, @@ -254,29 +256,21 @@ export function ChunksPanel({ citationListViewRequestId ? "list" : chunkDisplayModeState.mode; - const headerTitle = isPageAssetSource || visibleView === "original" + const headerTitle = activeVisibleView === "original" ? "Original File" : focusedChunkId ? "Referenced Chunks" : "Parsed Chunks"; const shouldMountOriginalPreview = - visibleView === "original" || - (originalPreviewCacheKey !== null && - mountedOriginalPreviewKey === originalPreviewCacheKey); + !isPageAssetSource && + (activeVisibleView === "original" || + (originalPreviewCacheKey !== null && + mountedOriginalPreviewKey === originalPreviewCacheKey)); const isTreeModeVisible = - !isPageAssetSource && visibleView === "parsed" && chunkDisplayMode === "tree"; - const headerSubtitle = isPageAssetSource ? ( - selectedSource ? ( - <> - Showing page images for{" "} - - {selectedSource} - - - ) : ( - "Select a source to preview its page images." - ) - ) : visibleView === "original" ? ( + !isPageAssetSource && + activeVisibleView === "parsed" && + chunkDisplayMode === "tree"; + const headerSubtitle = activeVisibleView === "original" ? ( selectedSource ? ( <> Showing the original file for{" "} @@ -335,7 +329,9 @@ export function ChunksPanel({

- {!isPageAssetSource && visibleView === "parsed" && chunks.length > 0 ? ( + {!isPageAssetSource && + activeVisibleView === "parsed" && + chunks.length > 0 ? (
Parsed @@ -379,7 +375,7 @@ export function ChunksPanel({
- + - {isPageAssetSource && selectedSourceView ? ( - - ) : isLoading ? ( + {isLoading ? ( ) : chunks.length === 0 && processingMessage ? ( @@ -437,7 +427,9 @@ export function ChunksPanel({ isOriginalPreviewAvailable={isOriginalPreviewAvailable} measureElement={measureVirtualChunkElement} onChunkClick={ - hasOriginalFile ? handleChunkSelected : undefined + hasOriginalFile && !isPageAssetSource + ? handleChunkSelected + : undefined } onReferenceClick={requestChunkFocus} selectedSourceFile={selectedSourceFile} @@ -472,7 +464,7 @@ export function ChunksPanel({ ) : null} {shouldMountOriginalPreview ? ( - + ; - readonly isLoading: boolean; - readonly message: string | null; - readonly totalPages: number; -}; - -const pageAssetPageSize = 50; - -function PageAssetDocumentViewer({ - focusedPageNumber, - focusedPageRequestId, - source, -}: { - readonly focusedPageNumber: number | null; - readonly focusedPageRequestId: number; - readonly source: SourceView; -}): ReactNode { - const [pageState, setPageState] = useState(() => ({ - pages: [], - loadedPageIndexes: new Set(), - isLoading: false, - message: null, - totalPages: Math.max( - 1, - Math.ceil( - (source.documentPresentation?.kind === "page-assets" - ? source.documentPresentation.pageCount - : 0) / pageAssetPageSize, - ), - ), - })); - const requestedPageIndexesRef = useRef>(new Set()); - const pageElementsRef = useRef>(new Map()); - const sourceIdRef = useRef(source.id); - - const loadPageIndex = useCallback( - (pageIndex: number): void => { - if (requestedPageIndexesRef.current.has(pageIndex)) return; - requestedPageIndexesRef.current.add(pageIndex); - setPageState((current) => ({ ...current, isLoading: true })); - - void workspaceClient - .fetchPageAssetPage(source.id, pageIndex) - .then((response) => { - setPageState((current) => { - const pagesByPageNumber = new Map( - current.pages.map((page) => [page.pageNumber, page]), - ); - for (const page of response.pages ?? []) { - pagesByPageNumber.set(page.pageNumber, page); - } - const loadedPageIndexes = new Set(current.loadedPageIndexes); - loadedPageIndexes.add(pageIndex); - - return { - pages: [...pagesByPageNumber.values()].sort( - (left, right) => left.pageNumber - right.pageNumber, - ), - loadedPageIndexes, - isLoading: false, - message: response.message ?? null, - totalPages: - response.pagination?.totalPages ?? current.totalPages, - }; - }); - }) - .catch(() => { - requestedPageIndexesRef.current.delete(pageIndex); - setPageState((current) => ({ ...current, isLoading: false })); - }); - }, - [source.id], - ); - - useEffect(() => { - if (sourceIdRef.current === source.id) return; - - sourceIdRef.current = source.id; - requestedPageIndexesRef.current = new Set(); - pageElementsRef.current = new Map(); - setPageState({ - pages: [], - loadedPageIndexes: new Set(), - isLoading: false, - message: null, - totalPages: Math.max( - 1, - Math.ceil( - (source.documentPresentation?.kind === "page-assets" - ? source.documentPresentation.pageCount - : 0) / pageAssetPageSize, - ), - ), - }); - }, [source.documentPresentation, source.id]); - - useEffect(() => { - loadPageIndex(1); - }, [loadPageIndex]); - - useEffect(() => { - if (!focusedPageNumber) return; - - loadPageIndex(Math.max(1, Math.ceil(focusedPageNumber / pageAssetPageSize))); - }, [focusedPageNumber, focusedPageRequestId, loadPageIndex]); - - useEffect(() => { - if (!focusedPageNumber) return; - - const element = pageElementsRef.current.get(focusedPageNumber); - if (!element) return; - - element.scrollIntoView({ behavior: "smooth", block: "start" }); - }, [focusedPageNumber, focusedPageRequestId, pageState.pages]); - - const nextPageIndex = getNextPageAssetPageIndex( - pageState.loadedPageIndexes, - pageState.totalPages, - ); - const canLoadMore = - nextPageIndex !== null && - nextPageIndex <= pageState.totalPages && - !pageState.isLoading; - - if (pageState.pages.length === 0 && pageState.isLoading) { - return ; - } - - if (pageState.pages.length === 0 && pageState.message) { - return ; - } - - if (pageState.pages.length === 0) { - return ; - } - - return ( -
- {pageState.pages.map((page) => ( - { - if (element) { - pageElementsRef.current.set(page.pageNumber, element); - return; - } - pageElementsRef.current.delete(page.pageNumber); - }} - /> - ))} - {canLoadMore ? ( -
- -
- ) : null} - {pageState.isLoading ? ( -
- Loading page images... -
- ) : null} -
- ); -} - -function PageAssetImage({ - isFocused, - page, - refCallback, -}: { - readonly isFocused: boolean; - readonly page: SourcePageAssetView; - readonly refCallback: (element: HTMLDivElement | null) => void; -}): ReactNode { - const [failedAssetUrl, setFailedAssetUrl] = useState(null); - const hasImageError = failedAssetUrl === page.assetUrl; - const aspectRatio = - page.width && page.height ? `${page.width} / ${page.height}` : undefined; - - return ( -
-
- Page {page.pageNumber} - {page.contentType} -
-
- {hasImageError ? ( - - ) : ( - // eslint-disable-next-line @next/next/no-img-element -- Page assets can be short-lived Knowhere URLs outside Next image optimization. - {`Page setFailedAssetUrl(page.assetUrl)} - /> - )} -
-
- ); -} - -function PageAssetImageUnavailable({ - pageNumber, -}: { - readonly pageNumber: number; -}): ReactNode { - return ( -
-
- -
-

- Page image unavailable. -

-
- ); -} - -function getNextPageAssetPageIndex( - loadedPageIndexes: ReadonlySet, - totalPages: number, -): number | null { - for (let pageIndex = 1; pageIndex <= totalPages; pageIndex += 1) { - if (!loadedPageIndexes.has(pageIndex)) return pageIndex; - } - - return null; -} - -function EmptyPageAssets(): ReactNode { - return ( -
- -

- No page images available. -

-
- ); -} - function UnavailableSourceMessage({ message, }: { diff --git a/src/components/parsed-chunk-card.test.ts b/src/components/parsed-chunk-card.test.ts index 946cffb..ff4fbef 100644 --- a/src/components/parsed-chunk-card.test.ts +++ b/src/components/parsed-chunk-card.test.ts @@ -1,6 +1,6 @@ // @vitest-environment jsdom import React from "react"; -import { cleanup, render, screen } from "@testing-library/react"; +import { cleanup, fireEvent, render, screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { afterEach, describe, expect, it, vi } from "vitest"; @@ -81,6 +81,48 @@ describe("ParsedChunkCard", () => { ).toBeTruthy(); }); + it("renders page citation assets instead of page summary content", () => { + render( + React.createElement(ParsedChunkCard, { + chunk: { + chunkId: "page_4", + type: "page", + content: "The summary should not be primary when an image exists.", + readableContent: "The summary should not be primary when an image exists.", + sourceTitle: "manual.pdf", + pageNums: [4], + pageAssets: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-4.png", + contentType: "image/png", + width: 1200, + height: 1600, + }, + ], + }, + isFocused: false, + isOriginalPreviewAvailable: true, + onChunkClick: vi.fn(), + onReferenceClick: vi.fn(), + }), + ); + + const pageImage = screen.getByRole("img", { name: "Page 4" }); + + expect(pageImage.getAttribute("src")).toBe( + "https://assets.example/page-4.png", + ); + expect(screen.getByText("image/png")).toBeTruthy(); + expect(screen.queryByText(/summary should not be primary/i)).toBeNull(); + expect( + screen.queryByRole("button", { name: /original file/i }), + ).toBeNull(); + + fireEvent.error(pageImage); + expect(screen.getByTestId("page-asset-image-unavailable-4")).toBeTruthy(); + }); + it("routes resolved artifact reference clicks to the target chunk", async () => { const user = userEvent.setup(); const onReferenceClick = vi.fn(); diff --git a/src/components/parsed-chunk-card.tsx b/src/components/parsed-chunk-card.tsx index 1ba233d..d794687 100644 --- a/src/components/parsed-chunk-card.tsx +++ b/src/components/parsed-chunk-card.tsx @@ -1,7 +1,15 @@ "use client"; -import { useMemo, type MouseEvent, type ReactNode } from "react"; -import { FileSearch, FileText, ImageIcon, Table2, Tags, TextQuote } from "lucide-react"; +import { useMemo, useState, type MouseEvent, type ReactNode } from "react"; +import { + FileSearch, + FileText, + ImageIcon, + ImageOff, + Table2, + Tags, + TextQuote, +} from "lucide-react"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; @@ -185,7 +193,9 @@ function ChunkSourcePanel({ ) : null}
- {onChunkClick && firstPageNumber !== null ? ( + {onChunkClick && + firstPageNumber !== null && + !hasPageCitationAssets(chunk) ? ( void; }): ReactNode { + const pageAssets = chunk.pageAssets ?? []; + return ( - -

- {chunk.readableContent ?? chunk.content} -

- + {pageAssets.length > 0 ? ( + + + + ) : ( + +

+ {chunk.readableContent ?? chunk.content} +

+ + )} ); } +function PageCitationAssets({ + assets, +}: { + readonly assets: NonNullable; +}): ReactNode { + return ( +
+ {assets.map((asset) => ( + + ))} +
+ ); +} + +function PageCitationAssetImage({ + asset, +}: { + readonly asset: NonNullable[number]; +}): ReactNode { + const [failedAssetUrl, setFailedAssetUrl] = useState(null); + const hasImageError = failedAssetUrl === asset.assetUrl; + const aspectRatio = + asset.width && asset.height ? `${asset.width} / ${asset.height}` : undefined; + + return ( +
+
+ Page {asset.pageNumber} + {asset.contentType} +
+
+ {hasImageError ? ( + + ) : ( + // eslint-disable-next-line @next/next/no-img-element -- Page assets can be short-lived Knowhere URLs outside Next image optimization. + {`Page setFailedAssetUrl(asset.assetUrl)} + /> + )} +
+
+ ); +} + +function PageCitationAssetUnavailable({ + pageNumber, +}: { + readonly pageNumber: number; +}): ReactNode { + return ( +
+
+ +
+

+ Page image unavailable. +

+
+ ); +} + function ImageChunkCard({ chunk, isFocused, @@ -467,6 +559,10 @@ function getImageChunkAssetUrl( return sourceOriginalFile.url; } +function hasPageCitationAssets(chunk: ParsedChunkView): boolean { + return (chunk.pageAssets?.length ?? 0) > 0; +} + function renderTextChunkContent( chunk: ParsedChunkView, onReferenceClick: (chunkId: string) => void, diff --git a/src/components/workspace-citation-focus.test.ts b/src/components/workspace-citation-focus.test.ts index bfc4654..21fce5a 100644 --- a/src/components/workspace-citation-focus.test.ts +++ b/src/components/workspace-citation-focus.test.ts @@ -169,8 +169,22 @@ describe("useWorkspaceCitationFocus", () => { expect(result.current.pendingCitationId).toBeNull(); }); - it("focuses page-asset citations without fetching full chunks", async () => { - const fetchChunks = vi.fn(async () => [prefetchedChunk]); + it("focuses page-asset citations through the loaded page chunk", async () => { + const pageChunk: ParsedChunkView = { + chunkId: "page_4", + documentId: "document_1", + type: "page", + content: "Page 4 summary", + sourceTitle: "Contract.pdf", + pageAssets: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-000004.png", + contentType: "image/png", + }, + ], + }; + const fetchChunks = vi.fn(async () => [pageChunk]); const selectSource = vi.fn(); const pageAssetSource: SourceView = { ...readySource, @@ -202,14 +216,14 @@ describe("useWorkspaceCitationFocus", () => { await result.current.handleCitationClick(pageCitation, "message_1:0"); }); - expect(fetchChunks).not.toHaveBeenCalled(); + expect(fetchChunks).toHaveBeenCalledWith("source_1"); expect(selectSource).toHaveBeenCalledWith("source_1"); - expect(result.current.focusedChunk.chunkId).toBeNull(); + expect(result.current.focusedChunk.chunkId).toBe("page_4"); expect(result.current.focusedPage).toEqual({ - pageNumber: 4, - requestId: 1, + pageNumber: null, + requestId: 0, }); - expect(result.current.citationListViewRequestId).toBe(0); + expect(result.current.citationListViewRequestId).toBe(1); }); it("reuses cached chunks for a different source without refetching", async () => { diff --git a/src/components/workspace-citation-focus.ts b/src/components/workspace-citation-focus.ts index c233671..5469cfa 100644 --- a/src/components/workspace-citation-focus.ts +++ b/src/components/workspace-citation-focus.ts @@ -211,17 +211,6 @@ export function useWorkspaceCitationFocus({ ) if (!source) return - if ( - workspaceCitationState.isPageAssetCitationTarget(source, citation) - ) { - if (selectedSourceId !== source.id) onSelectSource(source.id) - requestChunkFocus(null) - requestPageFocus( - workspaceCitationState.getCitationPageNumber(citation), - ) - return - } - setCitationListViewRequestId((current) => current + 1) const loadedChunkId = workspaceCitationState.getLoadedCitationChunkId({ @@ -291,7 +280,6 @@ export function useWorkspaceCitationFocus({ loadAllChunksForSource, onSelectSource, requestChunkFocus, - requestPageFocus, selectedChunks, selectedSourceId, sources, diff --git a/src/components/workspace-citation-state.test.ts b/src/components/workspace-citation-state.test.ts index ce5851e..cce0e6b 100644 --- a/src/components/workspace-citation-state.test.ts +++ b/src/components/workspace-citation-state.test.ts @@ -91,6 +91,44 @@ describe("workspaceCitationState", () => { ).toBeNull() }) + it("focuses loaded page chunks from page citation metadata", () => { + const citation: ChatCitationView = { + chunkType: "page", + score: 0.93, + pageCitationPageNumber: 4, + pageCitationAssetUrl: "https://assets.example/page-4.png", + source: { + documentId: "document_1", + sectionPath: "Page 4", + }, + } + + expect( + workspaceCitationState.getLoadedCitationChunkId({ + citation, + selectedSourceId: "source_1", + sourceId: "source_1", + selectedChunks: [ + { + chunkId: "page_4", + documentId: "document_1", + type: "page", + content: "Page 4 summary", + sourceTitle: "Contract.pdf", + pageAssets: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-4.png", + contentType: "image/png", + }, + ], + }, + ], + hasMoreSelectedChunks: true, + }), + ).toBe("page_4") + }) + describe("hasExactCitationTargetHint", () => { it("returns true when the citation has content text", () => { const citation: ChatCitationView = { @@ -120,6 +158,22 @@ describe("workspaceCitationState", () => { ).toBe(true) }) + it("returns true when the citation has page asset metadata", () => { + const citation: ChatCitationView = { + chunkType: "page", + score: 0.5, + pageCitationPageNumber: 4, + source: { + documentId: "document_1", + sectionPath: "Root", + }, + } + + expect( + workspaceCitationState.hasExactCitationTargetHint(citation), + ).toBe(true) + }) + it("returns false for source-only citations with no useful target hint", () => { const citation: ChatCitationView = { chunkType: "text", diff --git a/src/components/workspace-citation-state.ts b/src/components/workspace-citation-state.ts index cf33dad..0353409 100644 --- a/src/components/workspace-citation-state.ts +++ b/src/components/workspace-citation-state.ts @@ -27,10 +27,6 @@ type WorkspaceCitationStateModule = { readonly hasExactCitationTargetHint: ( citation: ChatCitationView, ) => boolean - readonly isPageAssetCitationTarget: ( - source: SourceView, - citation: ChatCitationView, - ) => boolean readonly getCitationPageNumber: (citation: ChatCitationView) => number | null readonly upsertPrefetchedChunks: ( current: PrefetchedChunksBySourceId, @@ -66,7 +62,10 @@ function getLoadedCitationChunkId( ? resolveCitationChunkByContent(input.citation, input.selectedChunks) : resolveCitationChunk(input.citation, input.selectedChunks) - return focusedChunk?.chunkId ?? null + return ( + focusedChunk?.chunkId ?? + resolveCitationPageChunkId(input.citation, input.selectedChunks) + ) } function hasExactCitationTargetHint(citation: ChatCitationView): boolean { @@ -74,6 +73,13 @@ function hasExactCitationTargetHint(citation: ChatCitationView): boolean { return true } + if ( + getCitationPageNumber(citation) !== null || + typeof citation.pageCitationAssetUrl === "string" + ) { + return true + } + const sectionPath = citation.source.sectionPath if (typeof sectionPath !== "string") return false @@ -84,17 +90,6 @@ function hasExactCitationTargetHint(citation: ChatCitationView): boolean { return true } -function isPageAssetCitationTarget( - source: SourceView, - citation: ChatCitationView, -): boolean { - return ( - source.documentPresentation?.kind === "page-assets" && - (getCitationPageNumber(citation) !== null || - typeof citation.pageCitationAssetUrl === "string") - ) -} - function getCitationPageNumber(citation: ChatCitationView): number | null { if ( typeof citation.pageCitationPageNumber === "number" && @@ -114,6 +109,43 @@ function getCitationPageNumber(citation: ChatCitationView): number | null { return Number.isSafeInteger(pageNumber) && pageNumber > 0 ? pageNumber : null } +function resolveCitationPageChunkId( + citation: ChatCitationView, + chunks: readonly ParsedChunkView[], +): string | null { + const pageNumber = getCitationPageNumber(citation) + if (pageNumber === null) return null + + const documentChunks = citation.source.documentId + ? chunks.filter((chunk) => chunk.documentId === citation.source.documentId) + : chunks + const matches = documentChunks.filter((chunk) => + isChunkForPageNumber(chunk, pageNumber), + ) + if (matches.length === 0) return null + + const pageAssetChunk = matches.find( + (chunk) => + chunk.type === "page" && + (chunk.pageAssets ?? []).some( + (pageAsset) => pageAsset.pageNumber === pageNumber, + ), + ) + + return pageAssetChunk?.chunkId ?? matches[0]?.chunkId ?? null +} + +function isChunkForPageNumber( + chunk: ParsedChunkView, + pageNumber: number, +): boolean { + return ( + (chunk.pageAssets ?? []).some( + (pageAsset) => pageAsset.pageNumber === pageNumber, + ) || (chunk.pageNums ?? []).includes(pageNumber) + ) +} + function upsertPrefetchedChunks( current: PrefetchedChunksBySourceId, sourceId: string, @@ -157,7 +189,6 @@ export const workspaceCitationState: WorkspaceCitationStateModule = { getLoadedCitationChunkId, getCitationPageNumber, hasExactCitationTargetHint, - isPageAssetCitationTarget, upsertPrefetchedChunks, removePrefetchedChunks, } diff --git a/src/components/workspace-selected-chunks.test.ts b/src/components/workspace-selected-chunks.test.ts index be575fb..f1edec1 100644 --- a/src/components/workspace-selected-chunks.test.ts +++ b/src/components/workspace-selected-chunks.test.ts @@ -10,12 +10,10 @@ import type { ParsedChunkView } from "@/domains/chunks/types"; import type { SourceView } from "@/domains/sources/types"; const fetchChunkPageMock = vi.hoisted(() => vi.fn()); -const fetchPageAssetPageMock = vi.hoisted(() => vi.fn()); vi.mock("@/domains/workspace/client", () => ({ workspaceClient: { fetchChunkPage: fetchChunkPageMock, - fetchPageAssetPage: fetchPageAssetPageMock, }, })); @@ -30,7 +28,6 @@ const readySource: SourceView = { describe("useWorkspaceSelectedChunks", () => { beforeEach(() => { fetchChunkPageMock.mockReset(); - fetchPageAssetPageMock.mockReset(); fetchChunkPageMock.mockResolvedValue({ chunks: [], pagination: { @@ -40,18 +37,9 @@ describe("useWorkspaceSelectedChunks", () => { totalPages: 1, }, }); - fetchPageAssetPageMock.mockResolvedValue({ - pages: [], - pagination: { - page: 1, - pageSize: 50, - total: 0, - totalPages: 0, - }, - }); }); - it("returns prefetched chunks while checking the visible page for media", async () => { + it("returns prefetched chunks while loading the visible chunk page", async () => { const { result } = renderHook( () => useWorkspaceSelectedChunks({ @@ -71,9 +59,6 @@ describe("useWorkspaceSelectedChunks", () => { { wrapper: createSWRWrapper }, ); - await waitFor(() => - expect(fetchPageAssetPageMock).toHaveBeenCalledWith("source_1", 1), - ); await waitFor(() => expect(fetchChunkPageMock).toHaveBeenCalledWith("source_1", 1), ); @@ -209,11 +194,34 @@ describe("useWorkspaceSelectedChunks", () => { expect(result.current.selectedChunks).toEqual([]); }); - it("does not fetch chunk pages for page-asset sources", () => { + it("loads chunk pages for page-asset sources", async () => { const pageAssetSource: SourceView = { ...readySource, documentPresentation: { kind: "page-assets", pageCount: 4 }, }; + fetchChunkPageMock.mockResolvedValue({ + chunks: [ + { + chunkId: "page_1", + type: "page", + content: "Page summary", + sourceTitle: "lecture.pdf", + pageAssets: [ + { + pageNumber: 1, + assetUrl: "https://blob.example/page-1.png", + contentType: "image/png", + }, + ], + }, + ], + pagination: { + page: 1, + pageSize: 50, + total: 1, + totalPages: 1, + }, + }); const { result } = renderHook( () => @@ -225,25 +233,35 @@ describe("useWorkspaceSelectedChunks", () => { { wrapper: createSWRWrapper }, ); - expect(result.current.selectedSource?.id).toBe("source_1"); - expect(result.current.selectedChunks).toEqual([]); + await waitFor(() => + expect(result.current.selectedChunks.map((chunk) => chunk.chunkId)).toEqual([ + "page_1", + ]), + ); expect(result.current.isSelectedChunksLoading).toBe(false); - expect(fetchPageAssetPageMock).not.toHaveBeenCalled(); - expect(fetchChunkPageMock).not.toHaveBeenCalled(); + expect(fetchChunkPageMock).toHaveBeenCalledWith("source_1", 1); }); - it("detects page assets for selected sources without presentation metadata", async () => { - fetchPageAssetPageMock.mockResolvedValue({ - pages: [ - { - pageNumber: 1, - assetUrl: "https://blob.example/page-1.png", - contentType: "image/png", - }, + it("detects page assets from selected source chunks without a second request", async () => { + fetchChunkPageMock.mockResolvedValue({ + chunks: [ { - pageNumber: 20, - assetUrl: "https://blob.example/page-20.png", - contentType: "image/png", + chunkId: "page_1", + type: "page", + content: "Page summary", + sourceTitle: "lecture.pdf", + pageAssets: [ + { + pageNumber: 1, + assetUrl: "https://blob.example/page-1.png", + contentType: "image/png", + }, + { + pageNumber: 20, + assetUrl: "https://blob.example/page-20.png", + contentType: "image/png", + }, + ], }, ], pagination: { @@ -270,9 +288,10 @@ describe("useWorkspaceSelectedChunks", () => { pageCount: 20, }), ); - expect(result.current.selectedChunks).toEqual([]); - expect(fetchPageAssetPageMock).toHaveBeenCalledWith("source_1", 1); - expect(fetchChunkPageMock).not.toHaveBeenCalled(); + expect(result.current.selectedChunks.map((chunk) => chunk.chunkId)).toEqual([ + "page_1", + ]); + expect(fetchChunkPageMock).toHaveBeenCalledWith("source_1", 1); }); it("requests a source refresh after loading an unlocalized remote source", async () => { diff --git a/src/components/workspace-selected-chunks.ts b/src/components/workspace-selected-chunks.ts index acbcd50..45b19a3 100644 --- a/src/components/workspace-selected-chunks.ts +++ b/src/components/workspace-selected-chunks.ts @@ -1,6 +1,6 @@ "use client" -import { useEffect, useMemo, useRef, useState } from "react" +import { useEffect, useMemo, useRef } from "react" import useSWRInfinite from "swr/infinite" import { workspaceClient } from "@/domains/workspace/client" @@ -8,7 +8,6 @@ import { workspaceClientCache, type SourceChunksKey, type SourceChunksResponse, - type SourcePageAssetsResponse, } from "@/domains/workspace/client-cache" import { resolveChunkConnectionTargets } from "@/domains/chunks" import type { ParsedChunkView } from "@/domains/chunks/types" @@ -31,10 +30,6 @@ type WorkspaceSelectedChunks = { readonly selectedSource: SourceView | undefined } -type PageAssetProbeState = - | { readonly status: "page-assets"; readonly pageCount: number } - | { readonly status: "parsed-chunks" } - export function useWorkspaceSelectedChunks({ selectedSourceId, sources, @@ -45,37 +40,12 @@ export function useWorkspaceSelectedChunks({ (source) => source.id === selectedSourceId, ) const remoteSourceRefreshRequestedIdsRef = useRef>(new Set()) - const requestedPageAssetProbeIdsRef = useRef>(new Set()) - const [pageAssetProbeBySourceId, setPageAssetProbeBySourceId] = useState< - Readonly> - >({}) - const pageAssetProbeState = rawSelectedSource - ? pageAssetProbeBySourceId[rawSelectedSource.id] - : undefined - const selectedSource = - rawSelectedSource && pageAssetProbeState?.status === "page-assets" - ? { - ...rawSelectedSource, - chunkCount: pageAssetProbeState.pageCount, - documentPresentation: { - kind: "page-assets" as const, - pageCount: pageAssetProbeState.pageCount, - }, - } - : rawSelectedSource + const selectedSource = rawSelectedSource const prefetchedSelectedChunks = selectedSourceId ? prefetchedChunksBySourceId[selectedSourceId] : undefined - const shouldProbePageAssets = - rawSelectedSource !== undefined && - rawSelectedSource.status === "ready" && - rawSelectedSource.documentPresentation === undefined && - pageAssetProbeState === undefined const selectedChunkSourceId = - selectedSource && - selectedSource.status === "ready" && - selectedSource.documentPresentation?.kind !== "page-assets" && - !shouldProbePageAssets + selectedSource && selectedSource.status === "ready" ? selectedSource.id : null const { @@ -118,9 +88,17 @@ export function useWorkspaceSelectedChunks({ : undefined, [pagedSelectedChunks, prefetchedSelectedChunks], ) - const selectedChunks = selectedSourceId - ? (resolvedPrefetchedChunks ?? pagedSelectedChunks) - : [] + const selectedChunks = useMemo( + () => + selectedSourceId + ? (resolvedPrefetchedChunks ?? pagedSelectedChunks) + : [], + [pagedSelectedChunks, resolvedPrefetchedChunks, selectedSourceId], + ) + const resolvedSelectedSource = useMemo( + () => getResolvedSelectedSource(selectedSource, selectedChunks), + [selectedChunks, selectedSource], + ) const hasMoreSelectedChunks = !prefetchedSelectedChunks && workspaceClientCache.hasMoreChunkPages(selectedChunkPages) @@ -132,40 +110,12 @@ export function useWorkspaceSelectedChunks({ typeof selectedChunkPages[selectedChunkPageCount - 1] === "undefined", ) const isSelectedChunksLoading = - shouldProbePageAssets || hasProcessingSelectedChunkPage || (selectedChunkSourceId !== null && !prefetchedSelectedChunks && !selectedChunkPages && isChunksLoading) - useEffect(() => { - const source = rawSelectedSource - if (!source || !shouldProbePageAssets) return - if (requestedPageAssetProbeIdsRef.current.has(source.id)) return - - requestedPageAssetProbeIdsRef.current.add(source.id) - void workspaceClient - .fetchPageAssetPage(source.id, 1) - .then((response) => { - setPageAssetProbeBySourceId((current) => ({ - ...current, - [source.id]: getPageAssetProbeState(response), - })) - - if (source.kind === "remote" && (response.pages?.length ?? 0) > 0) { - onRemoteSourceChunksLoaded?.(source.id) - } - }) - .catch(() => { - requestedPageAssetProbeIdsRef.current.delete(source.id) - setPageAssetProbeBySourceId((current) => ({ - ...current, - [source.id]: { status: "parsed-chunks" }, - })) - }) - }, [onRemoteSourceChunksLoaded, rawSelectedSource, shouldProbePageAssets]) - useEffect(() => { const sourceId = selectedSource?.id if (!sourceId || selectedSource.kind !== "remote") return @@ -188,23 +138,10 @@ export function useWorkspaceSelectedChunks({ isSelectedChunksLoadingMore, selectedChunksMessage, selectedChunks, - selectedSource, + selectedSource: resolvedSelectedSource, } } -function getPageAssetProbeState( - response: SourcePageAssetsResponse, -): PageAssetProbeState { - const pages = response.pages ?? [] - if (pages.length === 0) return { status: "parsed-chunks" } - - const maxPageNumber = Math.max( - ...pages.map((page) => page.pageNumber), - ) - const pageCount = Math.max(response.pagination?.total ?? 0, maxPageNumber) - return { status: "page-assets", pageCount } -} - function fetchChunksByKey([ , sourceId, @@ -213,6 +150,36 @@ function fetchChunksByKey([ return workspaceClient.fetchChunkPage(sourceId, page) } +function getResolvedSelectedSource( + source: SourceView | undefined, + chunks: readonly ParsedChunkView[], +): SourceView | undefined { + if (!source) return undefined + if (source.documentPresentation?.kind === "page-assets") return source + + const pageCount = getLoadedPageAssetCount(chunks) + if (!pageCount) return source + + return { + ...source, + documentPresentation: { + kind: "page-assets", + pageCount, + }, + } +} + +function getLoadedPageAssetCount( + chunks: readonly ParsedChunkView[], +): number | null { + const pageNumbers = chunks.flatMap((chunk) => + (chunk.pageAssets ?? []).map((pageAsset) => pageAsset.pageNumber), + ) + if (pageNumbers.length === 0) return null + + return Math.max(...pageNumbers) +} + function hasProcessingChunkPage( pages: readonly SourceChunksResponse[] | undefined, ): boolean { diff --git a/src/domains/chunks/index.test.ts b/src/domains/chunks/index.test.ts index 316bfd6..586d32b 100644 --- a/src/domains/chunks/index.test.ts +++ b/src/domains/chunks/index.test.ts @@ -169,6 +169,42 @@ describe("toParsedChunkView", () => { }); }); + it("maps usable page citation assets on page chunks", () => { + const chunk = makeDocumentChunk({ + id: "document_page_1", + chunkId: "parser_page_1", + chunkType: "page" as DocumentChunk["chunkType"], + metadata: { + pageAssets: [ + { + pageNum: 4, + assetUrl: "https://assets.example/page-4.png", + contentType: "image/png", + width: 1200, + height: 1600, + }, + { + pageNum: 5, + assetUrl: " ", + contentType: "image/png", + }, + ], + }, + }); + + expect(toParsedChunkView(chunk, "manual.pdf", "doc_123")).toMatchObject({ + pageAssets: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-4.png", + contentType: "image/png", + width: 1200, + height: 1600, + }, + ], + }); + }); + it("maps SDK-normalized page number metadata", () => { const chunk = makeDocumentChunk({ metadata: { diff --git a/src/domains/chunks/normalization.ts b/src/domains/chunks/normalization.ts index c6590b4..3ab6bae 100644 --- a/src/domains/chunks/normalization.ts +++ b/src/domains/chunks/normalization.ts @@ -35,6 +35,10 @@ function createParsedChunkView( metadata: input.metadata, }) const summary = getStringMetadata(input.metadata, "summary") + const pageAssets = + type === "page" + ? getPageAssetViews(input.metadata["pageAssets"], assetUrl) + : undefined return { chunkId: input.chunkId, @@ -47,6 +51,7 @@ function createParsedChunkView( readableContent: getReadableContent({ type, content, summary }), filePath, assetUrl, + pageAssets, summary, keywords: getStringArrayMetadata(input.metadata, "keywords"), pageNums: getPageNumbers( @@ -58,6 +63,38 @@ function createParsedChunkView( } } +function getPageAssetViews( + value: unknown, + fallbackAssetUrl?: string, +): ParsedChunkView["pageAssets"] | undefined { + if (!Array.isArray(value)) return undefined + + const pageAssets = value.flatMap((item): NonNullable => { + if (!isRecord(item)) return [] + const pageNumber = + getPositiveInteger(item["pageNum"]) ?? getPositiveInteger(item["pageNumber"]) + const assetUrl = getString(item["assetUrl"]) ?? fallbackAssetUrl + const contentType = getString(item["contentType"]) + if (!pageNumber || !assetUrl || !contentType) return [] + + return [ + { + pageNumber, + assetUrl, + contentType, + ...(getPositiveNumber(item["width"]) !== undefined + ? { width: getPositiveNumber(item["width"]) } + : {}), + ...(getPositiveNumber(item["height"]) !== undefined + ? { height: getPositiveNumber(item["height"]) } + : {}), + }, + ] + }) + + return pageAssets.length > 0 ? pageAssets : undefined +} + function resolveConnectionTargets( chunks: readonly ParsedChunkView[], ): ParsedChunkView[] { @@ -281,6 +318,20 @@ function getPageNumbers(value: unknown): number[] | undefined { return uniquePageNumbers.length > 0 ? uniquePageNumbers : undefined } +function getPositiveInteger(value: unknown): number | undefined { + return typeof value === "number" && + Number.isSafeInteger(value) && + value > 0 + ? value + : undefined +} + +function getPositiveNumber(value: unknown): number | undefined { + return typeof value === "number" && Number.isFinite(value) && value > 0 + ? value + : undefined +} + function isRecord(value: unknown): value is Readonly> { return typeof value === "object" && value !== null } diff --git a/src/domains/chunks/types.ts b/src/domains/chunks/types.ts index 2e84189..fedb17a 100644 --- a/src/domains/chunks/types.ts +++ b/src/domains/chunks/types.ts @@ -1,3 +1,5 @@ +import type { SourcePageAssetView } from "@/domains/sources/types" + export type ChunkType = "text" | "image" | "table" | "page" export type ParsedChunkConnection = { @@ -30,6 +32,7 @@ export type ParsedChunkView = { readonly filePath?: string /** Public Blob URL for parsed media/table artifacts when Notebook stored it. */ readonly assetUrl?: string + readonly pageAssets?: readonly SourcePageAssetView[] readonly summary?: string readonly keywords?: readonly string[] readonly pageNums?: readonly number[] diff --git a/src/domains/sources/page-assets.ts b/src/domains/sources/page-assets.ts deleted file mode 100644 index c23a702..0000000 --- a/src/domains/sources/page-assets.ts +++ /dev/null @@ -1,218 +0,0 @@ -import "server-only" - -import type { - DocumentChunk, - Knowledge, - KnowledgeReadResponse, -} from "@ontos-ai/knowhere-sdk" - -import type { ChunkPageParams } from "@/domains/chunks" -import type { SourcePageAssetView } from "./route-types" - -type ReadableSource = { - readonly documentId: string - readonly revisionKey?: string | null -} - -export type SourcePageAssetsPage = { - readonly pages: readonly SourcePageAssetView[] - readonly pagination: { - readonly page: number - readonly pageSize: number - readonly total: number - readonly totalPages: number - } -} - -type PageAssetReadClient = { - readonly documents: { - listChunks( - documentId: string, - params: { - readonly page: number - readonly pageSize: number - readonly chunkType: "page" - readonly includeAssetUrls: true - }, - ): Promise<{ - readonly chunks: readonly DocumentChunk[] - readonly pagination?: { - readonly page?: number - readonly pageSize?: number - readonly total?: number - readonly totalPages?: number - } - }> - } -} - -export async function readSourcePageAssets(input: { - readonly client: PageAssetReadClient - readonly knowledge: Knowledge - readonly source: ReadableSource - readonly params: ChunkPageParams -}): Promise { - const response = await input.knowledge.readChunks({ - documentId: input.source.documentId, - ...(input.source.revisionKey ? { revisionKey: input.source.revisionKey } : {}), - chunkType: "page", - page: input.params.page, - pageSize: input.params.pageSize, - }) - - const knowledgePage = toSourcePageAssetsPageFromKnowledgeResponse( - response, - input.params, - ) - if ( - isParsedStorageReadResponse(response) && - knowledgePage.pages.length > 0 - ) { - return knowledgePage - } - if (knowledgePage.pages.length > 0) return knowledgePage - - const remoteResponse = await input.client.documents.listChunks( - input.source.documentId, - { - page: input.params.page, - pageSize: input.params.pageSize, - chunkType: "page", - includeAssetUrls: true, - }, - ) - return toSourcePageAssetsPageFromRemoteResponse(remoteResponse, input.params) -} - -function toSourcePageAssetsPageFromKnowledgeResponse( - response: KnowledgeReadResponse, - params: ChunkPageParams, -): SourcePageAssetsPage { - const pages = response.chunks.flatMap((chunk): SourcePageAssetView[] => - readPageAssetViews(chunk.metadata.pageAssets, chunk.assetUrl), - ) - const maxPageNumber = getMaxPageNumber(pages) - const total = Math.max(response.totalChunks ?? 0, maxPageNumber ?? 0) - const resolvedTotal = total > 0 ? total : pages.length - const computedTotalPages = Math.max( - 1, - Math.ceil(resolvedTotal / params.pageSize), - ) - const totalPages = Math.max(response.totalPages ?? 0, computedTotalPages) - - return { - pages, - pagination: { - page: response.page ?? params.page, - pageSize: response.pageSize ?? params.pageSize, - total: resolvedTotal, - totalPages, - }, - } -} - -function toSourcePageAssetsPageFromRemoteResponse( - response: { - readonly chunks: readonly DocumentChunk[] - readonly pagination?: { - readonly page?: number - readonly pageSize?: number - readonly total?: number - readonly totalPages?: number - } - }, - params: ChunkPageParams, -): SourcePageAssetsPage { - const pages = response.chunks.flatMap((chunk): SourcePageAssetView[] => - readPageAssetViews(chunk.metadata.pageAssets, chunk.assetUrl ?? undefined), - ) - const maxPageNumber = getMaxPageNumber(pages) - const total = Math.max(response.pagination?.total ?? 0, maxPageNumber ?? 0) - const resolvedTotal = total > 0 ? total : pages.length - const computedTotalPages = Math.max( - 1, - Math.ceil(resolvedTotal / params.pageSize), - ) - const totalPages = Math.max( - response.pagination?.totalPages ?? 0, - computedTotalPages, - ) - - return { - pages, - pagination: { - page: response.pagination?.page ?? params.page, - pageSize: response.pagination?.pageSize ?? params.pageSize, - total: resolvedTotal, - totalPages, - }, - } -} - -function isParsedStorageReadResponse(response: KnowledgeReadResponse): boolean { - const resultDirectoryPath = response.document.resultDirectoryPath - return ( - typeof resultDirectoryPath === "string" && - resultDirectoryPath.startsWith("parsed-storage:") - ) -} - -function readPageAssetViews( - value: unknown, - fallbackAssetUrl?: string, -): SourcePageAssetView[] { - if (!Array.isArray(value)) return [] - - return value.flatMap((item): SourcePageAssetView[] => { - if (!isRecord(item)) return [] - const pageNumber = getPositiveInteger(item.pageNum) - const assetUrl = getTrimmedString(item.assetUrl) ?? fallbackAssetUrl - const contentType = getTrimmedString(item.contentType) - if (!pageNumber || !assetUrl || !contentType) return [] - - return [ - { - pageNumber, - assetUrl, - contentType, - ...(getPositiveNumber(item.width) !== undefined - ? { width: getPositiveNumber(item.width) } - : {}), - ...(getPositiveNumber(item.height) !== undefined - ? { height: getPositiveNumber(item.height) } - : {}), - }, - ] - }) -} - -function getTrimmedString(value: unknown): string | undefined { - if (typeof value !== "string") return undefined - const trimmed = value.trim() - return trimmed.length > 0 ? trimmed : undefined -} - -function getPositiveInteger(value: unknown): number | undefined { - return typeof value === "number" && - Number.isSafeInteger(value) && - value > 0 - ? value - : undefined -} - -function getPositiveNumber(value: unknown): number | undefined { - return typeof value === "number" && Number.isFinite(value) && value > 0 - ? value - : undefined -} - -function getMaxPageNumber( - pages: readonly SourcePageAssetView[], -): number | undefined { - if (pages.length === 0) return undefined - return Math.max(...pages.map((page) => page.pageNumber)) -} - -function isRecord(value: unknown): value is Readonly> { - return typeof value === "object" && value !== null -} diff --git a/src/domains/sources/route-page-assets.ts b/src/domains/sources/route-page-assets.ts deleted file mode 100644 index a2316fd..0000000 --- a/src/domains/sources/route-page-assets.ts +++ /dev/null @@ -1,208 +0,0 @@ -import { Effect } from "effect" - -import { readSourcePageAssets } from "./page-assets" -import { routeResult } from "@/lib/route-result" -import { displayReadUnavailable } from "./display-read-unavailable" -import { - decodeRemoteSourceId, - findRemoteLibraryDocumentBySourceId, -} from "./remote-library" -import { - getClientForWorkspace, - getKnowledgeResourcesForSource, -} from "./route-dependencies" -import { sourceRowRepository } from "./source-row-repository" -import type { - JsonRouteResult, - LoadSourcePageAssetsInput, - SourcePageAssetsBody, - SourceRouteServiceDependencies, -} from "./route-types" - -type RoutePageAssetsDependencies = Pick< - SourceRouteServiceDependencies, - | "ensureApiKeyForWorkspace" - | "ensureWorkspace" - | "getCurrentUser" - | "makeKnowhereClient" - | "sourceService" -> - -type RoutePageAssets = { - readonly loadSourcePageAssets: ( - input: LoadSourcePageAssetsInput, - ) => Promise> -} - -function createRoutePageAssets( - deps: RoutePageAssetsDependencies, -): RoutePageAssets { - return { - loadSourcePageAssets: (input: LoadSourcePageAssetsInput) => - Effect.runPromise(loadSourcePageAssetsEffect(input, deps)), - } -} - -const loadSourcePageAssetsEffect = ( - input: LoadSourcePageAssetsInput, - deps: RoutePageAssetsDependencies, -) => - Effect.gen(function* () { - if (!sourceRowRepository.isWorkspaceSourceId(input.sourceId)) { - const remoteResult = yield* loadRemotePageAssetsEffect(input, deps) - return remoteResult ?? sourceNotFound() - } - - const user = yield* Effect.tryPromise(() => deps.getCurrentUser()) - if (!user) return sourceNotFound() - - const workspace = yield* Effect.tryPromise(() => - deps.ensureWorkspace(user.id), - ) - const source = yield* Effect.tryPromise(() => - deps.sourceService.findInWorkspace(workspace.id, input.sourceId), - ) - if (!source) return sourceNotFound() - if (source.status !== "ready" || !source.knowhereDocumentId) { - return sourceNotReady() - } - - const documentId = source.knowhereDocumentId - const apiKey = yield* Effect.tryPromise(() => - deps.ensureApiKeyForWorkspace(workspace.id, input.cookieHeader), - ) - const readResources = getKnowledgeResourcesForSource({ - apiKey, - workspaceId: workspace.id, - sourceId: source.id, - documentId, - revisionKey: source.knowhereJobId, - }) - return yield* Effect.tryPromise(() => - readSourcePageAssets({ - client: readResources.client, - knowledge: readResources.knowledge, - source: { - documentId, - revisionKey: source.knowhereJobId, - }, - params: input.pageParams, - }), - ).pipe( - Effect.map((pageAssets) => routeResult.ok(pageAssets)), - Effect.catchAll((error) => - recoverUnavailablePageAssets(input, error), - ), - ) - }) - -const loadRemotePageAssetsEffect = ( - input: LoadSourcePageAssetsInput, - deps: RoutePageAssetsDependencies, -) => - Effect.gen(function* () { - if (!decodeRemoteSourceId(input.sourceId)) return null - - const user = yield* Effect.tryPromise(() => deps.getCurrentUser()) - if (!user) return null - - const workspace = yield* Effect.tryPromise(() => - deps.ensureWorkspace(user.id), - ) - const apiKey = yield* Effect.tryPromise(() => - deps.ensureApiKeyForWorkspace(workspace.id, input.cookieHeader), - ) - const client = yield* Effect.tryPromise(() => - getClientForWorkspace(workspace.id, input.cookieHeader, deps), - ) - const remoteDocument = yield* findRemoteLibraryDocumentBySourceId({ - sourceId: input.sourceId, - workspace, - client, - localSources: [], - }) - if (!remoteDocument) return null - - const source = yield* Effect.tryPromise(() => - deps.sourceService.localizeRemoteDocument(workspace.id, { - documentId: remoteDocument.documentId, - namespace: remoteDocument.namespace, - status: remoteDocument.status, - title: remoteDocument.title, - mimeType: remoteDocument.mimeType, - sizeBytes: remoteDocument.sizeBytes, - revisionKey: remoteDocument.revisionKey ?? null, - }), - ) - const documentId = source.knowhereDocumentId ?? remoteDocument.documentId - const revisionKey = - source.knowhereJobId ?? remoteDocument.revisionKey ?? null - const readResources = getKnowledgeResourcesForSource({ - apiKey, - workspaceId: workspace.id, - sourceId: source.id, - documentId, - revisionKey, - }) - return yield* Effect.tryPromise(() => - readSourcePageAssets({ - client: readResources.client, - knowledge: readResources.knowledge, - source: { documentId, revisionKey }, - params: input.pageParams, - }), - ).pipe( - Effect.map((pageAssets) => routeResult.ok(pageAssets)), - Effect.catchAll((error) => - recoverUnavailablePageAssets(input, error), - ), - ) - }) - -function sourceNotFound(): JsonRouteResult<{ readonly message: string }> { - return routeResult.error(404, "Source not found.") -} - -function sourceNotReady(): JsonRouteResult<{ readonly message: string }> { - return routeResult.error(409, "Source is not ready.") -} - -function sourcePageAssetsUnavailable( - input: LoadSourcePageAssetsInput, -): { - readonly pages: [] - readonly pagination: { - readonly page: number - readonly pageSize: number - readonly total: 0 - readonly totalPages: 0 - } - readonly message: string - readonly isUnavailable: true -} { - return { - pages: [], - pagination: { - page: input.pageParams.page, - pageSize: input.pageParams.pageSize, - total: 0, - totalPages: 0, - }, - message: displayReadUnavailable.message, - isUnavailable: true, - } -} - -function recoverUnavailablePageAssets( - input: LoadSourcePageAssetsInput, - error: unknown, -): Effect.Effect< - JsonRouteResult>, - unknown -> { - return displayReadUnavailable.isError(error) - ? Effect.succeed(routeResult.ok(sourcePageAssetsUnavailable(input))) - : Effect.fail(error) -} - -export { createRoutePageAssets } diff --git a/src/domains/sources/route-service.ts b/src/domains/sources/route-service.ts index 7e3b85f..9a19bef 100644 --- a/src/domains/sources/route-service.ts +++ b/src/domains/sources/route-service.ts @@ -4,14 +4,12 @@ import { createRouteArchive } from "./route-archive" import { createRouteChunks } from "./route-chunks" import { createSourceRouteDependencies } from "./route-dependencies" import { createRouteListing } from "./route-listing" -import { createRoutePageAssets } from "./route-page-assets" import { createRouteRetry } from "./route-retry" import { createRouteUpload } from "./route-upload" import type { ArchiveSourceInput, ListSourcesInput, LoadSourceChunksInput, - LoadSourcePageAssetsInput, RetrySourceInput, SourceRouteService, SourceRouteServiceOverrides, @@ -27,7 +25,6 @@ export function createSourceRouteService( const archive = createRouteArchive(deps) const retry = createRouteRetry(deps) const chunks = createRouteChunks(deps) - const pageAssets = createRoutePageAssets(deps) return { listSources: (input: ListSourcesInput) => listing.listSources(input), @@ -36,7 +33,5 @@ export function createSourceRouteService( retrySource: (input: RetrySourceInput) => retry.retrySource(input), loadSourceChunks: (input: LoadSourceChunksInput) => chunks.loadSourceChunks(input), - loadSourcePageAssets: (input: LoadSourcePageAssetsInput) => - pageAssets.loadSourcePageAssets(input), } } diff --git a/src/domains/sources/route-types.ts b/src/domains/sources/route-types.ts index 8ce806d..01d1883 100644 --- a/src/domains/sources/route-types.ts +++ b/src/domains/sources/route-types.ts @@ -5,7 +5,6 @@ import type { } from "@/domains/chunks" import type { ParsedChunkView } from "@/domains/chunks/types" import type { - SourcePageAssetView, SourceStatus, SourceView, } from "@/domains/sources/types" @@ -124,31 +123,6 @@ type SourceChunksBody = readonly message: string } -type SourcePageAssetsBody = - | { - readonly pages: readonly SourcePageAssetView[] - readonly pagination: { - readonly page: number - readonly pageSize: number - readonly total: number - readonly totalPages: number - } - } - | { - readonly pages: readonly [] - readonly pagination: { - readonly page: number - readonly pageSize: number - readonly total: 0 - readonly totalPages: 0 - } - readonly message: string - readonly isUnavailable: true - } - | { - readonly message: string - } - type ListSourcesInput = { readonly cookieHeader: string } @@ -176,12 +150,6 @@ type LoadSourceChunksInput = { readonly pageParams: ChunkPageParams } -type LoadSourcePageAssetsInput = { - readonly cookieHeader: string - readonly sourceId: string - readonly pageParams: ChunkPageParams -} - type SourceRouteService = { readonly listSources: ( input: ListSourcesInput, @@ -198,9 +166,6 @@ type SourceRouteService = { readonly loadSourceChunks: ( input: LoadSourceChunksInput, ) => Promise> - readonly loadSourcePageAssets: ( - input: LoadSourcePageAssetsInput, - ) => Promise> } type SourceWorkflowService = { @@ -304,12 +269,9 @@ export type { ListSourcesBody, ListSourcesInput, LoadSourceChunksInput, - LoadSourcePageAssetsInput, RetrySourceBody, RetrySourceInput, SourceChunksBody, - SourcePageAssetsBody, - SourcePageAssetView, SourceRouteDemoApi, SourceRouteKnowhereClient, SourceRouteService, diff --git a/src/domains/workspace/client-cache.ts b/src/domains/workspace/client-cache.ts index beaa41d..af5baae 100644 --- a/src/domains/workspace/client-cache.ts +++ b/src/domains/workspace/client-cache.ts @@ -11,9 +11,6 @@ import type { SourceView } from "@/domains/sources/types" type SourceChunksResponse = Awaited< ReturnType > -type SourcePageAssetsResponse = Awaited< - ReturnType -> type ChatThreadDetailResponse = Awaited< ReturnType > @@ -117,5 +114,4 @@ export type { ChatThreadKey, SourceChunksKey, SourceChunksResponse, - SourcePageAssetsResponse, } diff --git a/src/domains/workspace/client.test.ts b/src/domains/workspace/client.test.ts index 0c9df2b..47878bf 100644 --- a/src/domains/workspace/client.test.ts +++ b/src/domains/workspace/client.test.ts @@ -120,36 +120,6 @@ describe("workspaceClient", () => { }) }) - it("preserves page asset unavailable messages", async () => { - mockRouteClient.getJson.mockResolvedValue({ - pages: [], - message: - "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", - isUnavailable: true, - pagination: { - page: 1, - pageSize: 50, - total: 0, - totalPages: 0, - }, - }) - - const page = await workspaceClient.fetchPageAssetPage("source_1", 1) - - expect(page).toEqual({ - pages: [], - isUnavailable: true, - message: - "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", - pagination: { - page: 1, - pageSize: 50, - total: 0, - totalPages: 0, - }, - }) - }) - it("throws materialization route errors instead of treating them as empty sources", async () => { mockRouteClient.postJsonWithStatus.mockResolvedValue({ status: 502, diff --git a/src/domains/workspace/client.ts b/src/domains/workspace/client.ts index 5490d33..a415d6a 100644 --- a/src/domains/workspace/client.ts +++ b/src/domains/workspace/client.ts @@ -5,7 +5,6 @@ import type { } from "@/domains/chat/types" import type { ParsedChunkView } from "@/domains/chunks/types" import type { - SourcePageAssetView, SourceView, } from "@/domains/sources/types" import { workspaceRouteClient } from "./route-client" @@ -38,18 +37,6 @@ type SourceChunksResponse = { } } -type SourcePageAssetsResponse = { - pages?: SourcePageAssetView[] - isUnavailable?: boolean - message?: string - pagination?: { - page: number - pageSize: number - total: number - totalPages: number - } -} - type ChatThreadResponse = { thread?: ChatThreadView messages?: ChatMessageView[] @@ -107,7 +94,6 @@ export const workspaceClient = { keys: workspaceClientKeys, fetchChunks, fetchChunkPage, - fetchPageAssetPage, fetchSources, fetchChatThreads, fetchChatThread, @@ -156,26 +142,6 @@ async function fetchChunkPage( } } -async function fetchPageAssetPage( - sourceId: string, - page: number, -): Promise { - const searchParams = new URLSearchParams({ - page: String(page), - pageSize: String(workspaceClientConfig.sourceChunkPageSize), - }) - const body = await workspaceRouteClient.getJson( - `/api/sources/${encodeURIComponent(sourceId)}/page-assets?${searchParams.toString()}`, - ) - - return { - pages: Array.isArray(body.pages) ? body.pages : [], - ...(body.isUnavailable === true ? { isUnavailable: true } : {}), - ...(typeof body.message === "string" ? { message: body.message } : {}), - pagination: body.pagination, - } -} - async function fetchSources(): Promise { const body = await workspaceRouteClient.getJson( workspaceClientKeys.sources, From c83a0f38cbe1c0efc4a9a344d0b105a091c4fc42 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Wed, 8 Jul 2026 02:49:35 +0800 Subject: [PATCH 10/14] Fix remote document refresh and page dedupe --- src/components/chunks-panel-state.test.ts | 59 +++++++++++++ src/components/chunks-panel-state.ts | 34 ++++++++ src/components/chunks-panel.test.ts | 67 +++++++++++++++ src/components/chunks-panel.tsx | 17 ++-- src/components/workspace-citation-focus.ts | 3 - .../workspace-selected-chunks.test.ts | 27 +++--- src/components/workspace-selected-chunks.ts | 15 +--- src/components/workspace-shell.test.ts | 85 +++++++++++++++++++ src/components/workspace-shell.tsx | 1 - 9 files changed, 268 insertions(+), 40 deletions(-) diff --git a/src/components/chunks-panel-state.test.ts b/src/components/chunks-panel-state.test.ts index bcee0f0..53cb15a 100644 --- a/src/components/chunks-panel-state.test.ts +++ b/src/components/chunks-panel-state.test.ts @@ -139,6 +139,65 @@ describe("chunksPanelState", () => { ]) }) + it("deduplicates page-asset chunks by page number", () => { + const chunks: ParsedChunkView[] = [ + { + chunkId: "page_4_first", + type: "page", + content: "First page 4 summary.", + sourceTitle: "manual.pdf", + pageNums: [4], + pageAssets: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-4-a.png", + contentType: "image/png", + }, + ], + }, + { + chunkId: "page_4_duplicate", + type: "page", + content: "Duplicate page 4 summary.", + sourceTitle: "manual.pdf", + pageNums: [4], + pageAssets: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-4-b.png", + contentType: "image/png", + }, + ], + }, + { + chunkId: "page_5", + type: "page", + content: "Page 5 summary.", + sourceTitle: "manual.pdf", + pageAssets: [ + { + pageNumber: 5, + assetUrl: "https://assets.example/page-5.png", + contentType: "image/png", + }, + ], + }, + { + chunkId: "table_1", + type: "table", + content: "", + sourceTitle: "manual.pdf", + pageNums: [4], + }, + ] + + expect( + chunksPanelState + .getPageAssetChunksWithoutDuplicatePages(chunks) + .map((chunk) => chunk.chunkId), + ).toEqual(["page_4_first", "page_5", "table_1"]) + }) + it("formats Knowhere section paths and reference labels for display", () => { expect( chunksPanelState.formatChunkSectionPath( diff --git a/src/components/chunks-panel-state.ts b/src/components/chunks-panel-state.ts index dfd2ac7..0c64515 100644 --- a/src/components/chunks-panel-state.ts +++ b/src/components/chunks-panel-state.ts @@ -42,6 +42,9 @@ type ChunksPanelStateModule = { chunks: readonly ParsedChunkView[], focusedChunkId: string | null, ) => readonly ParsedChunkView[] + readonly getPageAssetChunksWithoutDuplicatePages: ( + chunks: readonly ParsedChunkView[], + ) => readonly ParsedChunkView[] readonly getReferenceLabel: (connection: ParsedChunkConnection) => string readonly getRenderableReferences: ( chunk: ParsedChunkView, @@ -157,6 +160,32 @@ function dedupeChunksById( return uniqueChunks } +function getPageAssetChunksWithoutDuplicatePages( + chunks: readonly ParsedChunkView[], +): readonly ParsedChunkView[] { + const seenPageNumbers = new Set() + + return chunks.filter((chunk) => { + if (chunk.type !== "page") return true + + const pageNumber = getPageAssetChunkPageNumber(chunk) + if (pageNumber === null) return true + if (seenPageNumbers.has(pageNumber)) return false + + seenPageNumbers.add(pageNumber) + return true + }) +} + +function getPageAssetChunkPageNumber(chunk: ParsedChunkView): number | null { + const pageAssetNumbers = (chunk.pageAssets ?? []) + .map((pageAsset) => pageAsset.pageNumber) + .filter(isPositivePageNumber) + if (pageAssetNumbers.length > 0) return Math.min(...pageAssetNumbers) + + return getFirstPageNumber(chunk) +} + function createMutableSectionTreeNode(input: { readonly id: string readonly kind: ChunkSectionTreeNodeKind @@ -336,6 +365,10 @@ function getFirstPageNumber(chunk: ParsedChunkView): number | null { return Math.min(...finitePageNumbers) } +function isPositivePageNumber(pageNumber: number): boolean { + return Number.isFinite(pageNumber) && pageNumber > 0 +} + function formatChunkSectionPath( sectionPath: ParsedChunkView["sectionPath"], ): string | null { @@ -454,6 +487,7 @@ export const chunksPanelState: ChunksPanelStateModule = { formatChunkSectionPath, formatReferenceLabel, getChunksWithFocusedFirst, + getPageAssetChunksWithoutDuplicatePages, getReferenceLabel, getRenderableReferences, } diff --git a/src/components/chunks-panel.test.ts b/src/components/chunks-panel.test.ts index 78a35f7..3031cb8 100644 --- a/src/components/chunks-panel.test.ts +++ b/src/components/chunks-panel.test.ts @@ -127,6 +127,73 @@ describe("ChunksPanel", () => { ).toBeNull(); }); + it("renders only the first page-asset chunk for each page number", async () => { + mockVisibleVirtualViewport(); + + render( + React.createElement(C, { + chunks: [ + { + chunkId: "page_4_first", + type: "page", + content: "First page 4 summary", + sourceTitle: "report.pdf", + pageNums: [4], + pageAssets: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-000004-a.png", + contentType: "image/png", + }, + ], + }, + { + chunkId: "page_4_duplicate", + type: "page", + content: "Duplicate page 4 summary", + sourceTitle: "report.pdf", + pageNums: [4], + pageAssets: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-000004-b.png", + contentType: "image/png", + }, + ], + }, + { + chunkId: "page_5", + type: "page", + content: "Page 5 summary", + sourceTitle: "report.pdf", + pageNums: [5], + pageAssets: [ + { + pageNumber: 5, + assetUrl: "https://assets.example/page-000005.png", + contentType: "image/png", + }, + ], + }, + ], + selectedSource: "report.pdf", + selectedSourceView: { + id: "source_1", + title: "report.pdf", + mimeType: "application/pdf", + status: "ready", + documentPresentation: { kind: "page-assets", pageCount: 5 }, + }, + }), + ); + + expect(await screen.findAllByRole("img", { name: "Page 4" })) + .toHaveLength(1); + expect(screen.getByRole("img", { name: "Page 5" })).toBeTruthy(); + expect(screen.queryByTestId("chunk-card-shell-page_4_duplicate")) + .toBeNull(); + }); + it("renders page chunks normally when no page assets exist", async () => { mockVisibleVirtualViewport(); diff --git a/src/components/chunks-panel.tsx b/src/components/chunks-panel.tsx index 71037c9..f49df7e 100644 --- a/src/components/chunks-panel.tsx +++ b/src/components/chunks-panel.tsx @@ -106,6 +106,13 @@ export function ChunksPanel({ const isPageAssetSource = selectedSourceView?.documentPresentation?.kind === "page-assets" || hasPageAssetChunks; + const displayChunks = useMemo( + () => + isPageAssetSource + ? chunksPanelState.getPageAssetChunksWithoutDuplicatePages(chunks) + : chunks, + [chunks, isPageAssetSource], + ); const effectiveVisibleView = isPageAssetSource ? "parsed" : undefined; const originalPreviewCacheKey = selectedSourceFile?.url ?? null; const isOriginalPreviewAvailable = @@ -144,7 +151,7 @@ export function ChunksPanel({ visibleChunks, visibleView, } = useChunksPanelWorkflow({ - chunks, + chunks: displayChunks, selectedSource, selectedSourceFile, focusedChunkId, @@ -331,7 +338,7 @@ export function ChunksPanel({
{!isPageAssetSource && activeVisibleView === "parsed" && - chunks.length > 0 ? ( + displayChunks.length > 0 ? (
{isLoading ? ( - ) : chunks.length === 0 && processingMessage ? ( + ) : displayChunks.length === 0 && processingMessage ? ( - ) : chunks.length === 0 ? ( + ) : displayChunks.length === 0 ? ( selectedSource ? ( ) : ( @@ -404,7 +411,7 @@ export function ChunksPanel({ ) : isTreeModeVisible ? ( Promise readonly initialPrefetchedChunksBySourceId?: PrefetchedChunksBySourceId - readonly onRemoteSourceChunksLoaded?: (sourceId: string) => void readonly onSelectSource: (sourceId: string | null) => void readonly selectedSourceId: string | null readonly sources: readonly SourceView[] @@ -58,7 +57,6 @@ type WorkspaceCitationFocus = { export function useWorkspaceCitationFocus({ fetchChunks, initialPrefetchedChunksBySourceId = {}, - onRemoteSourceChunksLoaded, onSelectSource, selectedSourceId, sources, @@ -100,7 +98,6 @@ export function useWorkspaceCitationFocus({ selectedSourceId, sources, prefetchedChunksBySourceId, - onRemoteSourceChunksLoaded, }) const requestChunkFocus = useCallback( diff --git a/src/components/workspace-selected-chunks.test.ts b/src/components/workspace-selected-chunks.test.ts index f1edec1..59cf03a 100644 --- a/src/components/workspace-selected-chunks.test.ts +++ b/src/components/workspace-selected-chunks.test.ts @@ -294,8 +294,7 @@ describe("useWorkspaceSelectedChunks", () => { expect(fetchChunkPageMock).toHaveBeenCalledWith("source_1", 1); }); - it("requests a source refresh after loading an unlocalized remote source", async () => { - const onRemoteSourceChunksLoaded = vi.fn(); + it("loads an unlocalized remote source without requesting a source refresh", async () => { const remoteSource: SourceView = { ...readySource, id: "knowhere-doc:default:doc_remote", @@ -320,30 +319,24 @@ describe("useWorkspaceSelectedChunks", () => { }, }); - const { rerender } = renderHook( - (input: { - readonly onRemoteSourceChunksLoaded: (sourceId: string) => void; - }) => + const { result } = renderHook( + () => useWorkspaceSelectedChunks({ selectedSourceId: "knowhere-doc:default:doc_remote", sources: [remoteSource], prefetchedChunksBySourceId: {}, - onRemoteSourceChunksLoaded: input.onRemoteSourceChunksLoaded, }), - { - initialProps: { onRemoteSourceChunksLoaded }, - wrapper: createSWRWrapper, - }, + { wrapper: createSWRWrapper }, ); await waitFor(() => - expect(onRemoteSourceChunksLoaded).toHaveBeenCalledWith( - "knowhere-doc:default:doc_remote", - ), + expect(result.current.selectedChunks.map((chunk) => chunk.chunkId)).toEqual([ + "chunk_1", + ]), + ); + expect(result.current.selectedSource?.id).toBe( + "knowhere-doc:default:doc_remote", ); - - rerender({ onRemoteSourceChunksLoaded }); - expect(onRemoteSourceChunksLoaded).toHaveBeenCalledTimes(1); }); it("returns an empty chunk list when no source is selected", () => { diff --git a/src/components/workspace-selected-chunks.ts b/src/components/workspace-selected-chunks.ts index 45b19a3..260a6f5 100644 --- a/src/components/workspace-selected-chunks.ts +++ b/src/components/workspace-selected-chunks.ts @@ -1,6 +1,6 @@ "use client" -import { useEffect, useMemo, useRef } from "react" +import { useMemo } from "react" import useSWRInfinite from "swr/infinite" import { workspaceClient } from "@/domains/workspace/client" @@ -17,7 +17,6 @@ type WorkspaceSelectedChunksInput = { readonly selectedSourceId: string | null readonly sources: readonly SourceView[] readonly prefetchedChunksBySourceId: Readonly> - readonly onRemoteSourceChunksLoaded?: (sourceId: string) => void } type WorkspaceSelectedChunks = { @@ -34,12 +33,10 @@ export function useWorkspaceSelectedChunks({ selectedSourceId, sources, prefetchedChunksBySourceId, - onRemoteSourceChunksLoaded, }: WorkspaceSelectedChunksInput): WorkspaceSelectedChunks { const rawSelectedSource = sources.find( (source) => source.id === selectedSourceId, ) - const remoteSourceRefreshRequestedIdsRef = useRef>(new Set()) const selectedSource = rawSelectedSource const prefetchedSelectedChunks = selectedSourceId ? prefetchedChunksBySourceId[selectedSourceId] @@ -116,16 +113,6 @@ export function useWorkspaceSelectedChunks({ !selectedChunkPages && isChunksLoading) - useEffect(() => { - const sourceId = selectedSource?.id - if (!sourceId || selectedSource.kind !== "remote") return - if (!selectedChunkPages || selectedChunkPages.length === 0) return - if (remoteSourceRefreshRequestedIdsRef.current.has(sourceId)) return - - remoteSourceRefreshRequestedIdsRef.current.add(sourceId) - onRemoteSourceChunksLoaded?.(sourceId) - }, [onRemoteSourceChunksLoaded, selectedChunkPages, selectedSource]) - function handleLoadMoreChunks(): void { if (!hasMoreSelectedChunks || isSelectedChunksLoadingMore) return void setSelectedChunkPageCount(selectedChunkPageCount + 1) diff --git a/src/components/workspace-shell.test.ts b/src/components/workspace-shell.test.ts index 651f4ba..308b60e 100644 --- a/src/components/workspace-shell.test.ts +++ b/src/components/workspace-shell.test.ts @@ -274,6 +274,91 @@ describe("WorkspaceShell", () => { expect(countFetches(fetch, "/api/sources/source_2/chunks")).toBe(0); }); + it("keeps a remote document open without refreshing sources after chunks load", async () => { + const remoteSourceId = "knowhere-doc:default:doc_remote"; + const encodedRemoteSourceId = encodeURIComponent(remoteSourceId); + const fetch = vi.fn(async (input) => { + const url = getRequestURL(input); + + if ( + url.pathname === `/api/sources/${encodedRemoteSourceId}/chunks` + ) { + return Response.json({ + chunks: [ + { + chunkId: "remote_page_1", + documentId: "doc_remote", + sectionPath: "Page 1", + type: "page", + content: "Remote page summary.", + sourceTitle: "remote.pdf", + pageNums: [1], + pageAssets: [ + { + pageNumber: 1, + assetUrl: "https://assets.example/page-1.png", + contentType: "image/png", + }, + ], + }, + ], + pagination: { + page: Number(url.searchParams.get("page") ?? "1"), + pageSize: 50, + total: 1, + totalPages: 1, + }, + }); + } + + if (url.pathname === "/api/sources") { + return Response.json({ + sources: [ + { + id: "source_localized", + kind: "workspace", + title: "remote.pdf", + status: "ready", + documentId: "doc_remote", + }, + ], + }); + } + + return Response.json({ message: "Unexpected request" }, { status: 404 }); + }); + vi.stubGlobal("fetch", fetch); + + render( + React.createElement(C, { + sources: [ + { + id: remoteSourceId, + kind: "remote", + title: "remote.pdf", + status: "ready", + documentId: "doc_remote", + excludedFromQuery: false, + }, + ], + }), + ); + + const desktopChunksPanel = within(screen.getByTestId("desktop-chunks-panel")); + await waitFor(() => { + expect(desktopChunksPanel.getByRole("img", { name: "Page 1" })) + .toBeTruthy(); + }); + + expect(countFetches(fetch, "/api/sources")).toBe(0); + expect( + countFetches( + fetch, + `/api/sources/${encodedRemoteSourceId}/chunks`, + ), + ).toBe(1); + }); + it("focuses guest citations on desktop using loaded demo chunks", async () => { const fetch = vi.fn(async (input) => { const url = getRequestURL(input); diff --git a/src/components/workspace-shell.tsx b/src/components/workspace-shell.tsx index bd45e2f..7dd43dc 100644 --- a/src/components/workspace-shell.tsx +++ b/src/components/workspace-shell.tsx @@ -114,7 +114,6 @@ function WorkspaceShellContent({ fetchChunks: workspaceClient.fetchChunks, initialPrefetchedChunksBySourceId: initialPrefetchedChunksBySourceId ?? undefined, - onRemoteSourceChunksLoaded: sourceWorkflow.handleSourcesRefresh, onSelectSource: handleCitationSourceSelected, selectedSourceId: sourceWorkflow.selectedSourceId, sources: sourceWorkflow.sources, From 15db165812fd1c8c38df77ee1871291e44d7bb90 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Wed, 8 Jul 2026 03:08:49 +0800 Subject: [PATCH 11/14] Fix inline parsed image viewing --- .../assets/[...assetPath]/route.test.ts | 51 ++++++ .../assets/[...assetPath]/route.ts | 1 + .../api/parsed-assets/inline/route.test.ts | 96 +++++++++++ src/app/api/parsed-assets/inline/route.ts | 159 ++++++++++++++++++ src/components/parsed-chunk-card.test.ts | 57 +++++++ src/components/parsed-chunk-card.tsx | 37 +++- 6 files changed, 396 insertions(+), 5 deletions(-) create mode 100644 src/app/api/demo-sources/[demoSourceId]/assets/[...assetPath]/route.test.ts create mode 100644 src/app/api/parsed-assets/inline/route.test.ts create mode 100644 src/app/api/parsed-assets/inline/route.ts diff --git a/src/app/api/demo-sources/[demoSourceId]/assets/[...assetPath]/route.test.ts b/src/app/api/demo-sources/[demoSourceId]/assets/[...assetPath]/route.test.ts new file mode 100644 index 0000000..aa1d09b --- /dev/null +++ b/src/app/api/demo-sources/[demoSourceId]/assets/[...assetPath]/route.test.ts @@ -0,0 +1,51 @@ +import { beforeEach, describe, expect, it, vi } from "vitest" + +const mocks = vi.hoisted(() => ({ + fetch: vi.fn(), +})) + +vi.mock("@/integrations/knowhere-demo", () => ({ + knowhereDemoApi: { + resolveApiURL: (pathname: string) => `https://demo.example${pathname}`, + }, +})) + +import { GET } from "./route" + +describe("GET /api/demo-sources/[demoSourceId]/assets/[...assetPath]", () => { + beforeEach(() => { + vi.clearAllMocks() + vi.stubGlobal("fetch", mocks.fetch) + }) + + it("serves demo image assets inline", async () => { + mocks.fetch.mockResolvedValue( + new Response("image bytes", { + status: 200, + headers: { + "content-type": "image/png", + }, + }), + ) + + const response = await GET( + new Request( + "http://localhost:3001/api/demo-sources/demo_1/assets/images/page-1.png", + ), + { + params: Promise.resolve({ + demoSourceId: "demo_1", + assetPath: ["images", "page-1.png"], + }), + }, + ) + + expect(response.status).toBe(200) + expect(response.headers.get("content-type")).toBe("image/png") + expect(response.headers.get("content-disposition")).toBe("inline") + expect(mocks.fetch).toHaveBeenCalledWith( + "https://demo.example/api/v1/demo/sources/demo_1/assets/images/page-1.png", + { cache: "no-store" }, + ) + }) +}) diff --git a/src/app/api/demo-sources/[demoSourceId]/assets/[...assetPath]/route.ts b/src/app/api/demo-sources/[demoSourceId]/assets/[...assetPath]/route.ts index 247f645..7dc89c3 100644 --- a/src/app/api/demo-sources/[demoSourceId]/assets/[...assetPath]/route.ts +++ b/src/app/api/demo-sources/[demoSourceId]/assets/[...assetPath]/route.ts @@ -33,6 +33,7 @@ export async function GET( status: 200, headers: { "content-type": response.headers.get("content-type") ?? "application/octet-stream", + "content-disposition": "inline", "cache-control": "public, max-age=3600", }, }) diff --git a/src/app/api/parsed-assets/inline/route.test.ts b/src/app/api/parsed-assets/inline/route.test.ts new file mode 100644 index 0000000..1bff9d3 --- /dev/null +++ b/src/app/api/parsed-assets/inline/route.test.ts @@ -0,0 +1,96 @@ +import { NextRequest } from "next/server" +import { beforeEach, describe, expect, it, vi } from "vitest" + +const mocks = vi.hoisted(() => ({ + getAuthenticated: vi.fn(), + getBlob: vi.fn(), +})) + +vi.mock("@/domains/workspace/request-context", () => ({ + notebookRequestContext: { + getAuthenticated: mocks.getAuthenticated, + }, +})) + +vi.mock("@vercel/blob", () => ({ + BlobNotFoundError: class BlobNotFoundError extends Error {}, + get: mocks.getBlob, +})) + +import { GET } from "./route" + +describe("GET /api/parsed-assets/inline", () => { + beforeEach(() => { + vi.clearAllMocks() + mocks.getAuthenticated.mockResolvedValue({ + workspace: { id: "workspace_1" }, + }) + }) + + it("serves Notebook parsed document image blobs inline", async () => { + const assetUrl = + "https://store.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_1/rev_1/page_citation_assets/page-1.png" + mocks.getBlob.mockResolvedValue({ + statusCode: 200, + stream: new Response("image bytes").body, + blob: { + contentType: "image/png", + }, + }) + + const response = await GET( + new NextRequest( + `http://localhost:3001/api/parsed-assets/inline?url=${encodeURIComponent( + assetUrl, + )}`, + ), + ) + + expect(response.status).toBe(200) + expect(response.headers.get("content-type")).toBe("image/png") + expect(response.headers.get("content-disposition")).toBe("inline") + await expect(response.text()).resolves.toBe("image bytes") + expect(mocks.getBlob).toHaveBeenCalledWith( + "workspaces/workspace_1/parsed-documents/doc_1/rev_1/page_citation_assets/page-1.png", + { access: "public" }, + ) + }) + + it("rejects blob URLs outside the authenticated workspace", async () => { + const assetUrl = + "https://store.public.blob.vercel-storage.com/workspaces/workspace_2/parsed-documents/doc_1/rev_1/page_citation_assets/page-1.png" + + const response = await GET( + new NextRequest( + `http://localhost:3001/api/parsed-assets/inline?url=${encodeURIComponent( + assetUrl, + )}`, + ), + ) + + expect(response.status).toBe(404) + expect(mocks.getBlob).not.toHaveBeenCalled() + }) + + it("does not serve non-image parsed assets through the inline image route", async () => { + const assetUrl = + "https://store.public.blob.vercel-storage.com/workspaces/workspace_1/sources/source_1/parsed-result/tables/table-1.html" + mocks.getBlob.mockResolvedValue({ + statusCode: 200, + stream: new Response("
").body, + blob: { + contentType: "text/html; charset=utf-8", + }, + }) + + const response = await GET( + new NextRequest( + `http://localhost:3001/api/parsed-assets/inline?url=${encodeURIComponent( + assetUrl, + )}`, + ), + ) + + expect(response.status).toBe(415) + }) +}) diff --git a/src/app/api/parsed-assets/inline/route.ts b/src/app/api/parsed-assets/inline/route.ts new file mode 100644 index 0000000..9e9beb3 --- /dev/null +++ b/src/app/api/parsed-assets/inline/route.ts @@ -0,0 +1,159 @@ +import { BlobNotFoundError, get } from "@vercel/blob" +import type { NextRequest } from "next/server" + +import { notebookRequestContext } from "@/domains/workspace/request-context" + +type ParsedAssetBlobPath = { + readonly pathname: string + readonly assetPath: string +} + +const notebookBlobHostSuffix = ".blob.vercel-storage.com" +const inlineImageContentTypes = new Set([ + "image/gif", + "image/jpeg", + "image/png", + "image/webp", +]) +const parsedDocumentsPathPattern = + /^\/?workspaces\/([^/]+)\/parsed-documents\/[^/]+\/[^/]+\/(.+)$/u +const parsedResultPathPattern = + /^\/?workspaces\/([^/]+)\/sources\/[^/]+\/parsed-result\/(.+)$/u + +export async function GET(request: NextRequest): Promise { + const { workspace } = await notebookRequestContext.getAuthenticated() + const parsedAssetBlobPath = getParsedAssetBlobPath( + request.nextUrl.searchParams.get("url"), + workspace.id, + ) + if (!parsedAssetBlobPath) { + return Response.json({ message: "Parsed asset not found." }, { status: 404 }) + } + + try { + const blob = await get(parsedAssetBlobPath.pathname, { access: "public" }) + if (!blob) { + return Response.json( + { message: "Parsed asset not found." }, + { status: 404 }, + ) + } + if (blob.statusCode !== 200) { + return new Response(null, { status: 304 }) + } + + const contentType = getInlineImageContentType( + blob.blob.contentType, + parsedAssetBlobPath.assetPath, + ) + if (!contentType) { + return Response.json( + { message: "Parsed asset is not an image." }, + { status: 415 }, + ) + } + + return new Response(blob.stream, { + status: 200, + headers: { + "content-type": contentType, + "content-disposition": "inline", + "cache-control": "public, max-age=3600", + }, + }) + } catch (error) { + if (error instanceof BlobNotFoundError) { + return Response.json( + { message: "Parsed asset not found." }, + { status: 404 }, + ) + } + throw error + } +} + +function getParsedAssetBlobPath( + assetUrl: string | null, + workspaceId: string, +): ParsedAssetBlobPath | null { + if (!assetUrl) return null + + const url = parseAbsoluteUrl(assetUrl) + if (!url) return null + if (url.protocol !== "https:") return null + if (!url.hostname.toLowerCase().endsWith(notebookBlobHostSuffix)) return null + + const pathname = decodePathname(url.pathname) + if (!pathname) return null + + const documentsMatch = parsedDocumentsPathPattern.exec(pathname) + if (documentsMatch) { + return getWorkspaceScopedAssetPath(documentsMatch, workspaceId, pathname) + } + + const parsedResultMatch = parsedResultPathPattern.exec(pathname) + if (parsedResultMatch) { + return getWorkspaceScopedAssetPath(parsedResultMatch, workspaceId, pathname) + } + + return null +} + +function getWorkspaceScopedAssetPath( + match: RegExpExecArray, + workspaceId: string, + pathname: string, +): ParsedAssetBlobPath | null { + const matchedWorkspaceId = match[1] + const assetPath = match[2] + if (matchedWorkspaceId !== workspaceId || !assetPath) return null + + return { + pathname: pathname.replace(/^\/+/u, ""), + assetPath, + } +} + +function getInlineImageContentType( + contentType: string | null, + assetPath: string, +): string | null { + const normalizedContentType = getBaseContentType(contentType) + if (normalizedContentType && inlineImageContentTypes.has(normalizedContentType)) { + return contentType + } + + return inferImageContentType(assetPath) +} + +function getBaseContentType(contentType: string | null): string | null { + const baseContentType = contentType?.split(";")[0]?.trim().toLowerCase() + return baseContentType && baseContentType.length > 0 ? baseContentType : null +} + +function inferImageContentType(assetPath: string): string | null { + const pathname = assetPath.toLowerCase().split("?")[0] ?? assetPath + if (pathname.endsWith(".png")) return "image/png" + if (pathname.endsWith(".jpg") || pathname.endsWith(".jpeg")) { + return "image/jpeg" + } + if (pathname.endsWith(".gif")) return "image/gif" + if (pathname.endsWith(".webp")) return "image/webp" + return null +} + +function parseAbsoluteUrl(value: string): URL | null { + try { + return new URL(value) + } catch { + return null + } +} + +function decodePathname(pathname: string): string | null { + try { + return decodeURIComponent(pathname) + } catch { + return null + } +} diff --git a/src/components/parsed-chunk-card.test.ts b/src/components/parsed-chunk-card.test.ts index ff4fbef..8bf5a42 100644 --- a/src/components/parsed-chunk-card.test.ts +++ b/src/components/parsed-chunk-card.test.ts @@ -123,6 +123,63 @@ describe("ParsedChunkCard", () => { expect(screen.getByTestId("page-asset-image-unavailable-4")).toBeTruthy(); }); + it("routes Notebook Blob page assets through the inline image endpoint", () => { + const assetUrl = + "https://store.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_1/rev_1/page_citation_assets/page-4.png"; + + render( + React.createElement(ParsedChunkCard, { + chunk: { + chunkId: "page_4", + type: "page", + content: "The summary should not be primary when an image exists.", + readableContent: "The summary should not be primary when an image exists.", + sourceTitle: "manual.pdf", + pageNums: [4], + pageAssets: [ + { + pageNumber: 4, + assetUrl, + contentType: "image/png", + }, + ], + }, + isFocused: false, + onReferenceClick: vi.fn(), + }), + ); + + const pageImage = screen.getByRole("img", { name: "Page 4" }); + expect(pageImage.getAttribute("src")).toBe( + `/api/parsed-assets/inline?url=${encodeURIComponent(assetUrl)}`, + ); + }); + + it("routes Notebook Blob image chunks through the inline image endpoint", () => { + const assetUrl = + "https://store.public.blob.vercel-storage.com/workspaces/workspace_1/sources/source_1/parsed-result/images/image-1.jpg"; + + render( + React.createElement(ParsedChunkCard, { + chunk: { + chunkId: "image_1", + type: "image", + content: "", + sourceTitle: "manual.pdf", + assetUrl, + summary: "A scanned diagram", + }, + isFocused: false, + onReferenceClick: vi.fn(), + }), + ); + + const image = screen.getByRole("img", { name: "A scanned diagram" }); + expect(image.getAttribute("src")).toBe( + `/api/parsed-assets/inline?url=${encodeURIComponent(assetUrl)}`, + ); + }); + it("routes resolved artifact reference clicks to the target chunk", async () => { const user = userEvent.setup(); const onReferenceClick = vi.fn(); diff --git a/src/components/parsed-chunk-card.tsx b/src/components/parsed-chunk-card.tsx index d794687..a8ddd6d 100644 --- a/src/components/parsed-chunk-card.tsx +++ b/src/components/parsed-chunk-card.tsx @@ -444,7 +444,8 @@ function PageCitationAssetImage({ readonly asset: NonNullable[number]; }): ReactNode { const [failedAssetUrl, setFailedAssetUrl] = useState(null); - const hasImageError = failedAssetUrl === asset.assetUrl; + const imageAssetUrl = getInlineImageAssetUrl(asset.assetUrl); + const hasImageError = failedAssetUrl === imageAssetUrl; const aspectRatio = asset.width && asset.height ? `${asset.width} / ${asset.height}` : undefined; @@ -463,11 +464,11 @@ function PageCitationAssetImage({ ) : ( // eslint-disable-next-line @next/next/no-img-element -- Page assets can be short-lived Knowhere URLs outside Next image optimization. {`Page setFailedAssetUrl(asset.assetUrl)} + onError={() => setFailedAssetUrl(imageAssetUrl)} /> )} @@ -509,6 +510,9 @@ function ImageChunkCard({ readonly sourceOriginalFile: SourceOriginalFileView | null; }): ReactNode { const imageAssetUrl = getImageChunkAssetUrl(chunk, sourceOriginalFile); + const inlineImageAssetUrl = imageAssetUrl + ? getInlineImageAssetUrl(imageAssetUrl) + : null; return ( - {imageAssetUrl ? ( + {inlineImageAssetUrl ? (
{/* eslint-disable-next-line @next/next/no-img-element -- Parsed artifact dimensions are not known before render. */} {chunk.summary @@ -559,6 +563,29 @@ function getImageChunkAssetUrl( return sourceOriginalFile.url; } +function getInlineImageAssetUrl(assetUrl: string): string { + if (!isNotebookBlobAssetUrl(assetUrl)) return assetUrl; + + return `/api/parsed-assets/inline?url=${encodeURIComponent(assetUrl)}`; +} + +function isNotebookBlobAssetUrl(assetUrl: string): boolean { + try { + const url = new URL(assetUrl); + if (!url.hostname.toLowerCase().endsWith(".blob.vercel-storage.com")) { + return false; + } + + const pathname = decodeURIComponent(url.pathname).toLowerCase(); + return ( + pathname.includes("/parsed-documents/") || + pathname.includes("/parsed-result/") + ); + } catch { + return false; + } +} + function hasPageCitationAssets(chunk: ParsedChunkView): boolean { return (chunk.pageAssets?.length ?? 0) > 0; } From 6d01583527e34ffe189f7c6629ec61a29d02ddcf Mon Sep 17 00:00:00 2001 From: suguanYang Date: Wed, 8 Jul 2026 03:31:45 +0800 Subject: [PATCH 12/14] Refactor chat to Knowhere tool runtime --- src/agent-harness/index.ts | 1 + src/agent-harness/knowhere-text.test.ts | 258 ++++++++++++ src/agent-harness/knowhere-text.ts | 300 ++++++++++++++ src/agent-harness/ledger.test.ts | 124 +++++- src/agent-harness/ledger.ts | 103 +++++ src/agent-harness/runtime.test.ts | 153 ++++--- src/agent-harness/runtime.ts | 504 ++++++++++++++++++------ src/agent-harness/types.ts | 52 ++- src/domains/chat/contracts.ts | 11 +- src/domains/chat/index.test.ts | 236 ++++++++++- src/domains/chat/index.ts | 9 + src/domains/chat/knowhere-tools.ts | 155 ++++++++ src/domains/chat/prompt.ts | 46 +-- src/domains/chat/route-answer.ts | 6 + src/domains/chat/service.test.ts | 2 + src/domains/chat/service.ts | 4 + 16 files changed, 1702 insertions(+), 262 deletions(-) create mode 100644 src/agent-harness/knowhere-text.test.ts create mode 100644 src/agent-harness/knowhere-text.ts create mode 100644 src/domains/chat/knowhere-tools.ts diff --git a/src/agent-harness/index.ts b/src/agent-harness/index.ts index bca503f..a1feff1 100644 --- a/src/agent-harness/index.ts +++ b/src/agent-harness/index.ts @@ -1,4 +1,5 @@ export * from "./ledger" +export * from "./knowhere-text" export * from "./runtime" export * from "./types" export * from "./validator" diff --git a/src/agent-harness/knowhere-text.test.ts b/src/agent-harness/knowhere-text.test.ts new file mode 100644 index 0000000..88159f5 --- /dev/null +++ b/src/agent-harness/knowhere-text.test.ts @@ -0,0 +1,258 @@ +import { describe, expect, it } from "vitest" +import type { + KnowledgeGrepResponse, + KnowledgeOutline, + KnowledgeReadResponse, + RetrievalQueryResponse, +} from "@ontos-ai/knowhere-sdk" + +import { createEvidenceLedger } from "./ledger" +import { knowhereToolText } from "./knowhere-text" + +describe("knowhereToolText", () => { + it("formats search results as tagged text with refs and no raw asset URLs", () => { + const ledger = createEvidenceLedger() + const response = makeSearchResponse() + const snapshot = ledger.addRetrievalResponse(response) + + const text = knowhereToolText.formatSearch({ + response, + retrievalCount: snapshot.retrievalCount, + chunks: snapshot.chunks, + assets: snapshot.assets, + }) + + expect(text).toContain('') + expect(text).toContain('ref="r1:result:1"') + expect(text).toContain('ref="asset:r1:result:1"') + expect(text).toContain("Call inspectImage") + expect(text).not.toContain("https://assets.example/page-1.png") + }) + + it("formats list and outline responses", () => { + const listText = knowhereToolText.formatListDocuments({ + documents: [ + { + documentId: "doc_1", + revisionKey: "job_1", + namespace: "notebook", + sourceFileName: "report.pdf", + title: "Report", + status: "ready", + }, + ], + }) + const outlineText = knowhereToolText.formatOutline(makeOutlineResponse()) + + expect(listText).toContain('') + expect(listText).toContain('documentId="doc_1"') + expect(outlineText).toContain( + '', + ) + expect(outlineText).toContain('sectionPath="Root / Revenue"') + }) + + it("formats full read chunk bodies without truncating large content", () => { + const ledger = createEvidenceLedger() + const largeContent = `BEGIN ${"full chunk body ".repeat(700)} END` + const response = makeReadResponse(largeContent) + const snapshot = ledger.addReadChunksResponse(response) + + const text = knowhereToolText.formatReadChunks({ + response, + chunks: snapshot.chunks, + assets: snapshot.assets, + }) + + expect(text).toContain('') + expect(text).toContain('ref="read1:chunk:1"') + expect(text).toContain(largeContent) + expect(text).not.toContain("...[truncated]") + expect(text).not.toContain("https://assets.example/page-1.png") + }) + + it("formats grep matches with continuation metadata", () => { + const ledger = createEvidenceLedger() + const response = makeGrepResponse() + const snapshot = ledger.addGrepChunksResponse(response) + + const text = knowhereToolText.formatGrepChunks({ + response, + chunks: snapshot.chunks, + assets: snapshot.assets, + }) + + expect(text).toContain('') + expect(text).toContain('truncated="true"') + expect(text).toContain('continuationCursor="cursor_2"') + expect(text).toContain('ref="grep1:match:1"') + expect(text).toContain("matched penalty snippet") + }) + + it("formats errors as tagged text", () => { + expect( + knowhereToolText.formatError({ + operation: "read_chunks", + message: "A documentId is required.", + }), + ).toBe( + [ + '', + "", + "A documentId is required.", + "", + "", + ].join("\n"), + ) + }) +}) + +function makeSearchResponse(): RetrievalQueryResponse { + return { + namespace: "notebook", + query: "page one", + routerUsed: "workflow_single_step", + answerText: null, + evidenceText: "Page one evidence.", + stopReason: "answer_done", + failureReason: null, + results: [ + { + content: "Page one summary.", + chunkType: "page", + score: 0.91, + assetUrl: "https://assets.example/page-1.png", + metadata: { + pageNums: [1], + pageAssets: [ + { + pageNum: 1, + artifactRef: "page_citation_assets/page-1.png", + assetUrl: "https://assets.example/page-1.png", + contentType: "image/png", + }, + ], + }, + source: { + documentId: "doc_1", + sourceFileName: "report.pdf", + sectionPath: "Root / Page 1", + }, + }, + ], + referencedChunks: [], + } +} + +function makeReadResponse(content: string): KnowledgeReadResponse { + return { + document: { + localDocumentId: "doc_1", + documentId: "doc_1", + jobId: "job_1", + namespace: "notebook", + sourceFileName: "report.pdf", + chunkCount: 1, + typeCounts: { text: 0, image: 0, table: 0, page: 1 }, + resultDirectoryPath: "parsed-storage:doc_1/job_1", + createdAt: new Date("2026-01-01T00:00:00Z"), + updatedAt: new Date("2026-01-01T00:00:00Z"), + }, + chunks: [ + { + position: 1, + chunkId: "chunk_page_1", + chunkType: "page", + contentSource: "content", + content, + readableContent: content, + sectionPath: "Root / Page 1", + sourceChunkPath: "pages/page-1.md", + filePath: "pages/page-1.png", + assetUrl: "https://assets.example/page-1.png", + pageNumbers: [1], + metadata: { + pageNums: [1], + pageAssets: [ + { + pageNum: 1, + artifactRef: "page_citation_assets/page-1.png", + assetUrl: "https://assets.example/page-1.png", + contentType: "image/png", + }, + ], + }, + }, + ], + page: 1, + pageSize: 1, + totalChunks: 1, + totalPages: 1, + } +} + +function makeOutlineResponse(): KnowledgeOutline { + return { + document: { + localDocumentId: "doc_1", + documentId: "doc_1", + jobId: "job_1", + namespace: "notebook", + sourceFileName: "report.pdf", + chunkCount: 1, + typeCounts: { text: 1, image: 0, table: 0, page: 0 }, + resultDirectoryPath: "parsed-storage:doc_1/job_1", + createdAt: new Date("2026-01-01T00:00:00Z"), + updatedAt: new Date("2026-01-01T00:00:00Z"), + }, + totalChunks: 1, + typeCounts: { text: 1, image: 0, table: 0, page: 0 }, + sections: [ + { + sectionPath: "Root / Revenue", + sectionTitle: "Revenue", + sectionLevel: 2, + summary: "Revenue summary.", + startChunk: 1, + endChunk: 1, + chunkCount: 1, + typeCounts: { text: 1, image: 0, table: 0, page: 0 }, + children: [], + }, + ], + sectionTree: [], + } +} + +function makeGrepResponse(): KnowledgeGrepResponse { + return { + document: { + localDocumentId: "doc_1", + documentId: "doc_1", + jobId: "job_1", + namespace: "notebook", + sourceFileName: "contract.pdf", + chunkCount: 4, + typeCounts: { text: 4, image: 0, table: 0, page: 0 }, + resultDirectoryPath: "parsed-storage:doc_1/job_1", + createdAt: new Date("2026-01-01T00:00:00Z"), + updatedAt: new Date("2026-01-01T00:00:00Z"), + }, + matches: [ + { + position: 3, + chunkId: "chunk_3", + chunkType: "text", + sectionPath: "Root / Penalties", + sourceChunkPath: "chunks/chunk-3.md", + filePath: "contract.pdf", + startOffset: 10, + endOffset: 17, + snippet: "matched penalty snippet", + }, + ], + scannedChunks: 4, + truncated: true, + continuationCursor: "cursor_2", + } +} diff --git a/src/agent-harness/knowhere-text.ts b/src/agent-harness/knowhere-text.ts new file mode 100644 index 0000000..69de846 --- /dev/null +++ b/src/agent-harness/knowhere-text.ts @@ -0,0 +1,300 @@ +import type { + KnowledgeGrepResponse, + KnowledgeOutline, + KnowledgeReadResponse, + RetrievalQueryResponse, +} from "@ontos-ai/knowhere-sdk" + +import type { + EvidenceAsset, + EvidenceChunk, + KnowhereListDocumentsResponse, +} from "./types" + +type EvidenceDelta = { + readonly chunks: readonly EvidenceChunk[] + readonly assets: readonly EvidenceAsset[] +} + +type SearchTextInput = EvidenceDelta & { + readonly response: RetrievalQueryResponse + readonly retrievalCount: number +} + +type ReadChunksTextInput = EvidenceDelta & { + readonly response: KnowledgeReadResponse +} + +type GrepChunksTextInput = EvidenceDelta & { + readonly response: KnowledgeGrepResponse +} + +type ErrorTextInput = { + readonly operation: KnowhereOperation + readonly message: string +} + +type KnowhereOperation = + | "search" + | "list_documents" + | "get_document_outline" + | "read_chunks" + | "grep_chunks" + +const assetInstruction = + "Notebook returned image/page asset refs. Call inspectImage with the asset refs when OCR, visual details, or verification are needed. Do not expose raw asset URLs." + +export const knowhereToolText = { + formatSearch(input: SearchTextInput): string { + return wrapKnowhereBlock("search", [ + formatTag("summary", { + retrievalCount: String(input.retrievalCount), + namespace: input.response.namespace, + query: input.response.query, + resultCount: String(input.response.results.length), + referencedChunkCount: String(input.response.referencedChunks.length), + stopReason: input.response.stopReason ?? undefined, + failureReason: input.response.failureReason ?? undefined, + }), + formatOptionalTextTag("evidence", input.response.evidenceText), + formatEvidenceChunks(input.chunks), + formatEvidenceAssets(input.assets), + formatAssetInstruction(input.assets), + ]) + }, + + formatListDocuments(response: KnowhereListDocumentsResponse): string { + return wrapKnowhereBlock("list_documents", [ + formatTag("summary", { documentCount: String(response.documents.length) }), + ...response.documents.map((document, index) => + formatSelfClosingTag("document", { + index: String(index + 1), + documentId: document.documentId, + localDocumentId: document.localDocumentId, + revisionKey: document.revisionKey, + namespace: document.namespace, + sourceFileName: document.sourceFileName, + title: document.title, + status: document.status, + chunkCount: + typeof document.chunkCount === "number" + ? String(document.chunkCount) + : undefined, + }), + ), + ]) + }, + + formatOutline(response: KnowledgeOutline): string { + return wrapKnowhereBlock("get_document_outline", [ + formatTag("document", { + documentId: response.document.documentId, + localDocumentId: response.document.localDocumentId, + revisionKey: response.document.jobId, + sourceFileName: response.document.sourceFileName, + totalChunks: String(response.totalChunks), + truncated: response.truncated === true ? "true" : undefined, + continuationCursor: response.continuationCursor, + }), + ...response.sections.map((section) => formatSection(section, 0)), + ]) + }, + + formatReadChunks(input: ReadChunksTextInput): string { + return wrapKnowhereBlock("read_chunks", [ + formatTag("document", { + documentId: input.response.document.documentId, + localDocumentId: input.response.document.localDocumentId, + revisionKey: input.response.document.jobId, + sourceFileName: input.response.document.sourceFileName, + page: + typeof input.response.page === "number" + ? String(input.response.page) + : undefined, + pageSize: + typeof input.response.pageSize === "number" + ? String(input.response.pageSize) + : undefined, + totalChunks: + typeof input.response.totalChunks === "number" + ? String(input.response.totalChunks) + : undefined, + totalPages: + typeof input.response.totalPages === "number" + ? String(input.response.totalPages) + : undefined, + nextChunk: + typeof input.response.nextChunk === "number" + ? String(input.response.nextChunk) + : undefined, + }), + formatEvidenceChunks(input.chunks), + formatEvidenceAssets(input.assets), + formatAssetInstruction(input.assets), + ]) + }, + + formatGrepChunks(input: GrepChunksTextInput): string { + return wrapKnowhereBlock("grep_chunks", [ + formatTag("document", { + documentId: input.response.document.documentId, + localDocumentId: input.response.document.localDocumentId, + revisionKey: input.response.document.jobId, + sourceFileName: input.response.document.sourceFileName, + matchCount: String(input.response.matches.length), + scannedChunks: String(input.response.scannedChunks), + truncated: input.response.truncated ? "true" : "false", + continuationCursor: input.response.continuationCursor, + }), + formatEvidenceChunks(input.chunks), + ]) + }, + + formatError(input: ErrorTextInput): string { + return [ + formatOpenTag("knowhere", { + operation: input.operation, + status: "error", + }), + formatTextTag("message", input.message), + "", + ].join("\n") + }, +} as const + +function wrapKnowhereBlock( + operation: KnowhereOperation, + parts: readonly string[], +): string { + return [ + formatOpenTag("knowhere", { operation, status: "ok" }), + ...parts.filter((part) => part.trim().length > 0), + "", + ].join("\n") +} + +function formatEvidenceChunks(chunks: readonly EvidenceChunk[]): string { + if (chunks.length === 0) return "" + + return [ + "", + ...chunks.map((chunk) => + [ + formatOpenTag("chunk", { + ref: chunk.ref, + kind: chunk.kind, + chunkId: chunk.chunkId, + chunkType: chunk.chunkType, + score: chunk.score === null ? undefined : String(chunk.score), + documentId: chunk.source.documentId ?? undefined, + sourceFileName: chunk.source.sourceFileName ?? undefined, + sectionPath: chunk.source.sectionPath ?? undefined, + sourceChunkPath: chunk.sourceChunkPath ?? undefined, + filePath: chunk.filePath ?? undefined, + assetRef: chunk.assetRef, + }), + formatTextTag("content", chunk.content), + "", + ].join("\n"), + ), + "", + ].join("\n") +} + +function formatEvidenceAssets(assets: readonly EvidenceAsset[]): string { + if (assets.length === 0) return "" + + return [ + "", + ...assets.map((asset) => + formatSelfClosingTag("asset", { + ref: asset.ref, + chunkRef: asset.chunkRef, + type: asset.type, + label: asset.label, + sourcePath: asset.sourcePath, + documentId: asset.source.documentId ?? undefined, + sectionPath: asset.source.sectionPath ?? undefined, + }), + ), + "", + ].join("\n") +} + +function formatAssetInstruction(assets: readonly EvidenceAsset[]): string { + if (!assets.some((asset) => asset.type === "image")) return "" + return formatTextTag("asset_instruction", assetInstruction) +} + +function formatSection( + section: KnowledgeOutline["sections"][number], + depth: number, +): string { + return [ + formatOpenTag("section", { + depth: String(depth), + sectionPath: section.sectionPath, + sectionTitle: section.sectionTitle, + sectionLevel: String(section.sectionLevel), + startChunk: + typeof section.startChunk === "number" + ? String(section.startChunk) + : undefined, + endChunk: + typeof section.endChunk === "number" + ? String(section.endChunk) + : undefined, + chunkCount: String(section.chunkCount), + }), + formatOptionalTextTag("summary", section.summary), + ...section.children.map((child) => formatSection(child, depth + 1)), + "", + ] + .filter((part) => part.trim().length > 0) + .join("\n") +} + +function formatOptionalTextTag( + tagName: string, + value: string | null | undefined, +): string { + const trimmedValue = value?.trim() + return trimmedValue ? formatTextTag(tagName, trimmedValue) : "" +} + +function formatTextTag(tagName: string, value: string): string { + return [`<${tagName}>`, value, ``].join("\n") +} + +function formatTag( + tagName: string, + attrs: Readonly>, +): string { + return `${formatOpenTag(tagName, attrs)}` +} + +function formatSelfClosingTag( + tagName: string, + attrs: Readonly>, +): string { + return `${formatOpenTag(tagName, attrs).slice(0, -1)} />` +} + +function formatOpenTag( + tagName: string, + attrs: Readonly>, +): string { + const serializedAttrs = Object.entries(attrs) + .filter((entry): entry is [string, string] => typeof entry[1] === "string") + .map(([key, value]) => `${key}="${escapeAttribute(value)}"`) + .join(" ") + return serializedAttrs ? `<${tagName} ${serializedAttrs}>` : `<${tagName}>` +} + +function escapeAttribute(value: string): string { + return value + .replaceAll("&", "&") + .replaceAll('"', """) + .replaceAll("<", "<") + .replaceAll(">", ">") +} diff --git a/src/agent-harness/ledger.test.ts b/src/agent-harness/ledger.test.ts index 22e0f6a..061a892 100644 --- a/src/agent-harness/ledger.test.ts +++ b/src/agent-harness/ledger.test.ts @@ -1,5 +1,9 @@ import { describe, expect, it } from "vitest" -import type { RetrievalQueryResponse } from "@ontos-ai/knowhere-sdk" +import type { + KnowledgeGrepResponse, + KnowledgeReadResponse, + RetrievalQueryResponse, +} from "@ontos-ai/knowhere-sdk" import { createEvidenceLedger } from "./ledger" @@ -63,6 +67,50 @@ describe("createEvidenceLedger", () => { }), ) }) + + it("adds read chunk refs and page image assets", () => { + const ledger = createEvidenceLedger() + + const snapshot = ledger.addReadChunksResponse(makeReadResponse()) + + expect(snapshot.chunks).toEqual([ + expect.objectContaining({ + ref: "read1:chunk:1", + kind: "read_chunk", + content: "Full page content.", + contentPreview: "Full page content.", + assetRef: "asset:read1:chunk:1", + }), + ]) + expect(snapshot.assets).toEqual([ + expect.objectContaining({ + ref: "asset:read1:chunk:1", + chunkRef: "read1:chunk:1", + type: "image", + sourcePath: "page_citation_assets/page-1.png", + }), + ]) + }) + + it("adds grep match refs", () => { + const ledger = createEvidenceLedger() + + const snapshot = ledger.addGrepChunksResponse(makeGrepResponse()) + + expect(snapshot.chunks).toEqual([ + expect.objectContaining({ + ref: "grep1:match:1", + kind: "grep_match", + chunkId: "chunk_3", + content: "matched penalty snippet", + source: expect.objectContaining({ + documentId: "doc_contract", + sourceFileName: "contract.pdf", + sectionPath: "Root / Penalties", + }), + }), + ]) + }) }) function makeRetrievalResponse(): RetrievalQueryResponse { @@ -135,3 +183,77 @@ function makePageAssetUrlRetrievalResponse(): RetrievalQueryResponse { ], } } + +function makeReadResponse(): KnowledgeReadResponse { + return { + document: { + localDocumentId: "doc_contract", + documentId: "doc_contract", + jobId: "job_contract", + namespace: "notebook", + sourceFileName: "contract.pdf", + chunkCount: 1, + typeCounts: { text: 0, image: 0, table: 0, page: 1 }, + resultDirectoryPath: "parsed-storage:doc_contract/job_contract", + createdAt: new Date("2026-01-01T00:00:00Z"), + updatedAt: new Date("2026-01-01T00:00:00Z"), + }, + chunks: [ + { + position: 1, + chunkId: "chunk_page_1", + chunkType: "page", + content: "Full page content.", + readableContent: "Full page content.", + sectionPath: "Root / Page 1", + sourceChunkPath: "pages/page-1.md", + filePath: "pages/page-1.png", + assetUrl: "https://assets.example/page-1.png", + pageNumbers: [1], + metadata: { + pageNums: [1], + pageAssets: [ + { + pageNum: 1, + artifactRef: "page_citation_assets/page-1.png", + assetUrl: "https://assets.example/page-1.png", + contentType: "image/png", + }, + ], + }, + }, + ], + } +} + +function makeGrepResponse(): KnowledgeGrepResponse { + return { + document: { + localDocumentId: "doc_contract", + documentId: "doc_contract", + jobId: "job_contract", + namespace: "notebook", + sourceFileName: "contract.pdf", + chunkCount: 4, + typeCounts: { text: 4, image: 0, table: 0, page: 0 }, + resultDirectoryPath: "parsed-storage:doc_contract/job_contract", + createdAt: new Date("2026-01-01T00:00:00Z"), + updatedAt: new Date("2026-01-01T00:00:00Z"), + }, + matches: [ + { + position: 3, + chunkId: "chunk_3", + chunkType: "text", + sectionPath: "Root / Penalties", + sourceChunkPath: "chunks/chunk-3.md", + filePath: "contract.pdf", + startOffset: 10, + endOffset: 17, + snippet: "matched penalty snippet", + }, + ], + scannedChunks: 4, + truncated: false, + } +} diff --git a/src/agent-harness/ledger.ts b/src/agent-harness/ledger.ts index f9d859d..d782ae3 100644 --- a/src/agent-harness/ledger.ts +++ b/src/agent-harness/ledger.ts @@ -1,4 +1,8 @@ import type { + KnowledgeGrepMatch, + KnowledgeGrepResponse, + KnowledgeReadChunk, + KnowledgeReadResponse, RetrievalQueryResponse, RetrievalResult, } from "@ontos-ai/knowhere-sdk" @@ -14,6 +18,8 @@ const imageExtensions = [".jpg", ".jpeg", ".png", ".gif", ".webp", ".svg"] as co type MutableLedger = { retrievalCount: number + readCount: number + grepCount: number chunks: EvidenceChunk[] assets: EvidenceAsset[] evidenceText: string[] @@ -41,6 +47,8 @@ export type EvidenceLedger = ReturnType export function createEvidenceLedger() { const ledger: MutableLedger = { retrievalCount: 0, + readCount: 0, + grepCount: 0, chunks: [], assets: [], evidenceText: [], @@ -104,6 +112,38 @@ export function createEvidenceLedger() { return snapshot(ledger) }, + addReadChunksResponse(response: KnowledgeReadResponse): EvidenceLedgerSnapshot { + ledger.readCount += 1 + const readIndex = ledger.readCount + + response.chunks.forEach((chunk, index) => { + addChunkFromReadChunk({ + ledger, + response, + chunk, + ref: `read${readIndex}:chunk:${index + 1}`, + }) + }) + + return snapshot(ledger) + }, + + addGrepChunksResponse(response: KnowledgeGrepResponse): EvidenceLedgerSnapshot { + ledger.grepCount += 1 + const grepIndex = ledger.grepCount + + response.matches.forEach((match, index) => { + addChunkFromGrepMatch({ + ledger, + response, + match, + ref: `grep${grepIndex}:match:${index + 1}`, + }) + }) + + return snapshot(ledger) + }, + read(ref: string, offset = 0, limit = 4_000) { const chunk = ledger.chunks.find((candidate) => candidate.ref === ref) if (!chunk) { @@ -179,6 +219,69 @@ function addChunkFromResult(input: { }) } +function addChunkFromReadChunk(input: { + readonly ledger: MutableLedger + readonly response: KnowledgeReadResponse + readonly chunk: KnowledgeReadChunk + readonly ref: string +}): void { + addChunk({ + ledger: input.ledger, + chunk: { + ref: input.ref, + kind: "read_chunk", + chunkId: input.chunk.chunkId, + content: input.chunk.content, + contentPreview: buildContentPreview(input.chunk.content), + chunkType: input.chunk.chunkType, + score: null, + sourceChunkPath: input.chunk.sourceChunkPath, + filePath: input.chunk.filePath, + metadata: input.chunk.metadata, + source: { + documentId: input.response.document.documentId, + sourceFileName: input.response.document.sourceFileName, + sectionPath: input.chunk.sectionPath, + }, + revisionKey: input.response.document.jobId, + ...(input.chunk.assetUrl ? { assetUrl: input.chunk.assetUrl } : {}), + }, + }) +} + +function addChunkFromGrepMatch(input: { + readonly ledger: MutableLedger + readonly response: KnowledgeGrepResponse + readonly match: KnowledgeGrepMatch + readonly ref: string +}): void { + addChunk({ + ledger: input.ledger, + chunk: { + ref: input.ref, + kind: "grep_match", + chunkId: input.match.chunkId, + content: input.match.snippet, + contentPreview: buildContentPreview(input.match.snippet), + chunkType: input.match.chunkType, + score: null, + sourceChunkPath: input.match.sourceChunkPath, + filePath: input.match.filePath, + metadata: { + position: input.match.position, + startOffset: input.match.startOffset, + endOffset: input.match.endOffset, + }, + source: { + documentId: input.response.document.documentId, + sourceFileName: input.response.document.sourceFileName, + sectionPath: input.match.sectionPath, + }, + revisionKey: input.response.document.jobId, + }, + }) +} + function addChunk(input: { readonly ledger: MutableLedger readonly chunk: Omit diff --git a/src/agent-harness/runtime.test.ts b/src/agent-harness/runtime.test.ts index 3e3c5fa..e78de33 100644 --- a/src/agent-harness/runtime.test.ts +++ b/src/agent-harness/runtime.test.ts @@ -14,8 +14,8 @@ import type { ContextPolicy, HarnessToolCallTrace, IntentFrame, + KnowhereToolRuntime, OutputManifest, - RetrievalCapability, } from "./types" describe("agent harness runtime", () => { @@ -29,7 +29,7 @@ describe("agent harness runtime", () => { }) it("passes only outer retrieval parameters to KNOWHERE after intent and context policy are declared", async () => { - const query = vi.fn().mockResolvedValue( + const query = vi.fn().mockResolvedValue( makeRetrievalResponse(), ) const state: { @@ -40,14 +40,12 @@ describe("agent harness runtime", () => { const tools = createHarnessTools({ state, ledger: createEvidenceLedger(), - retrieval: { query }, + knowhereTools: makeKnowhereTools(query), recentTurns: [], }) - expect(await executeTool(tools.retrieve, { query: "q4 chart" })).toEqual({ - ok: false, - message: "declareIntent must be called before retrieve.", - }) + expect(await executeTool(tools.knowhere_search, { query: "q4 chart" })) + .toContain('status="error"') await executeTool(tools.declareIntent, { task: "show_media", @@ -57,30 +55,27 @@ describe("agent harness runtime", () => { constraints: { desiredCount: 2, maxCount: 2 }, groundingPolicy: "must_use_sources", }) - expect(await executeTool(tools.retrieve, { query: "q4 chart" })).toEqual({ - ok: false, - message: "setContextPolicy must be called before retrieve.", - }) + expect(await executeTool(tools.knowhere_search, { query: "q4 chart" })) + .toContain("setContextPolicy must be called before knowhere_search.") await executeTool(tools.setContextPolicy, { carryHistory: "none", reason: "The current request is unrelated to previous turns.", activePriorTurnIds: [], }) - const result = await executeTool(tools.retrieve, { + const result = await executeTool(tools.knowhere_search, { query: "q4 chart", - modalities: ["image"], + targetContent: "image", topK: 2, purpose: "Find the two requested charts.", }) - expect(result).toMatchObject({ - ok: true, - retrievalCount: 1, - }) + expect(result).toContain('') + expect(result).toContain('ref="r1:result:1"') + expect(result).toContain('ref="asset:r1:result:1"') expect(query).toHaveBeenCalledWith({ query: "q4 chart", - modalities: ["image"], + targetContent: "image", topK: 2, purpose: "Find the two requested charts.", signalPaths: undefined, @@ -91,17 +86,17 @@ describe("agent harness runtime", () => { "LegalAction", ) expect(state.toolCalls?.map((call) => [call.tool, call.ok])).toEqual([ - ["retrieve", false], + ["knowhere_search", false], ["declareIntent", true], - ["retrieve", false], + ["knowhere_search", false], ["setContextPolicy", true], - ["retrieve", true], + ["knowhere_search", true], ]) }) - it("returns only newly retrieved evidence in each retrieve tool result", async () => { + it("returns only newly searched evidence in each knowhere_search tool result", async () => { const query = vi - .fn() + .fn() .mockResolvedValueOnce(makeRetrievalResponse()) .mockResolvedValueOnce({ ...makeRetrievalResponse(), @@ -143,22 +138,17 @@ describe("agent harness runtime", () => { const tools = createHarnessTools({ state, ledger, - retrieval: { query }, + knowhereTools: makeKnowhereTools(query), recentTurns: [], }) - const firstResult = await executeTool(tools.retrieve, { query: "first" }) - const secondResult = await executeTool(tools.retrieve, { query: "second" }) + const firstResult = await executeTool(tools.knowhere_search, { query: "first" }) + const secondResult = await executeTool(tools.knowhere_search, { query: "second" }) - expect(firstResult).toMatchObject({ - retrievalCount: 1, - chunks: [{ ref: "r1:result:1" }], - }) - expect(secondResult).toMatchObject({ - retrievalCount: 2, - evidenceText: "Second evidence", - chunks: [{ ref: "r2:result:1" }], - }) + expect(firstResult).toContain('ref="r1:result:1"') + expect(secondResult).toContain('retrievalCount="2"') + expect(secondResult).toContain("Second evidence") + expect(secondResult).toContain('ref="r2:result:1"') expect(JSON.stringify(secondResult)).not.toContain("r1:result:1") expect(ledger.snapshot().chunks.map((chunk) => chunk.ref)).toEqual([ "r1:result:1", @@ -171,7 +161,7 @@ describe("agent harness runtime", () => { const tools = createHarnessTools({ state: {}, ledger: createEvidenceLedger(), - retrieval: { query: vi.fn() }, + knowhereTools: makeKnowhereTools(), inspectImages, recentTurns: [], }) @@ -183,12 +173,13 @@ describe("agent harness runtime", () => { expect(result).toEqual({ ok: false, - message: "retrieve must be called before inspectImage.", + message: + "Knowhere evidence tools must return image assets before inspectImage.", inspected: [], skipped: [ { ref: "asset:r1:result:1", - reason: "No retrieval evidence is available yet.", + reason: "No Knowhere evidence is available yet.", }, ], }) @@ -202,7 +193,7 @@ describe("agent harness runtime", () => { const tools = createHarnessTools({ state: {}, ledger, - retrieval: { query: vi.fn() }, + knowhereTools: makeKnowhereTools(), inspectImages, recentTurns: [], }) @@ -223,7 +214,7 @@ describe("agent harness runtime", () => { }, { ref: "missing", - reason: "Ref was not returned by retrieve as an asset.", + reason: "Ref was not returned by Knowhere as an asset.", }, ], }) @@ -236,7 +227,7 @@ describe("agent harness runtime", () => { const tools = createHarnessTools({ state: {}, ledger, - retrieval: { query: vi.fn() }, + knowhereTools: makeKnowhereTools(), inspectImages: vi.fn(), recentTurns: [], }) @@ -268,7 +259,7 @@ describe("agent harness runtime", () => { const tools = createHarnessTools({ state: {}, ledger, - retrieval: { query: vi.fn() }, + knowhereTools: makeKnowhereTools(), inspectImages, recentTurns: [], }) @@ -324,7 +315,7 @@ describe("agent harness runtime", () => { const tools = createHarnessTools({ state: {}, ledger, - retrieval: { query: vi.fn() }, + knowhereTools: makeKnowhereTools(), inspectImages, recentTurns: [], }) @@ -368,7 +359,7 @@ describe("agent harness runtime", () => { const tools = createHarnessTools({ state, ledger: createEvidenceLedger(), - retrieval: { query: vi.fn() }, + knowhereTools: makeKnowhereTools(), recentTurns: [], }) @@ -427,7 +418,7 @@ describe("agent harness runtime", () => { const tools = createHarnessTools({ state, ledger: createEvidenceLedger(), - retrieval: { query: vi.fn() }, + knowhereTools: makeKnowhereTools(), recentTurns: [ { id: "turn_1", @@ -459,7 +450,7 @@ describe("agent harness runtime", () => { const tools = createHarnessTools({ state, ledger: createEvidenceLedger(), - retrieval: { query: vi.fn() }, + knowhereTools: makeKnowhereTools(), recentTurns: [ { id: "turn_1", @@ -521,7 +512,7 @@ describe("agent harness runtime", () => { { type: "tool-call", toolCallId: "call_1", - toolName: "retrieve", + toolName: "knowhere_search", input: { query: "q4" }, providerOptions: { google: { @@ -537,12 +528,10 @@ describe("agent harness runtime", () => { { type: "tool-result", toolCallId: "call_1", - toolName: "retrieve", + toolName: "knowhere_search", output: { - type: "json", - value: { - ok: true, - }, + type: "text", + value: '', }, providerOptions: { google: { @@ -574,12 +563,10 @@ describe("agent harness runtime", () => { { type: "tool-result", toolCallId: "call_1", - toolName: "retrieve", + toolName: "knowhere_search", output: { - type: "json", - value: { - ok: true, - }, + type: "text", + value: '', }, }, ], @@ -619,13 +606,11 @@ describe("agent harness runtime", () => { { type: "tool-result", toolCallId: "call_1", - toolName: "retrieve", + toolName: "knowhere_search", output: { - type: "json", - value: { - ok: true, - assets: [{ ref: "asset:r1:referenced:1", type: "image" }], - }, + type: "text", + value: + '', }, }, ], @@ -645,13 +630,11 @@ describe("agent harness runtime", () => { { type: "tool-result", toolCallId: "call_1", - toolName: "retrieve", + toolName: "knowhere_search", output: { - type: "json", - value: { - ok: true, - assets: [{ ref: "asset:r1:referenced:1", type: "image" }], - }, + type: "text", + value: + '', }, }, ], @@ -673,13 +656,11 @@ describe("agent harness runtime", () => { { type: "tool-result", toolCallId: "call_1", - toolName: "retrieve", + toolName: "knowhere_search", output: { - type: "json", - value: { - ok: true, - chunks: [{ ref: "r1:result:1" }], - }, + type: "text", + value: + '', }, providerOptions: { google: { @@ -704,13 +685,11 @@ describe("agent harness runtime", () => { { type: "tool-result", toolCallId: "call_1", - toolName: "retrieve", + toolName: "knowhere_search", output: { - type: "json", - value: { - ok: true, - chunks: [{ ref: "r1:result:1" }], - }, + type: "text", + value: + '', }, }, ], @@ -729,6 +708,18 @@ function executeTool(tool: unknown, input: unknown): Promise { return (tool as { execute: (input: unknown) => Promise }).execute(input) } +function makeKnowhereTools( + search: KnowhereToolRuntime["search"] = vi.fn(), +): KnowhereToolRuntime { + return { + search, + listDocuments: vi.fn().mockResolvedValue({ documents: [] }), + getDocumentOutline: vi.fn().mockRejectedValue(new Error("Not configured.")), + readChunks: vi.fn().mockRejectedValue(new Error("Not configured.")), + grepChunks: vi.fn().mockRejectedValue(new Error("Not configured.")), + } +} + function makeTurnInput(overrides: Partial = {}): AgentTurnInput { return { surface: "notebook_chat", diff --git a/src/agent-harness/runtime.ts b/src/agent-harness/runtime.ts index d38028e..e2b2783 100644 --- a/src/agent-harness/runtime.ts +++ b/src/agent-harness/runtime.ts @@ -9,6 +9,7 @@ import { import { z } from "zod" import { createEvidenceLedger } from "./ledger" +import { knowhereToolText } from "./knowhere-text" import type { AgentTurn, AgentTurnInput, @@ -20,9 +21,10 @@ import type { ImageInspectionResponse, InspectImages, IntentFrame, + KnowhereSearchRequest, + KnowhereSearchTargetContent, + KnowhereToolRuntime, OutputManifest, - RetrievalCapability, - TargetModality, } from "./types" import { validateOutputManifest } from "./validator" @@ -39,7 +41,7 @@ export type AgentHarnessModel = ToolLoopAgentSettings["model"] export type RunAgentHarnessInput = { readonly model: AgentHarnessModel readonly turn: AgentTurnInput - readonly retrieval: RetrievalCapability + readonly knowhereTools: KnowhereToolRuntime readonly inspectImages?: InspectImages readonly maxSteps?: number /** @@ -71,6 +73,53 @@ type HarnessStepPreparation = { } const targetModalitySchema = z.enum(["text", "image", "table"]) +const knowhereSearchTargetContentSchema = z.enum([ + "all", + "text", + "image", + "table", + "text_image", + "text_table", +]) +const knowledgeChunkTypeSchema = z.enum(["text", "image", "table", "page"]) + +const knowhereDocumentReferenceSchema = z.object({ + localDocumentId: z.string().min(1).optional(), + documentId: z.string().min(1).optional(), + jobId: z.string().min(1).optional(), + revisionKey: z.string().min(1).optional(), +}) + +const knowhereSearchSchema = z.object({ + query: z.string().min(1), + targetContent: knowhereSearchTargetContentSchema.default("all"), + purpose: z.string().optional(), + topK: z.number().int().min(1).max(12).optional(), + signalPaths: z.array(z.string().min(1)).max(8).optional(), + filterMode: z.enum(["keep", "delete"]).optional(), + threshold: z.number().min(0).max(1).optional(), +}) + +const knowhereReadChunksSchema = knowhereDocumentReferenceSchema.extend({ + page: z.number().int().min(1).optional(), + pageSize: z.number().int().min(1).max(50).optional(), + sectionPath: z.string().min(1).optional(), + startChunk: z.number().int().min(0).optional(), + endChunk: z.number().int().min(0).optional(), + chunkId: z.string().min(1).optional(), + chunkType: knowledgeChunkTypeSchema.optional(), +}) + +const knowhereGrepChunksSchema = knowhereDocumentReferenceSchema.extend({ + pattern: z.string().min(1), + continuationCursor: z.string().min(1).optional(), + isRegex: z.boolean().optional(), + isCaseSensitive: z.boolean().optional(), + maxResults: z.number().int().min(1).max(50).optional(), + chunkType: knowledgeChunkTypeSchema.optional(), + sectionPathPrefix: z.string().min(1).optional(), + contextChars: z.number().int().min(0).max(2_000).optional(), +}) const intentFrameSchema = z.object({ task: z.enum([ @@ -166,7 +215,7 @@ export async function runAgentHarness( const tools = createHarnessTools({ state, ledger, - retrieval: input.retrieval, + knowhereTools: input.knowhereTools, inspectImages: input.inspectImages, recentTurns: input.turn.recentTurns, }) @@ -364,7 +413,7 @@ function buildRevisionFeedback(errors: readonly string[]): string { function buildForcedFinalizationFeedback(): string { return [ "The retrieval step budget has been reached.", - "Do not search again or call any evidence-reading tools.", + "Do not search again or call any Knowhere evidence-reading tools.", "Use only the evidence and tool results already available in this turn.", "Call finalize now with the best supported answer.", "If the existing evidence is insufficient, explain the gap in unresolved", @@ -396,7 +445,7 @@ function hasUninspectedImageAssets(input: { export function createHarnessTools(input: { readonly state: HarnessToolState readonly ledger: ReturnType - readonly retrieval: RetrievalCapability + readonly knowhereTools: KnowhereToolRuntime readonly inspectImages?: InspectImages readonly recentTurns: readonly AgentTurn[] }) { @@ -433,82 +482,109 @@ export function createHarnessTools(input: { }), }), - retrieve: tool({ + knowhere_search: tool({ description: - "Ask KNOWHERE for evidence context. KNOWHERE handles internal navigation; this tool only submits a concise query and records returned evidence.", - inputSchema: z.object({ - query: z.string().min(1), - modalities: z.array(targetModalitySchema).default(["text"]), - purpose: z.string().optional(), - topK: z.number().int().min(1).max(12).optional(), - signalPaths: z.array(z.string().min(1)).max(8).optional(), - filterMode: z.enum(["keep", "delete"]).optional(), - threshold: z.number().min(0).max(1).optional(), - }), + "Search Knowhere for relevant Notebook evidence. Returns tagged text with evidence refs such as r1:result:1 and asset refs such as asset:r1:result:1.", + inputSchema: knowhereSearchSchema, execute: async (request) => traceToolCall(input.state, { - toolName: "retrieve", - inputSummary: summarizeRetrievalRequest(request), - execute: async () => { - if (!input.state.intent) { - return { - ok: false, - message: "declareIntent must be called before retrieve.", - } - } - if (!input.state.contextPolicy) { - return { - ok: false, - message: "setContextPolicy must be called before retrieve.", - } - } + toolName: "knowhere_search", + inputSummary: summarizeKnowhereSearchRequest(request), + execute: async () => + executeKnowhereSearch({ + state: input.state, + ledger: input.ledger, + knowhereTools: input.knowhereTools, + request, + }), + summarizeOutput: summarizeKnowhereTextOutput, + }), + }), - const beforeSnapshot = input.ledger.snapshot() - const response = await input.retrieval.query({ - query: request.query, - modalities: request.modalities as TargetModality[], - purpose: request.purpose, - topK: request.topK, - signalPaths: request.signalPaths, - filterMode: request.filterMode, - threshold: request.threshold, - }) - const snapshot = input.ledger.addRetrievalResponse(response) - const currentChunks = snapshot.chunks.slice( - beforeSnapshot.chunks.length, - ) - const currentAssets = snapshot.assets.slice( - beforeSnapshot.assets.length, - ) - return { - ok: true, - retrievalCount: snapshot.retrievalCount, - evidenceText: response.evidenceText ?? "", - stopReason: response.stopReason ?? null, - failureReason: response.failureReason ?? null, - chunks: currentChunks.map((chunk) => ({ - ref: chunk.ref, - kind: chunk.kind, - type: chunk.chunkType, - preview: chunk.contentPreview, - source: chunk.source, - assetRef: chunk.assetRef, - })), - assets: currentAssets.map((asset) => ({ - ref: asset.ref, - type: asset.type, - label: asset.label, - source: asset.source, - })), - } - }, - summarizeOutput: summarizeRetrieveOutput, + knowhere_list_documents: tool({ + description: + "List ready visible Notebook/Knowhere documents available for this chat turn. Use this to discover documentId and revisionKey before outline/read/grep.", + inputSchema: z.object({}), + execute: async () => + traceToolCall(input.state, { + toolName: "knowhere_list_documents", + inputSummary: {}, + execute: async () => + executeKnowhereTextTool({ + operation: "list_documents", + state: input.state, + execute: async () => + knowhereToolText.formatListDocuments( + await input.knowhereTools.listDocuments(), + ), + }), + summarizeOutput: summarizeKnowhereTextOutput, + }), + }), + + knowhere_get_document_outline: tool({ + description: + "Read a document outline from Knowhere parsed storage. Use documentId/revisionKey from knowhere_list_documents or search refs.", + inputSchema: knowhereDocumentReferenceSchema, + execute: async (request) => + traceToolCall(input.state, { + toolName: "knowhere_get_document_outline", + inputSummary: summarizeDocumentReference(request), + execute: async () => + executeKnowhereTextTool({ + operation: "get_document_outline", + state: input.state, + validate: () => validateDocumentReference(request), + execute: async () => + knowhereToolText.formatOutline( + await input.knowhereTools.getDocumentOutline(request), + ), + }), + summarizeOutput: summarizeKnowhereTextOutput, + }), + }), + + knowhere_read_chunks: tool({ + description: + "Read complete chunk bodies from Knowhere parsed storage. This tool never slices individual chunk content; control read size with page/pageSize, sectionPath, startChunk/endChunk, chunkId, and chunkType.", + inputSchema: knowhereReadChunksSchema, + execute: async (request) => + traceToolCall(input.state, { + toolName: "knowhere_read_chunks", + inputSummary: summarizeReadChunksRequest(request), + execute: async () => + executeKnowhereReadChunks({ + state: input.state, + ledger: input.ledger, + knowhereTools: input.knowhereTools, + request, + }), + summarizeOutput: summarizeKnowhereTextOutput, + }), + }), + + knowhere_grep_chunks: tool({ + description: + "Search chunk text with a literal or regex pattern. Returns bounded match snippets as grep refs such as grep1:match:1 and may include truncated=true with a continuationCursor.", + inputSchema: knowhereGrepChunksSchema, + execute: async (request) => + traceToolCall(input.state, { + toolName: "knowhere_grep_chunks", + inputSummary: summarizeGrepChunksRequest(request), + execute: async () => + executeKnowhereGrepChunks({ + state: input.state, + ledger: input.ledger, + knowhereTools: input.knowhereTools, + request, + }), + summarizeOutput: summarizeKnowhereTextOutput, }), }), inspectImage: tool({ description: - "Inspect retrieved image asset refs visually for OCR, visual details, comparisons, or verification. Use only after retrieve has returned image assets.", + "Inspect Knowhere image asset refs visually for OCR, visual details, comparisons, or verification. Use only after a Knowhere tool has returned image assets.", inputSchema: z.object({ refs: z.array(z.string().min(1)).min(1).max(imageInspectionRefLimit), question: z.string().min(1), @@ -529,32 +605,6 @@ export function createHarnessTools(input: { }), }), - readEvidence: tool({ - description: - "Read more text from an evidence chunk already returned by KNOWHERE.", - inputSchema: z.object({ - ref: z.string().min(1), - offset: z.number().int().min(0).optional(), - limit: z.number().int().min(1).max(8_000).optional(), - }), - execute: async (request) => - traceToolCall(input.state, { - toolName: "readEvidence", - inputSummary: { - ref: request.ref, - offset: request.offset ?? 0, - limit: request.limit ?? 4_000, - }, - execute: async () => - input.ledger.read( - request.ref, - request.offset ?? 0, - request.limit ?? 4_000, - ), - summarizeOutput: summarizeReadEvidenceOutput, - }), - }), - readPriorTurn: tool({ description: "Read the full text and citation labels of a specific prior turn by id " + @@ -689,14 +739,15 @@ async function inspectRetrievedImages(input: { } const snapshot = input.ledger.snapshot() - if (snapshot.retrievalCount === 0) { + if (snapshot.chunks.length === 0 && snapshot.assets.length === 0) { return { ok: false, - message: "retrieve must be called before inspectImage.", + message: + "Knowhere evidence tools must return image assets before inspectImage.", inspected: [], skipped: refs.map((ref) => ({ ref, - reason: "No retrieval evidence is available yet.", + reason: "No Knowhere evidence is available yet.", })), } } @@ -726,7 +777,7 @@ async function inspectRetrievedImages(input: { if (!asset) { skipped.push({ ref, - reason: "Ref was not returned by retrieve as an asset.", + reason: "Ref was not returned by Knowhere as an asset.", }) continue } @@ -804,6 +855,161 @@ async function inspectRetrievedImages(input: { } } +type KnowhereToolOperation = + | "search" + | "list_documents" + | "get_document_outline" + | "read_chunks" + | "grep_chunks" + +type KnowhereSearchToolRequest = z.infer +type KnowhereDocumentReferenceRequest = z.infer< + typeof knowhereDocumentReferenceSchema +> +type KnowhereReadChunksToolRequest = z.infer +type KnowhereGrepChunksToolRequest = z.infer +type DocumentReferenceSummary = { + readonly documentId?: string + readonly localDocumentId?: string + readonly hasJobId: boolean + readonly hasRevisionKey: boolean +} + +async function executeKnowhereSearch(input: { + readonly state: HarnessToolState + readonly ledger: ReturnType + readonly knowhereTools: KnowhereToolRuntime + readonly request: KnowhereSearchToolRequest +}): Promise { + return executeKnowhereTextTool({ + operation: "search", + state: input.state, + execute: async () => { + const beforeSnapshot = input.ledger.snapshot() + const response = await input.knowhereTools.search({ + query: input.request.query, + targetContent: input.request.targetContent, + purpose: input.request.purpose, + topK: input.request.topK, + signalPaths: input.request.signalPaths, + filterMode: input.request.filterMode, + threshold: input.request.threshold, + }) + const snapshot = input.ledger.addRetrievalResponse(response) + return knowhereToolText.formatSearch({ + response, + retrievalCount: snapshot.retrievalCount, + chunks: snapshot.chunks.slice(beforeSnapshot.chunks.length), + assets: snapshot.assets.slice(beforeSnapshot.assets.length), + }) + }, + }) +} + +async function executeKnowhereReadChunks(input: { + readonly state: HarnessToolState + readonly ledger: ReturnType + readonly knowhereTools: KnowhereToolRuntime + readonly request: KnowhereReadChunksToolRequest +}): Promise { + return executeKnowhereTextTool({ + operation: "read_chunks", + state: input.state, + validate: () => validateDocumentReference(input.request), + execute: async () => { + const beforeSnapshot = input.ledger.snapshot() + const response = await input.knowhereTools.readChunks(input.request) + const snapshot = input.ledger.addReadChunksResponse(response) + return knowhereToolText.formatReadChunks({ + response, + chunks: snapshot.chunks.slice(beforeSnapshot.chunks.length), + assets: snapshot.assets.slice(beforeSnapshot.assets.length), + }) + }, + }) +} + +async function executeKnowhereGrepChunks(input: { + readonly state: HarnessToolState + readonly ledger: ReturnType + readonly knowhereTools: KnowhereToolRuntime + readonly request: KnowhereGrepChunksToolRequest +}): Promise { + return executeKnowhereTextTool({ + operation: "grep_chunks", + state: input.state, + validate: () => validateDocumentReference(input.request), + execute: async () => { + const beforeSnapshot = input.ledger.snapshot() + const response = await input.knowhereTools.grepChunks(input.request) + const snapshot = input.ledger.addGrepChunksResponse(response) + return knowhereToolText.formatGrepChunks({ + response, + chunks: snapshot.chunks.slice(beforeSnapshot.chunks.length), + assets: snapshot.assets.slice(beforeSnapshot.assets.length), + }) + }, + }) +} + +async function executeKnowhereTextTool(input: { + readonly operation: KnowhereToolOperation + readonly state: HarnessToolState + readonly validate?: () => string | null + readonly execute: () => Promise +}): Promise { + const workflowError = validateKnowhereWorkflow(input.state, input.operation) + if (workflowError) { + return knowhereToolText.formatError({ + operation: input.operation, + message: workflowError, + }) + } + + const validationError = input.validate?.() + if (validationError) { + return knowhereToolText.formatError({ + operation: input.operation, + message: validationError, + }) + } + + try { + return await input.execute() + } catch (error) { + return knowhereToolText.formatError({ + operation: input.operation, + message: formatUnknownError(error), + }) + } +} + +function validateKnowhereWorkflow( + state: HarnessToolState, + operation: KnowhereToolOperation, +): string | null { + const toolName = `knowhere_${operation}` + if (!state.intent) return `declareIntent must be called before ${toolName}.` + if (!state.contextPolicy) { + return `setContextPolicy must be called before ${toolName}.` + } + return null +} + +function validateDocumentReference( + request: KnowhereDocumentReferenceRequest, +): string | null { + if ( + request.documentId || + request.localDocumentId || + request.jobId + ) { + return null + } + + return "A documentId, localDocumentId, or jobId is required." +} + function getUniqueTrimmedRefs(refs: readonly string[]): string[] { const normalizedRefs: string[] = [] for (const ref of refs) { @@ -860,6 +1066,7 @@ function recordToolCall( } function getToolTraceOk(output: unknown): boolean { + if (typeof output === "string") return !output.includes('status="error"') if (!isRecord(output)) return true if (typeof output.ok === "boolean") return output.ok if (typeof output.found === "boolean") return output.found @@ -884,9 +1091,9 @@ function summarizeContextPolicy(policy: ContextPolicy): unknown { } } -function summarizeRetrievalRequest(request: { +function summarizeKnowhereSearchRequest(request: { readonly query: string - readonly modalities?: readonly TargetModality[] + readonly targetContent?: KnowhereSearchTargetContent readonly purpose?: string readonly topK?: number readonly signalPaths?: readonly string[] @@ -895,7 +1102,7 @@ function summarizeRetrievalRequest(request: { }): unknown { return { query: request.query, - modalities: request.modalities ?? ["text"], + targetContent: request.targetContent ?? "all", purpose: request.purpose, topK: request.topK, signalPathCount: request.signalPaths?.length ?? 0, @@ -904,27 +1111,67 @@ function summarizeRetrievalRequest(request: { } } -function summarizeRetrieveOutput(output: unknown): unknown { - if (!isRecord(output)) return output +function summarizeDocumentReference( + request: KnowhereDocumentReferenceRequest, +): DocumentReferenceSummary { return { - ok: output.ok, - retrievalCount: output.retrievalCount, - stopReason: output.stopReason, - failureReason: output.failureReason, - chunkCount: Array.isArray(output.chunks) ? output.chunks.length : 0, - assetCount: Array.isArray(output.assets) ? output.assets.length : 0, + documentId: request.documentId, + localDocumentId: request.localDocumentId, + hasJobId: typeof request.jobId === "string", + hasRevisionKey: typeof request.revisionKey === "string", } } -function summarizeReadEvidenceOutput(output: unknown): unknown { - if (!isRecord(output)) return output +function summarizeReadChunksRequest( + request: KnowhereReadChunksToolRequest, +): unknown { return { - found: output.found, - ref: output.ref, - contentLength: output.contentLength, - offset: output.offset, - limit: output.limit, - hasMoreContent: output.hasMoreContent, + ...summarizeDocumentReference(request), + page: request.page, + pageSize: request.pageSize, + sectionPath: request.sectionPath, + startChunk: request.startChunk, + endChunk: request.endChunk, + chunkId: request.chunkId, + chunkType: request.chunkType, + } +} + +function summarizeGrepChunksRequest( + request: KnowhereGrepChunksToolRequest, +): unknown { + return { + ...summarizeDocumentReference(request), + patternLength: request.pattern.trim().length, + continuationCursor: request.continuationCursor, + isRegex: request.isRegex, + isCaseSensitive: request.isCaseSensitive, + maxResults: request.maxResults, + chunkType: request.chunkType, + sectionPathPrefix: request.sectionPathPrefix, + contextChars: request.contextChars, + } +} + +function summarizeKnowhereTextOutput(output: unknown): unknown { + if (typeof output !== "string") return output + return { + ok: !output.includes('status="error"'), + textLength: output.length, + chunkCount: countOccurrences(output, " { return typeof value === "object" && value !== null } +function formatUnknownError(error: unknown): string { + if (error instanceof Error) return error.message + return String(error) +} + export function buildHarnessSystemPrompt(turn: AgentTurnInput): string { return [ "You are the outer Knowhere Agent Harness.", @@ -1007,10 +1259,10 @@ export function buildHarnessSystemPrompt(turn: AgentTurnInput): string { "1. Call declareIntent first. Capture constraints like a requested image/table count in constraints.desiredCount.", "2. Call setContextPolicy next, deciding how prior turns should influence this turn.", "3. When the policy needs prior-turn detail (references or corrections), call readPriorTurn for the relevant ids.", - "4. Call retrieve only when evidence is needed. The query must be concise and self-contained.", - "5. For pixel-level details, OCR, visual comparison, image verification, or when the likely answer is only visible on a retrieved page/image asset, call inspectImage only after retrieve returned image asset refs.", + "4. Call knowhere_search when relevance search is needed. Use knowhere_list_documents, knowhere_get_document_outline, knowhere_read_chunks, and knowhere_grep_chunks for focused document reads.", + "5. For pixel-level details, OCR, visual comparison, image verification, or when the likely answer is only visible on a returned page/image asset, call inspectImage only after a Knowhere tool returned image asset refs.", `6. inspectImage accepts at most ${imageInspectionRefLimit} image asset refs per call and per turn.`, - "7. Use readEvidence only for chunk refs already in the evidence ledger.", + "7. knowhere_read_chunks returns complete chunk bodies; control size with page/pageSize, sectionPath, startChunk/endChunk, chunkId, and chunkType.", "8. Call finalize with text, citations, artifacts, and unresolved issues. finalize requires declareIntent and setContextPolicy first.", "", "Context rules:", @@ -1022,7 +1274,7 @@ export function buildHarnessSystemPrompt(turn: AgentTurnInput): string { "- Final output is the OutputManifest passed to finalize, not freeform tool JSON or trailing text.", "- artifacts with display=true are the exact images/tables shown. Never display every candidate; honor constraints.desiredCount / maxCount.", "- Use type=derived_table only for tables you create from evidence; every derived_table.sourceRefs entry must reference evidence in the ledger.", - "- citations and selected image/table artifact refs may only reference refs returned by retrieve (in the evidence ledger).", + "- citations and selected image/table artifact refs may only reference refs returned by Knowhere tools in the evidence ledger.", "- inspectImage observations are inspection notes, not new source refs. Final citations and displayed image artifacts must use the original retrieved image asset refs.", "- If text evidence identifies a relevant page/image but does not include the exact fact, inspect the returned image asset for OCR/detail before saying the answer is unavailable.", "- If evidence is insufficient, list it in unresolved instead of fabricating facts.", diff --git a/src/agent-harness/types.ts b/src/agent-harness/types.ts index 4601f6f..83bb08d 100644 --- a/src/agent-harness/types.ts +++ b/src/agent-harness/types.ts @@ -1,4 +1,10 @@ import type { + KnowledgeDocumentReference, + KnowledgeGrepParams, + KnowledgeGrepResponse, + KnowledgeOutline, + KnowledgeReadParams, + KnowledgeReadResponse, RetrievalQueryParams, RetrievalQueryResponse, } from "@ontos-ai/knowhere-sdk" @@ -77,23 +83,57 @@ export type AgentTurnInput = { } } -export type HarnessRetrievalRequest = Pick< +export type KnowhereSearchTargetContent = + | "all" + | "text" + | "image" + | "table" + | "text_image" + | "text_table" + +export type KnowhereSearchRequest = Pick< RetrievalQueryParams, "query" | "topK" | "signalPaths" | "filterMode" | "threshold" > & { - readonly modalities: readonly TargetModality[] + readonly targetContent?: KnowhereSearchTargetContent readonly purpose?: string } -export type RetrievalCapability = { - readonly query: ( - input: HarnessRetrievalRequest, +export type KnowhereDocumentSummary = { + readonly documentId?: string + readonly localDocumentId?: string + readonly revisionKey?: string + readonly namespace?: string + readonly sourceFileName: string + readonly title?: string + readonly status?: string + readonly chunkCount?: number + readonly typeCounts?: Readonly> +} + +export type KnowhereListDocumentsResponse = { + readonly documents: readonly KnowhereDocumentSummary[] +} + +export type KnowhereToolRuntime = { + readonly search: ( + input: KnowhereSearchRequest, ) => Promise + readonly listDocuments: () => Promise + readonly getDocumentOutline: ( + input: KnowledgeDocumentReference, + ) => Promise + readonly readChunks: ( + input: KnowledgeReadParams, + ) => Promise + readonly grepChunks: ( + input: KnowledgeGrepParams, + ) => Promise } export type EvidenceChunk = { readonly ref: string - readonly kind: "result" | "referenced_chunk" + readonly kind: "result" | "referenced_chunk" | "read_chunk" | "grep_match" readonly chunkId?: string readonly content: string readonly contentPreview: string diff --git a/src/domains/chat/contracts.ts b/src/domains/chat/contracts.ts index 3400af5..e43b232 100644 --- a/src/domains/chat/contracts.ts +++ b/src/domains/chat/contracts.ts @@ -1,16 +1,22 @@ import type { + Knowledge, RetrievalQueryParams, RetrievalQueryResponse, } from "@ontos-ai/knowhere-sdk" import type { Source } from "@/infrastructure/db/schema" -import type { HarnessRunResult, InspectImages } from "@/agent-harness" +import type { + HarnessRunResult, + InspectImages, + KnowhereToolRuntime, +} from "@/agent-harness" import type { ChatArtifactView, ChatCitationView, } from "@/domains/chat/types" import type { HardenMediaAssetUrls } from "./media-asset-hardening" import type { HardenChatAssetUrl } from "./media-assets" +import type { NotebookKnowhereRemoteDocumentClient } from "./knowhere-tools" export type RetrievalClient = { query(params: RetrievalQueryParams): Promise @@ -57,6 +63,7 @@ export type GenerateAnswer = (input: { sources: readonly Source[] excludedSourceIds: readonly string[] searchSources: SearchSources + knowhereTools?: KnowhereToolRuntime inspectImages?: InspectImages }) => Promise @@ -67,6 +74,8 @@ export type AnswerQuestionInput = { sources: readonly Source[] excludedSourceIds: readonly string[] retrieval: RetrievalClient + knowledge?: Knowledge + remoteDocumentClient?: NotebookKnowhereRemoteDocumentClient generateAnswer: GenerateAnswer hardenChatAssetUrl?: HardenChatAssetUrl hardenMediaAssetUrls?: HardenMediaAssetUrls diff --git a/src/domains/chat/index.test.ts b/src/domains/chat/index.test.ts index 9a34a4b..313d3ff 100644 --- a/src/domains/chat/index.test.ts +++ b/src/domains/chat/index.test.ts @@ -1,5 +1,11 @@ import { afterEach, describe, expect, it, vi } from "vitest" -import type { RetrievalResult } from "@ontos-ai/knowhere-sdk" +import type { + Knowledge, + KnowledgeGrepResponse, + KnowledgeOutline, + KnowledgeReadResponse, + RetrievalResult, +} from "@ontos-ai/knowhere-sdk" import { Effect } from "effect" import { ToolLoopAgent } from "ai" import type { HarnessRunResult } from "@/agent-harness" @@ -8,6 +14,7 @@ import { answerQuestionWithRetrieval, generateAgenticOutputManifest, parseChatRequestBody, + type GenerateAnswer, type SearchSources, } from "." import type { @@ -88,6 +95,7 @@ describe("answerQuestionWithRetrieval", () => { sources, excludedSourceIds: ["source_2", "knowhere-doc:default:doc_remote"], searchSources: expect.any(Function), + knowhereTools: expect.any(Object), }); expect(answer).toEqual({ answer: "The answer is grounded.", @@ -96,6 +104,132 @@ describe("answerQuestionWithRetrieval", () => { }); }); + it("exposes search, list, outline, read, and grep through the Knowhere tool runtime", async () => { + const result = makeRetrievalResult({ + chunkType: "image", + source: { + documentId: "doc_included", + sourceFileName: "notes.txt", + sectionPath: "images/diagram.png", + }, + }); + const retrieval = { + query: vi.fn().mockResolvedValue({ + results: [result], + evidenceText: "Diagram evidence.", + referencedChunks: [], + namespace: "notebook-workspace", + query: "diagram", + routerUsed: "workflow_single_step", + answerText: null, + }), + }; + const getDocumentOutline = vi.fn().mockResolvedValue(makeKnowledgeOutline()); + const readChunks = vi.fn().mockResolvedValue( + makeKnowledgeReadResponse("Full diagram chunk body."), + ); + const grepChunks = vi.fn().mockResolvedValue(makeKnowledgeGrepResponse()); + const knowledge = { + getDocumentOutline, + readChunks, + grepChunks, + } as unknown as Knowledge; + const listDocuments = vi.fn().mockResolvedValue({ + documents: [ + { + documentId: "doc_remote", + namespace: "default", + status: "ready", + currentJobResultId: "job_remote", + sourceFileName: "remote.pdf", + }, + ], + }); + const generateAnswer = vi.fn( + async ({ knowhereTools }: Parameters[0]) => { + if (!knowhereTools) throw new Error("Knowhere tools were not provided."); + + const searchResponse = await knowhereTools.search({ + query: "diagram", + targetContent: "image", + topK: 2, + }); + const documents = await knowhereTools.listDocuments(); + await knowhereTools.getDocumentOutline({ + documentId: "doc_included", + revisionKey: "job_123", + }); + await knowhereTools.readChunks({ + documentId: "doc_included", + revisionKey: "job_123", + page: 1, + pageSize: 2, + }); + await knowhereTools.grepChunks({ + documentId: "doc_included", + revisionKey: "job_123", + pattern: "diagram", + maxResults: 3, + }); + + expect(searchResponse.results).toEqual([result]); + expect( + documents.documents.map((document) => document.documentId), + ).toEqual(["doc_included", "doc_remote"]); + return makeHarnessRunResult("Runtime answer."); + }, + ); + const sources = [ + makeSource({ id: "source_included", knowhereDocumentId: "doc_included" }), + makeSource({ id: "source_excluded", knowhereDocumentId: "doc_excluded" }), + ]; + + const answer = await Effect.runPromise( + answerQuestionWithRetrieval({ + question: "Show the diagram.", + namespace: "notebook-workspace", + sources, + excludedSourceIds: ["source_excluded"], + retrieval, + knowledge, + remoteDocumentClient: { documents: { list: listDocuments } }, + generateAnswer, + messages: [], + }), + ); + + expect(retrieval.query).toHaveBeenCalledWith({ + namespace: "notebook-workspace", + query: "diagram", + topK: 2, + useAgentic: false, + dataType: 3, + excludeDocumentIds: ["doc_excluded"], + }); + expect(listDocuments).toHaveBeenCalledWith({ + namespace: "default", + page: 1, + pageSize: 200, + }); + expect(getDocumentOutline).toHaveBeenCalledWith({ + documentId: "doc_included", + revisionKey: "job_123", + }); + expect(readChunks).toHaveBeenCalledWith({ + documentId: "doc_included", + revisionKey: "job_123", + page: 1, + pageSize: 2, + }); + expect(grepChunks).toHaveBeenCalledWith({ + documentId: "doc_included", + revisionKey: "job_123", + pattern: "diagram", + maxResults: 3, + }); + expect(answer.answer).toBe("Runtime answer."); + }); + it("does not carry no-evidence metadata from default into a successful legacy namespace result", async () => { const legacyResult = makeRetrievalResult({ source: { @@ -501,6 +635,7 @@ describe("answerQuestionWithRetrieval", () => { sources, excludedSourceIds: [], searchSources: expect.any(Function), + knowhereTools: expect.any(Object), }); const expectedResult = { ...result, @@ -1568,6 +1703,7 @@ describe("answerQuestionWithRetrieval", () => { sources, excludedSourceIds: [], searchSources: expect.any(Function), + knowhereTools: expect.any(Object), }); expect(retrieval.query).toHaveBeenCalledWith({ namespace: "notebook-workspace", @@ -1673,6 +1809,7 @@ describe("answerQuestionWithRetrieval", () => { sources: [makeSource({ title: "TSLA-Q4-2025-Update.pdf" })], excludedSourceIds: [], searchSources: expect.any(Function), + knowhereTools: expect.any(Object), }); }); @@ -1821,9 +1958,9 @@ describe("generateAgenticOutputManifest", () => { reason: "The current request is self-contained.", activePriorTurnIds: [], }); - await tools.retrieve?.execute({ + await tools.knowhere_search?.execute({ query: "冯荣洲 身份证 图片", - modalities: ["text", "image"], + targetContent: "text_image", topK: 2, purpose: "Find exactly the requested identity-card images.", }); @@ -1954,9 +2091,9 @@ describe("generateAgenticOutputManifest", () => { reason: "The current request is self-contained.", activePriorTurnIds: [], }); - await tools.retrieve?.execute({ + await tools.knowhere_search?.execute({ query: "identity card front image", - modalities: ["image"], + targetContent: "image", topK: 1, purpose: "Find the ID card image to inspect.", }); @@ -2093,9 +2230,9 @@ describe("generateAgenticOutputManifest", () => { reason: "The current request is self-contained.", activePriorTurnIds: [], }); - await tools.retrieve?.execute({ + await tools.knowhere_search?.execute({ query: "进度计划 违约金 承包人", - modalities: ["text"], + targetContent: "text", topK: 6, purpose: "Find the contract clause and page for the liquidated damages amount.", }); @@ -2248,9 +2385,9 @@ describe("generateAgenticOutputManifest", () => { reason: "Self-contained request.", activePriorTurnIds: [], }); - await tools.retrieve?.execute({ + await tools.knowhere_search?.execute({ query: "身份证 图片", - modalities: ["image"], + targetContent: "image", topK: 3, purpose: "Find requested identity images.", }); @@ -2421,6 +2558,87 @@ function makeHarnessRunResult(text: string): HarnessRunResult { }; } +function makeKnowledgeOutline(): KnowledgeOutline { + return { + document: makeLocalKnowledgeDocument(), + totalChunks: 1, + typeCounts: { text: 1, image: 0, table: 0, page: 0 }, + sections: [ + { + sectionPath: "Root / Diagram", + sectionTitle: "Diagram", + sectionLevel: 2, + summary: "Diagram section.", + startChunk: 1, + endChunk: 1, + chunkCount: 1, + typeCounts: { text: 1, image: 0, table: 0, page: 0 }, + children: [], + }, + ], + sectionTree: [], + }; +} + +function makeKnowledgeReadResponse(content: string): KnowledgeReadResponse { + return { + document: makeLocalKnowledgeDocument(), + chunks: [ + { + position: 1, + chunkId: "chunk_1", + chunkType: "text", + content, + readableContent: content, + sectionPath: "Root / Diagram", + sourceChunkPath: "chunks/chunk-1.md", + filePath: "notes.txt", + metadata: {}, + }, + ], + page: 1, + pageSize: 1, + totalChunks: 1, + totalPages: 1, + }; +} + +function makeKnowledgeGrepResponse(): KnowledgeGrepResponse { + return { + document: makeLocalKnowledgeDocument(), + matches: [ + { + position: 1, + chunkId: "chunk_1", + chunkType: "text", + sectionPath: "Root / Diagram", + sourceChunkPath: "chunks/chunk-1.md", + filePath: "notes.txt", + startOffset: 0, + endOffset: 7, + snippet: "diagram", + }, + ], + scannedChunks: 1, + truncated: false, + }; +} + +function makeLocalKnowledgeDocument() { + return { + localDocumentId: "doc_included", + documentId: "doc_included", + jobId: "job_123", + namespace: "notebook-workspace", + sourceFileName: "notes.txt", + chunkCount: 1, + typeCounts: { text: 1, image: 0, table: 0, page: 0 }, + resultDirectoryPath: "parsed-storage:doc_included/job_123", + createdAt: new Date("2026-01-01T00:00:00Z"), + updatedAt: new Date("2026-01-01T00:00:00Z"), + }; +} + type KnowhereQueryResponseLogMeta = { readonly query: string readonly resultCount: number diff --git a/src/domains/chat/index.ts b/src/domains/chat/index.ts index 4f5dee6..c0dc51b 100644 --- a/src/domains/chat/index.ts +++ b/src/domains/chat/index.ts @@ -40,6 +40,7 @@ import { } from "./media-assets" import { enrichRetrievalResultsWithPageCitationAssetUrls } from "./page-citation-assets" import type { HardenableRetrievalResult } from "./media-asset-hardening" +import { notebookKnowhereTools } from "./knowhere-tools" const DEFAULT_TOP_K = 8 const NOTEBOOK_USE_AGENTIC_RETRIEVAL: NonNullable< @@ -217,6 +218,14 @@ export const answerQuestionWithRetrieval = ( sources: input.sources, excludedSourceIds: input.excludedSourceIds, searchSources, + knowhereTools: notebookKnowhereTools.createRuntime({ + namespace: input.namespace, + sources: input.sources, + excludedSourceIds: input.excludedSourceIds, + searchSources, + knowledge: input.knowledge, + remoteDocumentClient: input.remoteDocumentClient, + }), ...(input.inspectImages ? { inspectImages: input.inspectImages } : {}), }), ) diff --git a/src/domains/chat/knowhere-tools.ts b/src/domains/chat/knowhere-tools.ts new file mode 100644 index 0000000..210b487 --- /dev/null +++ b/src/domains/chat/knowhere-tools.ts @@ -0,0 +1,155 @@ +import { Effect } from "effect" +import type { Knowledge } from "@ontos-ai/knowhere-sdk" + +import type { + KnowhereDocumentSummary, + KnowhereToolRuntime, +} from "@/agent-harness" +import type { Source } from "@/infrastructure/db/schema" +import { + listRemoteLibraryDocuments, + type RemoteLibraryDocument, +} from "@/domains/sources/remote-library" +import type { SearchSources } from "./contracts" +import { excludeDocuments } from "./retrieval" + +type RemoteDocumentClient = Parameters< + typeof listRemoteLibraryDocuments +>[0]["client"] + +export type NotebookKnowhereRemoteDocumentClient = RemoteDocumentClient + +type NotebookKnowhereToolsInput = { + readonly namespace: string + readonly sources: readonly Source[] + readonly excludedSourceIds: readonly string[] + readonly searchSources: SearchSources + readonly knowledge?: Knowledge + readonly remoteDocumentClient?: RemoteDocumentClient +} + +type SearchOnlyRuntimeInput = { + readonly searchSources: SearchSources +} + +export const notebookKnowhereTools = { + createRuntime(input: NotebookKnowhereToolsInput): KnowhereToolRuntime { + return { + search: (request) => input.searchSources(request), + listDocuments: async () => ({ + documents: await listVisibleDocuments(input), + }), + getDocumentOutline: async (request) => { + const knowledge = requireKnowledge(input.knowledge) + return knowledge.getDocumentOutline(request) + }, + readChunks: async (request) => { + const knowledge = requireKnowledge(input.knowledge) + return knowledge.readChunks(request) + }, + grepChunks: async (request) => { + const knowledge = requireKnowledge(input.knowledge) + return knowledge.grepChunks(request) + }, + } + }, + + createSearchOnlyRuntime(input: SearchOnlyRuntimeInput): KnowhereToolRuntime { + return { + search: (request) => input.searchSources(request), + listDocuments: async () => ({ documents: [] }), + getDocumentOutline: async () => { + throw new Error("Knowhere document outline is not configured.") + }, + readChunks: async () => { + throw new Error("Knowhere chunk reads are not configured.") + }, + grepChunks: async () => { + throw new Error("Knowhere grep is not configured.") + }, + } + }, +} as const + +async function listVisibleDocuments( + input: NotebookKnowhereToolsInput, +): Promise { + const excludedSourceIds = new Set(input.excludedSourceIds) + const excludedDocumentIds = new Set( + excludeDocuments(input.sources, input.excludedSourceIds) + .excludeDocumentIds ?? [], + ) + const localDocuments = input.sources + .filter( + (source): source is Source & { readonly knowhereDocumentId: string } => + source.status === "ready" && + Boolean(source.knowhereDocumentId) && + !excludedSourceIds.has(source.id) && + !excludedDocumentIds.has(source.knowhereDocumentId ?? ""), + ) + .map((source): KnowhereDocumentSummary => ({ + documentId: source.knowhereDocumentId, + revisionKey: source.knowhereJobId ?? undefined, + namespace: input.namespace, + sourceFileName: source.title, + title: source.title, + status: source.status, + })) + + const remoteDocuments = await listVisibleRemoteDocuments({ + input, + localDocuments, + excludedDocumentIds, + }) + + return [...localDocuments, ...remoteDocuments] +} + +async function listVisibleRemoteDocuments(input: { + readonly input: NotebookKnowhereToolsInput + readonly localDocuments: readonly KnowhereDocumentSummary[] + readonly excludedDocumentIds: ReadonlySet +}): Promise { + if (!input.input.remoteDocumentClient) return [] + + const localDocumentIds = new Set( + input.localDocuments.flatMap((document): string[] => + document.documentId ? [document.documentId] : [], + ), + ) + const documents = await Effect.runPromise( + listRemoteLibraryDocuments({ + workspace: { namespace: input.input.namespace }, + client: input.input.remoteDocumentClient, + localSources: input.input.sources, + }), + ) + + return documents + .filter( + (document) => + document.status === "ready" && + !localDocumentIds.has(document.documentId) && + !input.excludedDocumentIds.has(document.documentId), + ) + .map(toRemoteDocumentSummary) +} + +function toRemoteDocumentSummary( + document: RemoteLibraryDocument, +): KnowhereDocumentSummary { + return { + documentId: document.documentId, + revisionKey: document.revisionKey, + namespace: document.namespace, + sourceFileName: + document.sourceFileName ?? document.title ?? document.documentId, + title: document.title, + status: document.status, + } +} + +function requireKnowledge(knowledge: Knowledge | undefined): Knowledge { + if (knowledge) return knowledge + throw new Error("Knowhere parsed-document reads are not configured.") +} diff --git a/src/domains/chat/prompt.ts b/src/domains/chat/prompt.ts index efee008..f424252 100644 --- a/src/domains/chat/prompt.ts +++ b/src/domains/chat/prompt.ts @@ -8,17 +8,15 @@ import { runAgentHarness, type AgentTurn, type AgentTurnInput, - type HarnessRetrievalRequest, type HarnessRunResult, type InspectImages, - type TargetModality, + type KnowhereToolRuntime, } from "@/agent-harness" import type { - AgenticRetrievalQuery, - AgenticRetrievalTargetContent, ChatHistoryMessage, SearchSources, } from "./contracts" +import { notebookKnowhereTools } from "./knowhere-tools" const RECENT_CONTEXT_MESSAGE_LIMIT = 8 const CONTEXT_CONTENT_CHAR_LIMIT = 900 @@ -30,6 +28,7 @@ type GenerateAgenticOutputManifestInput = { sources: readonly Source[] excludedSourceIds: readonly string[] searchSources: SearchSources + knowhereTools?: KnowhereToolRuntime inspectImages?: InspectImages } @@ -59,10 +58,11 @@ export const generateAgenticOutputManifestEffect = ( runAgentHarness({ model: CHAT_MODEL, turn, - retrieval: { - query: (request) => - input.searchSources(toAgenticRetrievalQuery(request)), - }, + knowhereTools: + input.knowhereTools ?? + notebookKnowhereTools.createSearchOnlyRuntime({ + searchSources: input.searchSources, + }), ...(input.inspectImages ? { inspectImages: input.inspectImages } : {}), }), ) @@ -124,36 +124,6 @@ function getCitationLabels( .filter((label) => label.length > 0) } -function toAgenticRetrievalQuery( - request: HarnessRetrievalRequest, -): AgenticRetrievalQuery { - return { - query: request.query, - targetContent: toAgenticRetrievalTargetContent(request.modalities), - purpose: request.purpose, - topK: request.topK, - signalPaths: request.signalPaths, - filterMode: request.filterMode, - threshold: request.threshold, - } -} - -function toAgenticRetrievalTargetContent( - modalities: readonly TargetModality[], -): AgenticRetrievalTargetContent { - const requestedModalities = new Set(modalities) - if (requestedModalities.has("image") && requestedModalities.has("text")) { - return "text_image" - } - if (requestedModalities.has("table") && requestedModalities.has("text")) { - return "text_table" - } - if (requestedModalities.has("image")) return "image" - if (requestedModalities.has("table")) return "table" - if (requestedModalities.has("text")) return "text" - return "all" -} - function formatSourceContext( sources: readonly Source[], excludedSourceIds: readonly string[], diff --git a/src/domains/chat/route-answer.ts b/src/domains/chat/route-answer.ts index 2f8ae6b..3f9fc9d 100644 --- a/src/domains/chat/route-answer.ts +++ b/src/domains/chat/route-answer.ts @@ -27,6 +27,7 @@ import type { Source } from "@/infrastructure/db/schema" import { CHAT_MODEL } from "@/lib/ai" import type { HardenChatAssetUrl } from "./media-assets" import { isAuthError } from "@/integrations/dashboard/api-key-service" +import { makeKnowhereClientWithParsedStorage } from "@/integrations/knowhere" import { summarizeUnknownError } from "@/lib/format-log-value" import { logger } from "@/lib/logger" import { routeResult, type RouteResult } from "@/lib/route-result" @@ -98,6 +99,9 @@ const answerChatEffect = (input: AnswerChatInput) => const parsedStorage = new BlobParsedDocumentStorage({ workspaceId: workspace.id, }) + const knowhereResources = makeKnowhereClientWithParsedStorage(apiKey, { + workspaceId: workspace.id, + }) const hardenChatAssetUrl: HardenChatAssetUrl = async ({ source, sourcePath, @@ -140,6 +144,8 @@ const answerChatEffect = (input: AnswerChatInput) => threadId: body.value.threadId, excludedSourceIds: body.value.excludedSourceIds, retrieval: client.retrieval, + knowledge: knowhereResources.knowledge, + remoteDocumentClient: client, generateAnswer: generateAgenticOutputManifest, hardenChatAssetUrl, hardenMediaAssetUrls: ({ results, artifacts }) => diff --git a/src/domains/chat/service.test.ts b/src/domains/chat/service.test.ts index 12df745..441a12c 100644 --- a/src/domains/chat/service.test.ts +++ b/src/domains/chat/service.test.ts @@ -67,6 +67,7 @@ describe("handleChatTurn", () => { sources, excludedSourceIds: ["source_excluded"], searchSources: expect.any(Function), + knowhereTools: expect.any(Object), }); expect(repository.appendMessageToThread).toHaveBeenNthCalledWith(1, "workspace_1", { threadId: "thread_1", @@ -216,6 +217,7 @@ describe("handleChatTurn", () => { sources, excludedSourceIds: [], searchSources: expect.any(Function), + knowhereTools: expect.any(Object), }); expect(retrieval.query).toHaveBeenCalledWith({ namespace: "notebook-namespace", diff --git a/src/domains/chat/service.ts b/src/domains/chat/service.ts index bfa7e42..7a7b2b5 100644 --- a/src/domains/chat/service.ts +++ b/src/domains/chat/service.ts @@ -66,6 +66,8 @@ type ChatTurnInput = { threadId?: string excludedSourceIds: readonly string[] retrieval: RetrievalClient + knowledge?: AnswerQuestionInput["knowledge"] + remoteDocumentClient?: AnswerQuestionInput["remoteDocumentClient"] generateAnswer: GenerateAnswer hardenChatAssetUrl?: AnswerQuestionInput["hardenChatAssetUrl"] hardenMediaAssetUrls?: AnswerQuestionInput["hardenMediaAssetUrls"] @@ -126,6 +128,8 @@ export const handleChatTurnEffect = (input: ChatTurnInput) => sources: readySources, excludedSourceIds: input.excludedSourceIds, retrieval: input.retrieval, + knowledge: input.knowledge, + remoteDocumentClient: input.remoteDocumentClient, generateAnswer: input.generateAnswer, hardenChatAssetUrl: input.hardenChatAssetUrl, hardenMediaAssetUrls: input.hardenMediaAssetUrls, From b59db5aaacf259b6a2e4a35008fe61be10f1cca2 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Wed, 8 Jul 2026 04:10:00 +0800 Subject: [PATCH 13/14] fix(chat): validate citation source refs --- src/agent-harness/knowhere-text.ts | 1 + src/agent-harness/runtime.test.ts | 11 + src/agent-harness/runtime.ts | 10 +- src/agent-harness/validator.test.ts | 136 +++++++++++ src/agent-harness/validator.ts | 79 ++++++ src/domains/chat/index.test.ts | 364 ++++++++++++++++++++++++++-- src/domains/chat/index.ts | 55 +---- src/domains/chat/service.test.ts | 4 +- 8 files changed, 580 insertions(+), 80 deletions(-) diff --git a/src/agent-harness/knowhere-text.ts b/src/agent-harness/knowhere-text.ts index 69de846..1c19368 100644 --- a/src/agent-harness/knowhere-text.ts +++ b/src/agent-harness/knowhere-text.ts @@ -214,6 +214,7 @@ function formatEvidenceAssets(assets: readonly EvidenceAsset[]): string { label: asset.label, sourcePath: asset.sourcePath, documentId: asset.source.documentId ?? undefined, + sourceFileName: asset.source.sourceFileName ?? undefined, sectionPath: asset.source.sectionPath ?? undefined, }), ), diff --git a/src/agent-harness/runtime.test.ts b/src/agent-harness/runtime.test.ts index e78de33..eec6064 100644 --- a/src/agent-harness/runtime.test.ts +++ b/src/agent-harness/runtime.test.ts @@ -28,6 +28,17 @@ describe("agent harness runtime", () => { expect(prompt).not.toContain("navigation action") }) + it("tells the agent to cite only refs with matching source metadata", () => { + const prompt = buildHarnessSystemPrompt(makeTurnInput()) + + expect(prompt).toContain( + "source.documentId, sourceFileName, and sectionPath must match the selected evidence ref exactly", + ) + expect(prompt).toContain( + "Omit citations when you cannot identify a supporting evidence ref", + ) + }) + it("passes only outer retrieval parameters to KNOWHERE after intent and context policy are declared", async () => { const query = vi.fn().mockResolvedValue( makeRetrievalResponse(), diff --git a/src/agent-harness/runtime.ts b/src/agent-harness/runtime.ts index e2b2783..e272340 100644 --- a/src/agent-harness/runtime.ts +++ b/src/agent-harness/runtime.ts @@ -21,7 +21,6 @@ import type { ImageInspectionResponse, InspectImages, IntentFrame, - KnowhereSearchRequest, KnowhereSearchTargetContent, KnowhereToolRuntime, OutputManifest, @@ -406,7 +405,9 @@ function buildRevisionFeedback(errors: readonly string[]): string { "answer contract.", "Do not exceed the user's requested artifact count, only cite or display", "evidence refs that exist in the evidence ledger, and do not fabricate", - "facts when evidence is missing.", + "facts when evidence is missing. Citation source.documentId,", + "sourceFileName, and sectionPath must match the selected evidence ref", + "exactly.", ].join("\n") } @@ -668,7 +669,8 @@ export function createHarnessTools(input: { description: "Finalize the user-facing output manifest. This is the only final answer " + "contract. Artifacts listed here with display=true are the exact set of " + - "images/tables shown to the user; cite only refs from the evidence ledger.", + "images/tables shown to the user; cite only refs from the evidence ledger " + + "and copy citation source metadata exactly from the selected ref.", inputSchema: outputManifestSchema, execute: async (manifest) => traceToolCall(input.state, { @@ -1275,6 +1277,8 @@ export function buildHarnessSystemPrompt(turn: AgentTurnInput): string { "- artifacts with display=true are the exact images/tables shown. Never display every candidate; honor constraints.desiredCount / maxCount.", "- Use type=derived_table only for tables you create from evidence; every derived_table.sourceRefs entry must reference evidence in the ledger.", "- citations and selected image/table artifact refs may only reference refs returned by Knowhere tools in the evidence ledger.", + "- For every citation, source.documentId, sourceFileName, and sectionPath must match the selected evidence ref exactly. Use null or omit the field only when the ref omits it.", + "- Omit citations when you cannot identify a supporting evidence ref with matching source metadata.", "- inspectImage observations are inspection notes, not new source refs. Final citations and displayed image artifacts must use the original retrieved image asset refs.", "- If text evidence identifies a relevant page/image but does not include the exact fact, inspect the returned image asset for OCR/detail before saying the answer is unavailable.", "- If evidence is insufficient, list it in unresolved instead of fabricating facts.", diff --git a/src/agent-harness/validator.test.ts b/src/agent-harness/validator.test.ts index be8de3e..89bf0ba 100644 --- a/src/agent-harness/validator.test.ts +++ b/src/agent-harness/validator.test.ts @@ -115,6 +115,142 @@ describe("validateOutputManifest", () => { ) }) + it("rejects citation source metadata that conflicts with the resolved evidence ref", () => { + const validation = validateOutputManifest({ + manifest: makeManifest({ + text: "The source discusses information hiding.", + citations: [ + { + ref: "r1:result:1", + label: "wrong source", + source: { + documentId: "doc_claimed", + sourceFileName: "claimed-source.pdf", + sectionPath: "Claimed Section", + }, + }, + ], + }), + intent: makeIntent({}), + contextPolicy: unrelatedContextPolicy, + ledger: { + ...emptyLedger, + chunks: [ + { + ref: "r1:result:1", + kind: "result", + content: "Information hiding is a module design principle.", + contentPreview: "Information hiding is a module design principle.", + chunkType: "text", + score: 0.9, + source: { + documentId: "doc_information_hiding", + sourceFileName: "information_hiding.pdf", + sectionPath: "Root / Module Design", + }, + }, + ], + }, + surface: "notebook_chat", + }) + + expect(validation.errors).toContain( + "Citation ref 'r1:result:1' source.documentId must match resolved evidence source. Expected 'doc_information_hiding', received 'doc_claimed'.", + ) + expect(validation.errors).toContain( + "Citation ref 'r1:result:1' source.sourceFileName must match resolved evidence source. Expected 'information_hiding.pdf', received 'claimed-source.pdf'.", + ) + expect(validation.errors).toContain( + "Citation ref 'r1:result:1' source.sectionPath must match resolved evidence source. Expected 'Root / Module Design', received 'Claimed Section'.", + ) + }) + + it("accepts citation source metadata that exactly matches the resolved evidence ref", () => { + const validation = validateOutputManifest({ + manifest: makeManifest({ + text: "Revenue increased.", + citations: [ + { + ref: "r1:result:1", + label: "report.pdf / Q4", + source: { + documentId: "doc_1", + sourceFileName: "report.pdf", + sectionPath: "Q4", + }, + }, + ], + }), + intent: makeIntent({}), + contextPolicy: unrelatedContextPolicy, + ledger: { + ...emptyLedger, + chunks: [ + { + ref: "r1:result:1", + kind: "result", + content: "Revenue increased.", + contentPreview: "Revenue increased.", + chunkType: "text", + score: 0.9, + source: { + documentId: "doc_1", + sourceFileName: "report.pdf", + sectionPath: "Q4", + }, + }, + ], + }, + surface: "notebook_chat", + }) + + expect(validation.ok).toBe(true) + expect(validation.errors).toEqual([]) + }) + + it("rejects missing citation documentId when the resolved evidence has one", () => { + const validation = validateOutputManifest({ + manifest: makeManifest({ + text: "Revenue increased.", + citations: [ + { + ref: "r1:result:1", + label: "report.pdf / Q4", + source: { + sourceFileName: "report.pdf", + sectionPath: "Q4", + }, + }, + ], + }), + intent: makeIntent({}), + contextPolicy: unrelatedContextPolicy, + ledger: { + ...emptyLedger, + chunks: [ + { + ref: "r1:result:1", + kind: "result", + content: "Revenue increased.", + contentPreview: "Revenue increased.", + chunkType: "text", + score: 0.9, + source: { + documentId: "doc_1", + sourceFileName: "report.pdf", + sectionPath: "Q4", + }, + }, + ], + }, + surface: "notebook_chat", + }) + + expect(validation.errors).toContain( + "Citation ref 'r1:result:1' source.documentId must match resolved evidence source. Expected 'doc_1', received missing.", + ) + }) + it("accepts source-backed derived tables and rejects missing source refs", () => { const validation = validateOutputManifest({ manifest: makeManifest({ diff --git a/src/agent-harness/validator.ts b/src/agent-harness/validator.ts index dd67343..fcd8c0e 100644 --- a/src/agent-harness/validator.ts +++ b/src/agent-harness/validator.ts @@ -3,6 +3,7 @@ import type { EvidenceLedgerSnapshot, HarnessToolCallTrace, IntentFrame, + OutputCitation, OutputManifest, } from "./types" @@ -21,6 +22,15 @@ export type ManifestValidationResult = { readonly errors: readonly string[] } +const citationSourceFields = [ + "documentId", + "sourceFileName", + "sectionPath", +] as const + +type CitationSourceField = (typeof citationSourceFields)[number] +type EvidenceSource = EvidenceLedgerSnapshot["chunks"][number]["source"] + export function validateOutputManifest( input: ManifestValidationInput, ): ManifestValidationResult { @@ -71,6 +81,14 @@ function validateArtifactRefs( ...input.ledger.chunks.map((chunk) => chunk.ref), ...input.ledger.assets.map((asset) => asset.ref), ]) + const sourcesByRef = new Map([ + ...input.ledger.chunks.map( + (chunk): readonly [string, EvidenceSource] => [chunk.ref, chunk.source], + ), + ...input.ledger.assets.map( + (asset): readonly [string, EvidenceSource] => [asset.ref, asset.source], + ), + ]) for (const artifact of input.manifest.artifacts) { if (artifact.type === "derived_table") { @@ -102,10 +120,71 @@ function validateArtifactRefs( for (const citation of input.manifest.citations) { if (!knownRefs.has(citation.ref)) { errors.push(`Citation ref '${citation.ref}' was not found in the evidence ledger.`) + continue } + + const source = sourcesByRef.get(citation.ref) + if (!source) continue + validateCitationSource( + { + ref: citation.ref, + declared: citation.source, + resolved: source, + }, + errors, + ) + } +} + +function validateCitationSource( + input: { + readonly ref: string + readonly declared: OutputCitation["source"] + readonly resolved: EvidenceSource + }, + errors: string[], +): void { + for (const field of citationSourceFields) { + validateCitationSourceField( + { + ref: input.ref, + field, + declaredValue: input.declared[field], + resolvedValue: input.resolved[field], + }, + errors, + ) } } +function validateCitationSourceField( + input: { + readonly ref: string + readonly field: CitationSourceField + readonly declaredValue: string | null | undefined + readonly resolvedValue: string | null | undefined + }, + errors: string[], +): void { + const declaredValue = normalizeSourceValue(input.declaredValue) + const resolvedValue = normalizeSourceValue(input.resolvedValue) + if (declaredValue === resolvedValue) return + + errors.push( + `Citation ref '${input.ref}' source.${input.field} must match resolved evidence source. Expected ${formatSourceValue( + resolvedValue, + )}, received ${formatSourceValue(declaredValue)}.`, + ) +} + +function normalizeSourceValue(value: string | null | undefined): string | null { + return value ?? null +} + +function formatSourceValue(value: string | null): string { + return value === null ? "missing" : `'${value}'` +} + function validateArtifactCounts( input: ManifestValidationInput, errors: string[], diff --git a/src/domains/chat/index.test.ts b/src/domains/chat/index.test.ts index 313d3ff..4281520 100644 --- a/src/domains/chat/index.test.ts +++ b/src/domains/chat/index.test.ts @@ -99,7 +99,51 @@ describe("answerQuestionWithRetrieval", () => { }); expect(answer).toEqual({ answer: "The answer is grounded.", - citations: [result], + citations: [], + artifacts: [], + }); + }); + + it("does not create source chips from retrieval results when the manifest has no citations", async () => { + const unrelatedResult = makeRetrievalResult({ + content: "Information hiding is unrelated to the requested source.", + source: { + documentId: "doc_information_hiding", + sourceFileName: "information_hiding.pdf", + sectionPath: "Root", + }, + }); + const retrieval = { + query: vi.fn().mockResolvedValue({ + results: [unrelatedResult], + evidenceText: "Information hiding evidence.", + referencedChunks: [], + namespace: "notebook-workspace", + query: "requested fact", + routerUsed: "workflow_single_step", + answerText: null, + }), + }; + const generateAnswer = vi.fn(async ({ searchSources }) => { + await searchSources({ query: "requested fact" }); + return makeHarnessRunResult("The answer omits citations."); + }); + + const answer = await Effect.runPromise( + answerQuestionWithRetrieval({ + question: "What does the selected source say?", + namespace: "notebook-workspace", + sources: [makeSource()], + excludedSourceIds: [], + retrieval, + generateAnswer, + messages: [], + }), + ); + + expect(answer).toEqual({ + answer: "The answer omits citations.", + citations: [], artifacts: [], }); }); @@ -304,7 +348,7 @@ describe("answerQuestionWithRetrieval", () => { ); expect(answer).toEqual({ answer: "The legacy answer is grounded.", - citations: [legacyResult], + citations: [], artifacts: [], }); }); @@ -562,8 +606,18 @@ describe("answerQuestionWithRetrieval", () => { }; const generateAnswer = vi.fn(async ({ searchSources }) => { await searchSources({ query: "What improved?" }); - return makeHarnessRunResult( + return makeHarnessRunResultWithLedger( "Revenue improved [Source 1: revenue growth]. Margins expanded [Source 2: margin expansion].", + { + citations: [ + makeOutputCitation("r1:result:1", firstResult), + makeOutputCitation("r1:result:2", secondResult), + ], + chunks: [ + makeEvidenceChunkFromRetrievalResult("r1:result:1", firstResult), + makeEvidenceChunkFromRetrievalResult("r1:result:2", secondResult), + ], + }, ); }); @@ -606,8 +660,9 @@ describe("answerQuestionWithRetrieval", () => { }; const generateAnswer = vi.fn(async ({ searchSources }) => { await searchSources({ query: "Tesla xAI investment" }); - return makeHarnessRunResult( + return makeCitedHarnessRunResult( "Tesla invested in xAI [Source 1: xAI investment].", + result, ); }); const sources = [ @@ -678,7 +733,10 @@ describe("answerQuestionWithRetrieval", () => { targetContent: "image", purpose: "Find visual rocket launch chunks.", }); - return makeHarnessRunResult(`Use this launch photo. ${upstreamAssetUrl}`); + return makeCitedHarnessRunResult( + `Use this launch photo. ${upstreamAssetUrl}`, + result, + ); }); const hardenChatAssetUrl = vi .fn() @@ -938,7 +996,10 @@ describe("answerQuestionWithRetrieval", () => { }; const generateAnswer = vi.fn(async ({ searchSources }) => { await searchSources({ query: "page four evidence" }); - return makeHarnessRunResult(`This page has the answer. ${storedPageAssetUrl}`); + return makeCitedHarnessRunResult( + `This page has the answer. ${storedPageAssetUrl}`, + result, + ); }); const hardenMediaAssetUrls = vi.fn( async ({ @@ -1022,6 +1083,35 @@ describe("answerQuestionWithRetrieval", () => { "https://blob.example/workspaces/workspace_1/sources/source_pages/parsed-result/page_citation_assets/page-6.png"; const hardenedPageAssetUrl = "https://blob.example/workspaces/workspace_1/chat-assets/source-source_pages/page-6.png"; + const pageMetadata = { + pageNums: [6], + pageAssets: [ + { + pageNum: 6, + artifactRef: "page_citation_assets/page-6.png", + assetUrl: rawPageAssetUrl, + contentType: "image/png", + }, + ], + }; + const referencedPageChunk: HarnessRunResult["trace"]["ledger"]["chunks"][number] = { + ref: "r1:referenced:1", + kind: "referenced_chunk", + chunkId: "chunk_page_6", + content: "", + contentPreview: "", + chunkType: "page", + score: null, + filePath: null, + metadata: pageMetadata, + source: { + documentId: "doc_pages", + sourceFileName: null, + sectionPath: "Page 6", + }, + revisionKey: "job_1", + assetUrl: rawPageAssetUrl, + }; const retrieval = { query: vi.fn().mockResolvedValue({ results: [], @@ -1035,17 +1125,7 @@ describe("answerQuestionWithRetrieval", () => { filePath: null, jobId: "job_1", assetUrl: rawPageAssetUrl, - metadata: { - pageNums: [6], - pageAssets: [ - { - pageNum: 6, - artifactRef: "page_citation_assets/page-6.png", - assetUrl: rawPageAssetUrl, - contentType: "image/png", - }, - ], - }, + metadata: pageMetadata, }, ], namespace: "notebook-workspace", @@ -1056,7 +1136,16 @@ describe("answerQuestionWithRetrieval", () => { }; const generateAnswer = vi.fn(async ({ searchSources }) => { await searchSources({ query: "page six evidence" }); - return makeHarnessRunResult("This page has referenced evidence."); + return makeHarnessRunResultWithLedger("This page has referenced evidence.", { + citations: [ + { + ref: "r1:referenced:1", + label: "Page 6", + source: referencedPageChunk.source, + }, + ], + chunks: [referencedPageChunk], + }); }); const hardenMediaAssetUrls = vi.fn( async ({ @@ -1403,6 +1492,103 @@ describe("answerQuestionWithRetrieval", () => { }); }); + it("returns a safe fallback when a manifest citation ref declares the wrong source", async () => { + process.env.AI_GATEWAY_API_KEY = "test_gateway_key"; + const result = makeRetrievalResult({ + content: "Information hiding is a module design principle.", + source: { + documentId: "doc_information_hiding", + sourceFileName: "information_hiding.pdf", + sectionPath: "Root / Module Design", + }, + }); + const retrieval = { + query: vi.fn().mockResolvedValue({ + results: [result], + evidenceText: "Information hiding evidence.", + referencedChunks: [], + namespace: "notebook-workspace", + query: "information hiding", + routerUsed: "workflow_single_step", + answerText: null, + }), + }; + let generateCallCount = 0; + vi.spyOn(ToolLoopAgent.prototype, "generate").mockImplementation( + async function mockGenerate( + this: ToolLoopAgent, + ): ReturnType { + generateCallCount += 1; + const tools = this.tools as unknown as Record< + string, + { execute: (input: unknown) => Promise } + >; + + if (generateCallCount === 1) { + await tools.declareIntent?.execute({ + task: "answer", + dependsOnPreviousTurn: false, + retrievalNeeded: "yes", + targetModalities: ["text"], + constraints: { citationRequired: true }, + groundingPolicy: "must_use_sources", + }); + await tools.setContextPolicy?.execute({ + carryHistory: "none", + reason: "Self-contained request.", + activePriorTurnIds: [], + }); + await tools.knowhere_search?.execute({ + query: "information hiding", + targetContent: "text", + }); + } + + await tools.finalize?.execute({ + text: "Information hiding is a module design principle.", + citations: [ + { + ref: "r1:result:1", + label: "claimed-source.pdf / Claimed", + source: { + documentId: "doc_claimed", + sourceFileName: "claimed-source.pdf", + sectionPath: "Claimed", + }, + }, + ], + artifacts: [], + unresolved: [], + }); + + return { + text: "ignored", + response: { messages: [] }, + } as unknown as Awaited>; + }, + ); + + const answer = await Effect.runPromise( + answerQuestionWithRetrieval({ + question: "What is information hiding?", + namespace: "notebook-workspace", + sources: [makeSource()], + excludedSourceIds: [], + retrieval, + generateAnswer: generateAgenticOutputManifest, + messages: [], + }), + ); + + expect(generateCallCount).toBe(2); + expect(answer).toEqual({ + answer: + "I couldn't safely finish that response because the agent output did not pass Notebook's validation checks. Please try again.", + citations: [], + artifacts: [], + }); + }); + it("keeps image-only harness output instead of treating it as no results", async () => { const assetUrl = "https://blob.example/images/diagram.png"; const retrieval = { @@ -1866,6 +2052,23 @@ describe("answerQuestionWithRetrieval", () => { }); it("uses structured referenced chunks from RetrievalQueryResponse as citations", async () => { + const referencedImageChunk: HarnessRunResult["trace"]["ledger"]["chunks"][number] = { + ref: "r1:referenced:1", + kind: "referenced_chunk", + chunkId: "chunk_1", + content: "", + contentPreview: "", + chunkType: "image", + score: null, + filePath: "images/launch.jpg", + source: { + documentId: "doc_spacex", + sourceFileName: null, + sectionPath: "Assets / images / launch.jpg", + }, + revisionKey: "job_1", + assetUrl: "https://blob.example/images/launch.jpg", + }; const retrieval = { query: vi.fn().mockResolvedValue({ results: [], @@ -1892,7 +2095,16 @@ describe("answerQuestionWithRetrieval", () => { query: "SpaceX launch image", targetContent: "image", }); - return makeHarnessRunResult("Here is the launch image."); + return makeHarnessRunResultWithLedger("Here is the launch image.", { + citations: [ + { + ref: "r1:referenced:1", + label: "launch image", + source: referencedImageChunk.source, + }, + ], + chunks: [referencedImageChunk], + }); }); const answer = await Effect.runPromise( @@ -1972,7 +2184,7 @@ describe("generateAgenticOutputManifest", () => { label: "商务标文件.pdf / 身份证正面", source: { documentId: "doc_identity", - sourceFileName: "商务标文件.pdf", + sourceFileName: "document-generated.pdf", sectionPath: "身份证正面", }, }, @@ -2109,7 +2321,7 @@ describe("generateAgenticOutputManifest", () => { label: "identity.pdf / images/id-front.png", source: { documentId: "doc_identity", - sourceFileName: "identity.pdf", + sourceFileName: "generated.pdf", sectionPath: "images/id-front.png", }, }, @@ -2249,7 +2461,7 @@ describe("generateAgenticOutputManifest", () => { label: "投标书 / （6）现场工期进度管理方面的违约责任", source: { documentId: "doc_contract", - sourceFileName: "投标书.pdf", + sourceFileName: null, sectionPath: "Root / （6）现场工期进度管理方面的违约责任", }, }, @@ -2393,7 +2605,17 @@ describe("generateAgenticOutputManifest", () => { }); await tools.finalize?.execute({ text: "见下方图片。", - citations: [{ ref: "r1:result:1", label: "id" }], + citations: [ + { + ref: "r1:result:1", + label: "ids.pdf / 身份证 1", + source: { + documentId: "doc_identity", + sourceFileName: "ids.pdf", + sectionPath: "身份证 1", + }, + }, + ], artifacts: [1, 2, 3].map((index) => ({ type: "image", ref: `asset:r1:result:${index}`, @@ -2405,7 +2627,17 @@ describe("generateAgenticOutputManifest", () => { } else { await tools.finalize?.execute({ text: "见下方图片。", - citations: [{ ref: "r1:result:1", label: "id" }], + citations: [ + { + ref: "r1:result:1", + label: "ids.pdf / 身份证 1", + source: { + documentId: "doc_identity", + sourceFileName: "ids.pdf", + sectionPath: "身份证 1", + }, + }, + ], artifacts: [1, 2].map((index) => ({ type: "image", ref: `asset:r1:result:${index}`, @@ -2558,6 +2790,90 @@ function makeHarnessRunResult(text: string): HarnessRunResult { }; } +function makeCitedHarnessRunResult( + text: string, + result: RetrievalResult, + ref = "r1:result:1", +): HarnessRunResult { + return makeHarnessRunResultWithLedger(text, { + citations: [makeOutputCitation(ref, result)], + chunks: [makeEvidenceChunkFromRetrievalResult(ref, result)], + }); +} + +function makeHarnessRunResultWithLedger( + text: string, + input: { + readonly citations?: HarnessRunResult["manifest"]["citations"] + readonly chunks?: HarnessRunResult["trace"]["ledger"]["chunks"] + readonly assets?: HarnessRunResult["trace"]["ledger"]["assets"] + readonly artifacts?: HarnessRunResult["manifest"]["artifacts"] + }, +): HarnessRunResult { + const chunks = input.chunks ?? []; + return { + manifest: { + text, + citations: input.citations ?? [], + artifacts: input.artifacts ?? [], + unresolved: [], + }, + trace: { + ...makeHarnessRunResult("").trace, + ledger: { + retrievalCount: chunks.length > 0 ? 1 : 0, + chunks, + assets: input.assets ?? [], + evidenceText: [], + stopReasons: [], + failureReasons: [], + decisionTraces: [], + }, + }, + }; +} + +function makeOutputCitation( + ref: string, + result: RetrievalResult, +): HarnessRunResult["manifest"]["citations"][number] { + return { + ref, + label: [result.source.sourceFileName, result.source.sectionPath] + .filter(Boolean) + .join(" / "), + source: { + documentId: result.source.documentId, + sourceFileName: result.source.sourceFileName, + sectionPath: result.source.sectionPath, + }, + }; +} + +function makeEvidenceChunkFromRetrievalResult( + ref: string, + result: RetrievalResult, +): HarnessRunResult["trace"]["ledger"]["chunks"][number] { + return { + ref, + kind: "result", + ...(result.chunkId ? { chunkId: result.chunkId } : {}), + content: result.content, + contentPreview: result.content, + chunkType: result.chunkType, + score: result.score, + ...(result.sourceChunkPath ? { sourceChunkPath: result.sourceChunkPath } : {}), + ...(result.filePath ? { filePath: result.filePath } : {}), + ...(result.metadata ? { metadata: result.metadata } : {}), + source: { + documentId: result.source.documentId, + sourceFileName: result.source.sourceFileName, + sectionPath: result.source.sectionPath, + }, + ...(result.assetUrl ? { assetUrl: result.assetUrl } : {}), + }; +} + function makeKnowledgeOutline(): KnowledgeOutline { return { document: makeLocalKnowledgeDocument(), diff --git a/src/domains/chat/index.ts b/src/domains/chat/index.ts index c0dc51b..d5cd82c 100644 --- a/src/domains/chat/index.ts +++ b/src/domains/chat/index.ts @@ -254,8 +254,6 @@ export const answerQuestionWithRetrieval = ( const rawResults = selectCitationRawResults({ generatedAnswer, - retrievalResponses, - sources: input.sources, }) if ( rawResults.length === 0 && @@ -858,14 +856,12 @@ function normalizeTopK(value: number | undefined): number { /** * Display citations come from the agent-curated manifest (the refs it chose to - * cite), resolved against the evidence ledger. Only when the agent cited - * nothing do we fall back to the full set of retrieved results, so a grounded - * answer still shows its sources instead of appearing unsupported. + * cite), resolved against the evidence ledger. If the manifest has no citations, + * only displayed artifacts produce citation chips; arbitrary retrieval results + * stay hidden so Notebook does not imply support from an unselected source. */ function selectCitationRawResults(input: { readonly generatedAnswer: HarnessRunResult - readonly retrievalResponses: readonly RetrievalQueryResponse[] - readonly sources: readonly AnswerQuestionInput["sources"][number][] }): RetrievalResult[] { const curated = mapManifestCitationsToResults(input.generatedAnswer) if (curated.length > 0) return curated @@ -873,7 +869,7 @@ function selectCitationRawResults(input: { input.generatedAnswer, ) if (displayedArtifacts.length > 0) return displayedArtifacts - return collectRetrievalResults(input.retrievalResponses, input.sources) + return [] } function mapManifestCitationsToResults( @@ -1008,49 +1004,6 @@ function hasDisplayedManifestArtifacts(result: HarnessRunResult): boolean { return result.manifest.artifacts.some((artifact) => artifact.display) } -function collectRetrievalResults( - responses: readonly RetrievalQueryResponse[], - sources: readonly AnswerQuestionInput["sources"][number][], -): RetrievalResult[] { - const results: RetrievalResult[] = [] - const seenKeys = new Set() - const sourceTitlesByDocumentId = new Map( - sources.flatMap((source): readonly [string, string][] => - source.knowhereDocumentId ? [[source.knowhereDocumentId, source.title]] : [], - ), - ) - - for (const response of responses) { - for (const result of [ - ...response.results, - ...response.referencedChunks.map((chunk): RetrievalResult => ({ - chunkId: chunk.chunkId, - content: "", - chunkType: chunk.chunkType, - score: null, - ...(chunk.assetUrl ? { assetUrl: chunk.assetUrl } : {}), - ...(chunk.sourceChunkPath ? { sourceChunkPath: chunk.sourceChunkPath } : {}), - ...(chunk.filePath ? { filePath: chunk.filePath } : {}), - ...(chunk.metadata ? { metadata: chunk.metadata } : {}), - source: { - documentId: chunk.documentId, - sourceFileName: sourceTitlesByDocumentId.get(chunk.documentId), - sectionPath: chunk.sectionPath, - }, - })), - ]) { - const key = getRetrievalResultKey(result) - if (seenKeys.has(key)) continue - - seenKeys.add(key) - results.push(result) - if (results.length >= MAX_CITATION_RESULTS) return results - } - } - - return results -} - function formatRetrievalEvidenceText( responses: readonly RetrievalQueryResponse[], ): string | undefined { diff --git a/src/domains/chat/service.test.ts b/src/domains/chat/service.test.ts index 441a12c..3a1ab82 100644 --- a/src/domains/chat/service.test.ts +++ b/src/domains/chat/service.test.ts @@ -48,7 +48,7 @@ describe("handleChatTurn", () => { { role: "assistant", content: "Grounded answer.", - citations: [makeRetrievalResult()], + citations: undefined, }, ], }); @@ -78,7 +78,7 @@ describe("handleChatTurn", () => { threadId: "thread_1", role: "assistant", content: "Grounded answer.", - citations: [makeRetrievalResult()], + citations: [], artifacts: [], }); }); From e4bc2b6da96a47ed7a6ac47d4149417b8e7a0c29 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Wed, 8 Jul 2026 04:53:31 +0800 Subject: [PATCH 14/14] Trust agent chat output --- src/agent-harness/index.ts | 1 - src/agent-harness/runtime.test.ts | 63 +-- src/agent-harness/runtime.ts | 150 ++------ src/agent-harness/types.ts | 4 +- src/agent-harness/validator.test.ts | 570 ---------------------------- src/agent-harness/validator.ts | 363 ------------------ src/domains/chat/index.test.ts | 32 +- src/domains/chat/index.ts | 22 +- src/domains/chat/prompt.ts | 2 +- 9 files changed, 52 insertions(+), 1155 deletions(-) delete mode 100644 src/agent-harness/validator.test.ts delete mode 100644 src/agent-harness/validator.ts diff --git a/src/agent-harness/index.ts b/src/agent-harness/index.ts index a1feff1..5fbdf7a 100644 --- a/src/agent-harness/index.ts +++ b/src/agent-harness/index.ts @@ -2,4 +2,3 @@ export * from "./ledger" export * from "./knowhere-text" export * from "./runtime" export * from "./types" -export * from "./validator" diff --git a/src/agent-harness/runtime.test.ts b/src/agent-harness/runtime.test.ts index eec6064..04614a7 100644 --- a/src/agent-harness/runtime.test.ts +++ b/src/agent-harness/runtime.test.ts @@ -28,24 +28,23 @@ describe("agent harness runtime", () => { expect(prompt).not.toContain("navigation action") }) - it("tells the agent to cite only refs with matching source metadata", () => { + it("tells the agent citation metadata is optional and ledger-resolved", () => { const prompt = buildHarnessSystemPrompt(makeTurnInput()) expect(prompt).toContain( - "source.documentId, sourceFileName, and sectionPath must match the selected evidence ref exactly", + "Citation label and source metadata are optional", ) expect(prompt).toContain( - "Omit citations when you cannot identify a supporting evidence ref", + "Notebook resolves citation metadata from evidence refs when possible", ) + expect(prompt).not.toContain("must match the selected evidence ref exactly") }) - it("passes only outer retrieval parameters to KNOWHERE after intent and context policy are declared", async () => { + it("passes only outer retrieval parameters to KNOWHERE without planning-tool gating", async () => { const query = vi.fn().mockResolvedValue( makeRetrievalResponse(), ) const state: { - intent?: IntentFrame - contextPolicy?: ContextPolicy toolCalls?: HarnessToolCallTrace[] } = {} const tools = createHarnessTools({ @@ -55,25 +54,6 @@ describe("agent harness runtime", () => { recentTurns: [], }) - expect(await executeTool(tools.knowhere_search, { query: "q4 chart" })) - .toContain('status="error"') - - await executeTool(tools.declareIntent, { - task: "show_media", - dependsOnPreviousTurn: false, - retrievalNeeded: "yes", - targetModalities: ["image"], - constraints: { desiredCount: 2, maxCount: 2 }, - groundingPolicy: "must_use_sources", - }) - expect(await executeTool(tools.knowhere_search, { query: "q4 chart" })) - .toContain("setContextPolicy must be called before knowhere_search.") - - await executeTool(tools.setContextPolicy, { - carryHistory: "none", - reason: "The current request is unrelated to previous turns.", - activePriorTurnIds: [], - }) const result = await executeTool(tools.knowhere_search, { query: "q4 chart", targetContent: "image", @@ -97,10 +77,6 @@ describe("agent harness runtime", () => { "LegalAction", ) expect(state.toolCalls?.map((call) => [call.tool, call.ok])).toEqual([ - ["knowhere_search", false], - ["declareIntent", true], - ["knowhere_search", false], - ["setContextPolicy", true], ["knowhere_search", true], ]) }) @@ -360,10 +336,8 @@ describe("agent harness runtime", () => { }) }) - it("blocks finalize until intent and context policy are declared", async () => { + it("accepts finalize output without planning-tool gating", async () => { const state: { - intent?: IntentFrame - contextPolicy?: ContextPolicy finalizedManifest?: OutputManifest finalized?: boolean } = {} @@ -381,31 +355,6 @@ describe("agent harness runtime", () => { unresolved: [], } - expect(await executeTool(tools.finalize, manifest)).toEqual({ - ok: false, - message: "declareIntent must be called before finalize.", - }) - expect(state.finalizedManifest).toBeUndefined() - - await executeTool(tools.declareIntent, { - task: "answer_question", - dependsOnPreviousTurn: false, - retrievalNeeded: "no", - targetModalities: ["text"], - constraints: {}, - groundingPolicy: "may_use_sources", - }) - expect(await executeTool(tools.finalize, manifest)).toEqual({ - ok: false, - message: "setContextPolicy must be called before finalize.", - }) - expect(state.finalizedManifest).toBeUndefined() - - await executeTool(tools.setContextPolicy, { - carryHistory: "none", - reason: "Self-contained request.", - activePriorTurnIds: [], - }) expect(await executeTool(tools.finalize, manifest)).toMatchObject({ ok: true, text: "Answer.", diff --git a/src/agent-harness/runtime.ts b/src/agent-harness/runtime.ts index e272340..8963fdf 100644 --- a/src/agent-harness/runtime.ts +++ b/src/agent-harness/runtime.ts @@ -25,10 +25,8 @@ import type { KnowhereToolRuntime, OutputManifest, } from "./types" -import { validateOutputManifest } from "./validator" const defaultMaxSteps = 14 -const defaultMaxRevisions = 1 const imageInspectionReminderStepNumber = 12 const forcedFinalizationStepNumber = 13 const imageInspectionRefLimit = 6 @@ -43,11 +41,6 @@ export type RunAgentHarnessInput = { readonly knowhereTools: KnowhereToolRuntime readonly inspectImages?: InspectImages readonly maxSteps?: number - /** - * How many times the agent may revise after a failed validation pass before - * the harness gives up and returns the last manifest with recorded errors. - */ - readonly maxRevisions?: number } type HarnessToolState = { @@ -164,12 +157,14 @@ const contextPolicySchema = z.object({ const outputCitationSchema = z.object({ ref: z.string().min(1), - label: z.string().min(1), - source: z.object({ - documentId: z.string().nullable().optional(), - sourceFileName: z.string().nullable().optional(), - sectionPath: z.string().nullable().optional(), - }), + label: z.string().min(1).optional(), + source: z + .object({ + documentId: z.string().nullable().optional(), + sourceFileName: z.string().nullable().optional(), + sectionPath: z.string().nullable().optional(), + }) + .optional(), }) const selectedOutputArtifactSchema = z.object({ @@ -236,46 +231,11 @@ export async function runAgentHarness( ], }) - const maxRevisions = input.maxRevisions ?? defaultMaxRevisions - let messages = buildHarnessMessages(input.turn) - let manifest = buildFallbackManifest("") - let validationErrors: readonly string[] = [] - let revisionsUsed = 0 - - for (let attempt = 0; ; attempt += 1) { - const response = await agent.generate({ messages }) - manifest = - state.finalizedManifest ?? buildFallbackManifest(response.text.trim()) - - const validation = validateOutputManifest({ - manifest, - intent: state.intent, - contextPolicy: state.contextPolicy, - finalized: state.finalized === true, - ledger: ledger.snapshot(), - toolCalls: state.toolCalls, - surface: input.turn.surface, - }) - validationErrors = validation.errors - - if (validation.ok || attempt >= maxRevisions) break - - // Self-correction (reflexion): continue the same conversation with - // structured validator feedback and require a fresh finalize so the agent - // can repair its own contract violations instead of shipping them. - revisionsUsed += 1 - state.finalizedManifest = undefined - state.finalized = false - messages = [ - ...messages, - ...(response.response.messages as ModelMessage[]), - { - role: "user", - content: buildRevisionFeedback(validation.errors), - }, - ] - } - + const response = await agent.generate({ + messages: buildHarnessMessages(input.turn), + }) + const manifest = + state.finalizedManifest ?? buildFallbackManifest(response.text.trim()) const ledgerSnapshot = ledger.snapshot() return { manifest, @@ -286,8 +246,8 @@ export async function runAgentHarness( finalized: state.finalized === true, priorTurnReads: [...(state.priorTurnReads ?? [])], toolCalls: [...(state.toolCalls ?? [])], - validationErrors, - revisionsUsed, + validationErrors: [], + revisionsUsed: 0, }, } } @@ -395,22 +355,6 @@ type ModelMessageForRole = Extract< { readonly role: TRole } > -function buildRevisionFeedback(errors: readonly string[]): string { - return [ - "Your finalize output did not satisfy the output contract:", - ...errors.map((error) => `- ${error}`), - "", - "Fix every issue and call finalize again with a corrected manifest.", - "You must call finalize; freeform assistant text is not a valid final", - "answer contract.", - "Do not exceed the user's requested artifact count, only cite or display", - "evidence refs that exist in the evidence ledger, and do not fabricate", - "facts when evidence is missing. Citation source.documentId,", - "sourceFileName, and sectionPath must match the selected evidence ref", - "exactly.", - ].join("\n") -} - function buildForcedFinalizationFeedback(): string { return [ "The retrieval step budget has been reached.", @@ -453,7 +397,7 @@ export function createHarnessTools(input: { return { declareIntent: tool({ description: - "Declare the user's intent before any other action. This is working memory, not a final answer.", + "Declare the user's intent when it helps plan the response. This is working memory, not a final answer.", inputSchema: intentFrameSchema, execute: async (intent): Promise => traceToolCall(input.state, { @@ -493,7 +437,6 @@ export function createHarnessTools(input: { inputSummary: summarizeKnowhereSearchRequest(request), execute: async () => executeKnowhereSearch({ - state: input.state, ledger: input.ledger, knowhereTools: input.knowhereTools, request, @@ -513,7 +456,6 @@ export function createHarnessTools(input: { execute: async () => executeKnowhereTextTool({ operation: "list_documents", - state: input.state, execute: async () => knowhereToolText.formatListDocuments( await input.knowhereTools.listDocuments(), @@ -534,7 +476,6 @@ export function createHarnessTools(input: { execute: async () => executeKnowhereTextTool({ operation: "get_document_outline", - state: input.state, validate: () => validateDocumentReference(request), execute: async () => knowhereToolText.formatOutline( @@ -555,7 +496,6 @@ export function createHarnessTools(input: { inputSummary: summarizeReadChunksRequest(request), execute: async () => executeKnowhereReadChunks({ - state: input.state, ledger: input.ledger, knowhereTools: input.knowhereTools, request, @@ -574,7 +514,6 @@ export function createHarnessTools(input: { inputSummary: summarizeGrepChunksRequest(request), execute: async () => executeKnowhereGrepChunks({ - state: input.state, ledger: input.ledger, knowhereTools: input.knowhereTools, request, @@ -669,26 +608,13 @@ export function createHarnessTools(input: { description: "Finalize the user-facing output manifest. This is the only final answer " + "contract. Artifacts listed here with display=true are the exact set of " + - "images/tables shown to the user; cite only refs from the evidence ledger " + - "and copy citation source metadata exactly from the selected ref.", + "images/tables shown to the user; cite evidence refs when available.", inputSchema: outputManifestSchema, execute: async (manifest) => traceToolCall(input.state, { toolName: "finalize", inputSummary: summarizeManifest(manifest), execute: async () => { - if (!input.state.intent) { - return { - ok: false as const, - message: "declareIntent must be called before finalize.", - } - } - if (!input.state.contextPolicy) { - return { - ok: false as const, - message: "setContextPolicy must be called before finalize.", - } - } input.state.finalizedManifest = manifest input.state.finalized = true return { ok: true as const, ...manifest } @@ -878,14 +804,12 @@ type DocumentReferenceSummary = { } async function executeKnowhereSearch(input: { - readonly state: HarnessToolState readonly ledger: ReturnType readonly knowhereTools: KnowhereToolRuntime readonly request: KnowhereSearchToolRequest }): Promise { return executeKnowhereTextTool({ operation: "search", - state: input.state, execute: async () => { const beforeSnapshot = input.ledger.snapshot() const response = await input.knowhereTools.search({ @@ -909,14 +833,12 @@ async function executeKnowhereSearch(input: { } async function executeKnowhereReadChunks(input: { - readonly state: HarnessToolState readonly ledger: ReturnType readonly knowhereTools: KnowhereToolRuntime readonly request: KnowhereReadChunksToolRequest }): Promise { return executeKnowhereTextTool({ operation: "read_chunks", - state: input.state, validate: () => validateDocumentReference(input.request), execute: async () => { const beforeSnapshot = input.ledger.snapshot() @@ -932,14 +854,12 @@ async function executeKnowhereReadChunks(input: { } async function executeKnowhereGrepChunks(input: { - readonly state: HarnessToolState readonly ledger: ReturnType readonly knowhereTools: KnowhereToolRuntime readonly request: KnowhereGrepChunksToolRequest }): Promise { return executeKnowhereTextTool({ operation: "grep_chunks", - state: input.state, validate: () => validateDocumentReference(input.request), execute: async () => { const beforeSnapshot = input.ledger.snapshot() @@ -956,18 +876,9 @@ async function executeKnowhereGrepChunks(input: { async function executeKnowhereTextTool(input: { readonly operation: KnowhereToolOperation - readonly state: HarnessToolState readonly validate?: () => string | null readonly execute: () => Promise }): Promise { - const workflowError = validateKnowhereWorkflow(input.state, input.operation) - if (workflowError) { - return knowhereToolText.formatError({ - operation: input.operation, - message: workflowError, - }) - } - const validationError = input.validate?.() if (validationError) { return knowhereToolText.formatError({ @@ -986,18 +897,6 @@ async function executeKnowhereTextTool(input: { } } -function validateKnowhereWorkflow( - state: HarnessToolState, - operation: KnowhereToolOperation, -): string | null { - const toolName = `knowhere_${operation}` - if (!state.intent) return `declareIntent must be called before ${toolName}.` - if (!state.contextPolicy) { - return `setContextPolicy must be called before ${toolName}.` - } - return null -} - function validateDocumentReference( request: KnowhereDocumentReferenceRequest, ): string | null { @@ -1257,15 +1156,15 @@ export function buildHarnessSystemPrompt(turn: AgentTurnInput): string { "KNOWHERE is only an evidence provider. Do not infer or control its internal navigation algorithm.", "Your job is to understand intent, decide context use, optionally retrieve evidence, select evidence/artifacts, create source-backed derived tables when useful, and finalize an output manifest.", "", - "Required workflow:", - "1. Call declareIntent first. Capture constraints like a requested image/table count in constraints.desiredCount.", - "2. Call setContextPolicy next, deciding how prior turns should influence this turn.", + "Recommended workflow:", + "1. Call declareIntent when it helps you plan the response. Capture constraints like a requested image/table count in constraints.desiredCount.", + "2. Call setContextPolicy when prior turns may influence this turn.", "3. When the policy needs prior-turn detail (references or corrections), call readPriorTurn for the relevant ids.", "4. Call knowhere_search when relevance search is needed. Use knowhere_list_documents, knowhere_get_document_outline, knowhere_read_chunks, and knowhere_grep_chunks for focused document reads.", "5. For pixel-level details, OCR, visual comparison, image verification, or when the likely answer is only visible on a returned page/image asset, call inspectImage only after a Knowhere tool returned image asset refs.", `6. inspectImage accepts at most ${imageInspectionRefLimit} image asset refs per call and per turn.`, "7. knowhere_read_chunks returns complete chunk bodies; control size with page/pageSize, sectionPath, startChunk/endChunk, chunkId, and chunkType.", - "8. Call finalize with text, citations, artifacts, and unresolved issues. finalize requires declareIntent and setContextPolicy first.", + "8. Call finalize with text, citations, artifacts, and unresolved issues when you are ready to answer.", "", "Context rules:", "- If the current user request is unrelated to prior turns, set carryHistory to none and do not reuse prior topics.", @@ -1276,13 +1175,12 @@ export function buildHarnessSystemPrompt(turn: AgentTurnInput): string { "- Final output is the OutputManifest passed to finalize, not freeform tool JSON or trailing text.", "- artifacts with display=true are the exact images/tables shown. Never display every candidate; honor constraints.desiredCount / maxCount.", "- Use type=derived_table only for tables you create from evidence; every derived_table.sourceRefs entry must reference evidence in the ledger.", - "- citations and selected image/table artifact refs may only reference refs returned by Knowhere tools in the evidence ledger.", - "- For every citation, source.documentId, sourceFileName, and sectionPath must match the selected evidence ref exactly. Use null or omit the field only when the ref omits it.", - "- Omit citations when you cannot identify a supporting evidence ref with matching source metadata.", + "- Prefer citation and selected image/table artifact refs returned by Knowhere tools in the evidence ledger.", + "- Citation label and source metadata are optional. Notebook resolves citation metadata from evidence refs when possible.", + "- If evidence is relevant but you cannot identify a supporting evidence ref, answer with unresolved issues instead of fabricating a ref.", "- inspectImage observations are inspection notes, not new source refs. Final citations and displayed image artifacts must use the original retrieved image asset refs.", "- If text evidence identifies a relevant page/image but does not include the exact fact, inspect the returned image asset for OCR/detail before saying the answer is unavailable.", "- If evidence is insufficient, list it in unresolved instead of fabricating facts.", - "- After a validation-feedback message, fix all listed issues and call finalize again.", `Surface: ${turn.surface}`, `Output capabilities: ${JSON.stringify(turn.outputCapabilities)}`, turn.sourceContext ? `Searchable source context:\n${turn.sourceContext}` : "", diff --git a/src/agent-harness/types.ts b/src/agent-harness/types.ts index 83bb08d..36b6a27 100644 --- a/src/agent-harness/types.ts +++ b/src/agent-harness/types.ts @@ -209,8 +209,8 @@ export type EvidenceLedgerSnapshot = { export type OutputCitation = { readonly ref: string - readonly label: string - readonly source: EvidenceChunk["source"] + readonly label?: string + readonly source?: EvidenceChunk["source"] } export type OutputArtifact = { diff --git a/src/agent-harness/validator.test.ts b/src/agent-harness/validator.test.ts deleted file mode 100644 index 89bf0ba..0000000 --- a/src/agent-harness/validator.test.ts +++ /dev/null @@ -1,570 +0,0 @@ -import { describe, expect, it } from "vitest" - -import { validateOutputManifest } from "./validator" -import type { - ContextPolicy, - EvidenceLedgerSnapshot, - IntentFrame, - OutputManifest, -} from "./types" - -describe("validateOutputManifest", () => { - it("requires finalize to be the successful output path", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ text: "Freeform answer." }), - intent: makeIntent({ groundingPolicy: "no_retrieval" }), - contextPolicy: unrelatedContextPolicy, - finalized: false, - ledger: emptyLedger, - surface: "notebook_chat", - }) - - expect(validation.errors).toContain( - "Agent must call finalize to produce the output manifest.", - ) - }) - - it("requires the agent to declare intent and context policy before finalizing", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ text: "Answer." }), - ledger: emptyLedger, - surface: "notebook_chat", - }) - - expect(validation.errors).toContain( - "Agent must declare intent before finalizing.", - ) - expect(validation.errors).toContain( - "Agent must set context policy before finalizing.", - ) - }) - - it("limits displayed artifacts using the declared intent instead of hard-coded media rules", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ - artifacts: [ - { - type: "image", - ref: "asset:r1:result:1", - display: true, - reason: "front", - }, - { - type: "image", - ref: "asset:r1:result:2", - display: true, - reason: "back", - }, - { - type: "image", - ref: "asset:r1:result:3", - display: true, - reason: "extra candidate", - }, - ], - }), - intent: makeIntent({ desiredCount: 2, maxCount: 2 }), - contextPolicy: unrelatedContextPolicy, - ledger: { - ...emptyLedger, - assets: [ - makeAsset("asset:r1:result:1"), - makeAsset("asset:r1:result:2"), - makeAsset("asset:r1:result:3"), - ], - }, - surface: "notebook_chat", - }) - - expect(validation.errors).toContain( - "Displayed artifact count 3 exceeds desired count 2.", - ) - expect(validation.errors).toContain( - "Displayed artifact count 3 exceeds maximum count 2.", - ) - }) - - it("rejects grounded answers that use evidence without citations or selected artifacts", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ text: "Revenue increased." }), - intent: makeIntent({}), - contextPolicy: unrelatedContextPolicy, - ledger: { - ...emptyLedger, - chunks: [ - { - ref: "r1:result:1", - kind: "result", - content: "Revenue increased.", - contentPreview: "Revenue increased.", - chunkType: "text", - score: 0.9, - source: { - documentId: "doc_1", - sourceFileName: "report.pdf", - sectionPath: "Q4", - }, - }, - ], - }, - surface: "notebook_chat", - }) - - expect(validation.errors).toContain( - "Grounded output used evidence but did not cite or display any selected evidence.", - ) - }) - - it("rejects citation source metadata that conflicts with the resolved evidence ref", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ - text: "The source discusses information hiding.", - citations: [ - { - ref: "r1:result:1", - label: "wrong source", - source: { - documentId: "doc_claimed", - sourceFileName: "claimed-source.pdf", - sectionPath: "Claimed Section", - }, - }, - ], - }), - intent: makeIntent({}), - contextPolicy: unrelatedContextPolicy, - ledger: { - ...emptyLedger, - chunks: [ - { - ref: "r1:result:1", - kind: "result", - content: "Information hiding is a module design principle.", - contentPreview: "Information hiding is a module design principle.", - chunkType: "text", - score: 0.9, - source: { - documentId: "doc_information_hiding", - sourceFileName: "information_hiding.pdf", - sectionPath: "Root / Module Design", - }, - }, - ], - }, - surface: "notebook_chat", - }) - - expect(validation.errors).toContain( - "Citation ref 'r1:result:1' source.documentId must match resolved evidence source. Expected 'doc_information_hiding', received 'doc_claimed'.", - ) - expect(validation.errors).toContain( - "Citation ref 'r1:result:1' source.sourceFileName must match resolved evidence source. Expected 'information_hiding.pdf', received 'claimed-source.pdf'.", - ) - expect(validation.errors).toContain( - "Citation ref 'r1:result:1' source.sectionPath must match resolved evidence source. Expected 'Root / Module Design', received 'Claimed Section'.", - ) - }) - - it("accepts citation source metadata that exactly matches the resolved evidence ref", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ - text: "Revenue increased.", - citations: [ - { - ref: "r1:result:1", - label: "report.pdf / Q4", - source: { - documentId: "doc_1", - sourceFileName: "report.pdf", - sectionPath: "Q4", - }, - }, - ], - }), - intent: makeIntent({}), - contextPolicy: unrelatedContextPolicy, - ledger: { - ...emptyLedger, - chunks: [ - { - ref: "r1:result:1", - kind: "result", - content: "Revenue increased.", - contentPreview: "Revenue increased.", - chunkType: "text", - score: 0.9, - source: { - documentId: "doc_1", - sourceFileName: "report.pdf", - sectionPath: "Q4", - }, - }, - ], - }, - surface: "notebook_chat", - }) - - expect(validation.ok).toBe(true) - expect(validation.errors).toEqual([]) - }) - - it("rejects missing citation documentId when the resolved evidence has one", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ - text: "Revenue increased.", - citations: [ - { - ref: "r1:result:1", - label: "report.pdf / Q4", - source: { - sourceFileName: "report.pdf", - sectionPath: "Q4", - }, - }, - ], - }), - intent: makeIntent({}), - contextPolicy: unrelatedContextPolicy, - ledger: { - ...emptyLedger, - chunks: [ - { - ref: "r1:result:1", - kind: "result", - content: "Revenue increased.", - contentPreview: "Revenue increased.", - chunkType: "text", - score: 0.9, - source: { - documentId: "doc_1", - sourceFileName: "report.pdf", - sectionPath: "Q4", - }, - }, - ], - }, - surface: "notebook_chat", - }) - - expect(validation.errors).toContain( - "Citation ref 'r1:result:1' source.documentId must match resolved evidence source. Expected 'doc_1', received missing.", - ) - }) - - it("accepts source-backed derived tables and rejects missing source refs", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ - artifacts: [ - { - type: "derived_table", - ref: "derived:table:1", - title: "Revenue comparison", - columns: ["Metric", "Value"], - rows: [["Revenue", "$10M"]], - sourceRefs: ["r1:result:1"], - display: true, - reason: "Structured comparison requested by the user.", - }, - { - type: "derived_table", - ref: "derived:table:2", - title: "Invalid table", - columns: ["Metric", "Value"], - rows: [["Revenue"]], - sourceRefs: ["missing"], - display: true, - reason: "Demonstrates validation.", - }, - ], - }), - intent: makeIntent({}), - contextPolicy: unrelatedContextPolicy, - ledger: { - ...emptyLedger, - chunks: [ - { - ref: "r1:result:1", - kind: "result", - content: "Revenue was $10M.", - contentPreview: "Revenue was $10M.", - chunkType: "text", - score: 0.9, - source: { - documentId: "doc_1", - sourceFileName: "report.pdf", - sectionPath: "Revenue", - }, - }, - ], - }, - surface: "notebook_chat", - }) - - expect(validation.errors).toContain( - "Derived table source ref 'missing' was not found in the evidence ledger.", - ) - expect(validation.errors).toContain( - "Derived table row 1 has 1 cells but expected 2.", - ) - }) - - it("requires compare outputs to cite at least two evidence refs", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ - text: "A is stronger than B.", - citations: [ - { - ref: "r1:result:1", - label: "report.pdf / A", - source: { - documentId: "doc_1", - sourceFileName: "report.pdf", - sectionPath: "A", - }, - }, - ], - }), - intent: { - ...makeIntent({}), - task: "compare", - }, - contextPolicy: unrelatedContextPolicy, - ledger: { - ...emptyLedger, - chunks: [ - { - ref: "r1:result:1", - kind: "result", - content: "A is strong.", - contentPreview: "A is strong.", - chunkType: "text", - score: 0.9, - source: { - documentId: "doc_1", - sourceFileName: "report.pdf", - sectionPath: "A", - }, - }, - ], - }, - surface: "notebook_chat", - }) - - expect(validation.errors).toContain( - "Compare outputs that must use sources require at least two evidence refs or an explicit unresolved reason.", - ) - }) - - it("keeps typing compose output insertion-ready", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ text: "- bullet\n- list" }), - intent: makeIntent({ groundingPolicy: "no_retrieval" }), - contextPolicy: unrelatedContextPolicy, - ledger: emptyLedger, - surface: "typing_compose", - }) - - expect(validation.errors).toContain( - "Typing compose output must be insertion-ready plain text.", - ) - }) - - it("rejects image inspection failure claims when inspectImage was not called", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ - text: "图像检查工具未能成功读取该页面的具体条款内容。", - citations: [ - { - ref: "r1:referenced:1", - label: "contract.pdf / page 8", - source: { - documentId: "doc_contract", - sourceFileName: "contract.pdf", - sectionPath: "Root / （6）现场工期进度管理方面的违约责任", - }, - }, - ], - }), - intent: makeIntent({}), - contextPolicy: unrelatedContextPolicy, - ledger: { - ...emptyLedger, - chunks: [ - { - ref: "r1:referenced:1", - kind: "referenced_chunk", - content: "", - contentPreview: "", - chunkType: "page", - score: null, - source: { - documentId: "doc_contract", - sourceFileName: "contract.pdf", - sectionPath: "Root / （6）现场工期进度管理方面的违约责任", - }, - assetRef: "asset:r1:referenced:1", - }, - ], - assets: [makeAsset("asset:r1:referenced:1")], - }, - toolCalls: [], - surface: "notebook_chat", - }) - - expect(validation.errors).toContain( - "Final output must not claim image/OCR inspection succeeded or failed unless inspectImage was called.", - ) - }) - - it("allows image inspection failure claims after inspectImage was called", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ - text: "对该页面的图像识别未成功提取具体数值。", - citations: [ - { - ref: "r1:referenced:1", - label: "contract.pdf / page 8", - source: { - documentId: "doc_contract", - sourceFileName: "contract.pdf", - sectionPath: "Root / （6）现场工期进度管理方面的违约责任", - }, - }, - ], - }), - intent: makeIntent({}), - contextPolicy: unrelatedContextPolicy, - ledger: { - ...emptyLedger, - chunks: [ - { - ref: "r1:referenced:1", - kind: "referenced_chunk", - content: "", - contentPreview: "", - chunkType: "page", - score: null, - source: { - documentId: "doc_contract", - sourceFileName: "contract.pdf", - sectionPath: "Root / （6）现场工期进度管理方面的违约责任", - }, - assetRef: "asset:r1:referenced:1", - }, - ], - assets: [makeAsset("asset:r1:referenced:1")], - }, - toolCalls: [ - { - tool: "inspectImage", - ok: false, - inputSummary: { refs: ["asset:r1:referenced:1"] }, - outputSummary: { inspectedCount: 0 }, - startedAt: "2026-07-06T00:00:00.000Z", - durationMs: 1, - }, - ], - surface: "notebook_chat", - }) - - expect(validation.errors).not.toContain( - "Final output must not claim image/OCR inspection succeeded or failed unless inspectImage was called.", - ) - }) - - it("rejects unreadable page claims when image assets are available but uninspected", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ - text: "相关的页面图像无法通过目前的检测工具获取。", - unresolved: ["由于无法直接读取该页面的详细条款内容，目前无法给出确切的赔偿金额。"], - }), - intent: makeIntent({}), - contextPolicy: unrelatedContextPolicy, - ledger: { - ...emptyLedger, - chunks: [ - { - ref: "r1:referenced:1", - kind: "referenced_chunk", - content: "", - contentPreview: "", - chunkType: "page", - score: null, - source: { - documentId: "doc_contract", - sourceFileName: "contract.pdf", - sectionPath: "Root / （6）现场工期进度管理方面的违约责任", - }, - assetRef: "asset:r1:referenced:1", - }, - ], - assets: [makeAsset("asset:r1:referenced:1")], - }, - toolCalls: [], - surface: "notebook_chat", - }) - - expect(validation.errors).toContain( - "Final output must inspect available image assets before claiming retrieved page/image content cannot be read.", - ) - expect(validation.errors).toContain( - "Final output must not claim image/OCR inspection succeeded or failed unless inspectImage was called.", - ) - }) -}) - -const emptyLedger: EvidenceLedgerSnapshot = { - retrievalCount: 0, - chunks: [], - assets: [], - evidenceText: [], - stopReasons: [], - failureReasons: [], - decisionTraces: [], -} - -const unrelatedContextPolicy: ContextPolicy = { - carryHistory: "none", - reason: "The current turn is self-contained.", - activePriorTurnIds: [], -} - -function makeIntent( - constraints: IntentFrame["constraints"] & { - readonly groundingPolicy?: IntentFrame["groundingPolicy"] - }, -): IntentFrame { - return { - task: "answer", - dependsOnPreviousTurn: false, - retrievalNeeded: constraints.groundingPolicy === "no_retrieval" ? "no" : "yes", - targetModalities: ["text"], - constraints, - groundingPolicy: constraints.groundingPolicy ?? "must_use_sources", - } -} - -function makeManifest(overrides: Partial): OutputManifest { - return { - text: "", - citations: [], - artifacts: [], - unresolved: [], - ...overrides, - } -} - -function makeAsset(ref: string): EvidenceLedgerSnapshot["assets"][number] { - return { - ref, - chunkRef: ref.replace("asset:", ""), - type: "image", - assetUrl: `https://assets.example/${ref}.png`, - label: ref, - source: { - documentId: "doc_1", - sourceFileName: "report.pdf", - sectionPath: ref, - }, - } -} diff --git a/src/agent-harness/validator.ts b/src/agent-harness/validator.ts deleted file mode 100644 index fcd8c0e..0000000 --- a/src/agent-harness/validator.ts +++ /dev/null @@ -1,363 +0,0 @@ -import type { - ContextPolicy, - EvidenceLedgerSnapshot, - HarnessToolCallTrace, - IntentFrame, - OutputCitation, - OutputManifest, -} from "./types" - -export type ManifestValidationInput = { - readonly manifest: OutputManifest - readonly intent?: IntentFrame - readonly contextPolicy?: ContextPolicy - readonly finalized?: boolean - readonly ledger: EvidenceLedgerSnapshot - readonly toolCalls?: readonly HarnessToolCallTrace[] - readonly surface: "notebook_chat" | "typing_compose" | "typing_quick_ask" -} - -export type ManifestValidationResult = { - readonly ok: boolean - readonly errors: readonly string[] -} - -const citationSourceFields = [ - "documentId", - "sourceFileName", - "sectionPath", -] as const - -type CitationSourceField = (typeof citationSourceFields)[number] -type EvidenceSource = EvidenceLedgerSnapshot["chunks"][number]["source"] - -export function validateOutputManifest( - input: ManifestValidationInput, -): ManifestValidationResult { - const errors: string[] = [] - const text = input.manifest.text.trim() - - if (!text && input.manifest.artifacts.every((artifact) => !artifact.display)) { - errors.push("Final output must contain text or at least one displayed artifact.") - } - - validateWorkflow(input, errors) - validateArtifactRefs(input, errors) - validateArtifactCounts(input, errors) - validateGrounding(input, errors) - validateTaskEvidence(input, errors) - validateImageInspectionClaims(input, errors) - validateUnavailableImageContentClaims(input, errors) - validateTypingText(input, errors) - - return { - ok: errors.length === 0, - errors, - } -} - -function validateWorkflow( - input: ManifestValidationInput, - errors: string[], -): void { - if (input.finalized === false) { - errors.push("Agent must call finalize to produce the output manifest.") - } - - if (!input.intent) { - errors.push("Agent must declare intent before finalizing.") - } - - if (!input.contextPolicy) { - errors.push("Agent must set context policy before finalizing.") - } -} - -function validateArtifactRefs( - input: ManifestValidationInput, - errors: string[], -): void { - const knownRefs = new Set([ - ...input.ledger.chunks.map((chunk) => chunk.ref), - ...input.ledger.assets.map((asset) => asset.ref), - ]) - const sourcesByRef = new Map([ - ...input.ledger.chunks.map( - (chunk): readonly [string, EvidenceSource] => [chunk.ref, chunk.source], - ), - ...input.ledger.assets.map( - (asset): readonly [string, EvidenceSource] => [asset.ref, asset.source], - ), - ]) - - for (const artifact of input.manifest.artifacts) { - if (artifact.type === "derived_table") { - artifact.rows.forEach((row, index) => { - if (row.length !== artifact.columns.length) { - errors.push( - `Derived table row ${index + 1} has ${row.length} cells but expected ${artifact.columns.length}.`, - ) - } - }) - - for (const ref of artifact.sourceRefs) { - if (!knownRefs.has(ref)) { - errors.push( - `Derived table source ref '${ref}' was not found in the evidence ledger.`, - ) - } - } - continue - } - - if (!knownRefs.has(artifact.ref)) { - errors.push( - `Artifact ref '${artifact.ref}' was not found in the evidence ledger.`, - ) - } - } - - for (const citation of input.manifest.citations) { - if (!knownRefs.has(citation.ref)) { - errors.push(`Citation ref '${citation.ref}' was not found in the evidence ledger.`) - continue - } - - const source = sourcesByRef.get(citation.ref) - if (!source) continue - validateCitationSource( - { - ref: citation.ref, - declared: citation.source, - resolved: source, - }, - errors, - ) - } -} - -function validateCitationSource( - input: { - readonly ref: string - readonly declared: OutputCitation["source"] - readonly resolved: EvidenceSource - }, - errors: string[], -): void { - for (const field of citationSourceFields) { - validateCitationSourceField( - { - ref: input.ref, - field, - declaredValue: input.declared[field], - resolvedValue: input.resolved[field], - }, - errors, - ) - } -} - -function validateCitationSourceField( - input: { - readonly ref: string - readonly field: CitationSourceField - readonly declaredValue: string | null | undefined - readonly resolvedValue: string | null | undefined - }, - errors: string[], -): void { - const declaredValue = normalizeSourceValue(input.declaredValue) - const resolvedValue = normalizeSourceValue(input.resolvedValue) - if (declaredValue === resolvedValue) return - - errors.push( - `Citation ref '${input.ref}' source.${input.field} must match resolved evidence source. Expected ${formatSourceValue( - resolvedValue, - )}, received ${formatSourceValue(declaredValue)}.`, - ) -} - -function normalizeSourceValue(value: string | null | undefined): string | null { - return value ?? null -} - -function formatSourceValue(value: string | null): string { - return value === null ? "missing" : `'${value}'` -} - -function validateArtifactCounts( - input: ManifestValidationInput, - errors: string[], -): void { - const displayedCount = input.manifest.artifacts.filter( - (artifact) => artifact.display, - ).length - const desiredCount = input.intent?.constraints.desiredCount - const maxCount = input.intent?.constraints.maxCount - - if (typeof desiredCount === "number" && displayedCount > desiredCount) { - errors.push( - `Displayed artifact count ${displayedCount} exceeds desired count ${desiredCount}.`, - ) - } - - if (typeof maxCount === "number" && displayedCount > maxCount) { - errors.push( - `Displayed artifact count ${displayedCount} exceeds maximum count ${maxCount}.`, - ) - } -} - -function validateGrounding( - input: ManifestValidationInput, - errors: string[], -): void { - if (input.intent?.groundingPolicy !== "must_use_sources") return - - const hasLedgerEvidence = - input.ledger.chunks.length > 0 || input.ledger.evidenceText.length > 0 - const hasOutputEvidence = - input.manifest.citations.length > 0 || - input.manifest.artifacts.some((artifact) => artifact.display) - const hasUnresolved = input.manifest.unresolved.length > 0 - - if (!hasLedgerEvidence && !hasUnresolved) { - errors.push( - "Grounded output requires evidence or an explicit unresolved reason.", - ) - } - - if (hasLedgerEvidence && !hasOutputEvidence && !hasUnresolved) { - errors.push( - "Grounded output used evidence but did not cite or display any selected evidence.", - ) - } -} - -function validateTaskEvidence( - input: ManifestValidationInput, - errors: string[], -): void { - if (input.intent?.groundingPolicy !== "must_use_sources") return - if (input.manifest.unresolved.length > 0) return - - const refs = getOutputEvidenceRefs(input.manifest) - - if (input.intent.task === "compare" && refs.size < 2) { - errors.push( - "Compare outputs that must use sources require at least two evidence refs or an explicit unresolved reason.", - ) - } - - if (input.intent.task === "summarize" && refs.size < 1) { - errors.push( - "Summaries that must use sources require at least one evidence ref or an explicit unresolved reason.", - ) - } -} - -function getOutputEvidenceRefs(manifest: OutputManifest): Set { - const refs = new Set() - - for (const citation of manifest.citations) refs.add(citation.ref) - - for (const artifact of manifest.artifacts) { - if (!artifact.display) continue - if (artifact.type === "derived_table") { - artifact.sourceRefs.forEach((ref) => refs.add(ref)) - } else { - refs.add(artifact.ref) - } - } - - return refs -} - -function validateTypingText( - input: ManifestValidationInput, - errors: string[], -): void { - if (input.surface !== "typing_compose") return - - const text = input.manifest.text - if (/```|^\s*#{1,6}\s|^\s*[-*]\s/m.test(text)) { - errors.push("Typing compose output must be insertion-ready plain text.") - } -} - -function validateImageInspectionClaims( - input: ManifestValidationInput, - errors: string[], -): void { - const claimText = [input.manifest.text, ...input.manifest.unresolved].join("\n") - if (!mentionsImageInspectionResult(claimText)) return - - if (hasImageInspectionToolCall(input)) return - - errors.push( - "Final output must not claim image/OCR inspection succeeded or failed unless inspectImage was called.", - ) -} - -function validateUnavailableImageContentClaims( - input: ManifestValidationInput, - errors: string[], -): void { - const hasImageAssets = input.ledger.assets.some((asset) => asset.type === "image") - if (!hasImageAssets) return - if (hasImageInspectionToolCall(input)) return - - const claimText = [input.manifest.text, ...input.manifest.unresolved].join("\n") - if (!mentionsUnavailableImageContent(claimText)) return - - errors.push( - "Final output must inspect available image assets before claiming retrieved page/image content cannot be read.", - ) -} - -function hasImageInspectionToolCall(input: ManifestValidationInput): boolean { - return input.toolCalls?.some((call) => call.tool === "inspectImage") === true -} - -function mentionsImageInspectionResult(value: string): boolean { - const normalized = value.replace(/\s+/g, " ").trim() - if (!normalized) return false - - return ( - /\b(?:image|visual)\s+(?:inspection|recognition|analysis)\s+(?:failed|did not|could not|was unable|found|showed|confirmed)/iu.test( - normalized, - ) || - /\bOCR\s+(?:failed|did not|could not|was unable|found|showed|confirmed|extracted)/iu.test( - normalized, - ) || - /(?:图像|图片|视觉).{0,8}(?:识别|检查|检视|查看|分析).{0,12}(?:未|没|无法|不能|不成功|失败|成功|显示|发现|提取|读取)/u.test( - normalized, - ) || - /(?:图像|图片|视觉).{0,12}(?:未|没|无法|不能|不成功|失败).{0,16}(?:识别|检查|检视|查看|分析|检测|提取|读取|获取)/u.test( - normalized, - ) || - /(?:未|没|无法|不能|不成功|失败).{0,12}(?:图像|图片|视觉|OCR).{0,12}(?:识别|检查|检视|查看|分析|检测|提取|读取|获取)/u.test( - normalized, - ) - ) -} - -function mentionsUnavailableImageContent(value: string): boolean { - const normalized = value.replace(/\s+/g, " ").trim() - if (!normalized) return false - - return ( - /\b(?:cannot|can't|could not|unable to|was unable to|did not)\s+(?:directly\s+)?(?:read|inspect|access|extract|see|view)\b.{0,48}\b(?:page|image|visual|OCR|content|clause|details|text)\b/iu.test( - normalized, - ) || - /\b(?:page|image|visual|OCR|content|clause|details|text)\b.{0,48}\b(?:cannot|can't|could not|unable to|was unable to|did not)\s+(?:directly\s+)?(?:read|inspect|access|extract|see|view)\b/iu.test( - normalized, - ) || - /(?:无法|不能|未能|没法|没有办法).{0,8}(?:直接)?(?:读取|查看|识别|检测|提取|看清|访问|获取).{0,16}(?:页面|页|图片|图像|条款|内容|细节|文字)/u.test( - normalized, - ) || - /(?:页面|页|图片|图像|条款|内容|细节|文字).{0,16}(?:无法|不能|未能|没法|没有办法).{0,16}(?:直接)?(?:读取|查看|识别|检测|提取|看清|访问|获取)/u.test( - normalized, - ) - ) -} diff --git a/src/domains/chat/index.test.ts b/src/domains/chat/index.test.ts index 4281520..697fe77 100644 --- a/src/domains/chat/index.test.ts +++ b/src/domains/chat/index.test.ts @@ -1446,7 +1446,7 @@ describe("answerQuestionWithRetrieval", () => { ]); }); - it("returns a safe fallback when the harness still has validation errors", async () => { + it("returns agent output when a legacy harness trace has validation errors", async () => { const retrieval = { query: vi.fn().mockResolvedValue({ results: [makeRetrievalResult()], @@ -1485,14 +1485,13 @@ describe("answerQuestionWithRetrieval", () => { ); expect(answer).toEqual({ - answer: - "I couldn't safely finish that response because the agent output did not pass Notebook's validation checks. Please try again.", + answer: "This invalid answer should not ship.", citations: [], artifacts: [], }); }); - it("returns a safe fallback when a manifest citation ref declares the wrong source", async () => { + it("renders answer and resolves citations when manifest source metadata is wrong", async () => { process.env.AI_GATEWAY_API_KEY = "test_gateway_key"; const result = makeRetrievalResult({ content: "Information hiding is a module design principle.", @@ -1580,13 +1579,16 @@ describe("answerQuestionWithRetrieval", () => { }), ); - expect(generateCallCount).toBe(2); - expect(answer).toEqual({ - answer: - "I couldn't safely finish that response because the agent output did not pass Notebook's validation checks. Please try again.", - citations: [], - artifacts: [], - }); + expect(generateCallCount).toBe(1); + expect(answer.answer).toBe("Information hiding is a module design principle."); + expect(answer.artifacts).toEqual([]); + expect(answer.citations.map((citation) => citation.source)).toEqual([ + { + documentId: "doc_information_hiding", + sourceFileName: "information_hiding.pdf", + sectionPath: "Root / Module Design", + }, + ]); }); it("keeps image-only harness output instead of treating it as no results", async () => { @@ -2570,7 +2572,7 @@ describe("generateAgenticOutputManifest", () => { expect(result.trace.validationErrors).toEqual([]); }); - it("self-corrects an over-budget manifest via a validation-feedback revision", async () => { + it("keeps an over-budget manifest after the first generation", async () => { process.env.AI_GATEWAY_API_KEY = "test_gateway_key"; let generateCallCount = 0; vi.spyOn(ToolLoopAgent.prototype, "generate").mockImplementation( @@ -2687,12 +2689,12 @@ describe("generateAgenticOutputManifest", () => { searchSources, }); - expect(generateCallCount).toBe(2); - expect(result.trace.revisionsUsed).toBe(1); + expect(generateCallCount).toBe(1); + expect(result.trace.revisionsUsed).toBe(0); expect(result.trace.validationErrors).toEqual([]); expect( result.manifest.artifacts.filter((artifact) => artifact.display).length, - ).toBe(2); + ).toBe(3); }); }); diff --git a/src/domains/chat/index.ts b/src/domains/chat/index.ts index d5cd82c..d4dbcc5 100644 --- a/src/domains/chat/index.ts +++ b/src/domains/chat/index.ts @@ -16,7 +16,6 @@ import type { EvidenceChunk, HarnessRunResult, OutputArtifact, - OutputCitation, } from "@/agent-harness" import { toChatCitationViews, @@ -55,8 +54,6 @@ const KNOWHERE_RESPONSE_TEXT_LOG_LIMIT = 200 const KNOWHERE_CHUNK_LOG_LIMIT = 100 const KNOWHERE_RESPONSE_LOG_ITEM_LIMIT = 20 const NO_RESULTS_ANSWER = "I couldn't find that in your sources." -const HARNESS_VALIDATION_FAILURE_ANSWER = - "I couldn't safely finish that response because the agent output did not pass Notebook's validation checks. Please try again." const RAW_URL_PATTERN = /https?:\/\/[^\s)\]}>"']+/g const REDACTED_MEDIA_URL = "[media asset URL hidden]" const RETRIEVAL_TARGET_CONTENT_DATA_TYPES: Readonly< @@ -234,23 +231,8 @@ export const answerQuestionWithRetrieval = ( answerLength: generatedAnswer.manifest.text.length, retrievalCallCount: retrievalResponses.length, citationCount: generatedAnswer.manifest.citations.length, - harnessValidationErrorCount: generatedAnswer.trace.validationErrors.length, - revisionsUsed: generatedAnswer.trace.revisionsUsed, + finalized: generatedAnswer.trace.finalized, }) - if (generatedAnswer.trace.validationErrors.length > 0) { - logger.warn("chat-agent: validation failed; returning safe fallback", { - validationErrors: generatedAnswer.trace.validationErrors, - revisionsUsed: generatedAnswer.trace.revisionsUsed, - finalized: generatedAnswer.trace.finalized, - intentTask: generatedAnswer.trace.intent?.task ?? null, - retrievalCallCount: retrievalResponses.length, - }) - return { - answer: HARNESS_VALIDATION_FAILURE_ANSWER, - citations: [] as ChatCitationView[], - artifacts: [] as ChatArtifactView[], - } - } const rawResults = selectCitationRawResults({ generatedAnswer, @@ -437,7 +419,7 @@ function toChatArtifactView(input: { } function normalizeHarnessSource( - source: OutputCitation["source"], + source: EvidenceChunk["source"], sources: readonly AnswerQuestionInput["sources"][number][], ): ChatCitationView["source"] { const sourceTitle = source.documentId diff --git a/src/domains/chat/prompt.ts b/src/domains/chat/prompt.ts index f424252..a4ae326 100644 --- a/src/domains/chat/prompt.ts +++ b/src/domains/chat/prompt.ts @@ -74,7 +74,7 @@ export const generateAgenticOutputManifestEffect = ( citationCount: result.manifest.citations.length, artifactCount: result.manifest.artifacts.length, unresolvedCount: result.manifest.unresolved.length, - validationErrorCount: result.trace.validationErrors.length, + finalized: result.trace.finalized, intentTask: result.trace.intent?.task ?? null, carryHistory: result.trace.contextPolicy?.carryHistory ?? null, })