diff --git a/docs/parsed-results-citation-ux.md b/docs/parsed-results-citation-ux.md new file mode 100644 index 00000000..4730107c --- /dev/null +++ b/docs/parsed-results-citation-ux.md @@ -0,0 +1,72 @@ +# Parsed Results + citation UX + +Recorded 2026-08-18 so we do not relitigate the Figma / issue discussion. + +Sources: + +- Figma [Knowhere (Copy) / brain](https://www.figma.com/design/0jjj7SQ1N0zo4D9WzJoNJh/Knowhere--Copy-?node-id=4000-37742) (`4000:37742`) +- GitHub [knowhere#222](https://github.com/Ontos-AI/knowhere/issues/222), [knowhere#223](https://github.com/Ontos-AI/knowhere/issues/223) +- Product lock from that conversation (issues and Figma notes are partly stale) + +This work lands in **Knowhere Notebook**. Region geometry later lands in **Knowhere retrieval** (`Ontos-AI/knowhere`) and is only consumed here when the payload exists. + +## Locked product + +Vision / page-image mode only. There is **no original PDF preview** as a citation destination. + +| Topic | Decision | +| --- | --- | +| Layout | Sources \| Parsed Results \| Assistant (today’s Notebook columns). Ignore the Figma frame that put Assistant in the middle. | +| Rename | Knowhere Notebook → **Knowhere Brain** | +| Canvas | Title **Parsed Results**, subtitle **From {source title}** | +| Page card | **Page N + parse path**, page image, keywords. No keywords → hide the keyword row, still show the image. Screen 2 “text” is a crop overlay on the page image, not a text chunk. | +| List / Tree | Keep **both**. Tree is the **existing** section tree, also in page-image mode. **Default Tree**. Citation click may switch to List so the page card is visible. | +| Inline chips | One chip **per citation**, not per page. Label `{source title}/pN` (same name as the Sources row). Max width 250px, truncate. Same page can be two chips; later they get different crop boxes. | +| Chip / `pN` click | Switch the **viewed** source if needed (not query checkboxes), open that file’s parsed page, scroll to it. Chat may cite any workspace file while the canvas shows another. | +| Footer SOURCES | Numbered list, one row per file, `{title}` + cited page links `p25 p26 p27`. Those links do the **same jump** as chips. | +| Copy / download | Copy answer / code / tables. Download **this assistant turn** as **Markdown and PDF**, **without citations** (no chips, no SOURCES list). | +| Region box | **Required product**, **not this slice**. When retrieval sends a crop, draw `#8E51FF` / 25% / 8px radius on that page image (flash twice, 3s, fade). Until then, page jump only. Do not fake a full-page box. | + +Out of this slice: original-file viewer for citations, region overlay, faking geometry. + +## Current code (gaps) + +Almost all of this is UI in this repo. Page images, keywords, page numbers, and source-title chips already exist in pieces. + +| Area | Today | Need | +| --- | --- | --- | +| Branding | “Knowhere Notebook” in `top-nav.tsx`, `lib/app-metadata.ts`, login, `global-error.tsx` | Brain | +| Canvas copy | “Parsed Chunks” / “Showing all parsed chunks from …” | Parsed Results / From {title} | +| Original preview | Hidden for `page-assets` sources (`chunks-panel.tsx`) | Keep hidden. Citations must not reopen it. | +| List / Tree | Tree **hidden** for page-asset sources (`!isPageAssetSource`). Default mode is already `"tree"` but unused in page mode. | Show the toggle for page mode; Tree uses existing `ChunkSectionTree`. | +| Page cards | `PageChunkCard` already shows page image + `ChunkKeywords` | Match Figma title (`Page N` + path). Image-first. Hide keywords when empty (already returns null). | +| Inline chips | `[Source N: …]` tokens are **stripped** (`buildCitationContentMarkdown`). Chips live in a footer `AssistantSources`. | Replace tokens with inline `{title}/pN` chips in the answer. | +| Chip label | Filename / source title only, **no `/pN`**. | `{title}/pN` | +| Dedupe | `getDisplayCitations` and `selectCitationRawResults` drop same-key citations. | Keep one chip per citation so two regions on one page stay two chips. | +| Footer SOURCES | Flat chips | Numbered per-file list + page links | +| Click target | `handleCitationClick` focuses a **chunk**. `focusedPage` is plumbed to `ChunksPanel` but **not used**. | Focus/scroll the **page** card; `onSelectSource` for file B while viewing A already exists. | +| Export | None | Per-turn copy + MD/PDF without citations | + +Later API (do not block Notebook UI): + +- Retrieval / agent citation payload has page number and page image today (`pageCitationPageNumber`, `pageCitationAssetUrl`). +- It does **not** have a citation crop box. Worker bboxes are for asset crops, not “this quote lives here.” +- When Knowhere adds a region on the citation, extend `ChatCitationView` + persistence, then overlay on `PageCitationAssetImage`. + +## Implementation slices (Notebook) + +Do these in order. Tests follow existing ownership: model tests for labels/grouping, component tests for chips/cards/export, panel tests for click → source + page. + +1. **Chrome** — rename Brain; Parsed Results header/subtitle; enable List/Tree on page-asset sources; default Tree. +2. **Page cards** — Figma card header (page + path); image + optional keywords; citation click forces List (already does via `citationListViewRequestId`) and scrolls to the page chunk. +3. **Citations** — stop stripping without replacement; inline chips `{title}/pN`; stop collapsing same-page citations; numbered SOURCES with `pN` links; click switches viewed source and page-focuses. +4. **Copy / download** — message actions: clipboard of visible answer (code/tables included as rendered); MD + PDF of this turn with the same citation-stripped body we already produce for display (`buildCitationContentMarkdown` / equivalent). PDF is client-side from that Markdown, not the original document. +5. **Later** — region overlay when retrieval sends a box. Keep citation identity stable (do not key chips by `title+page` only). + +## Do not + +- Route citation clicks to original PDF preview. +- Deduplicate inline chips by page. +- Draw a highlight without retrieval geometry. +- Change which sources are **checked for the next query** when jumping the canvas. +- Put this feature in archived `knowhere-api`; new retrieval fields go in `Ontos-AI/knowhere`. diff --git a/e2e/chunk-original-preview.e2e.ts b/e2e/chunk-original-preview.e2e.ts index 1a6502c4..c7e0426c 100644 --- a/e2e/chunk-original-preview.e2e.ts +++ b/e2e/chunk-original-preview.e2e.ts @@ -7,7 +7,7 @@ test("reopens the demo PDF preview after returning from original to parsed chunk }) => { await page.goto("/") await expect( - page.getByRole("heading", { name: "Parsed Chunks" }), + page.getByRole("heading", { name: "Parsed Results" }), ).toBeVisible() await openFirstChunk(page) @@ -18,7 +18,7 @@ test("reopens the demo PDF preview after returning from original to parsed chunk .getByRole("button", { name: "Parsed", exact: true }) .click() await expect( - page.getByRole("heading", { name: "Parsed Chunks" }), + page.getByRole("heading", { name: "Parsed Results" }), ).toBeVisible() await openFirstChunk(page) diff --git a/e2e/citation-dedupe.e2e.ts b/e2e/citation-dedupe.e2e.ts index 78de15e8..071a986f 100644 --- a/e2e/citation-dedupe.e2e.ts +++ b/e2e/citation-dedupe.e2e.ts @@ -87,3 +87,24 @@ test("keeps duplicate source labels clickable for separate documents", async ({ await expect(page.getByText("First report source content.")).toBeVisible() expect(firstSourceChunkRequests).toBeGreaterThan(0) }) + +test("keeps two chips to the same title/pN as separate buttons", async ({ + context, + page, +}) => { + await context.addCookies([ + { + name: "better-auth.session_token", + value: "playwright", + url: "http://localhost:3000", + }, + ]) + + await page.goto("/e2e/citation-same-page") + + const chatPanel = page.getByTestId("desktop-chat-panel") + await expect(chatPanel.getByTestId("citation-chip")).toHaveCount(2) + await expect( + chatPanel.getByRole("button", { name: "Open source spacex-s1.pdf/p26" }), + ).toHaveCount(2) +}) diff --git a/src/agent-harness/ledger.test.ts b/src/agent-harness/ledger.test.ts index 061a8928..15f1b20e 100644 --- a/src/agent-harness/ledger.test.ts +++ b/src/agent-harness/ledger.test.ts @@ -1,5 +1,6 @@ import { describe, expect, it } from "vitest" import type { + KnowledgeGrepMatch, KnowledgeGrepResponse, KnowledgeReadResponse, RetrievalQueryResponse, @@ -111,6 +112,76 @@ describe("createEvidenceLedger", () => { }), ]) }) + + it("copies page metadata onto grep matches from the same chunk id", () => { + const ledger = createEvidenceLedger() + ledger.addReadChunksResponse(makeReadResponse()) + + const snapshot = ledger.addGrepChunksResponse({ + ...makeGrepResponse(), + matches: [ + { + position: 1, + chunkId: "chunk_page_1", + chunkType: "page", + sectionPath: "Root / Page 1", + sourceChunkPath: "pages/page-1.md", + startOffset: 0, + endOffset: 12, + snippet: "Full page snip", + }, + ], + }) + + expect(snapshot.chunks[1]).toEqual( + expect.objectContaining({ + ref: "grep1:match:1", + kind: "grep_match", + chunkId: "chunk_page_1", + metadata: expect.objectContaining({ + pageNums: [1], + position: 1, + startOffset: 0, + endOffset: 12, + }), + }), + ) + }) + + it("copies pageNumbers from the grep match when the SDK provides them", () => { + const ledger = createEvidenceLedger() + + const snapshot = ledger.addGrepChunksResponse({ + ...makeGrepResponse(), + matches: [ + { + position: 1, + chunkId: "chunk_page_4", + chunkType: "page", + sectionPath: "FINANCIAL SUMMARY", + sourceChunkPath: "pages/page-4.md", + startOffset: 0, + endOffset: 12, + snippet: "automotive revenues", + pageNumbers: [4], + } as KnowledgeGrepMatch, + ], + }) + + expect(snapshot.chunks[0]).toEqual( + expect.objectContaining({ + ref: "grep1:match:1", + kind: "grep_match", + chunkId: "chunk_page_4", + metadata: expect.objectContaining({ + pageNums: [4], + position: 1, + startOffset: 0, + endOffset: 12, + }), + }), + ) + }) }) function makeRetrievalResponse(): RetrievalQueryResponse { diff --git a/src/agent-harness/ledger.ts b/src/agent-harness/ledger.ts index d782ae3e..aa9d58e5 100644 --- a/src/agent-harness/ledger.ts +++ b/src/agent-harness/ledger.ts @@ -255,6 +255,19 @@ function addChunkFromGrepMatch(input: { readonly match: KnowledgeGrepMatch readonly ref: string }): void { + const matchWithPages = input.match as KnowledgeGrepMatch & { + readonly pageNumbers?: readonly number[] + readonly metadata?: Readonly> + } + const donor = input.ledger.chunks.find( + (chunk) => + chunk.chunkId === input.match.chunkId && hasPageMetadata(chunk.metadata), + ) + const pageNums = + matchWithPages.pageNumbers && matchWithPages.pageNumbers.length > 0 + ? [...matchWithPages.pageNumbers] + : undefined + addChunk({ ledger: input.ledger, chunk: { @@ -268,6 +281,9 @@ function addChunkFromGrepMatch(input: { sourceChunkPath: input.match.sourceChunkPath, filePath: input.match.filePath, metadata: { + ...(donor?.metadata ?? {}), + ...(matchWithPages.metadata ?? {}), + ...(pageNums ? { pageNums } : {}), position: input.match.position, startOffset: input.match.startOffset, endOffset: input.match.endOffset, @@ -282,6 +298,26 @@ function addChunkFromGrepMatch(input: { }) } +function hasPageMetadata( + metadata: Readonly> | undefined, +): boolean { + if (!metadata) return false + const values = [ + metadata.pageNums, + metadata.page_nums, + metadata.pageNum, + metadata.page_num, + metadata.pageAssets, + metadata.page_assets, + ] + return values.some((value) => { + if (Array.isArray(value)) return value.length > 0 + if (typeof value === "number") return Number.isSafeInteger(value) && value > 0 + if (typeof value === "string") return value.trim().length > 0 + return false + }) +} + 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 04614a77..12873d39 100644 --- a/src/agent-harness/runtime.test.ts +++ b/src/agent-harness/runtime.test.ts @@ -40,6 +40,15 @@ describe("agent harness runtime", () => { expect(prompt).not.toContain("must match the selected evidence ref exactly") }) + it("tells the agent to emit [[cite:n]] markers instead of title/pN or [1]", () => { + const prompt = buildHarnessSystemPrompt(makeTurnInput()) + + expect(prompt).toContain("[[cite:n]]") + expect(prompt).toContain("1-based index into the citations array") + expect(prompt).toContain("Do not write title/pN, [1], Markdown footnotes") + expect(prompt).toContain("Do not collapse same-page citations") + }) + it("passes only outer retrieval parameters to KNOWHERE without planning-tool gating", async () => { const query = vi.fn().mockResolvedValue( makeRetrievalResponse(), diff --git a/src/agent-harness/runtime.ts b/src/agent-harness/runtime.ts index 8963fdf5..0ad95d8d 100644 --- a/src/agent-harness/runtime.ts +++ b/src/agent-harness/runtime.ts @@ -1176,6 +1176,9 @@ 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.", "- Prefer citation and selected image/table artifact refs returned by Knowhere tools in the evidence ledger.", + "- Place [[cite:n]] immediately after the supported claim. n is the 1-based index into the citations array passed to finalize.", + "- Do not write title/pN, [1], Markdown footnotes, or [Source N: ...] in the answer text. Notebook renders chips from [[cite:n]] and citation metadata.", + "- Repeat [[cite:n]] when another claim uses the same page. Do not collapse same-page citations to one row.", "- 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.", diff --git a/src/app/e2e/citation-same-page/page.tsx b/src/app/e2e/citation-same-page/page.tsx new file mode 100644 index 00000000..c825f918 --- /dev/null +++ b/src/app/e2e/citation-same-page/page.tsx @@ -0,0 +1,78 @@ +import { notFound } from "next/navigation" + +import { WorkspaceShell } from "@/components/workspace-shell" +import type { ChatMessageView, ChatThreadView } from "@/domains/chat/types" +import type { SourceView } from "@/domains/sources/types" + +const sources: SourceView[] = [ + { + id: "source_spacex", + title: "spacex-s1.pdf", + mimeType: "application/pdf", + status: "ready", + documentId: "doc_spacex", + chunkCount: 1, + documentPresentation: { kind: "page-assets", pageCount: 26 }, + }, +] + +const chatThreads: ChatThreadView[] = [ + { + id: "thread_1", + title: "Same page citations", + createdAt: "2026-05-17T00:00:00.000Z", + updatedAt: "2026-05-17T00:00:00.000Z", + }, +] + +const chatMessages: ChatMessageView[] = [ + { + id: "assistant_1", + role: "assistant", + content: "Revenue grew [[cite:1]] and later expanded [[cite:2]].", + citations: [ + { + chunkType: "page", + score: 0.91, + pageCitationPageNumber: 26, + source: { + documentId: "doc_spacex", + sourceFileName: "spacex-s1.pdf", + sectionPath: "Page 26", + }, + }, + { + chunkType: "page", + score: 0.89, + pageCitationPageNumber: 26, + source: { + documentId: "doc_spacex", + sourceFileName: "spacex-s1.pdf", + sectionPath: "Page 26", + }, + }, + ], + }, +] + +export default function CitationSamePageTestPage() { + if (process.env.NODE_ENV === "production") notFound() + + return ( + + ) +} diff --git a/src/app/global-error.tsx b/src/app/global-error.tsx index 5fd5cc79..409a3af5 100644 --- a/src/app/global-error.tsx +++ b/src/app/global-error.tsx @@ -15,11 +15,11 @@ export default function GlobalError({ return ( - Something went wrong - Knowhere Notebook + Something went wrong - Knowhere Brain

