diff --git a/src/domains/demo/view.test.ts b/src/domains/demo/view.test.ts new file mode 100644 index 00000000..593ff12b --- /dev/null +++ b/src/domains/demo/view.test.ts @@ -0,0 +1,140 @@ +import { describe, expect, it } from "vitest" + +import { chatCitationModel } from "@/components/chat-citation-model" +import { demoView } from "@/domains/demo/view" +import type { DemoCatalog, DemoCitation, DemoSource } from "@/integrations/knowhere-demo" + +describe("demoView.toChatMessages", () => { + it("uses the live citation mapper so demo answers render title/pN chips", () => { + const messages = demoView.toChatMessages( + makeCatalog({ + answer: + "Tesla entered an agreement to invest about $2 billion in xAI. [[cite:1]]", + citations: [makeCitation({ pageCitationPageNumber: 12 })], + }), + ) + const assistant = messages.find((message) => message.role === "assistant") + + expect(assistant?.content).toContain("[[cite:1]]") + expect(assistant?.citations).toEqual([ + expect.objectContaining({ + pageCitationPageNumber: 12, + pageCitationAssetUrl: + "/api/demo-sources/demo-tsla-q4-2025/assets/page_citation_assets/page-12.png", + source: expect.objectContaining({ + sourceFileName: "TSLA-Q4-2025-Update.pdf", + sectionPath: "TSLA-Q4-2025-Update.pdf/OTHER UPDATES", + }), + }), + ]) + expect( + chatCitationModel.getCitationChipLabel(assistant!.citations![0]!, {}), + ).toBe("TSLA-Q4-2025-Update.pdf/p12") + }) + + it("resolves page chips from catalog page_nums when the explicit page field is omitted", () => { + const messages = demoView.toChatMessages( + makeCatalog({ + title: "spacex-s1.pdf", + answer: + "The filing says SpaceX operates about 9,600 Starlink satellites. [[cite:1]]", + citations: [ + makeCitation({ + demoSourceId: "demo-spacex-s1", + canonicalDocumentId: "demo-doc-spacex-s1", + pageCitationPageNumber: undefined, + pageCitationAssetUrl: undefined, + pageNums: [28], + source: { + documentId: "demo-doc-spacex-s1", + sourceFileName: "spacex-s1.pdf", + sectionPath: "spacex-s1.pdf/Root", + }, + }), + ], + }), + ) + const assistant = messages.find((message) => message.role === "assistant") + + expect(assistant?.citations?.[0]?.pageCitationPageNumber).toBe(28) + expect( + chatCitationModel.getCitationChipLabel(assistant!.citations![0]!, {}), + ).toBe("spacex-s1.pdf/p28") + }) + + it("appends cite markers when the catalog answer still omits them", () => { + const messages = demoView.toChatMessages( + makeCatalog({ + answer: "Tesla entered an agreement to invest about $2 billion in xAI.", + citations: [makeCitation({ pageCitationPageNumber: 12 })], + }), + ) + const assistant = messages.find((message) => message.role === "assistant") + + expect(assistant?.content).toBe( + "Tesla entered an agreement to invest about $2 billion in xAI. [[cite:1]]", + ) + }) +}) + +function makeCatalog(input: { + readonly title?: string + readonly answer: string + readonly citations: readonly DemoCitation[] +}): DemoCatalog { + return { + officialLibrary: { categories: [], sources: [] }, + sources: [makeDemoSource(input)], + } +} + +function makeDemoSource(input: { + readonly title?: string + readonly answer: string + readonly citations: readonly DemoCitation[] +}): DemoSource { + return { + demoSourceId: "demo-tsla-q4-2025", + canonicalDocumentId: "demo-doc-tsla-q4-2025", + title: input.title ?? "TSLA-Q4-2025-Update.pdf", + mimeType: "application/pdf", + sizeBytes: 1024, + status: "ready", + chunkCount: 71, + originalFile: { + url: "/api/v1/demo/sources/demo-tsla-q4-2025/original", + mimeType: "application/pdf", + sizeBytes: 1024, + canDownload: false, + }, + examples: [ + { + id: "demo-tsla-q4-2025-xai", + question: "What does the document say about Tesla's xAI investment?", + answer: input.answer, + citations: input.citations, + }, + ], + } +} + +function makeCitation( + overrides: Partial = {}, +): DemoCitation { + return { + demoSourceId: "demo-tsla-q4-2025", + canonicalDocumentId: "demo-doc-tsla-q4-2025", + canonicalChunkId: "demo-tsla-q4-2025:chunk", + chunkId: "chunk", + chunkType: "page", + content: "Tesla entered into an agreement to invest approximately", + pageCitationAssetUrl: + "/api/demo-sources/demo-tsla-q4-2025/assets/page_citation_assets/page-12.png", + source: { + documentId: "demo-doc-tsla-q4-2025", + sourceFileName: "TSLA-Q4-2025-Update.pdf", + sectionPath: "TSLA-Q4-2025-Update.pdf/OTHER UPDATES", + }, + ...overrides, + } +} diff --git a/src/domains/demo/view.ts b/src/domains/demo/view.ts index 0b9a39b9..129c7fe9 100644 --- a/src/domains/demo/view.ts +++ b/src/domains/demo/view.ts @@ -1,10 +1,13 @@ -import { parsedChunkNormalization } from "@/domains/chunks/normalization" +import { toChatCitationViews } from "@/domains/chat/citations" +import type { PageCitationAssetRetrievalResult } from "@/domains/chat/page-citation-assets" import type { ChatMessageView } from "@/domains/chat/types" +import { parsedChunkNormalization } from "@/domains/chunks/normalization" import type { ParsedChunkView } from "@/domains/chunks/types" import { demoOriginalFile } from "@/domains/demo/original-file" import type { SourceView } from "@/domains/sources/types" import type { DemoCatalog, + DemoCitation, DemoChunk, DemoSource, } from "@/integrations/knowhere-demo" @@ -42,34 +45,65 @@ function toSourceView(source: DemoSource): SourceView { function toChatMessages(catalog: DemoCatalog): ChatMessageView[] { return catalog.sources.flatMap((source) => - source.examples.flatMap((example): ChatMessageView[] => [ - { - id: `${example.id}-user`, - role: "user", - content: example.question, - }, - { - id: `${example.id}-assistant`, - role: "assistant", - content: example.answer, - citations: example.citations.map((citation) => ({ - chunkType: citation.chunkType, - score: 0.95, - content: citation.content, - ...(citation.description - ? { description: citation.description } - : {}), - source: { - documentId: citation.canonicalDocumentId, - sourceFileName: citation.source.sourceFileName, - sectionPath: citation.source.sectionPath, - }, - })), - }, - ]), + source.examples.flatMap((example): ChatMessageView[] => { + const answer = withCitationMarkers( + example.answer, + example.citations.length, + ) + return [ + { + id: `${example.id}-user`, + role: "user", + content: example.question, + }, + { + id: `${example.id}-assistant`, + role: "assistant", + content: answer, + citations: toDemoChatCitationViews(example.citations, answer), + }, + ] + }), ) } +function toDemoChatCitationViews( + citations: readonly DemoCitation[], + answer: string, +) { + return toChatCitationViews( + citations.map(toDemoRetrievalResult), + answer, + ).map((citation, index) => { + const description = citations[index]?.description + return description ? { ...citation, description } : citation + }) +} + +function toDemoRetrievalResult( + citation: DemoCitation, +): PageCitationAssetRetrievalResult { + return { + content: citation.content, + chunkType: citation.chunkType, + score: 0.95, + ...(citation.pageCitationAssetUrl + ? { pageCitationAssetUrl: citation.pageCitationAssetUrl } + : {}), + ...(citation.pageCitationPageNumber + ? { pageCitationPageNumber: citation.pageCitationPageNumber } + : {}), + ...(citation.pageNums && citation.pageNums.length > 0 + ? { metadata: { page_nums: [...citation.pageNums] } } + : {}), + source: { + documentId: citation.canonicalDocumentId, + sourceFileName: citation.source.sourceFileName, + sectionPath: citation.source.sectionPath, + }, + } +} + function toParsedChunkView( source: SourceView, chunk: DemoChunk, @@ -87,3 +121,14 @@ function toParsedChunkView( sourceTitle: source.title, }) } + +const citeMarkerPattern = /\[\[cite:\d+\]\]/ + +function withCitationMarkers(answer: string, citationCount: number): string { + if (citationCount < 1 || citeMarkerPattern.test(answer)) return answer + const markers = Array.from( + { length: citationCount }, + (_, index) => `[[cite:${index + 1}]]`, + ).join(" ") + return `${answer.trimEnd()} ${markers}` +} diff --git a/src/domains/workspace/initial-state.test.ts b/src/domains/workspace/initial-state.test.ts index 9fc14505..0ed05e23 100644 --- a/src/domains/workspace/initial-state.test.ts +++ b/src/domains/workspace/initial-state.test.ts @@ -85,7 +85,7 @@ describe("loadWorkspaceShellInitialState", () => { { id: "demo-example-1-assistant", role: "assistant", - content: "Tesla delivered higher revenue.", + content: "Tesla delivered higher revenue. [[cite:1]]", citations: [ { chunkType: "text", diff --git a/src/integrations/knowhere-demo.test.ts b/src/integrations/knowhere-demo.test.ts index fe4f7e91..aa560418 100644 --- a/src/integrations/knowhere-demo.test.ts +++ b/src/integrations/knowhere-demo.test.ts @@ -252,6 +252,70 @@ describe("knowhereDemoApi", () => { chunkCount: 922, }) }) + + it("maps demo example citation page metadata onto Notebook citation fields", async () => { + globalThis.fetch = vi.fn().mockResolvedValue( + new Response( + JSON.stringify({ + sources: [ + { + 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", + size_bytes: 1024, + status: "ready", + chunk_count: 71, + original_file: { + url: "/api/v1/demo/sources/demo-tsla-q4-2025/original", + mime_type: "application/pdf", + size_bytes: 1024, + can_download: false, + }, + examples: [ + { + id: "demo-tsla-q4-2025-xai", + question: "What does the document say about Tesla's xAI investment?", + answer: + "Tesla entered an agreement to invest about $2 billion. [[cite:1]]", + citations: [ + { + demo_source_id: "demo-tsla-q4-2025", + canonical_document_id: "demo-doc-tsla-q4-2025", + canonical_chunk_id: "demo-tsla-q4-2025:chunk", + chunk_id: "chunk", + chunk_type: "page", + content: "Tesla entered into an agreement", + page_citation_page_number: 12, + page_nums: [12], + page_citation_asset_url: + "/api/v1/demo/sources/demo-tsla-q4-2025/assets/page_citation_assets/page-12.png", + source: { + document_id: "demo-doc-tsla-q4-2025", + source_file_name: "TSLA-Q4-2025-Update.pdf", + section_path: "TSLA-Q4-2025-Update.pdf/OTHER UPDATES", + }, + }, + ], + }, + ], + }, + ], + official_library: { categories: [], sources: [] }, + }), + { status: 200, headers: { "content-type": "application/json" } }, + ), + ) + + const catalog = await knowhereDemoApi.fetchCatalog() + + expect(catalog.sources[0]?.examples[0]?.citations[0]).toMatchObject({ + pageCitationPageNumber: 12, + pageNums: [12], + pageCitationAssetUrl: + "/api/demo-sources/demo-tsla-q4-2025/assets/page_citation_assets/page-12.png", + }) + }) }) function restoreEnv(key: string, value: string | undefined): void { diff --git a/src/integrations/knowhere-demo.ts b/src/integrations/knowhere-demo.ts index f4342a59..83c56eee 100644 --- a/src/integrations/knowhere-demo.ts +++ b/src/integrations/knowhere-demo.ts @@ -11,6 +11,9 @@ export type DemoCitation = { readonly chunkType: string readonly content: string readonly description?: string + readonly pageCitationPageNumber?: number + readonly pageCitationAssetUrl?: string + readonly pageNums?: readonly number[] readonly source: { readonly documentId: string readonly sourceFileName: string @@ -155,6 +158,9 @@ type DemoCitationResponse = { readonly chunk_type?: unknown readonly content?: unknown readonly description?: unknown + readonly page_citation_page_number?: unknown + readonly page_citation_asset_url?: unknown + readonly page_nums?: unknown readonly source?: { readonly document_id?: unknown readonly source_file_name?: unknown @@ -442,6 +448,13 @@ function toDemoExample(example: DemoExampleResponse): DemoExample { function toDemoCitation(citation: DemoCitationResponse): DemoCitation { const source = citation.source ?? {} const description = optionalString(citation.description) + const pageNums = toPositiveIntegers(citation.page_nums) + const pageCitationPageNumber = + optionalPositiveInteger(citation.page_citation_page_number) ?? pageNums[0] + const pageCitationAssetUrl = toDemoAssetUrl( + requireString(citation.demo_source_id), + optionalString(citation.page_citation_asset_url), + ) return { demoSourceId: requireString(citation.demo_source_id), canonicalDocumentId: requireString(citation.canonical_document_id), @@ -450,6 +463,11 @@ function toDemoCitation(citation: DemoCitationResponse): DemoCitation { chunkType: requireString(citation.chunk_type), content: requireString(citation.content), ...(description ? { description } : {}), + ...(pageCitationPageNumber !== undefined + ? { pageCitationPageNumber } + : {}), + ...(pageCitationAssetUrl ? { pageCitationAssetUrl } : {}), + ...(pageNums.length > 0 ? { pageNums } : {}), source: { documentId: requireString(source.document_id), sourceFileName: requireString(source.source_file_name), @@ -673,6 +691,20 @@ function optionalPositiveNumber(value: unknown): number | undefined { : undefined } +function optionalPositiveInteger(value: unknown): number | undefined { + return typeof value === "number" && Number.isSafeInteger(value) && value > 0 + ? value + : undefined +} + +function toPositiveIntegers(value: unknown): readonly number[] { + if (!Array.isArray(value)) return [] + return value.flatMap((item) => { + const page = optionalPositiveInteger(item) + return page === undefined ? [] : [page] + }) +} + function toDemoAssetUrl( demoSourceId: string, assetUrl: string | undefined,