- Notebook error + Brain error

Something went wrong diff --git a/src/app/login/page.tsx b/src/app/login/page.tsx index c802af3f..8ae3bc73 100644 --- a/src/app/login/page.tsx +++ b/src/app/login/page.tsx @@ -32,7 +32,7 @@ export async function LoginContent() {

- Knowhere Notebook + Knowhere Brain

{ + it("labels chips as title/pN and keeps two same-page citations as two identities", () => { + expect( + chatCitationModel.getCitationChipLabel(pageCitation, { + doc_1: "spacex-s1.pdf", + }), + ).toBe("spacex-s1.pdf/p26") + + const groups = chatCitationModel.groupCitationsByFile( + "assistant_1", + [pageCitation, samePageCitation, otherFileCitation], + { + doc_1: "spacex-s1.pdf", + doc_2: "report.pdf", + }, + ) + + expect(groups).toHaveLength(2) + expect(groups[0]?.entries.map((entry) => entry.citationId)).toEqual([ + "assistant_1:0", + "assistant_1:1", + ]) + expect(groups[0]?.entries.map((entry) => entry.pageNumber)).toEqual([26, 26]) + expect( + chatCitationModel + .uniquePageLinkEntries(groups[0]!.entries) + .map((entry) => entry.pageNumber), + ).toEqual([26]) + expect(groups[1]?.title).toBe("report.pdf") + }) + + it("embeds cite markers as links and leaves fenced code unchanged", () => { + const markdown = chatCitationModel.embedCitationMarkersAsLinks( + [ + "Revenue grew [[cite:1]] and later [[cite:2]].", + "", + "```ts", + "const value = '[[cite:1]]'", + "```", + "", + "Legacy [Source 1: spacex-s1.pdf / Page 26] remains.", + ].join("\n"), + [pageCitation, samePageCitation], + { doc_1: "spacex-s1.pdf" }, + ) + + expect(markdown).toContain( + "[spacex-s1.pdf/p26](knowhere-cite://1)", + ) + expect(markdown).toContain( + "[spacex-s1.pdf/p26](knowhere-cite://2)", + ) + expect(markdown).toContain("const value = '[[cite:1]]'") + expect(markdown).not.toContain("[Source 1:") + expect(markdown).not.toContain("[[cite:1]] and later") + }) + + it("drops unknown cite markers from display markdown", () => { + const markdown = chatCitationModel.embedCitationMarkersAsLinks( + "Unsupported claim [[cite:9]].", + [pageCitation], + { doc_1: "spacex-s1.pdf" }, + ) + + expect(markdown).toBe("Unsupported claim.") + expect(markdown).not.toContain("knowhere-cite://9") + }) + + it("strips markers from export markdown and keeps copy labels", () => { + const message: ChatMessageView = { + id: "assistant_1", + role: "assistant", + content: [ + "Revenue grew [[cite:1]].", + "", + "```ts", + "const value = 1;", + "```", + ].join("\n"), + citations: [pageCitation], + artifacts: [ + { + type: "derived_table", + display: true, + title: "Revenue", + columns: ["Year", "Amount"], + rows: [["2025", "4.4B"]], + }, + ], + } + + expect(chatCitationModel.getExportMarkdown(message)).toBe( + [ + "Revenue grew.", + "", + "```ts", + "const value = 1;", + "```", + "", + "### Revenue", + "| Year | Amount |", + "| --- | --- |", + "| 2025 | 4.4B |", + ].join("\n"), + ) + expect(chatCitationModel.getCopyMarkdown(message, { doc_1: "spacex-s1.pdf" })) + .toContain("Revenue grew spacex-s1.pdf/p26.") + expect(chatCitationModel.getCopyMarkdown(message, { doc_1: "spacex-s1.pdf" })) + .toContain("const value = 1;") + }) + + it("keeps knowhere cite hrefs out of the default markdown URL transform", () => { + expect( + chatCitationModel.transformMarkdownUrl("knowhere-cite://2", () => ""), + ).toBe("knowhere-cite://2") + expect(chatCitationModel.parseKnowhereCiteIndex("knowhere-cite://2")).toBe(2) + expect(chatCitationModel.isKnowhereCiteHref("https://example.com")).toBe( + false, + ) + }) +}) diff --git a/src/components/chat-citation-model.ts b/src/components/chat-citation-model.ts new file mode 100644 index 00000000..c3bac75c --- /dev/null +++ b/src/components/chat-citation-model.ts @@ -0,0 +1,284 @@ +import { workspaceCitationState } from "@/components/workspace-citation-state" +import type { + ChatArtifactView, + ChatCitationView, + ChatMessageView, +} from "@/domains/chat/types" + +export const knowhereCiteHrefPrefix = "knowhere-cite://" + +export type CitationSourceGroupEntry = { + readonly citation: ChatCitationView + readonly citationId: string + readonly citationIndex: number + readonly chipLabel: string + readonly pageNumber: number | null +} + +export type CitationSourceGroup = { + readonly key: string + readonly title: string + readonly documentId: string | null + readonly entries: readonly CitationSourceGroupEntry[] +} + +export const chatCitationModel = { + knowhereCiteHrefPrefix, + embedCitationMarkersAsLinks, + getCitationChipLabel, + getCopyMarkdown, + getExportMarkdown, + getSourceTitle, + groupCitationsByFile, + isKnowhereCiteHref, + parseKnowhereCiteIndex, + stripCitationMarkers, + transformMarkdownUrl, + uniquePageLinkEntries, +} as const + +function embedCitationMarkersAsLinks( + content: string, + citations: readonly ChatCitationView[], + sourceTitlesByDocumentId: Readonly>, +): string { + return mapOutsideCodeFences(content, (text) => + collapseRemovedCitationSpace( + replaceCitationTokens(text, (index) => { + const citation = citations[index - 1] + if (!citation) return "" + return toCiteMarkdownLink( + getCitationChipLabel(citation, sourceTitlesByDocumentId), + index, + ) + }), + ), + ) +} + +function getCopyMarkdown( + message: ChatMessageView, + sourceTitlesByDocumentId: Readonly>, +): string { + const copiedBody = mapOutsideCodeFences(message.content, (text) => + collapseRemovedCitationSpace( + replaceCitationTokens(text, (index) => { + const citation = message.citations?.[index - 1] + if (!citation) return "" + return ` ${getCitationChipLabel(citation, sourceTitlesByDocumentId)}` + }), + ), + ) + + return appendDerivedTables(copiedBody, message.artifacts) +} + +function getExportMarkdown(message: ChatMessageView): string { + return appendDerivedTables(stripCitationMarkers(message.content), message.artifacts) +} + +function stripCitationMarkers(content: string): string { + return mapOutsideCodeFences(content, (text) => + collapseRemovedCitationSpace(replaceCitationTokens(text, () => "")), + ) +} + +function groupCitationsByFile( + messageId: string, + citations: readonly ChatCitationView[], + sourceTitlesByDocumentId: Readonly>, +): readonly CitationSourceGroup[] { + const groups: CitationSourceGroup[] = [] + const groupIndexByKey = new Map() + + for (const [citationIndex, citation] of citations.entries()) { + const title = getSourceTitle(citation, sourceTitlesByDocumentId) + const documentId = getTrimmedField(citation.source.documentId) + const key = documentId + ? `document:${documentId}` + : `file:${citationIndex}:${title}` + const entry: CitationSourceGroupEntry = { + citation, + citationId: `${messageId}:${citationIndex}`, + citationIndex, + chipLabel: getCitationChipLabel(citation, sourceTitlesByDocumentId), + pageNumber: workspaceCitationState.getCitationPageNumber(citation), + } + const existingIndex = groupIndexByKey.get(key) + if (existingIndex === undefined) { + groupIndexByKey.set(key, groups.length) + groups.push({ + key, + title, + documentId, + entries: [entry], + }) + continue + } + + const existing = groups[existingIndex] + if (!existing) continue + groups[existingIndex] = { + ...existing, + entries: [...existing.entries, entry], + } + } + + return groups +} + +function uniquePageLinkEntries( + entries: readonly CitationSourceGroupEntry[], +): readonly CitationSourceGroupEntry[] { + const seenPageNumbers = new Set() + const uniqueEntries: CitationSourceGroupEntry[] = [] + + for (const entry of entries) { + if (entry.pageNumber === null || seenPageNumbers.has(entry.pageNumber)) { + continue + } + seenPageNumbers.add(entry.pageNumber) + uniqueEntries.push(entry) + } + + return uniqueEntries +} + +function getCitationChipLabel( + citation: ChatCitationView, + sourceTitlesByDocumentId: Readonly>, +): string { + const title = getSourceTitle(citation, sourceTitlesByDocumentId) + const pageNumber = workspaceCitationState.getCitationPageNumber(citation) + return pageNumber === null ? title : `${title}/p${pageNumber}` +} + +function getSourceTitle( + citation: ChatCitationView, + sourceTitlesByDocumentId: Readonly>, +): string { + const documentId = getTrimmedField(citation.source.documentId) + const sourceTitle = documentId + ? getTrimmedField(sourceTitlesByDocumentId[documentId]) + : null + if (sourceTitle) return sourceTitle + + const sourceFileName = getTrimmedField(citation.source.sourceFileName) + if (sourceFileName && !isGeneratedKnowhereFileName(sourceFileName)) { + return sourceFileName + } + + return "Source" +} + +function isKnowhereCiteHref(href: string | undefined): boolean { + return typeof href === "string" && href.startsWith(knowhereCiteHrefPrefix) +} + +function parseKnowhereCiteIndex(href: string | undefined): number | null { + if (!href || !href.startsWith(knowhereCiteHrefPrefix)) return null + const value = href.slice(knowhereCiteHrefPrefix.length) + const index = Number.parseInt(value, 10) + if (!Number.isSafeInteger(index) || index < 1) return null + return index +} + +function transformMarkdownUrl(value: string, transformDefault: (next: string) => string): string { + if (isKnowhereCiteHref(value)) return value + return transformDefault(value) +} + +function replaceCitationTokens( + text: string, + replaceIndex: (index: number) => string, +): string { + return text + .replace(/\[\[cite:(\d+)\]\]/g, (_match, rawIndex: string) => { + const index = Number.parseInt(rawIndex, 10) + if (!Number.isSafeInteger(index) || index < 1) return "" + return replaceIndex(index) + }) + .replace(/\[Source\s+(\d+)\s*:\s*[^\]]*\]/g, (_match, rawIndex: string) => { + const index = Number.parseInt(rawIndex, 10) + if (!Number.isSafeInteger(index) || index < 1) return "" + return replaceIndex(index) + }) +} + +function toCiteMarkdownLink(label: string, index: number): string { + return `[${escapeMarkdownLinkLabel(label)}](${knowhereCiteHrefPrefix}${index})` +} + +function escapeMarkdownLinkLabel(label: string): string { + return label.replace(/\\/g, "\\\\").replace(/\[/g, "\\[").replace(/\]/g, "\\]") +} + +function collapseRemovedCitationSpace(text: string): string { + return text.replaceAll(" ", " ").replace(/ +([.,;:])/g, "$1").trimEnd() +} + +function mapOutsideCodeFences( + content: string, + mapText: (text: string) => string, +): string { + const lines = content.split("\n") + const output: string[] = [] + let fenceMarker: string | null = null + + for (const line of lines) { + if (fenceMarker) { + output.push(line) + if (line.startsWith(fenceMarker)) fenceMarker = null + continue + } + + const fenceMatch = /^(```+|~~~+)/.exec(line) + if (fenceMatch) { + fenceMarker = fenceMatch[1] ?? null + output.push(line) + continue + } + + output.push(mapText(line)) + } + + return output.join("\n") +} + +function appendDerivedTables( + markdown: string, + artifacts: readonly ChatArtifactView[] | undefined, +): string { + const tables = (artifacts ?? []).flatMap((artifact, index) => { + if (artifact.display === false || artifact.type !== "derived_table") return [] + if (!artifact.title || !artifact.columns || !artifact.rows) return [] + return [formatDerivedTable(artifact, index)] + }) + if (tables.length === 0) return markdown + const body = markdown.trimEnd() + return body.length > 0 ? `${body}\n\n${tables.join("\n\n")}` : tables.join("\n\n") +} + +function formatDerivedTable( + artifact: ChatArtifactView, + index: number, +): string { + const columns = artifact.columns ?? [] + const header = `| ${columns.join(" | ")} |` + const separator = `| ${columns.map(() => "---").join(" | ")} |` + const rows = (artifact.rows ?? []).map( + (row) => `| ${columns.map((_, columnIndex) => row[columnIndex] ?? "").join(" | ")} |`, + ) + return [`### ${artifact.title ?? `Table ${index + 1}`}`, header, separator, ...rows].join( + "\n", + ) +} + +function getTrimmedField(value: string | null | undefined): string | null { + const trimmed = value?.trim() ?? "" + return trimmed.length > 0 ? trimmed : null +} + +function isGeneratedKnowhereFileName(value: string): boolean { + return /^document-[A-Za-z0-9_-]{16,}\.[A-Za-z0-9]+$/u.test(value) +} diff --git a/src/components/chat-message-export-pdf.test.ts b/src/components/chat-message-export-pdf.test.ts new file mode 100644 index 00000000..8bbe53cf --- /dev/null +++ b/src/components/chat-message-export-pdf.test.ts @@ -0,0 +1,16 @@ +import { describe, expect, it } from "vitest" + +import { createPdfBlobFromMarkdown } from "./chat-message-export-pdf" + +describe("createPdfBlobFromMarkdown", () => { + it("builds a PDF without citation markers", async () => { + const blob = createPdfBlobFromMarkdown("Revenue grew in the quarter.") + const text = await blob.text() + + expect(blob.type).toBe("application/pdf") + expect(text.startsWith("%PDF-1.4")).toBe(true) + expect(text).toContain("Revenue grew in the quarter.") + expect(text).not.toContain("[[cite:") + expect(text).not.toContain("SOURCES") + }) +}) diff --git a/src/components/chat-message-export-pdf.ts b/src/components/chat-message-export-pdf.ts new file mode 100644 index 00000000..4c4aa3ba --- /dev/null +++ b/src/components/chat-message-export-pdf.ts @@ -0,0 +1,136 @@ +const pdfPageWidth = 612 +const pdfPageHeight = 792 +const pdfMargin = 54 +const pdfFontSize = 11 +const pdfLineHeight = 14 +const pdfCharsPerLine = 90 + +export async function downloadAnswerPdf( + filename: string, + markdown: string, +): Promise { + const blob = createPdfBlobFromMarkdown(markdown) + downloadBlob(filename, blob) +} + +export function createPdfBlobFromMarkdown(markdown: string): Blob { + const lines = wrapPdfLines(toWinAnsi(markdown), pdfCharsPerLine) + const linesPerPage = Math.max( + 1, + Math.floor((pdfPageHeight - pdfMargin * 2) / pdfLineHeight), + ) + const pageCount = Math.max(1, Math.ceil(lines.length / linesPerPage)) + const objects: string[] = [] + const pageObjectNumbers: number[] = [] + const fontObjectNumber = 3 + pageCount * 2 + + for (let pageIndex = 0; pageIndex < pageCount; pageIndex += 1) { + const pageLines = lines.slice( + pageIndex * linesPerPage, + (pageIndex + 1) * linesPerPage, + ) + const content = buildPageContent(pageLines) + const contentObjectNumber = 3 + pageIndex * 2 + const pageObjectNumber = contentObjectNumber + 1 + objects.push(pdfObject(contentObjectNumber, `<< /Length ${content.length} >>\nstream\n${content}\nendstream`)) + objects.push( + pdfObject( + pageObjectNumber, + `<< /Type /Page /Parent 2 0 R /MediaBox [0 0 ${pdfPageWidth} ${pdfPageHeight}] /Contents ${contentObjectNumber} 0 R /Resources << /Font << /F1 ${fontObjectNumber} 0 R >> >> >>`, + ), + ) + pageObjectNumbers.push(pageObjectNumber) + } + + objects.unshift( + pdfObject(1, "<< /Type /Catalog /Pages 2 0 R >>"), + pdfObject( + 2, + `<< /Type /Pages /Count ${pageCount} /Kids [${pageObjectNumbers + .map((objectNumber) => `${objectNumber} 0 R`) + .join(" ")}] >>`, + ), + ) + objects.push( + pdfObject(fontObjectNumber, "<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>"), + ) + + return buildPdf(objects) +} + +function buildPageContent(lines: readonly string[]): string { + const commands = [ + "BT", + `/F1 ${pdfFontSize} Tf`, + `${pdfMargin} ${pdfPageHeight - pdfMargin} Td`, + `${pdfLineHeight} TL`, + ] + for (const [index, line] of lines.entries()) { + commands.push(`(${escapePdfLiteral(line)}) Tj`) + if (index < lines.length - 1) commands.push("T*") + } + commands.push("ET") + return commands.join("\n") +} + +function wrapPdfLines(text: string, maxChars: number): string[] { + const lines: string[] = [] + for (const paragraph of text.split(/\r?\n/)) { + if (paragraph.length === 0) { + lines.push("") + continue + } + + let remaining = paragraph + while (remaining.length > maxChars) { + let breakAt = remaining.lastIndexOf(" ", maxChars) + if (breakAt < 1) breakAt = maxChars + lines.push(remaining.slice(0, breakAt)) + remaining = remaining.slice(breakAt).trimStart() + } + lines.push(remaining) + } + return lines.length > 0 ? lines : [""] +} + +function toWinAnsi(text: string): string { + return Array.from(text, (character) => + character.charCodeAt(0) < 256 ? character : "?", + ).join("") +} + +function escapePdfLiteral(text: string): string { + return text.replace(/\\/g, "\\\\").replace(/\(/g, "\\(").replace(/\)/g, "\\)") +} + +function pdfObject(objectNumber: number, body: string): string { + return `${objectNumber} 0 obj\n${body}\nendobj` +} + +function buildPdf(objects: readonly string[]): Blob { + const xrefOffsets: number[] = [0] + let body = "%PDF-1.4\n" + for (const object of objects) { + xrefOffsets.push(body.length) + body += `${object}\n` + } + const xrefStart = body.length + const xrefEntries = [ + "xref", + `0 ${objects.length + 1}`, + "0000000000 65535 f ", + ...xrefOffsets.slice(1).map((offset) => `${String(offset).padStart(10, "0")} 00000 n `), + ] + body += `${xrefEntries.join("\n")}\n` + body += `trailer\n<< /Size ${objects.length + 1} /Root 1 0 R >>\nstartxref\n${xrefStart}\n%%EOF\n` + return new Blob([body], { type: "application/pdf" }) +} + +function downloadBlob(filename: string, blob: Blob): void { + const objectUrl = URL.createObjectURL(blob) + const anchor = document.createElement("a") + anchor.href = objectUrl + anchor.download = filename + anchor.click() + URL.revokeObjectURL(objectUrl) +} diff --git a/src/components/chat-message-list.test.ts b/src/components/chat-message-list.test.ts index 94441282..e59fd540 100644 --- a/src/components/chat-message-list.test.ts +++ b/src/components/chat-message-list.test.ts @@ -70,13 +70,14 @@ describe("ChatMessageList", () => { id: "assistant_1", role: "assistant", content: [ - "Capital expenditure appears in the appendix. [Source 1: spacex-s1.pdf / Assets / tables / table-25 Capital Expenditures.html]", - "Drivers are discussed elsewhere. [Source 3: spacex-s1.pdf / MD&A / Drivers of Our Performance]", + "Capital expenditure appears in the appendix. [[cite:1]]", + "Drivers are discussed elsewhere. [[cite:3]]", ].join("\n\n"), citations: [ { chunkType: "table", score: 0.9, + pageCitationPageNumber: 25, source: { documentId: "doc_1", sourceFileName: "spacex-s1.pdf", @@ -87,6 +88,7 @@ describe("ChatMessageList", () => { { chunkType: "table", score: 0.91, + pageCitationPageNumber: 25, source: { documentId: "doc_1", sourceFileName: "spacex-s1.pdf", @@ -97,6 +99,7 @@ describe("ChatMessageList", () => { { chunkType: "text", score: 0.8, + pageCitationPageNumber: 40, source: { documentId: "doc_1", sourceFileName: "spacex-s1.pdf", @@ -110,16 +113,27 @@ describe("ChatMessageList", () => { }), ); - expect(screen.getByText("Capital expenditure appears in the appendix.")) + expect(screen.getByText(/Capital expenditure appears in the appendix./u)) .toBeTruthy(); expect(screen.queryByText(/Source 1/u)).toBeNull(); - expect(screen.queryByText(/Source 3/u)).toBeNull(); + expect(screen.queryByText(/\[\[cite:/u)).toBeNull(); expect(screen.getByText("Sources")).toBeTruthy(); - const sourceChips = screen.getAllByRole("button", { - name: "Open source spacex-s1.pdf", - }); + const sourceChips = screen.getAllByTestId("citation-chip"); expect(sourceChips).toHaveLength(2); + expect(sourceChips[0]?.textContent).toBe("spacex-s1.pdf/p25"); + expect(sourceChips[1]?.textContent).toBe("spacex-s1.pdf/p40"); + expect(sourceChips[0]?.className).toContain("bg-muted"); + expect(sourceChips[0]?.className).toContain("h-5"); + expect( + screen.getByRole("button", { name: "Open source spacex-s1.pdf" }), + ).toBeTruthy(); + expect( + screen.getAllByRole("button", { name: "Open page 25 of spacex-s1.pdf" }), + ).toHaveLength(1); + expect( + screen.getByRole("button", { name: "Open page 40 of spacex-s1.pdf" }), + ).toBeTruthy(); await user.hover(sourceChips[0]!); @@ -161,7 +175,7 @@ describe("ChatMessageList", () => { ); const citationButton = screen.getByRole("button", { - name: "Open source report.pdf", + name: "Open page 4 of report.pdf", }); expect( screen.queryByRole("link", { @@ -179,7 +193,7 @@ describe("ChatMessageList", () => { ); }); - it("removes description-only source labels without changing other markdown whitespace", () => { + it("renders inline chips for leftover Source tokens without rewriting fenced code", () => { render( React.createElement(ChatMessageList, { messages: [ @@ -198,6 +212,7 @@ describe("ChatMessageList", () => { chunkType: "text", score: 0.9, description: "revenue growth", + pageCitationPageNumber: 2, source: { documentId: "doc_1", sourceFileName: "notes.pdf", @@ -210,13 +225,56 @@ describe("ChatMessageList", () => { }), ); - expect(screen.getByText("Revenue improved.")).toBeTruthy(); + expect(screen.getByText(/Revenue improved/u)).toBeTruthy(); expect(screen.queryByText(/Source 1/u)).toBeNull(); + expect(screen.getByRole("button", { name: "Open source notes.pdf/p2" })) + .toBeTruthy(); expect(document.querySelector("code.language-ts")?.textContent).toContain( "const value = 1;", ); }); + it("keeps two same-page citation chips as separate buttons", () => { + render( + React.createElement(ChatMessageList, { + messages: [ + { + id: "assistant_1", + role: "assistant", + content: "First claim [[cite:1]] and second claim [[cite:2]].", + citations: [ + { + chunkType: "page", + score: 0.9, + pageCitationPageNumber: 2, + source: { + documentId: "doc_1", + sourceFileName: "spacex-s1.pdf", + sectionPath: "Page 2", + }, + }, + { + chunkType: "page", + score: 0.88, + pageCitationPageNumber: 2, + source: { + documentId: "doc_1", + sourceFileName: "spacex-s1.pdf", + sectionPath: "Page 2", + }, + }, + ], + }, + ], + }), + ); + + expect( + screen.getAllByRole("button", { name: "Open source spacex-s1.pdf/p2" }), + ).toHaveLength(2); + expect(screen.getAllByTestId("citation-chip")).toHaveLength(2); + }); + it("preserves repeated spaces when there are no citation tokens to remove", () => { render( React.createElement(ChatMessageList, { @@ -533,4 +591,50 @@ describe("ChatMessageList", () => { within(screen.getByTestId("chat-scroll")).getByText("What changed?"), ).toBeTruthy(); }); + + it("copies visible answer text and keeps export markdown citation-stripped", async () => { + const user = userEvent.setup(); + const writeText = vi.fn().mockResolvedValue(undefined); + Object.defineProperty(navigator, "clipboard", { + configurable: true, + value: { writeText }, + }); + + render( + React.createElement(ChatMessageList, { + messages: [ + { + id: "assistant_1", + role: "assistant", + content: "Revenue grew [[cite:1]].", + citations: [ + { + chunkType: "page", + score: 0.9, + pageCitationPageNumber: 26, + source: { + documentId: "doc_1", + sourceFileName: "spacex-s1.pdf", + sectionPath: "Page 26", + }, + }, + ], + }, + ], + sourceTitlesByDocumentId: { doc_1: "spacex-s1.pdf" }, + }), + ); + + await user.click(screen.getByRole("button", { name: "Copy answer" })); + expect(writeText).toHaveBeenCalledWith( + expect.stringContaining("Revenue grew spacex-s1.pdf/p26."), + ); + expect(writeText.mock.calls[0]?.[0]).not.toContain("[[cite:"); + expect( + screen.getByRole("button", { name: "Download answer as Markdown" }), + ).toBeTruthy(); + expect( + screen.getByRole("button", { name: "Download answer as PDF" }), + ).toBeTruthy(); + }); }); diff --git a/src/components/chat-message-list.tsx b/src/components/chat-message-list.tsx index 38c50593..64ab794f 100644 --- a/src/components/chat-message-list.tsx +++ b/src/components/chat-message-list.tsx @@ -2,16 +2,18 @@ import { type CSSProperties, type ReactElement } from "react"; import { type VirtualItem } from "@tanstack/react-virtual"; -import { ImageIcon, MessageCircle } from "lucide-react"; +import { Copy, Download, FileText, ImageIcon, MessageCircle } from "lucide-react"; import ReactMarkdown, { defaultUrlTransform, type Components, } from "react-markdown"; import remarkGfm from "remark-gfm"; +import { chatCitationModel } from "@/components/chat-citation-model"; import { ChatDiagramCard } from "@/components/chat-diagram-card"; import { useChatMessageListWorkflow } from "@/components/chat-message-list-workflow"; import { chatPanelModel } from "@/components/chat-panel-model"; +import { Button } from "@/components/ui/button"; import { ScrollArea } from "@/components/ui/scroll-area"; import { Spinner } from "@/components/ui/spinner"; import { @@ -27,14 +29,11 @@ import type { ChatMessageView, } from "@/domains/chat/types"; -type DisplayCitation = { +type DisplayImageCitation = { readonly citation: ChatCitationView; readonly citationId: string; readonly label: string; readonly tooltipLabel: string; -}; - -type DisplayImageCitation = DisplayCitation & { readonly assetUrl: string; }; @@ -306,10 +305,6 @@ function MessageBubble({ ); } - const displayCitations = getDisplayCitations( - message, - sourceTitlesByDocumentId, - ); const displayImageArtifacts = getDisplayImageArtifacts( message, sourceTitlesByDocumentId, @@ -319,7 +314,7 @@ function MessageBubble({ ? displayImageArtifacts : getDisplayImageCitations(message, sourceTitlesByDocumentId); const displayDerivedTables = getDisplayDerivedTableArtifacts(message); - const citationContentMarkdown = buildCitationContentMarkdown( + const citationContentMarkdown = chatCitationModel.embedCitationMarkersAsLinks( message.content, message.citations ?? [], sourceTitlesByDocumentId, @@ -327,9 +322,14 @@ function MessageBubble({ return (
+
{displayDerivedTables.length > 0 && (
@@ -371,11 +371,17 @@ function MessageBubble({
)}
+ +
); } @@ -418,78 +424,6 @@ function AssistantDiagram({ ); } -function buildCitationContentMarkdown( - content: string, - citations: readonly ChatCitationView[], - sourceTitlesByDocumentId: Readonly>, -): string { - let rewrittenContent = content; - - for (const [index, citation] of citations.entries()) { - const displayCitation = { - citation, - citationId: "", - label: getCitationSourceChipLabel(citation, sourceTitlesByDocumentId), - tooltipLabel: chatPanelModel.getCitationLabel( - citation, - sourceTitlesByDocumentId, - ), - }; - for (const token of getInlineCitationTokens(displayCitation, index)) { - if (!rewrittenContent.includes(token)) continue; - - rewrittenContent = removeInlineCitationToken(rewrittenContent, token); - } - } - - return rewrittenContent; -} - -function removeInlineCitationToken(content: string, token: string): string { - return content - .replaceAll(` ${token}`, "") - .replaceAll(`${token} `, "") - .replaceAll(token, ""); -} - -function getInlineCitationTokens( - displayCitation: DisplayCitation, - index: number, -): readonly string[] { - const label = displayCitation.label; - const slashLabel = label.replace(/\s+·\s+/gu, " / "); - const sourceName = getCitationSourceName(label); - const sectionPath = getTrimmedCitationField( - displayCitation.citation.source.sectionPath, - ); - const description = getTrimmedCitationField(displayCitation.citation.description); - const sectionLabel = sectionPath ? `${sourceName} / ${sectionPath}` : null; - const descriptionLabel = description ? `${sourceName} / ${description}` : null; - const citationNumber = index + 1; - const tokens = [ - `[${label}]`, - `[${slashLabel}]`, - `[Source ${citationNumber}: ${label}]`, - `[Source ${citationNumber}: ${slashLabel}]`, - sectionLabel ? `[${sectionLabel}]` : null, - sectionLabel ? `[Source ${citationNumber}: ${sectionLabel}]` : null, - descriptionLabel ? `[${descriptionLabel}]` : null, - descriptionLabel ? `[Source ${citationNumber}: ${descriptionLabel}]` : null, - description ? `[${description}]` : null, - description ? `[Source ${citationNumber}: ${description}]` : null, - ]; - - return Array.from( - new Set(tokens.filter((token): token is string => Boolean(token))), - ).sort((left, right): number => right.length - left.length); -} - -function getCitationSourceName(label: string): string { - const [sourceName] = label.split(/\s+·\s+/u); - const normalized = sourceName?.trim(); - return normalized && normalized.length > 0 ? normalized : label; -} - function DerivedTableArtifactView({ artifact, }: { @@ -540,12 +474,49 @@ function DerivedTableArtifactView({ function AssistantMessageContent({ content, + message, + onCitationClick, + pendingCitationId, + sourceTitlesByDocumentId, }: { readonly content: string; + readonly message: ChatMessageView; + readonly onCitationClick?: ( + citation: ChatCitationView, + citationId: string, + ) => void; + readonly pendingCitationId?: string | null; + readonly sourceTitlesByDocumentId: Readonly>; }): ReactElement { const markdownComponents: Components = { ...assistantMarkdownComponents, a: ({ href, children }) => { + const citationIndex = chatCitationModel.parseKnowhereCiteIndex(href); + if (citationIndex !== null) { + const citation = (message.citations ?? [])[citationIndex - 1]; + if (!citation) return null; + const citationId = chatPanelModel.getCitationId( + message.id, + citationIndex - 1, + ); + return ( + + ); + } + return ( void; readonly pendingCitationId?: string | null; + readonly sourceTitlesByDocumentId: Readonly>; }): ReactElement | null { - if (displayCitations.length === 0) return null; + const groups = chatCitationModel.groupCitationsByFile( + message.id, + message.citations ?? [], + sourceTitlesByDocumentId, + ); + if (groups.length === 0) return null; return (
@@ -593,26 +571,148 @@ function AssistantSources({ Sources

-
- {displayCitations.map((displayCitation) => ( - +
    + {groups.map((group, groupIndex) => ( +
  1. + + {groupIndex + 1}. + + entry.citationId === pendingCitationId, + )} + onCitationClick={onCitationClick} + /> + {chatCitationModel.uniquePageLinkEntries(group.entries).map( + (entry) => ( + + ), + )} +
    2. ))} -
+
); } +function CitationSourceTitle({ + group, + isPending, + onCitationClick, +}: { + readonly group: ReturnType[number]; + readonly isPending: boolean; + readonly onCitationClick?: ( + citation: ChatCitationView, + citationId: string, + ) => void; +}): ReactElement { + const firstEntry = group.entries[0]; + if (!firstEntry) { + return {group.title}; + } + + return ( + + ); +} + +function AssistantMessageActions({ + message, + sourceTitlesByDocumentId, +}: { + readonly message: ChatMessageView; + readonly sourceTitlesByDocumentId: Readonly>; +}): ReactElement { + const copyMarkdown = chatCitationModel.getCopyMarkdown( + message, + sourceTitlesByDocumentId, + ); + const exportMarkdown = chatCitationModel.getExportMarkdown(message); + + return ( +
+ + + +
+ ); +} + +function downloadTextFile(filename: string, contents: string): void { + const blob = new Blob([contents], { type: "text/markdown;charset=utf-8" }); + const objectUrl = URL.createObjectURL(blob); + const anchor = document.createElement("a"); + anchor.href = objectUrl; + anchor.download = filename; + anchor.click(); + URL.revokeObjectURL(objectUrl); +} + function transformAssistantMarkdownUrl(value: string): string { - return defaultUrlTransform(value); + return chatCitationModel.transformMarkdownUrl(value, defaultUrlTransform); } function CitationChip({ @@ -634,15 +734,17 @@ function CitationChip({ ) => void; }): ReactElement { return ( - + - ) : null} + {workspaceSources.length === 0 ? ( diff --git a/src/components/top-nav.tsx b/src/components/top-nav.tsx index 5c32df86..641fcf8b 100644 --- a/src/components/top-nav.tsx +++ b/src/components/top-nav.tsx @@ -34,7 +34,7 @@ export function TopNav({

- Knowhere Notebook + Knowhere Brain

{ await result.current.handleCitationClick(pageCitation, "message_1:0"); }); - expect(fetchChunks).toHaveBeenCalledWith("source_1"); + expect(fetchChunks).toHaveBeenCalledWith("source_1", { + chunkType: "page", + untilPageNumber: 4, + }); expect(selectSource).toHaveBeenCalledWith("source_1"); expect(result.current.focusedChunk.chunkId).toBe("page_4"); expect(result.current.focusedPage).toEqual({ - pageNumber: null, - requestId: 0, + pageNumber: 4, + requestId: 1, + citationId: "message_1:0", }); expect(result.current.citationListViewRequestId).toBe(1); }); + it("uses cached SWR pages for file B while viewing A without loading all chunk types", 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 viewingSource: SourceView = { + id: "source_2", + title: "Other.pdf", + mimeType: "application/pdf", + status: "ready", + documentId: "document_2", + }; + 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 cache = new Map([ + [ + unstable_serialize(["source-chunks", "source_1", 1]), + { + data: { + chunks: [pageChunk], + pagination: { + page: 1, + pageSize: 50, + total: 1, + totalPages: 1, + }, + }, + }, + ], + ]); + + const { result } = renderHook( + () => + useWorkspaceCitationFocus({ + fetchChunks, + onSelectSource: selectSource, + selectedSourceId: "source_2", + sources: [pageAssetSource, viewingSource], + }), + { + wrapper: ({ children }: { readonly children: ReactNode }) => + React.createElement( + SWRConfig, + { value: { provider: () => cache } }, + children, + ), + }, + ); + + 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).toBe("page_4"); + expect(result.current.focusedPage.pageNumber).toBe(4); + expect(result.current.focusedPage.citationId).toBe("message_1:0"); + }); + + it("loads only page-type chunks when the tree asks for a full page-asset source", 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 pageAssetSource: SourceView = { + ...readySource, + documentPresentation: { kind: "page-assets", pageCount: 8 }, + }; + + const { result } = renderHook( + () => + useWorkspaceCitationFocus({ + fetchChunks, + onSelectSource: vi.fn(), + selectedSourceId: "source_1", + sources: [pageAssetSource], + }), + { wrapper: createSWRWrapper }, + ); + + await act(async () => { + result.current.handleLoadAllChunks(); + }); + + await waitFor(() => { + expect(fetchChunks).toHaveBeenCalledWith("source_1", { chunkType: "page" }); + }); + expect(result.current.prefetchedChunksBySourceId).toEqual({ + source_1: [pageChunk], + }); + }); + it("reuses cached chunks for a different source without refetching", async () => { const fetchChunks = vi.fn(async () => [prefetchedChunk]); const selectSource = vi.fn(); @@ -261,6 +394,51 @@ describe("useWorkspaceCitationFocus", () => { "source_1", ); }); + + it("focuses a second citation on the selected source without refetching", async () => { + const secondChunk: ParsedChunkView = { + chunkId: "chunk_2", + documentId: "document_1", + sectionPath: "Costs", + type: "text", + content: "Costs rose in the quarter.", + sourceTitle: "Contract.pdf", + }; + const secondCitation: ChatCitationView = { + chunkType: "text", + score: 0.9, + content: "Costs rose in the quarter.", + source: { + documentId: "document_1", + sourceFileName: "Contract.pdf", + sectionPath: "Costs", + }, + }; + const fetchChunks = vi.fn(async () => [prefetchedChunk, secondChunk]); + const selectSource = vi.fn(); + + const { result } = renderHook(() => + useWorkspaceCitationFocus({ + fetchChunks, + onSelectSource: selectSource, + selectedSourceId: "source_1", + sources: [readySource], + }), + { wrapper: createSWRWrapper }, + ); + + await act(async () => { + await result.current.handleCitationClick(citation, "message_1:0"); + }); + await act(async () => { + await result.current.handleCitationClick(secondCitation, "message_1:1"); + }); + + expect(fetchChunks).toHaveBeenCalledTimes(1); + expect(fetchChunks).toHaveBeenCalledWith("source_1"); + expect(result.current.focusedChunk.chunkId).toBe("chunk_2"); + expect(result.current.focusedPage.citationId).toBe("message_1:1"); + }); }); function createSWRWrapper({ diff --git a/src/components/workspace-citation-focus.ts b/src/components/workspace-citation-focus.ts index 70b0acf9..379dc290 100644 --- a/src/components/workspace-citation-focus.ts +++ b/src/components/workspace-citation-focus.ts @@ -1,12 +1,17 @@ "use client" import { useCallback, useRef, useState } from "react" +import { useSWRConfig } from "swr" import { workspaceCitationState } from "@/components/workspace-citation-state" import { useWorkspaceSelectedChunks } from "@/components/workspace-selected-chunks" import type { ChatCitationView } from "@/domains/chat/types" import type { ParsedChunkView } from "@/domains/chunks/types" import type { SourceView } from "@/domains/sources/types" +import { + type FetchChunksOptions, +} from "@/domains/workspace/client" +import { workspaceClientCache } from "@/domains/workspace/client-cache" type FocusedChunkState = { readonly chunkId: string | null @@ -16,6 +21,7 @@ type FocusedChunkState = { type FocusedPageState = { readonly pageNumber: number | null readonly requestId: number + readonly citationId: string | null } type PrefetchedChunksBySourceId = Readonly> @@ -24,7 +30,10 @@ type PrefetchedChunksUpdater = ( ) => PrefetchedChunksBySourceId type WorkspaceCitationFocusInput = { - readonly fetchChunks: (sourceId: string) => Promise + readonly fetchChunks: ( + sourceId: string, + options?: FetchChunksOptions, + ) => Promise readonly initialPrefetchedChunksBySourceId?: PrefetchedChunksBySourceId readonly onSelectSource: (sourceId: string | null) => void readonly selectedSourceId: string | null @@ -68,6 +77,7 @@ export function useWorkspaceCitationFocus({ const [focusedPage, setFocusedPage] = useState({ pageNumber: null, requestId: 0, + citationId: null, }) const [pendingCitationId, setPendingCitationId] = useState( null, @@ -77,6 +87,8 @@ export function useWorkspaceCitationFocus({ const [fullChunkLoadingSourceId, setFullChunkLoadingSourceId] = useState< string | null >(null) + const { cache, mutate } = useSWRConfig() + const partialPrefetchSourceIdsRef = useRef>(new Set()) const fullChunkRequestsBySourceIdRef = useRef< Map> >(new Map()) @@ -109,12 +121,16 @@ export function useWorkspaceCitationFocus({ }, [], ) - const requestPageFocus = useCallback((pageNumber: number | null): void => { - setFocusedPage((current) => ({ - pageNumber, - requestId: current.requestId + 1, - })) - }, []) + const requestPageFocus = useCallback( + (pageNumber: number | null, citationId: string | null = null): void => { + setFocusedPage((current) => ({ + pageNumber, + citationId, + requestId: current.requestId + 1, + })) + }, + [], + ) const updatePrefetchedChunksBySourceId = useCallback( (updater: PrefetchedChunksUpdater): void => { @@ -147,14 +163,31 @@ export function useWorkspaceCitationFocus({ ) const loadAllChunksForSource = useCallback( - (sourceId: string): Promise => { + ( + sourceId: string, + options?: FetchChunksOptions, + ): Promise => { const existingRequest = fullChunkRequestsBySourceIdRef.current.get(sourceId) if (existingRequest) return existingRequest setFullChunkLoadingSourceId(sourceId) - const request = fetchChunks(sourceId) + const source = sources.find((candidate) => candidate.id === sourceId) + const chunkType = + options?.chunkType ?? + (source?.documentPresentation?.kind === "page-assets" + ? "page" + : undefined) + const fetchOptions: FetchChunksOptions | undefined = chunkType + ? { chunkType, ...options } + : options + const request = ( + fetchOptions + ? fetchChunks(sourceId, fetchOptions) + : fetchChunks(sourceId) + ) .then((chunks) => { + partialPrefetchSourceIdsRef.current.delete(sourceId) updatePrefetchedChunksBySourceId((current) => workspaceCitationState.upsertPrefetchedChunks( current, @@ -162,6 +195,13 @@ export function useWorkspaceCitationFocus({ chunks, ), ) + workspaceClientCache.hydrateSourceChunks( + (key, data) => { + void mutate(key, data, { revalidate: false }) + }, + sourceId, + chunks, + ) return chunks }) .finally(() => { @@ -174,25 +214,32 @@ export function useWorkspaceCitationFocus({ fullChunkRequestsBySourceIdRef.current.set(sourceId, request) return request }, - [fetchChunks, updatePrefetchedChunksBySourceId], + [fetchChunks, mutate, sources, updatePrefetchedChunksBySourceId], ) const handleLoadAllChunks = useCallback((): void => { + if (!selectedSourceId) return + + const hasPartialPrefetch = + partialPrefetchSourceIdsRef.current.has(selectedSourceId) if ( - !selectedSourceId || - prefetchedChunksBySourceIdRef.current[selectedSourceId] || - fullChunkRequestedSourceIdsRef.current.has(selectedSourceId) || - fullChunkRequestsBySourceIdRef.current.has(selectedSourceId) + !hasPartialPrefetch && + (prefetchedChunksBySourceIdRef.current[selectedSourceId] || + fullChunkRequestedSourceIdsRef.current.has(selectedSourceId) || + fullChunkRequestsBySourceIdRef.current.has(selectedSourceId)) ) { return } fullChunkRequestedSourceIdsRef.current.add(selectedSourceId) - void loadAllChunksForSource(selectedSourceId) - }, [ - loadAllChunksForSource, - selectedSourceId, - ]) + const source = sources.find((candidate) => candidate.id === selectedSourceId) + void loadAllChunksForSource( + selectedSourceId, + source?.documentPresentation?.kind === "page-assets" + ? { chunkType: "page" } + : undefined, + ) + }, [loadAllChunksForSource, selectedSourceId, sources]) const handleCitationClick = useCallback( async ( @@ -209,17 +256,62 @@ export function useWorkspaceCitationFocus({ if (!source) return setCitationListViewRequestId((current) => current + 1) + const pageNumber = workspaceCitationState.getCitationPageNumber(citation) + const isPageAssetSource = + source.documentPresentation?.kind === "page-assets" - const loadedChunkId = workspaceCitationState.getLoadedCitationChunkId({ - citation, - selectedSourceId, - sourceId: source.id, - selectedChunks, - hasMoreSelectedChunks, - }) - if (loadedChunkId) { - requestChunkFocus(loadedChunkId) - return + const focusFromChunks = ( + chunks: readonly ParsedChunkView[], + hasMore: boolean, + ): string | null => + workspaceCitationState.getLoadedCitationChunkId({ + citation, + selectedSourceId: source.id, + sourceId: source.id, + selectedChunks: chunks, + hasMoreSelectedChunks: hasMore, + }) + + const applyFocus = (chunkId: string | null): void => { + if (selectedSourceId !== source.id) onSelectSource(source.id) + requestChunkFocus(chunkId) + requestPageFocus(pageNumber, citationId) + } + + if (selectedSourceId === source.id) { + const loadedChunkId = workspaceCitationState.getLoadedCitationChunkId({ + citation, + selectedSourceId, + sourceId: source.id, + selectedChunks, + hasMoreSelectedChunks, + }) + if (loadedChunkId) { + requestChunkFocus(loadedChunkId) + requestPageFocus(pageNumber, citationId) + return + } + } + + const prefetchedChunks = prefetchedChunksBySourceIdRef.current[source.id] + if (prefetchedChunks) { + const prefetchedChunkId = focusFromChunks(prefetchedChunks, false) + if (prefetchedChunkId) { + applyFocus(prefetchedChunkId) + return + } + } + + const swrChunks = workspaceClientCache.getCachedSourceChunks( + cache, + source.id, + ) + if (swrChunks && swrChunks.length > 0) { + const cachedChunkId = focusFromChunks(swrChunks, true) + if (cachedChunkId) { + applyFocus(cachedChunkId) + return + } } if (!workspaceCitationState.hasExactCitationTargetHint(citation)) { @@ -229,43 +321,45 @@ export function useWorkspaceCitationFocus({ ) if (selectedSourceId !== source.id) onSelectSource(source.id) requestChunkFocus(null) + requestPageFocus(null) return } - const cachedChunks = prefetchedChunksBySourceIdRef.current[source.id] - if (cachedChunks) { - const cachedFocusId = - workspaceCitationState.getLoadedCitationChunkId({ - citation, - selectedSourceId: source.id, - sourceId: source.id, - selectedChunks: cachedChunks, - hasMoreSelectedChunks: false, - }) - updatePrefetchedChunksBySourceId((current) => - workspaceCitationState.upsertPrefetchedChunks( - current, + requestChunkFocus(null) + const fetchOptions: FetchChunksOptions | undefined = isPageAssetSource + ? { + chunkType: "page", + ...(pageNumber !== null ? { untilPageNumber: pageNumber } : {}), + } + : undefined + const chunks = + isPageAssetSource && pageNumber !== null + ? await fetchChunks(source.id, fetchOptions) + : await loadAllChunksForSource(source.id, fetchOptions) + + if (isPageAssetSource && pageNumber !== null) { + if (!fullChunkRequestsBySourceIdRef.current.has(source.id)) { + partialPrefetchSourceIdsRef.current.add(source.id) + updatePrefetchedChunksBySourceId((current) => + workspaceCitationState.upsertPrefetchedChunks( + current, + source.id, + chunks, + ), + ) + workspaceClientCache.hydrateSourceChunks( + (key, data) => { + void mutate(key, data, { revalidate: false }) + }, source.id, - cachedChunks, - ), - ) - if (selectedSourceId !== source.id) onSelectSource(source.id) - requestChunkFocus(cachedFocusId) - return + chunks, + ) + } } - requestChunkFocus(null) - const chunks = await loadAllChunksForSource(source.id) - const prefetchedChunkId = - workspaceCitationState.getLoadedCitationChunkId({ - citation, - selectedSourceId: source.id, - sourceId: source.id, - selectedChunks: chunks, - hasMoreSelectedChunks: false, - }) - onSelectSource(source.id) - requestChunkFocus(prefetchedChunkId) + applyFocus( + focusFromChunks(chunks, false), + ) } finally { setPendingCitationId((current) => current === citationId ? null : current, @@ -273,10 +367,14 @@ export function useWorkspaceCitationFocus({ } }, [ + cache, + fetchChunks, hasMoreSelectedChunks, loadAllChunksForSource, + mutate, onSelectSource, requestChunkFocus, + requestPageFocus, selectedChunks, selectedSourceId, sources, diff --git a/src/components/workspace-selected-chunks.test.ts b/src/components/workspace-selected-chunks.test.ts index 59cf03ac..7adb3ac8 100644 --- a/src/components/workspace-selected-chunks.test.ts +++ b/src/components/workspace-selected-chunks.test.ts @@ -239,7 +239,9 @@ describe("useWorkspaceSelectedChunks", () => { ]), ); expect(result.current.isSelectedChunksLoading).toBe(false); - expect(fetchChunkPageMock).toHaveBeenCalledWith("source_1", 1); + expect(fetchChunkPageMock).toHaveBeenCalledWith("source_1", 1, { + chunkType: "page", + }); }); it("detects page assets from selected source chunks without a second request", async () => { diff --git a/src/components/workspace-selected-chunks.ts b/src/components/workspace-selected-chunks.ts index 260a6f51..da357e2b 100644 --- a/src/components/workspace-selected-chunks.ts +++ b/src/components/workspace-selected-chunks.ts @@ -45,6 +45,10 @@ export function useWorkspaceSelectedChunks({ selectedSource && selectedSource.status === "ready" ? selectedSource.id : null + const pageChunkType = + selectedSource?.documentPresentation?.kind === "page-assets" + ? "page" + : undefined const { data: selectedChunkPages, isLoading: isChunksLoading, @@ -57,7 +61,12 @@ export function useWorkspaceSelectedChunks({ pageIndex, previousPageData, ), - fetchChunksByKey, + ([, sourceId, page]: SourceChunksKey) => + pageChunkType + ? workspaceClient.fetchChunkPage(sourceId, page, { + chunkType: pageChunkType, + }) + : workspaceClient.fetchChunkPage(sourceId, page), { revalidateIfStale: false, keepPreviousData: false, @@ -129,14 +138,6 @@ export function useWorkspaceSelectedChunks({ } } -function fetchChunksByKey([ - , - sourceId, - page, -]: SourceChunksKey): Promise { - return workspaceClient.fetchChunkPage(sourceId, page) -} - function getResolvedSelectedSource( source: SourceView | undefined, chunks: readonly ParsedChunkView[], diff --git a/src/components/workspace-shell-layout.test.ts b/src/components/workspace-shell-layout.test.ts index d414f1ad..f32d24d7 100644 --- a/src/components/workspace-shell-layout.test.ts +++ b/src/components/workspace-shell-layout.test.ts @@ -77,6 +77,11 @@ describe("WorkspaceShellLayout", () => { "350px", ) expect(screen.getByTestId("desktop-chat-panel").style.width).toBe("420px") + expect( + within(screen.getByTestId("desktop-sources-panel")).getByRole("button", { + name: "Open library", + }), + ).toBeTruthy() }) it("renders compact sidebars when the side panels are collapsed", () => { @@ -175,6 +180,11 @@ describe("WorkspaceShellLayout", () => { screen.getByRole("button", { name: "Show sources panel" }), ).toBeTruthy() expect(screen.getByRole("button", { name: "Show chat panel" })).toBeTruthy() + expect( + within(screen.getByTestId("desktop-sources-panel")).getByRole("button", { + name: "Open library", + }), + ).toBeTruthy() screen.getByRole("button", { name: "Open source Quarterly Report.pdf" }).click() expect(handleSourceSelected).toHaveBeenCalledWith("source_report") @@ -273,5 +283,10 @@ describe("WorkspaceShellLayout", () => { name: "Delete Very Long Quarterly Report Filename.pdf", }), ).toBeTruthy() + expect( + within(desktopSourcesPanel).getByRole("button", { + name: "Open library", + }), + ).toBeTruthy() }) }) diff --git a/src/components/workspace-shell-layout.tsx b/src/components/workspace-shell-layout.tsx index 72c67241..33c9394e 100644 --- a/src/components/workspace-shell-layout.tsx +++ b/src/components/workspace-shell-layout.tsx @@ -1,5 +1,6 @@ import { useCallback, type ReactElement } from "react" import { + BookOpen, Database, FileText, MessageSquare, @@ -44,6 +45,7 @@ type FocusedChunkState = { type FocusedPageState = { readonly pageNumber: number | null readonly requestId: number + readonly citationId?: string | null } type WorkspaceShellUser = { @@ -209,7 +211,9 @@ export function WorkspaceShellLayout( props.onDesktopPanelExpand("sources")} + onLibraryOpen={props.onLibraryOpen} onSourceSelected={props.onSourceSelected} /> ) : ( @@ -287,6 +291,7 @@ export function WorkspaceShellLayout( focusedChunkRequestId={props.focusedChunk.requestId} focusedPageNumber={focusedPage.pageNumber} focusedPageRequestId={focusedPage.requestId} + focusedCitationId={focusedPage.citationId ?? null} isLoading={props.isSelectedChunksLoading} isLoadingAllChunks={props.isSelectedAllChunksLoading} isLoadingMore={props.isSelectedChunksLoadingMore} @@ -432,6 +437,7 @@ export function WorkspaceShellLayout( focusedChunkRequestId={props.focusedChunk.requestId} focusedPageNumber={focusedPage.pageNumber} focusedPageRequestId={focusedPage.requestId} + focusedCitationId={focusedPage.citationId ?? null} isLoading={props.isSelectedChunksLoading} isLoadingAllChunks={props.isSelectedAllChunksLoading} isLoadingMore={props.isSelectedChunksLoadingMore} @@ -564,12 +570,16 @@ function DesktopPanelRestoreButton({ } function CompactSourcesSidebar({ + isLibraryOpen = false, onExpand, + onLibraryOpen, onSourceSelected, selectedSourceId, sources, }: { + readonly isLibraryOpen?: boolean readonly onExpand: () => void + readonly onLibraryOpen?: () => void readonly onSourceSelected: (sourceId: string | null) => void readonly selectedSourceId: string | null readonly sources: readonly SourceView[] @@ -582,6 +592,16 @@ function CompactSourcesSidebar({ onClick={onExpand} />
+ onLibraryOpen?.()} + > + + +
{sources.length === 0 ? ( { ); const desktopChunksPanel = within(screen.getByTestId("desktop-chunks-panel")); + await waitFor(() => { + expect(desktopChunksPanel.getByRole("button", { name: "List" })).toBeTruthy(); + }); + fireEvent.click(desktopChunksPanel.getByRole("button", { name: "List" })); await waitFor(() => { expect(desktopChunksPanel.getByRole("img", { name: "Page 1" })) .toBeTruthy(); @@ -537,7 +541,7 @@ describe("WorkspaceShell", () => { { id: "assistant_1", role: "assistant", - content: "The answer uses two sections.", + content: "The answer uses two sections. [[cite:1]] [[cite:2]]", citations: [ { content: "First cited section", @@ -620,19 +624,17 @@ describe("WorkspaceShell", () => { }); await user.click(sendButton); - await desktopChatPanel.findAllByRole("button", { - name: "Open source doc.pdf", - }); - const citationButtons = desktopChatPanel.getAllByRole( - "button", - { - name: "Open source doc.pdf", - }, - ); - expect(citationButtons).toHaveLength(2); - const firstCitation = citationButtons[0] as HTMLButtonElement; - const secondCitation = citationButtons[1] as HTMLButtonElement; - await user.click(firstCitation); + await desktopChatPanel.findAllByTestId("citation-chip"); + const getCitationChip = (citationId: string): HTMLButtonElement => { + const chip = desktopChatPanel + .getAllByTestId("citation-chip") + .find((element) => element.getAttribute("data-citation-id") === citationId); + expect(chip).toBeTruthy(); + return chip as HTMLButtonElement; + }; + expect(getCitationChip("assistant_1:0")).toBeTruthy(); + expect(getCitationChip("assistant_1:1")).toBeTruthy(); + await user.click(getCitationChip("assistant_1:0")); await waitFor(() => { expect( @@ -650,9 +652,9 @@ describe("WorkspaceShell", () => { expect(countFetches(fetch, "/api/sources/source_1/chunks")).toBe(1); await waitFor(() => { - expect(secondCitation.disabled).toBe(false); + expect(getCitationChip("assistant_1:1").disabled).toBe(false); }); - await user.click(secondCitation); + await user.click(getCitationChip("assistant_1:1")); await waitFor(() => { const topRow = screen diff --git a/src/domains/chat/citations.test.ts b/src/domains/chat/citations.test.ts index b6daa0a1..c07a69f6 100644 --- a/src/domains/chat/citations.test.ts +++ b/src/domains/chat/citations.test.ts @@ -44,6 +44,25 @@ describe("toChatCitationViews", () => { "https://blob.example/pages/page-000004.png", ) }) + + it("copies the page number from retrieval metadata onto the citation chip", () => { + const citations = toChatCitationViews( + [ + makeRetrievalResult({ + chunkType: "page", + metadata: { page_nums: [26] }, + source: { + documentId: "doc_1", + sourceFileName: "spacex-s1.pdf", + sectionPath: "spacex-s1.pdf / Overview", + }, + }), + ], + "Revenue grew.", + ) + + expect(citations[0]?.pageCitationPageNumber).toBe(26) + }) }) function makeRetrievalResult( diff --git a/src/domains/chat/citations.ts b/src/domains/chat/citations.ts index cdc9ddad..ea203f75 100644 --- a/src/domains/chat/citations.ts +++ b/src/domains/chat/citations.ts @@ -2,7 +2,10 @@ import type { RetrievalResult } from "@ontos-ai/knowhere-sdk" import type { Source } from "@/infrastructure/db/schema" import type { ChatCitationView } from "@/domains/chat/types" -import type { PageCitationAssetRetrievalResult } from "./page-citation-assets" +import { + resolvePageCitationPageNumber, + type PageCitationAssetRetrievalResult, +} from "./page-citation-assets" export function toChatCitationViews( results: readonly PageCitationAssetRetrievalResult[], @@ -12,6 +15,7 @@ export function toChatCitationViews( return results.map((result, index) => { const description = descriptionsBySourceNumber.get(index + 1) + const pageCitationPageNumber = resolvePageCitationPageNumber(result) return { content: result.content, chunkType: result.chunkType, @@ -20,8 +24,8 @@ export function toChatCitationViews( ...(result.pageCitationAssetUrl ? { pageCitationAssetUrl: result.pageCitationAssetUrl } : {}), - ...(result.pageCitationPageNumber - ? { pageCitationPageNumber: result.pageCitationPageNumber } + ...(pageCitationPageNumber + ? { pageCitationPageNumber } : {}), ...(description ? { description } : {}), source: { diff --git a/src/domains/chat/index.test.ts b/src/domains/chat/index.test.ts index f026dc69..e9556ac3 100644 --- a/src/domains/chat/index.test.ts +++ b/src/domains/chat/index.test.ts @@ -639,6 +639,61 @@ describe("answerQuestionWithRetrieval", () => { ]); }); + it("keeps two manifest citations to the same evidence chunk as two answer citations", async () => { + const result = makeRetrievalResult({ + content: "Revenue grew on the same page twice.", + chunkType: "page", + source: { + documentId: "doc_included", + sourceFileName: "spacex-s1.pdf", + sectionPath: "Page 26", + }, + }); + const retrieval = { + query: vi.fn().mockResolvedValue({ + results: [result], + evidenceText: "Revenue page evidence.", + referencedChunks: [], + namespace: "notebook-workspace", + query: "What grew?", + routerUsed: "workflow_single_step", + answerText: null, + }), + }; + const generateAnswer = vi.fn(async ({ searchSources }) => { + await searchSources({ query: "What grew?" }); + return makeHarnessRunResultWithLedger( + "Revenue grew [[cite:1]] and later expanded [[cite:2]].", + { + citations: [ + makeOutputCitation("r1:result:1", result), + makeOutputCitation("r1:result:1", result), + ], + chunks: [makeEvidenceChunkFromRetrievalResult("r1:result:1", result)], + }, + ); + }); + + const answer = await Effect.runPromise( + answerQuestionWithRetrieval({ + question: "What grew?", + namespace: "notebook-workspace", + sources: [makeSource()], + excludedSourceIds: [], + retrieval, + generateAnswer, + messages: [], + }), + ); + + expect(answer.answer).toBe( + "Revenue grew [[cite:1]] and later expanded [[cite:2]].", + ); + expect(answer.citations).toHaveLength(2); + expect(answer.citations[0]?.source.sectionPath).toBe("Page 26"); + expect(answer.citations[1]?.source.sectionPath).toBe("Page 26"); + }); + it("uses Notebook source titles instead of generated Knowhere filenames", async () => { const result = makeRetrievalResult({ source: { @@ -1215,6 +1270,87 @@ describe("answerQuestionWithRetrieval", () => { expect(answer.citations[0]?.pageCitationAssetUrl).toBe(hardenedPageAssetUrl); }); + it("hydrates page numbers for grep citations from the matching parsed chunk", async () => { + const grepChunk = { + ref: "grep1:match:1", + kind: "grep_match" as const, + chunkId: "chunk_financial_summary", + content: "ept percentages and per share data)\nTotal automotive revenues\n17,693", + contentPreview: "ept percentages and per share data)", + chunkType: "page", + score: null, + metadata: { + position: 1, + startOffset: 12, + endOffset: 80, + }, + source: { + documentId: "doc_tsla", + sourceFileName: "TSLA-Q4-2025-Update.pdf", + sectionPath: "FINANCIAL SUMMARY", + }, + }; + const retrieval = { + query: vi.fn(), + }; + const readChunks = vi.fn().mockResolvedValue({ + document: { + documentId: "doc_tsla", + sourceFileName: "TSLA-Q4-2025-Update.pdf", + }, + chunks: [ + { + position: 1, + chunkId: "chunk_financial_summary", + chunkType: "page", + content: "Full financial summary page.", + readableContent: "Full financial summary page.", + sectionPath: "FINANCIAL SUMMARY", + sourceChunkPath: "pages/page-4.md", + pageNumbers: [4, 5], + metadata: { + pageNums: [4, 5], + }, + }, + ], + }); + const generateAnswer = vi.fn(async () => + makeHarnessRunResultWithLedger( + "Automotive revenue was $17,693 million [[cite:1]].", + { + citations: [{ ref: "grep1:match:1" }], + chunks: [grepChunk], + }, + ), + ); + + const answer = await Effect.runPromise( + answerQuestionWithRetrieval({ + question: "Tesla automotive revenue", + namespace: "notebook-workspace", + sources: [ + makeSource({ + id: "source_tsla", + title: "TSLA-Q4-2025-Update.pdf", + knowhereDocumentId: "doc_tsla", + }), + ], + excludedSourceIds: [], + retrieval, + knowledge: { readChunks } as unknown as Knowledge, + generateAnswer, + messages: [], + }), + ); + + expect(readChunks).toHaveBeenCalledWith({ + documentId: "doc_tsla", + chunkId: "chunk_financial_summary", + }); + expect(answer.citations[0]?.pageCitationPageNumber).toBe(4); + expect(answer.citations[0]?.source.sectionPath).toBe("FINANCIAL SUMMARY"); + }); + it("returns only harness-selected artifacts when retrieval has extra media candidates", async () => { const frontAssetUrl = "https://blob.example/images/id-front.jpg"; const backAssetUrl = "https://blob.example/images/id-back.jpg"; diff --git a/src/domains/chat/index.ts b/src/domains/chat/index.ts index f45c88d0..a6990c46 100644 --- a/src/domains/chat/index.ts +++ b/src/domains/chat/index.ts @@ -37,7 +37,10 @@ import { enrichRetrievalResultsWithAssetUrls, removeRetrievedMediaAssetUrls, } from "./media-assets" -import { enrichRetrievalResultsWithPageCitationAssetUrls } from "./page-citation-assets" +import { + enrichRetrievalResultsWithPageCitationAssetUrls, + resolvePageCitationPageNumber, +} from "./page-citation-assets" import type { HardenableRetrievalResult } from "./media-asset-hardening" import { notebookKnowhereTools } from "./knowhere-tools" @@ -233,9 +236,15 @@ export const answerQuestionWithRetrieval = ( finalized: generatedAnswer.trace.finalized, }) - const rawResults = selectCitationRawResults({ - generatedAnswer, - }) + const rawResults = yield* Effect.tryPromise(() => + hydrateMissingCitationPageMetadata({ + results: selectCitationRawResults({ + generatedAnswer, + }), + ledgerChunks: generatedAnswer.trace.ledger.chunks, + knowledge: input.knowledge, + }), + ) if ( rawResults.length === 0 && generatedAnswer.manifest.text.trim().length === 0 && @@ -871,7 +880,6 @@ function mapManifestCitationsToResults( ) const results: RetrievalResult[] = [] - const seenKeys = new Set() for (const citation of result.manifest.citations) { const chunk = @@ -879,12 +887,11 @@ function mapManifestCitationsToResults( resolveChunkForAssetRef(citation.ref, assetsByRef, chunksByRef) if (!chunk) continue - const retrievalResult = toRetrievalResultFromEvidenceChunk(chunk) - const key = getRetrievalResultKey(retrievalResult) - if (seenKeys.has(key)) continue - - seenKeys.add(key) - results.push(retrievalResult) + results.push( + toRetrievalResultFromEvidenceChunk( + mergeChunkPageMetadata(chunk, result.trace.ledger.chunks), + ), + ) if (results.length >= MAX_CITATION_RESULTS) break } @@ -934,7 +941,9 @@ function mapDisplayedManifestArtifactsToResults( resolveChunkForAssetRef(sourceRef, assetsByRef, chunksByRef) if (!chunk) continue - const retrievalResult = toRetrievalResultFromEvidenceChunk(chunk) + const retrievalResult = toRetrievalResultFromEvidenceChunk( + mergeChunkPageMetadata(chunk, result.trace.ledger.chunks), + ) const key = getRetrievalResultKey(retrievalResult) if (seenKeys.has(key)) continue @@ -950,7 +959,9 @@ function mapDisplayedManifestArtifactsToResults( resolveChunkForAssetRef(artifact.ref, assetsByRef, chunksByRef) if (!chunk) continue - const retrievalResult = toRetrievalResultFromEvidenceChunk(chunk) + const retrievalResult = toRetrievalResultFromEvidenceChunk( + mergeChunkPageMetadata(chunk, result.trace.ledger.chunks), + ) const key = getRetrievalResultKey(retrievalResult) if (seenKeys.has(key)) continue @@ -962,6 +973,101 @@ function mapDisplayedManifestArtifactsToResults( return results } +function mergeChunkPageMetadata( + chunk: EvidenceChunk, + ledgerChunks: readonly EvidenceChunk[], +): EvidenceChunk { + if (hasResolvablePageNumber(chunk) || !chunk.chunkId) return chunk + + const donor = ledgerChunks.find( + (candidate) => + candidate.ref !== chunk.ref && + candidate.chunkId === chunk.chunkId && + hasResolvablePageNumber(candidate), + ) + if (!donor?.metadata) return chunk + + return { + ...chunk, + metadata: { ...donor.metadata, ...chunk.metadata }, + } +} + +async function hydrateMissingCitationPageMetadata(input: { + readonly results: readonly RetrievalResult[] + readonly ledgerChunks: readonly EvidenceChunk[] + readonly knowledge: AnswerQuestionInput["knowledge"] +}): Promise { + const results = input.results.map((result) => { + const donor = input.ledgerChunks.find( + (chunk) => + Boolean(result.chunkId) && + chunk.chunkId === result.chunkId && + hasResolvablePageNumber(chunk), + ) + if (!donor?.metadata || resolvePageCitationPageNumber(result)) { + return result + } + return { + ...result, + metadata: { ...donor.metadata, ...result.metadata }, + } + }) + + const knowledge = input.knowledge + if (!knowledge) return results + + return Promise.all( + results.map((result) => + hydrateResultPageMetadataFromKnowledge(result, knowledge), + ), + ) +} + +async function hydrateResultPageMetadataFromKnowledge( + result: RetrievalResult, + knowledge: NonNullable, +): Promise { + if (resolvePageCitationPageNumber(result)) return result + const documentId = result.source.documentId + const chunkId = result.chunkId + if (!documentId || !chunkId) return result + + try { + const response = await knowledge.readChunks({ documentId, chunkId }) + const chunk = response.chunks[0] + if (!chunk) return result + return { + ...result, + metadata: { + ...(chunk.metadata ?? {}), + ...(chunk.pageNumbers && chunk.pageNumbers.length > 0 + ? { pageNums: chunk.pageNumbers } + : {}), + ...result.metadata, + }, + } + } catch { + return result + } +} + +function hasResolvablePageNumber(chunk: EvidenceChunk): boolean { + return ( + resolvePageCitationPageNumber({ + content: chunk.content, + chunkType: chunk.chunkType, + score: chunk.score, + metadata: chunk.metadata, + source: { + documentId: chunk.source.documentId ?? undefined, + sourceFileName: chunk.source.sourceFileName ?? undefined, + sectionPath: chunk.source.sectionPath ?? undefined, + }, + }) !== undefined + ) +} + function toRetrievalResultFromEvidenceChunk( chunk: EvidenceChunk, ): RetrievalResult { diff --git a/src/domains/chat/page-citation-assets.test.ts b/src/domains/chat/page-citation-assets.test.ts index 4553d5d1..5b6c2b25 100644 --- a/src/domains/chat/page-citation-assets.test.ts +++ b/src/domains/chat/page-citation-assets.test.ts @@ -125,9 +125,52 @@ describe("enrichRetrievalResultsWithPageCitationAssetUrls", () => { }) expect(result?.pageCitationAssetUrl).toBeUndefined() + expect(result?.pageCitationPageNumber).toBe(4) }) - it("leaves non-page results unchanged even when they have page metadata", async () => { + it("sets the page number from page_nums even when page assets are missing", async () => { + const [result] = await enrichRetrievalResultsWithPageCitationAssetUrls({ + results: [ + makeRetrievalResult({ + chunkType: "page", + metadata: { page_nums: "25" }, + source: { + documentId: "doc_1", + sourceFileName: "spacex-s1.pdf", + sectionPath: "spacex-s1.pdf / MD&A", + }, + }), + ], + sources: [makeSource()], + }) + + expect(result?.pageCitationAssetUrl).toBeUndefined() + expect(result?.pageCitationPageNumber).toBe(25) + }) + + it("reads snake_case page_assets for the cited page number", async () => { + const [result] = await enrichRetrievalResultsWithPageCitationAssetUrls({ + results: [ + makeRetrievalResult({ + chunkType: "page", + metadata: { + page_nums: [11], + page_assets: [ + { + page_num: 11, + artifact_ref: "page_citation_assets/page-11.png", + }, + ], + }, + }), + ], + sources: [makeSource()], + }) + + expect(result?.pageCitationPageNumber).toBe(11) + }) + + it("attaches page numbers to text citations without turning them into page images", async () => { const [result] = await enrichRetrievalResultsWithPageCitationAssetUrls({ results: [ makeRetrievalResult({ @@ -148,6 +191,7 @@ describe("enrichRetrievalResultsWithPageCitationAssetUrls", () => { }) expect(result?.pageCitationAssetUrl).toBeUndefined() + expect(result?.pageCitationPageNumber).toBe(4) }) }) diff --git a/src/domains/chat/page-citation-assets.ts b/src/domains/chat/page-citation-assets.ts index 223498b1..7a388a75 100644 --- a/src/domains/chat/page-citation-assets.ts +++ b/src/domains/chat/page-citation-assets.ts @@ -48,32 +48,50 @@ export async function enrichRetrievalResultsWithPageCitationAssetUrls({ ) } +export function resolvePageCitationPageNumber( + result: RetrievalResult & { readonly pageCitationPageNumber?: number }, +): number | undefined { + const existing = getPositiveInteger(result.pageCitationPageNumber) + if (existing) return existing + + const pageNumbers = getPageNumbers(result.metadata) + const directAsset = getDirectPageCitationAsset(result, pageNumbers) + return ( + directAsset?.pageNum ?? + pageNumbers[0] ?? + parseSectionPathPageNumber(result.source.sectionPath) ?? + undefined + ) +} + async function enrichRetrievalResultWithPageCitationAssetUrl(input: { readonly result: RetrievalResult readonly sourcesByDocumentId: ReadonlyMap readonly hardenChatAssetUrl?: HardenChatAssetUrl }): Promise { - if (!isPageResult(input.result)) return input.result - const pageNumbers = getPageNumbers(input.result.metadata) const directAsset = getDirectPageCitationAsset(input.result, pageNumbers) - const sourceAssetUrl = await getStoredPageCitationAssetUrl({ - result: input.result, - directAsset, - sourcesByDocumentId: input.sourcesByDocumentId, - hardenChatAssetUrl: input.hardenChatAssetUrl, - }) - if (sourceAssetUrl || directAsset?.pageNum) { - return { - ...input.result, - ...(sourceAssetUrl ? { pageCitationAssetUrl: sourceAssetUrl } : {}), - ...(directAsset?.pageNum - ? { pageCitationPageNumber: directAsset.pageNum } - : {}), - } + const pageCitationPageNumber = resolvePageCitationPageNumber(input.result) + const sourceAssetUrl = isPageResult(input.result) + ? await getStoredPageCitationAssetUrl({ + result: input.result, + directAsset, + sourcesByDocumentId: input.sourcesByDocumentId, + hardenChatAssetUrl: input.hardenChatAssetUrl, + }) + : null + + if (!sourceAssetUrl && pageCitationPageNumber === undefined) { + return input.result } - return input.result + return { + ...input.result, + ...(sourceAssetUrl ? { pageCitationAssetUrl: sourceAssetUrl } : {}), + ...(pageCitationPageNumber + ? { pageCitationPageNumber } + : {}), + } } async function getStoredPageCitationAssetUrl(input: { @@ -105,7 +123,9 @@ function getDirectPageCitationAsset( result: RetrievalResult, pageNumbers: readonly number[], ): PageCitationAssetCandidate | null { - const candidates = parsePageCitationAssetCandidates(result.metadata?.pageAssets) + const candidates = parsePageCitationAssetCandidates( + result.metadata?.pageAssets ?? result.metadata?.page_assets, + ) if (pageNumbers.length > 0) { const matchingCandidates = candidates.filter((candidate) => @@ -124,15 +144,27 @@ function parsePageCitationAssetCandidates( return value.flatMap((item): PageCitationAssetCandidate[] => { if (!isRecord(item)) return [] - const pageNum = getPositiveInteger(item.pageNum) + const pageNum = + getPositiveInteger(item.pageNum) ?? + getPositiveInteger(item.page_num) ?? + getPositiveInteger(item.pageNumber) if (!pageNum) return [] return [ { pageNum, - artifactRef: getTrimmedString(item.artifactRef) ?? undefined, - assetUrl: getTrimmedString(item.assetUrl) ?? undefined, - contentType: getTrimmedString(item.contentType) ?? undefined, + artifactRef: + getTrimmedString(item.artifactRef) ?? + getTrimmedString(item.artifact_ref) ?? + undefined, + assetUrl: + getTrimmedString(item.assetUrl) ?? + getTrimmedString(item.asset_url) ?? + undefined, + contentType: + getTrimmedString(item.contentType) ?? + getTrimmedString(item.content_type) ?? + undefined, }, ] }) @@ -143,23 +175,38 @@ function getPageNumbers( ): readonly number[] { if (!metadata) return [] - const values = [metadata.pageNums, metadata.page_nums, metadata.pageNum] const pageNumbers = new Set() + collectPageNumbers(metadata.pageNums, pageNumbers) + collectPageNumbers(metadata.page_nums, pageNumbers) + collectPageNumbers(metadata.pageNum, pageNumbers) + collectPageNumbers(metadata.page_num, pageNumbers) + return [...pageNumbers].sort((left, right) => left - right) +} - for (const value of values) { - if (Array.isArray(value)) { - for (const item of value) { - const pageNum = getPositiveInteger(item) - if (pageNum) pageNumbers.add(pageNum) - } - continue - } - - const pageNum = getPositiveInteger(value) - if (pageNum) pageNumbers.add(pageNum) +function collectPageNumbers(value: unknown, pageNumbers: Set): void { + if (Array.isArray(value)) { + for (const item of value) collectPageNumbers(item, pageNumbers) + return } - return [...pageNumbers].sort((left, right) => left - right) + if (typeof value === "string" && value.includes(",")) { + for (const part of value.split(",")) collectPageNumbers(part.trim(), pageNumbers) + return + } + + const pageNum = getPositiveInteger(value) + if (pageNum) pageNumbers.add(pageNum) +} + +function parseSectionPathPageNumber( + sectionPath: string | null | undefined, +): number | null { + if (typeof sectionPath !== "string") return null + const match = + /\bpage\s+(\d+)\b/i.exec(sectionPath) ?? + /(?:^|[^\w])p(\d+)(?:[^\w]|$)/i.exec(sectionPath) + if (!match) return null + return getPositiveInteger(match[1]) } function getTrimmedString(value: unknown): string | null { @@ -169,11 +216,14 @@ function getTrimmedString(value: unknown): string | null { } function getPositiveInteger(value: unknown): number | null { - return typeof value === "number" && - Number.isSafeInteger(value) && - value > 0 - ? value - : null + if (typeof value === "number" && Number.isSafeInteger(value) && value > 0) { + return value + } + if (typeof value === "string") { + const parsed = Number.parseInt(value.trim(), 10) + if (Number.isSafeInteger(parsed) && parsed > 0) return parsed + } + return null } function isRecord(value: unknown): value is Record { diff --git a/src/domains/chunks/index.test.ts b/src/domains/chunks/index.test.ts index 586d32b1..e4ab72c4 100644 --- a/src/domains/chunks/index.test.ts +++ b/src/domains/chunks/index.test.ts @@ -205,6 +205,39 @@ describe("toParsedChunkView", () => { }); }); + it("maps demo API snake_case page citation assets on page chunks", () => { + const chunk = makeDocumentChunk({ + id: "document_page_1", + chunkId: "parser_page_1", + chunkType: "page" as DocumentChunk["chunkType"], + metadata: { + page_assets: [ + { + page_num: 8, + asset_url: + "/api/demo-sources/demo-tsla-q4-2025/assets/page_citation_assets/page-8.png", + content_type: "image/png", + width: 1200, + height: 1600, + }, + ], + }, + }); + + expect(toParsedChunkView(chunk, "manual.pdf", "doc_123")).toMatchObject({ + pageAssets: [ + { + pageNumber: 8, + assetUrl: + "/api/demo-sources/demo-tsla-q4-2025/assets/page_citation_assets/page-8.png", + contentType: "image/png", + width: 1200, + height: 1600, + }, + ], + }); + }); + it("maps SDK-normalized page number metadata", () => { const chunk = makeDocumentChunk({ metadata: { @@ -445,6 +478,22 @@ describe("loadChunkPageForSource", () => { }); }); + it("reads optional chunkType from the chunks query", () => { + expect( + getChunkPageParams( + new URLSearchParams({ + page: "2", + pageSize: "50", + chunkType: "page", + }), + ), + ).toEqual({ + page: 2, + pageSize: 50, + chunkType: "page", + }); + }); + it("resolves connection targets across already-loaded infinite pages", () => { const chunks = resolveChunkConnectionTargets([ makeParsedChunkView({ diff --git a/src/domains/chunks/index.ts b/src/domains/chunks/index.ts index e78293e8..7a38b886 100644 --- a/src/domains/chunks/index.ts +++ b/src/domains/chunks/index.ts @@ -40,9 +40,12 @@ export type LoadChunksOptions = { assetUrlsByFilePath?: Readonly> } +export type ChunkReadType = "text" | "image" | "table" | "page" + export type ChunkPageParams = { page: number pageSize: number + chunkType?: ChunkReadType } export type ChunkPagePagination = { @@ -60,6 +63,7 @@ export type ChunkPage = { export function getChunkPageParams( searchParams: URLSearchParams, ): ChunkPageParams { + const chunkType = parseChunkReadType(searchParams.get("chunkType")) return { page: normalizePositiveInteger(searchParams.get("page"), 1), pageSize: normalizePageSize( @@ -68,7 +72,15 @@ export function getChunkPageParams( defaultChunkPageSize, ), ), + ...(chunkType ? { chunkType } : {}), + } +} + +function parseChunkReadType(value: string | null): ChunkReadType | undefined { + if (value === "text" || value === "image" || value === "table" || value === "page") { + return value } + return undefined } export const loadChunksForSource = ( @@ -130,6 +142,7 @@ export const loadChunkPageForSource = ( client.documents.listChunks(source.knowhereDocumentId!, { page: params.page, pageSize: params.pageSize, + ...(params.chunkType ? { chunkType: params.chunkType } : {}), includeAssetUrls: true, }), ) diff --git a/src/domains/chunks/normalization.ts b/src/domains/chunks/normalization.ts index 3ab6bae7..53fe88b7 100644 --- a/src/domains/chunks/normalization.ts +++ b/src/domains/chunks/normalization.ts @@ -37,7 +37,10 @@ function createParsedChunkView( const summary = getStringMetadata(input.metadata, "summary") const pageAssets = type === "page" - ? getPageAssetViews(input.metadata["pageAssets"], assetUrl) + ? getPageAssetViews( + input.metadata["pageAssets"] ?? input.metadata["page_assets"], + assetUrl, + ) : undefined return { @@ -72,9 +75,15 @@ function getPageAssetViews( 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"]) + getPositiveInteger(item["pageNum"]) ?? + getPositiveInteger(item["pageNumber"]) ?? + getPositiveInteger(item["page_num"]) + const assetUrl = + getString(item["assetUrl"]) ?? + getString(item["asset_url"]) ?? + fallbackAssetUrl + const contentType = + getString(item["contentType"]) ?? getString(item["content_type"]) if (!pageNumber || !assetUrl || !contentType) return [] return [ diff --git a/src/domains/chunks/read.test.ts b/src/domains/chunks/read.test.ts index f63c8977..32bf8cd6 100644 --- a/src/domains/chunks/read.test.ts +++ b/src/domains/chunks/read.test.ts @@ -69,6 +69,43 @@ describe("readSourceChunkPage", () => { }) }) + it("passes chunkType through to knowledge.readChunks", async () => { + const listChunks = vi.fn() + const readChunks = vi.fn(async () => ({ + document: { + localDocumentId: "doc_1", + resultDirectoryPath: "parsed-storage:doc_1", + }, + chunks: [ + makeReadChunk({ + chunkType: "page", + filePath: "pages/page-000004.png", + }), + ], + page: 1, + pageSize: 50, + totalChunks: 8, + totalPages: 1, + })) + const knowledge = { readChunks } as unknown as Knowledge + + await readSourceChunkPage({ + client: { documents: { listChunks } }, + knowledge, + source: { documentId: "doc_1", title: "notes.pdf", revisionKey: "rev_1" }, + params: { page: 1, pageSize: 50, chunkType: "page" }, + }) + + expect(readChunks).toHaveBeenCalledWith({ + documentId: "doc_1", + revisionKey: "rev_1", + page: 1, + pageSize: 50, + chunkType: "page", + }) + expect(listChunks).not.toHaveBeenCalled() + }) + it("falls back to Knowhere asset URLs when the parsed-storage probe reads remote chunks", async () => { const readChunks = vi.fn(async () => ({ document: { diff --git a/src/domains/chunks/read.ts b/src/domains/chunks/read.ts index ae474cf3..552ba383 100644 --- a/src/domains/chunks/read.ts +++ b/src/domains/chunks/read.ts @@ -12,13 +12,12 @@ import { toParsedChunkViewFromReadChunk, type ChunkPage, type ChunkPageParams, + type ChunkReadType, } 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 @@ -32,7 +31,7 @@ type DisplayReadClient = { params: { readonly page: number readonly pageSize: number - readonly chunkType?: DisplayReadChunkType + readonly chunkType?: ChunkReadType readonly includeAssetUrls: true }, ): Promise<{ @@ -65,6 +64,7 @@ export async function readSourceChunkPage(input: { ...(input.source.revisionKey ? { revisionKey: input.source.revisionKey } : {}), page: input.params.page, pageSize: input.params.pageSize, + ...(input.params.chunkType ? { chunkType: input.params.chunkType } : {}), }) if (shouldUseKnowledgeChunkResponse(response)) { @@ -76,6 +76,7 @@ export async function readSourceChunkPage(input: { { page: input.params.page, pageSize: input.params.pageSize, + ...(input.params.chunkType ? { chunkType: input.params.chunkType } : {}), includeAssetUrls: true, }, ) @@ -130,6 +131,7 @@ export async function readAllSourceChunks(input: { readonly client: DisplayReadClient readonly knowledge: Knowledge readonly source: ReadableSource + readonly chunkType?: ChunkReadType }): Promise { const chunks: ParsedChunkView[] = [] let page = 1 @@ -143,6 +145,7 @@ export async function readAllSourceChunks(input: { : {}), page, pageSize: loadAllPageSize, + ...(input.chunkType ? { chunkType: input.chunkType } : {}), }) if (shouldUseKnowledgeChunkResponse(response)) { for (const chunk of response.chunks) { @@ -164,6 +167,7 @@ export async function readAllSourceChunks(input: { { page, pageSize: loadAllPageSize, + ...(input.chunkType ? { chunkType: input.chunkType } : {}), includeAssetUrls: true, }, ) diff --git a/src/domains/sources/route-chunks.ts b/src/domains/sources/route-chunks.ts index c636cb6a..9cfc99e2 100644 --- a/src/domains/sources/route-chunks.ts +++ b/src/domains/sources/route-chunks.ts @@ -2,7 +2,8 @@ import { Effect } from "effect" import { demoView } from "@/domains/demo/view" import { readAllSourceChunks, readSourceChunkPage } from "@/domains/chunks/read" -import { resolveChunkConnectionTargets } from "@/domains/chunks" +import { resolveChunkConnectionTargets, type ChunkReadType } from "@/domains/chunks" +import type { ParsedChunkView } from "@/domains/chunks/types" import type { DemoChunkPage } from "@/integrations/knowhere-demo" import { logger } from "@/lib/logger" import { routeResult } from "@/lib/route-result" @@ -118,6 +119,9 @@ const loadSourceChunksEffect = ( client: readResources.client, knowledge: readResources.knowledge, source: readableSource, + ...(input.pageParams.chunkType + ? { chunkType: input.pageParams.chunkType } + : {}), }), ).pipe( Effect.map((chunks) => @@ -199,6 +203,9 @@ const loadRemoteChunkPageEffect = ( client: readResources.client, knowledge: readResources.knowledge, source: readableSource, + ...(input.pageParams.chunkType + ? { chunkType: input.pageParams.chunkType } + : {}), }), ).pipe( Effect.map((chunks) => @@ -252,10 +259,13 @@ const loadDemoChunkPageEffect = ( status: "ready" as const, documentId: documentIdOverride ?? page.canonicalDocumentId, } - const chunks = pages.flatMap((demoChunkPage) => - demoChunkPage.chunks.map((chunk) => - demoView.toParsedChunkView(source, chunk), + const chunks = filterChunksByType( + pages.flatMap((demoChunkPage) => + demoChunkPage.chunks.map((chunk) => + demoView.toParsedChunkView(source, chunk), + ), ), + input.pageParams.chunkType, ) return routeResult.ok( @@ -319,6 +329,14 @@ function getErrorMessage(error: unknown): string { return String(error) } +function filterChunksByType( + chunks: readonly ParsedChunkView[], + chunkType: ChunkReadType | undefined, +): ParsedChunkView[] { + if (!chunkType) return [...chunks] + return chunks.filter((chunk) => chunk.type === chunkType) +} + function sourceNotFound(): JsonRouteResult<{ readonly message: string }> { return routeResult.error(404, "Source not found.") } diff --git a/src/domains/workspace/client-cache.test.ts b/src/domains/workspace/client-cache.test.ts index a04080a2..e2ee2cc8 100644 --- a/src/domains/workspace/client-cache.test.ts +++ b/src/domains/workspace/client-cache.test.ts @@ -52,4 +52,43 @@ describe("workspaceClientCache", () => { messages: [], }); }); + + it("reads cached source chunk pages for citation jumps", () => { + const cache = new Map([ + [ + unstable_serialize(["source-chunks", "source_1", 1]), + { + data: { + chunks: [ + { + chunkId: "page_4", + type: "page", + content: "Page 4", + sourceTitle: "report.pdf", + }, + ], + pagination: { + page: 1, + pageSize: 50, + total: 1, + totalPages: 1, + }, + }, + }, + ], + ]) as unknown as Cache; + + expect(workspaceClientCache.getCachedSourceChunks(cache, "source_1")).toEqual( + [ + { + chunkId: "page_4", + type: "page", + content: "Page 4", + sourceTitle: "report.pdf", + }, + ], + ); + expect(workspaceClientCache.getCachedSourceChunks(cache, "source_missing")) + .toBeNull(); + }); }); diff --git a/src/domains/workspace/client-cache.ts b/src/domains/workspace/client-cache.ts index af5baaef..1902db90 100644 --- a/src/domains/workspace/client-cache.ts +++ b/src/domains/workspace/client-cache.ts @@ -6,6 +6,7 @@ import type { ChatMessageView, ChatThreadView, } from "@/domains/chat/types" +import type { ParsedChunkView } from "@/domains/chunks/types" import type { SourceView } from "@/domains/sources/types" type SourceChunksResponse = Awaited< @@ -27,6 +28,10 @@ type WorkspaceClientCache = { cache: Cache, threadId: string, ) => ChatThreadDetailResponse | null + readonly getCachedSourceChunks: ( + cache: Cache, + sourceId: string, + ) => ParsedChunkView[] | null readonly getChatThreadKey: (threadId: string) => ChatThreadKey readonly getSourceChunksKey: ( sourceId: string | null, @@ -40,6 +45,11 @@ type WorkspaceClientCache = { pages: readonly SourceChunksResponse[] | undefined, ) => boolean readonly hasPendingSources: (sources: readonly SourceView[]) => boolean + readonly hydrateSourceChunks: ( + mutate: (key: SourceChunksKey, data: SourceChunksResponse) => unknown, + sourceId: string, + chunks: readonly ParsedChunkView[], + ) => void } function getSourceChunksKey( @@ -90,6 +100,57 @@ function hasMoreChunkPage(page: SourceChunksResponse): boolean { return page.pagination.page < page.pagination.totalPages } +function isSourceChunksResponse( + value: unknown, +): value is SourceChunksResponse { + if (!value || typeof value !== "object") return false + return Array.isArray((value as Partial).chunks) +} + +function getCachedSourceChunks( + cache: Cache, + sourceId: string, +): ParsedChunkView[] | null { + const chunks: ParsedChunkView[] = [] + let previous: SourceChunksResponse | null = null + + for (let pageIndex = 0; pageIndex < 200; pageIndex += 1) { + const key = getSourceChunksKey(sourceId, pageIndex, previous) + if (!key) break + + const cachedState = cache.get(unstable_serialize(key)) + const cachedData = cachedState?.data + if (!isSourceChunksResponse(cachedData)) { + return pageIndex === 0 ? null : chunks + } + + chunks.push(...(cachedData.chunks ?? [])) + previous = cachedData + } + + return chunks.length > 0 ? chunks : null +} + +function hydrateSourceChunks( + mutate: (key: SourceChunksKey, data: SourceChunksResponse) => unknown, + sourceId: string, + chunks: readonly ParsedChunkView[], +): void { + const key = getSourceChunksKey(sourceId, 0, null) + if (!key) return + + const total = chunks.length + mutate(key, { + chunks: [...chunks], + pagination: { + page: 1, + pageSize: Math.max(total, 1), + total, + totalPages: 1, + }, + }) +} + function isChatThreadDetailResponse( value: unknown, threadId: string, @@ -102,11 +163,13 @@ function isChatThreadDetailResponse( export const workspaceClientCache: WorkspaceClientCache = { getCachedChatThreadData, + getCachedSourceChunks, getChatThreadKey, getSourceChunksKey, hasLoadedChatThreadData, hasMoreChunkPages, hasPendingSources, + hydrateSourceChunks, } export type { diff --git a/src/domains/workspace/client.test.ts b/src/domains/workspace/client.test.ts index 47878bfb..f92cd32f 100644 --- a/src/domains/workspace/client.test.ts +++ b/src/domains/workspace/client.test.ts @@ -63,6 +63,48 @@ describe("workspaceClient", () => { }) }) + it("fetches all chunks filtered to page type until the cited page is found", async () => { + mockRouteClient.getJson.mockResolvedValueOnce({ + chunks: [ + { + chunkId: "page_1", + type: "page", + content: "Page 1", + sourceTitle: "report.pdf", + pageAssets: [{ pageNumber: 1, assetUrl: "https://a/1.png", contentType: "image/png" }], + }, + ], + pagination: { page: 1, pageSize: 50, total: 2, totalPages: 2 }, + }) + mockRouteClient.getJson.mockResolvedValueOnce({ + chunks: [ + { + chunkId: "page_4", + type: "page", + content: "Page 4", + sourceTitle: "report.pdf", + pageAssets: [{ pageNumber: 4, assetUrl: "https://a/4.png", contentType: "image/png" }], + }, + ], + pagination: { page: 2, pageSize: 50, total: 2, totalPages: 2 }, + }) + + const chunks = await workspaceClient.fetchChunks("source_1", { + chunkType: "page", + untilPageNumber: 4, + }) + + expect(mockRouteClient.getJson).toHaveBeenNthCalledWith( + 1, + "/api/sources/source_1/chunks?page=1&pageSize=50&chunkType=page", + ) + expect(mockRouteClient.getJson).toHaveBeenNthCalledWith( + 2, + "/api/sources/source_1/chunks?page=2&pageSize=50&chunkType=page", + ) + expect(chunks.map((chunk) => chunk.chunkId)).toEqual(["page_1", "page_4"]) + }) + it("preserves source chunk processing messages", async () => { mockRouteClient.getJson.mockResolvedValue({ chunks: [], diff --git a/src/domains/workspace/client.ts b/src/domains/workspace/client.ts index 77203832..1f6fbf4c 100644 --- a/src/domains/workspace/client.ts +++ b/src/domains/workspace/client.ts @@ -24,6 +24,11 @@ const workspaceClientConfig = { sourceChunkPageSize: 50, } as const +export type FetchChunksOptions = { + readonly chunkType?: "text" | "image" | "table" | "page" + readonly untilPageNumber?: number +} + type SourceChunksResponse = { chunks?: ParsedChunkView[] isProcessing?: boolean @@ -107,12 +112,23 @@ export const workspaceClient = { archiveChatThread, } as const -async function fetchChunks(sourceId: string): Promise { +async function fetchChunks( + sourceId: string, + options?: FetchChunksOptions, +): Promise { + if (options?.untilPageNumber != null) { + const { chunks } = await fetchChunkPagesUntil(sourceId, options) + return chunks + } + try { + const searchParams = new URLSearchParams() + if (options?.chunkType) searchParams.set("chunkType", options.chunkType) + const query = searchParams.toString() const body = await workspaceRouteClient.getJson<{ chunks?: ParsedChunkView[] }>( - `/api/sources/${encodeURIComponent(sourceId)}/chunks`, + `/api/sources/${encodeURIComponent(sourceId)}/chunks${query ? `?${query}` : ""}`, ) return Array.isArray(body.chunks) ? body.chunks : [] } catch { @@ -123,11 +139,13 @@ async function fetchChunks(sourceId: string): Promise { async function fetchChunkPage( sourceId: string, page: number, + options?: Pick, ): Promise { const searchParams = new URLSearchParams({ page: String(page), pageSize: String(workspaceClientConfig.sourceChunkPageSize), }) + if (options?.chunkType) searchParams.set("chunkType", options.chunkType) const body = await workspaceRouteClient.getJson( `/api/sources/${encodeURIComponent(sourceId)}/chunks?${searchParams.toString()}`, ) @@ -143,6 +161,49 @@ async function fetchChunkPage( } } +async function fetchChunkPagesUntil( + sourceId: string, + options: FetchChunksOptions, +): Promise<{ + readonly chunks: ParsedChunkView[] + readonly pages: SourceChunksResponse[] +}> { + const pages: SourceChunksResponse[] = [] + const chunks: ParsedChunkView[] = [] + let page = 1 + let totalPages = 1 + + do { + const result = await fetchChunkPage(sourceId, page, { + chunkType: options.chunkType, + }) + pages.push(result) + chunks.push(...(result.chunks ?? [])) + if ( + options.untilPageNumber != null && + chunksContainPage(chunks, options.untilPageNumber) + ) { + return { chunks, pages } + } + totalPages = result.pagination?.totalPages ?? 1 + page += 1 + } while (page <= totalPages) + + return { chunks, pages } +} + +function chunksContainPage( + chunks: readonly ParsedChunkView[], + pageNumber: number, +): boolean { + return chunks.some( + (chunk) => + (chunk.pageAssets ?? []).some( + (pageAsset) => pageAsset.pageNumber === pageNumber, + ) || (chunk.pageNums ?? []).includes(pageNumber), + ) +} + async function fetchSources(): Promise { const body = await workspaceRouteClient.getJson( workspaceClientKeys.sources, diff --git a/src/integrations/knowhere-demo.test.ts b/src/integrations/knowhere-demo.test.ts index 79268189..fe4f7e91 100644 --- a/src/integrations/knowhere-demo.test.ts +++ b/src/integrations/knowhere-demo.test.ts @@ -106,6 +106,80 @@ describe("knowhereDemoApi", () => { ) }) + it("rewrites demo page citation assets onto the Notebook asset proxy", async () => { + globalThis.fetch = vi.fn().mockResolvedValue( + new Response( + JSON.stringify({ + demo_source_id: "demo-tsla-q4-2025", + canonical_document_id: "demo-doc-tsla-q4-2025", + title: "TSLA-Q4-2025-Update.pdf", + mime_type: "application/pdf", + chunks: [ + { + id: "demo-tsla-q4-2025:page-1", + chunk_id: "page-1", + chunk_type: "page", + content: "Tesla energy storage deployments.", + section_path: "TSLA-Q4-2025-Update.pdf/SUMMARY", + source_chunk_path: "TSLA-Q4-2025-Update.pdf/SUMMARY", + file_path: "page_citation_assets/page-8.png", + sort_order: 8, + metadata: { + page_nums: [8], + page_assets: [ + { + page_num: 8, + artifact_ref: "page_citation_assets/page-8.png", + content_type: "image/png", + source: "knowhere-rendered-page-citation-source", + asset_url: + "/api/v1/demo/sources/demo-tsla-q4-2025/assets/page_citation_assets/page-8.png", + width: 1200, + height: 1600, + }, + ], + }, + asset_url: + "/api/v1/demo/sources/demo-tsla-q4-2025/assets/page_citation_assets/page-8.png", + }, + ], + pagination: { + page: 1, + page_size: 100, + total: 1, + total_pages: 1, + }, + }), + { status: 200, headers: { "content-type": "application/json" } }, + ), + ) + + const page = await knowhereDemoApi.fetchChunkPage({ + demoSourceId: "demo-tsla-q4-2025", + page: 1, + pageSize: 100, + }) + + expect(page.chunks[0]).toMatchObject({ + chunkType: "page", + filePath: "page_citation_assets/page-8.png", + assetUrl: + "/api/demo-sources/demo-tsla-q4-2025/assets/page_citation_assets/page-8.png", + metadata: { + pageAssets: [ + { + pageNum: 8, + contentType: "image/png", + assetUrl: + "/api/demo-sources/demo-tsla-q4-2025/assets/page_citation_assets/page-8.png", + width: 1200, + height: 1600, + }, + ], + }, + }) + }) + it("maps Official Library metadata from the demo catalog", async () => { globalThis.fetch = vi.fn().mockResolvedValue( new Response( diff --git a/src/integrations/knowhere-demo.ts b/src/integrations/knowhere-demo.ts index 8507dabc..f4342a59 100644 --- a/src/integrations/knowhere-demo.ts +++ b/src/integrations/knowhere-demo.ts @@ -481,6 +481,10 @@ function toDemoChunk( demoSourceId: string, chunk: DemoChunkResponse, ): DemoChunk { + const metadata = rewriteDemoPageAssetMetadata( + demoSourceId, + toRecord(chunk.metadata), + ) return { id: requireString(chunk.id), chunkId: requireString(chunk.chunk_id), @@ -490,8 +494,10 @@ function toDemoChunk( sourceChunkPath: optionalString(chunk.source_chunk_path) ?? null, filePath: optionalString(chunk.file_path) ?? null, sortOrder: requireNumber(chunk.sort_order), - metadata: toRecord(chunk.metadata), - assetUrl: toDemoAssetUrl(demoSourceId, optionalString(chunk.asset_url)), + metadata, + assetUrl: + toDemoAssetUrl(demoSourceId, optionalString(chunk.asset_url)) ?? + firstDemoPageAssetUrl(metadata), } } @@ -603,6 +609,70 @@ function toRecord(value: unknown): Readonly> { return value as Readonly> } +function rewriteDemoPageAssetMetadata( + demoSourceId: string, + metadata: Readonly>, +): Readonly> { + const rawAssets = metadata.pageAssets ?? metadata.page_assets + if (!Array.isArray(rawAssets)) return metadata + + const pageAssets = rawAssets.flatMap((item) => { + if (typeof item !== "object" || item === null || Array.isArray(item)) { + return [] + } + const asset = item as Readonly> + const pageNum = + optionalPositiveNumber(asset.pageNum) ?? + optionalPositiveNumber(asset.pageNumber) ?? + optionalPositiveNumber(asset.page_num) + const contentType = + optionalString(asset.contentType) ?? optionalString(asset.content_type) + const artifactRef = optionalString(asset.artifact_ref) + const remoteAssetUrl = + optionalString(asset.assetUrl) ?? + optionalString(asset.asset_url) ?? + (artifactRef + ? `/api/v1/demo/sources/${encodeURIComponent(demoSourceId)}/assets/${artifactRef}` + : undefined) + const assetUrl = toDemoAssetUrl(demoSourceId, remoteAssetUrl) + if (!pageNum || !contentType || !assetUrl) return [] + + return [ + { + pageNum, + assetUrl, + contentType, + ...(optionalPositiveNumber(asset.width) !== undefined + ? { width: optionalPositiveNumber(asset.width) } + : {}), + ...(optionalPositiveNumber(asset.height) !== undefined + ? { height: optionalPositiveNumber(asset.height) } + : {}), + }, + ] + }) + + const next: Record = { ...metadata, pageAssets } + delete next.page_assets + return next +} + +function firstDemoPageAssetUrl( + metadata: Readonly>, +): string | null { + const pageAssets = metadata.pageAssets + if (!Array.isArray(pageAssets) || pageAssets.length === 0) return null + const first = pageAssets[0] + if (typeof first !== "object" || first === null) return null + return optionalString((first as Readonly>).assetUrl) ?? null +} + +function optionalPositiveNumber(value: unknown): number | undefined { + return typeof value === "number" && Number.isFinite(value) && value > 0 + ? value + : undefined +} + function toDemoAssetUrl( demoSourceId: string, assetUrl: string | undefined, diff --git a/src/lib/app-metadata.ts b/src/lib/app-metadata.ts index 705b93d2..db9b41dc 100644 --- a/src/lib/app-metadata.ts +++ b/src/lib/app-metadata.ts @@ -1,7 +1,7 @@ import type { Metadata } from "next"; export const appMetadata: Metadata = { - title: "Knowhere Notebook", + title: "Knowhere Brain", description: "Upload documents, explore parsed content, and ask questions about your knowledge.", icons: {