diff --git a/package.json b/package.json index 94d1c3a..cf48c4a 100644 --- a/package.json +++ b/package.json @@ -28,7 +28,7 @@ "@effect/platform": "^0.96.1", "@napi-rs/canvas": "^1.0.2", "@neondatabase/serverless": "^1.1.0", - "@ontos-ai/knowhere-sdk": "^2.1.1", + "@ontos-ai/knowhere-sdk": "^2.1.2", "@radix-ui/react-alert-dialog": "^1.1.15", "@radix-ui/react-checkbox": "^1.3.3", "@radix-ui/react-dialog": "^1.1.15", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 48aabe5..12dc39f 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -24,8 +24,8 @@ importers: specifier: ^1.1.0 version: 1.1.0 '@ontos-ai/knowhere-sdk': - specifier: ^2.1.1 - version: 2.1.1 + specifier: ^2.1.2 + version: 2.1.2 '@radix-ui/react-alert-dialog': specifier: ^1.1.15 version: 1.1.15(@types/react-dom@19.2.3(@types/react@19.2.14))(@types/react@19.2.14)(react-dom@19.2.4(react@19.2.4))(react@19.2.4) @@ -1603,8 +1603,8 @@ packages: resolution: {integrity: sha512-nn5ozdjYQpUCZlWGuxcJY/KpxkWQs4DcbMCmKojjyrYDEAGy4Ce19NN4v5MduafTwJlbKc99UA8YhSVqq9yPZA==} engines: {node: '>=12.4.0'} - '@ontos-ai/knowhere-sdk@2.1.1': - resolution: {integrity: sha512-K33ylB/QjVYlLgow+/Hw+uNhQbWWrMTIqZ3IAlx5buHSI2EHPVP9+kDsfslwpfffYSn85j9Ab0/E9qTG0QotTw==} + '@ontos-ai/knowhere-sdk@2.1.2': + resolution: {integrity: sha512-m1wbQZNesExcg2yoEnmqFRlML/tprv11y/v7tlTLXEKTqKOY5zpG67idxXtbmw2qjOjQkR02HhgUwo/SdrFWPA==} engines: {node: '>=22.13.0', npm: '>=10.0.0', pnpm: '>=9.0.0'} '@open-draft/deferred-promise@2.2.0': @@ -6731,7 +6731,7 @@ snapshots: '@nolyfill/is-core-module@1.0.39': {} - '@ontos-ai/knowhere-sdk@2.1.1': + '@ontos-ai/knowhere-sdk@2.1.2': dependencies: axios: 1.18.1 jszip: 3.10.1 diff --git a/src/agent-harness/index.ts b/src/agent-harness/index.ts index bca503f..5fbdf7a 100644 --- a/src/agent-harness/index.ts +++ b/src/agent-harness/index.ts @@ -1,4 +1,4 @@ export * from "./ledger" +export * from "./knowhere-text" export * from "./runtime" export * from "./types" -export * from "./validator" diff --git a/src/agent-harness/knowhere-text.test.ts b/src/agent-harness/knowhere-text.test.ts new file mode 100644 index 0000000..88159f5 --- /dev/null +++ b/src/agent-harness/knowhere-text.test.ts @@ -0,0 +1,258 @@ +import { describe, expect, it } from "vitest" +import type { + KnowledgeGrepResponse, + KnowledgeOutline, + KnowledgeReadResponse, + RetrievalQueryResponse, +} from "@ontos-ai/knowhere-sdk" + +import { createEvidenceLedger } from "./ledger" +import { knowhereToolText } from "./knowhere-text" + +describe("knowhereToolText", () => { + it("formats search results as tagged text with refs and no raw asset URLs", () => { + const ledger = createEvidenceLedger() + const response = makeSearchResponse() + const snapshot = ledger.addRetrievalResponse(response) + + const text = knowhereToolText.formatSearch({ + response, + retrievalCount: snapshot.retrievalCount, + chunks: snapshot.chunks, + assets: snapshot.assets, + }) + + expect(text).toContain('') + expect(text).toContain('ref="r1:result:1"') + expect(text).toContain('ref="asset:r1:result:1"') + expect(text).toContain("Call inspectImage") + expect(text).not.toContain("https://assets.example/page-1.png") + }) + + it("formats list and outline responses", () => { + const listText = knowhereToolText.formatListDocuments({ + documents: [ + { + documentId: "doc_1", + revisionKey: "job_1", + namespace: "notebook", + sourceFileName: "report.pdf", + title: "Report", + status: "ready", + }, + ], + }) + const outlineText = knowhereToolText.formatOutline(makeOutlineResponse()) + + expect(listText).toContain('') + expect(listText).toContain('documentId="doc_1"') + expect(outlineText).toContain( + '', + ) + expect(outlineText).toContain('sectionPath="Root / Revenue"') + }) + + it("formats full read chunk bodies without truncating large content", () => { + const ledger = createEvidenceLedger() + const largeContent = `BEGIN ${"full chunk body ".repeat(700)} END` + const response = makeReadResponse(largeContent) + const snapshot = ledger.addReadChunksResponse(response) + + const text = knowhereToolText.formatReadChunks({ + response, + chunks: snapshot.chunks, + assets: snapshot.assets, + }) + + expect(text).toContain('') + expect(text).toContain('ref="read1:chunk:1"') + expect(text).toContain(largeContent) + expect(text).not.toContain("...[truncated]") + expect(text).not.toContain("https://assets.example/page-1.png") + }) + + it("formats grep matches with continuation metadata", () => { + const ledger = createEvidenceLedger() + const response = makeGrepResponse() + const snapshot = ledger.addGrepChunksResponse(response) + + const text = knowhereToolText.formatGrepChunks({ + response, + chunks: snapshot.chunks, + assets: snapshot.assets, + }) + + expect(text).toContain('') + expect(text).toContain('truncated="true"') + expect(text).toContain('continuationCursor="cursor_2"') + expect(text).toContain('ref="grep1:match:1"') + expect(text).toContain("matched penalty snippet") + }) + + it("formats errors as tagged text", () => { + expect( + knowhereToolText.formatError({ + operation: "read_chunks", + message: "A documentId is required.", + }), + ).toBe( + [ + '', + "", + "A documentId is required.", + "", + "", + ].join("\n"), + ) + }) +}) + +function makeSearchResponse(): RetrievalQueryResponse { + return { + namespace: "notebook", + query: "page one", + routerUsed: "workflow_single_step", + answerText: null, + evidenceText: "Page one evidence.", + stopReason: "answer_done", + failureReason: null, + results: [ + { + content: "Page one summary.", + chunkType: "page", + score: 0.91, + assetUrl: "https://assets.example/page-1.png", + metadata: { + pageNums: [1], + pageAssets: [ + { + pageNum: 1, + artifactRef: "page_citation_assets/page-1.png", + assetUrl: "https://assets.example/page-1.png", + contentType: "image/png", + }, + ], + }, + source: { + documentId: "doc_1", + sourceFileName: "report.pdf", + sectionPath: "Root / Page 1", + }, + }, + ], + referencedChunks: [], + } +} + +function makeReadResponse(content: string): KnowledgeReadResponse { + return { + document: { + localDocumentId: "doc_1", + documentId: "doc_1", + jobId: "job_1", + namespace: "notebook", + sourceFileName: "report.pdf", + chunkCount: 1, + typeCounts: { text: 0, image: 0, table: 0, page: 1 }, + resultDirectoryPath: "parsed-storage:doc_1/job_1", + createdAt: new Date("2026-01-01T00:00:00Z"), + updatedAt: new Date("2026-01-01T00:00:00Z"), + }, + chunks: [ + { + position: 1, + chunkId: "chunk_page_1", + chunkType: "page", + contentSource: "content", + content, + readableContent: content, + sectionPath: "Root / Page 1", + sourceChunkPath: "pages/page-1.md", + filePath: "pages/page-1.png", + assetUrl: "https://assets.example/page-1.png", + pageNumbers: [1], + metadata: { + pageNums: [1], + pageAssets: [ + { + pageNum: 1, + artifactRef: "page_citation_assets/page-1.png", + assetUrl: "https://assets.example/page-1.png", + contentType: "image/png", + }, + ], + }, + }, + ], + page: 1, + pageSize: 1, + totalChunks: 1, + totalPages: 1, + } +} + +function makeOutlineResponse(): KnowledgeOutline { + return { + document: { + localDocumentId: "doc_1", + documentId: "doc_1", + jobId: "job_1", + namespace: "notebook", + sourceFileName: "report.pdf", + chunkCount: 1, + typeCounts: { text: 1, image: 0, table: 0, page: 0 }, + resultDirectoryPath: "parsed-storage:doc_1/job_1", + createdAt: new Date("2026-01-01T00:00:00Z"), + updatedAt: new Date("2026-01-01T00:00:00Z"), + }, + totalChunks: 1, + typeCounts: { text: 1, image: 0, table: 0, page: 0 }, + sections: [ + { + sectionPath: "Root / Revenue", + sectionTitle: "Revenue", + sectionLevel: 2, + summary: "Revenue summary.", + startChunk: 1, + endChunk: 1, + chunkCount: 1, + typeCounts: { text: 1, image: 0, table: 0, page: 0 }, + children: [], + }, + ], + sectionTree: [], + } +} + +function makeGrepResponse(): KnowledgeGrepResponse { + return { + document: { + localDocumentId: "doc_1", + documentId: "doc_1", + jobId: "job_1", + namespace: "notebook", + sourceFileName: "contract.pdf", + chunkCount: 4, + typeCounts: { text: 4, image: 0, table: 0, page: 0 }, + resultDirectoryPath: "parsed-storage:doc_1/job_1", + createdAt: new Date("2026-01-01T00:00:00Z"), + updatedAt: new Date("2026-01-01T00:00:00Z"), + }, + matches: [ + { + position: 3, + chunkId: "chunk_3", + chunkType: "text", + sectionPath: "Root / Penalties", + sourceChunkPath: "chunks/chunk-3.md", + filePath: "contract.pdf", + startOffset: 10, + endOffset: 17, + snippet: "matched penalty snippet", + }, + ], + scannedChunks: 4, + truncated: true, + continuationCursor: "cursor_2", + } +} diff --git a/src/agent-harness/knowhere-text.ts b/src/agent-harness/knowhere-text.ts new file mode 100644 index 0000000..1c19368 --- /dev/null +++ b/src/agent-harness/knowhere-text.ts @@ -0,0 +1,301 @@ +import type { + KnowledgeGrepResponse, + KnowledgeOutline, + KnowledgeReadResponse, + RetrievalQueryResponse, +} from "@ontos-ai/knowhere-sdk" + +import type { + EvidenceAsset, + EvidenceChunk, + KnowhereListDocumentsResponse, +} from "./types" + +type EvidenceDelta = { + readonly chunks: readonly EvidenceChunk[] + readonly assets: readonly EvidenceAsset[] +} + +type SearchTextInput = EvidenceDelta & { + readonly response: RetrievalQueryResponse + readonly retrievalCount: number +} + +type ReadChunksTextInput = EvidenceDelta & { + readonly response: KnowledgeReadResponse +} + +type GrepChunksTextInput = EvidenceDelta & { + readonly response: KnowledgeGrepResponse +} + +type ErrorTextInput = { + readonly operation: KnowhereOperation + readonly message: string +} + +type KnowhereOperation = + | "search" + | "list_documents" + | "get_document_outline" + | "read_chunks" + | "grep_chunks" + +const assetInstruction = + "Notebook returned image/page asset refs. Call inspectImage with the asset refs when OCR, visual details, or verification are needed. Do not expose raw asset URLs." + +export const knowhereToolText = { + formatSearch(input: SearchTextInput): string { + return wrapKnowhereBlock("search", [ + formatTag("summary", { + retrievalCount: String(input.retrievalCount), + namespace: input.response.namespace, + query: input.response.query, + resultCount: String(input.response.results.length), + referencedChunkCount: String(input.response.referencedChunks.length), + stopReason: input.response.stopReason ?? undefined, + failureReason: input.response.failureReason ?? undefined, + }), + formatOptionalTextTag("evidence", input.response.evidenceText), + formatEvidenceChunks(input.chunks), + formatEvidenceAssets(input.assets), + formatAssetInstruction(input.assets), + ]) + }, + + formatListDocuments(response: KnowhereListDocumentsResponse): string { + return wrapKnowhereBlock("list_documents", [ + formatTag("summary", { documentCount: String(response.documents.length) }), + ...response.documents.map((document, index) => + formatSelfClosingTag("document", { + index: String(index + 1), + documentId: document.documentId, + localDocumentId: document.localDocumentId, + revisionKey: document.revisionKey, + namespace: document.namespace, + sourceFileName: document.sourceFileName, + title: document.title, + status: document.status, + chunkCount: + typeof document.chunkCount === "number" + ? String(document.chunkCount) + : undefined, + }), + ), + ]) + }, + + formatOutline(response: KnowledgeOutline): string { + return wrapKnowhereBlock("get_document_outline", [ + formatTag("document", { + documentId: response.document.documentId, + localDocumentId: response.document.localDocumentId, + revisionKey: response.document.jobId, + sourceFileName: response.document.sourceFileName, + totalChunks: String(response.totalChunks), + truncated: response.truncated === true ? "true" : undefined, + continuationCursor: response.continuationCursor, + }), + ...response.sections.map((section) => formatSection(section, 0)), + ]) + }, + + formatReadChunks(input: ReadChunksTextInput): string { + return wrapKnowhereBlock("read_chunks", [ + formatTag("document", { + documentId: input.response.document.documentId, + localDocumentId: input.response.document.localDocumentId, + revisionKey: input.response.document.jobId, + sourceFileName: input.response.document.sourceFileName, + page: + typeof input.response.page === "number" + ? String(input.response.page) + : undefined, + pageSize: + typeof input.response.pageSize === "number" + ? String(input.response.pageSize) + : undefined, + totalChunks: + typeof input.response.totalChunks === "number" + ? String(input.response.totalChunks) + : undefined, + totalPages: + typeof input.response.totalPages === "number" + ? String(input.response.totalPages) + : undefined, + nextChunk: + typeof input.response.nextChunk === "number" + ? String(input.response.nextChunk) + : undefined, + }), + formatEvidenceChunks(input.chunks), + formatEvidenceAssets(input.assets), + formatAssetInstruction(input.assets), + ]) + }, + + formatGrepChunks(input: GrepChunksTextInput): string { + return wrapKnowhereBlock("grep_chunks", [ + formatTag("document", { + documentId: input.response.document.documentId, + localDocumentId: input.response.document.localDocumentId, + revisionKey: input.response.document.jobId, + sourceFileName: input.response.document.sourceFileName, + matchCount: String(input.response.matches.length), + scannedChunks: String(input.response.scannedChunks), + truncated: input.response.truncated ? "true" : "false", + continuationCursor: input.response.continuationCursor, + }), + formatEvidenceChunks(input.chunks), + ]) + }, + + formatError(input: ErrorTextInput): string { + return [ + formatOpenTag("knowhere", { + operation: input.operation, + status: "error", + }), + formatTextTag("message", input.message), + "", + ].join("\n") + }, +} as const + +function wrapKnowhereBlock( + operation: KnowhereOperation, + parts: readonly string[], +): string { + return [ + formatOpenTag("knowhere", { operation, status: "ok" }), + ...parts.filter((part) => part.trim().length > 0), + "", + ].join("\n") +} + +function formatEvidenceChunks(chunks: readonly EvidenceChunk[]): string { + if (chunks.length === 0) return "" + + return [ + "", + ...chunks.map((chunk) => + [ + formatOpenTag("chunk", { + ref: chunk.ref, + kind: chunk.kind, + chunkId: chunk.chunkId, + chunkType: chunk.chunkType, + score: chunk.score === null ? undefined : String(chunk.score), + documentId: chunk.source.documentId ?? undefined, + sourceFileName: chunk.source.sourceFileName ?? undefined, + sectionPath: chunk.source.sectionPath ?? undefined, + sourceChunkPath: chunk.sourceChunkPath ?? undefined, + filePath: chunk.filePath ?? undefined, + assetRef: chunk.assetRef, + }), + formatTextTag("content", chunk.content), + "", + ].join("\n"), + ), + "", + ].join("\n") +} + +function formatEvidenceAssets(assets: readonly EvidenceAsset[]): string { + if (assets.length === 0) return "" + + return [ + "", + ...assets.map((asset) => + formatSelfClosingTag("asset", { + ref: asset.ref, + chunkRef: asset.chunkRef, + type: asset.type, + label: asset.label, + sourcePath: asset.sourcePath, + documentId: asset.source.documentId ?? undefined, + sourceFileName: asset.source.sourceFileName ?? undefined, + sectionPath: asset.source.sectionPath ?? undefined, + }), + ), + "", + ].join("\n") +} + +function formatAssetInstruction(assets: readonly EvidenceAsset[]): string { + if (!assets.some((asset) => asset.type === "image")) return "" + return formatTextTag("asset_instruction", assetInstruction) +} + +function formatSection( + section: KnowledgeOutline["sections"][number], + depth: number, +): string { + return [ + formatOpenTag("section", { + depth: String(depth), + sectionPath: section.sectionPath, + sectionTitle: section.sectionTitle, + sectionLevel: String(section.sectionLevel), + startChunk: + typeof section.startChunk === "number" + ? String(section.startChunk) + : undefined, + endChunk: + typeof section.endChunk === "number" + ? String(section.endChunk) + : undefined, + chunkCount: String(section.chunkCount), + }), + formatOptionalTextTag("summary", section.summary), + ...section.children.map((child) => formatSection(child, depth + 1)), + "", + ] + .filter((part) => part.trim().length > 0) + .join("\n") +} + +function formatOptionalTextTag( + tagName: string, + value: string | null | undefined, +): string { + const trimmedValue = value?.trim() + return trimmedValue ? formatTextTag(tagName, trimmedValue) : "" +} + +function formatTextTag(tagName: string, value: string): string { + return [`<${tagName}>`, value, ``].join("\n") +} + +function formatTag( + tagName: string, + attrs: Readonly>, +): string { + return `${formatOpenTag(tagName, attrs)}` +} + +function formatSelfClosingTag( + tagName: string, + attrs: Readonly>, +): string { + return `${formatOpenTag(tagName, attrs).slice(0, -1)} />` +} + +function formatOpenTag( + tagName: string, + attrs: Readonly>, +): string { + const serializedAttrs = Object.entries(attrs) + .filter((entry): entry is [string, string] => typeof entry[1] === "string") + .map(([key, value]) => `${key}="${escapeAttribute(value)}"`) + .join(" ") + return serializedAttrs ? `<${tagName} ${serializedAttrs}>` : `<${tagName}>` +} + +function escapeAttribute(value: string): string { + return value + .replaceAll("&", "&") + .replaceAll('"', """) + .replaceAll("<", "<") + .replaceAll(">", ">") +} diff --git a/src/agent-harness/ledger.test.ts b/src/agent-harness/ledger.test.ts index 22e0f6a..061a892 100644 --- a/src/agent-harness/ledger.test.ts +++ b/src/agent-harness/ledger.test.ts @@ -1,5 +1,9 @@ import { describe, expect, it } from "vitest" -import type { RetrievalQueryResponse } from "@ontos-ai/knowhere-sdk" +import type { + KnowledgeGrepResponse, + KnowledgeReadResponse, + RetrievalQueryResponse, +} from "@ontos-ai/knowhere-sdk" import { createEvidenceLedger } from "./ledger" @@ -63,6 +67,50 @@ describe("createEvidenceLedger", () => { }), ) }) + + it("adds read chunk refs and page image assets", () => { + const ledger = createEvidenceLedger() + + const snapshot = ledger.addReadChunksResponse(makeReadResponse()) + + expect(snapshot.chunks).toEqual([ + expect.objectContaining({ + ref: "read1:chunk:1", + kind: "read_chunk", + content: "Full page content.", + contentPreview: "Full page content.", + assetRef: "asset:read1:chunk:1", + }), + ]) + expect(snapshot.assets).toEqual([ + expect.objectContaining({ + ref: "asset:read1:chunk:1", + chunkRef: "read1:chunk:1", + type: "image", + sourcePath: "page_citation_assets/page-1.png", + }), + ]) + }) + + it("adds grep match refs", () => { + const ledger = createEvidenceLedger() + + const snapshot = ledger.addGrepChunksResponse(makeGrepResponse()) + + expect(snapshot.chunks).toEqual([ + expect.objectContaining({ + ref: "grep1:match:1", + kind: "grep_match", + chunkId: "chunk_3", + content: "matched penalty snippet", + source: expect.objectContaining({ + documentId: "doc_contract", + sourceFileName: "contract.pdf", + sectionPath: "Root / Penalties", + }), + }), + ]) + }) }) function makeRetrievalResponse(): RetrievalQueryResponse { @@ -135,3 +183,77 @@ function makePageAssetUrlRetrievalResponse(): RetrievalQueryResponse { ], } } + +function makeReadResponse(): KnowledgeReadResponse { + return { + document: { + localDocumentId: "doc_contract", + documentId: "doc_contract", + jobId: "job_contract", + namespace: "notebook", + sourceFileName: "contract.pdf", + chunkCount: 1, + typeCounts: { text: 0, image: 0, table: 0, page: 1 }, + resultDirectoryPath: "parsed-storage:doc_contract/job_contract", + createdAt: new Date("2026-01-01T00:00:00Z"), + updatedAt: new Date("2026-01-01T00:00:00Z"), + }, + chunks: [ + { + position: 1, + chunkId: "chunk_page_1", + chunkType: "page", + content: "Full page content.", + readableContent: "Full page content.", + sectionPath: "Root / Page 1", + sourceChunkPath: "pages/page-1.md", + filePath: "pages/page-1.png", + assetUrl: "https://assets.example/page-1.png", + pageNumbers: [1], + metadata: { + pageNums: [1], + pageAssets: [ + { + pageNum: 1, + artifactRef: "page_citation_assets/page-1.png", + assetUrl: "https://assets.example/page-1.png", + contentType: "image/png", + }, + ], + }, + }, + ], + } +} + +function makeGrepResponse(): KnowledgeGrepResponse { + return { + document: { + localDocumentId: "doc_contract", + documentId: "doc_contract", + jobId: "job_contract", + namespace: "notebook", + sourceFileName: "contract.pdf", + chunkCount: 4, + typeCounts: { text: 4, image: 0, table: 0, page: 0 }, + resultDirectoryPath: "parsed-storage:doc_contract/job_contract", + createdAt: new Date("2026-01-01T00:00:00Z"), + updatedAt: new Date("2026-01-01T00:00:00Z"), + }, + matches: [ + { + position: 3, + chunkId: "chunk_3", + chunkType: "text", + sectionPath: "Root / Penalties", + sourceChunkPath: "chunks/chunk-3.md", + filePath: "contract.pdf", + startOffset: 10, + endOffset: 17, + snippet: "matched penalty snippet", + }, + ], + scannedChunks: 4, + truncated: false, + } +} diff --git a/src/agent-harness/ledger.ts b/src/agent-harness/ledger.ts index f9d859d..d782ae3 100644 --- a/src/agent-harness/ledger.ts +++ b/src/agent-harness/ledger.ts @@ -1,4 +1,8 @@ import type { + KnowledgeGrepMatch, + KnowledgeGrepResponse, + KnowledgeReadChunk, + KnowledgeReadResponse, RetrievalQueryResponse, RetrievalResult, } from "@ontos-ai/knowhere-sdk" @@ -14,6 +18,8 @@ const imageExtensions = [".jpg", ".jpeg", ".png", ".gif", ".webp", ".svg"] as co type MutableLedger = { retrievalCount: number + readCount: number + grepCount: number chunks: EvidenceChunk[] assets: EvidenceAsset[] evidenceText: string[] @@ -41,6 +47,8 @@ export type EvidenceLedger = ReturnType export function createEvidenceLedger() { const ledger: MutableLedger = { retrievalCount: 0, + readCount: 0, + grepCount: 0, chunks: [], assets: [], evidenceText: [], @@ -104,6 +112,38 @@ export function createEvidenceLedger() { return snapshot(ledger) }, + addReadChunksResponse(response: KnowledgeReadResponse): EvidenceLedgerSnapshot { + ledger.readCount += 1 + const readIndex = ledger.readCount + + response.chunks.forEach((chunk, index) => { + addChunkFromReadChunk({ + ledger, + response, + chunk, + ref: `read${readIndex}:chunk:${index + 1}`, + }) + }) + + return snapshot(ledger) + }, + + addGrepChunksResponse(response: KnowledgeGrepResponse): EvidenceLedgerSnapshot { + ledger.grepCount += 1 + const grepIndex = ledger.grepCount + + response.matches.forEach((match, index) => { + addChunkFromGrepMatch({ + ledger, + response, + match, + ref: `grep${grepIndex}:match:${index + 1}`, + }) + }) + + return snapshot(ledger) + }, + read(ref: string, offset = 0, limit = 4_000) { const chunk = ledger.chunks.find((candidate) => candidate.ref === ref) if (!chunk) { @@ -179,6 +219,69 @@ function addChunkFromResult(input: { }) } +function addChunkFromReadChunk(input: { + readonly ledger: MutableLedger + readonly response: KnowledgeReadResponse + readonly chunk: KnowledgeReadChunk + readonly ref: string +}): void { + addChunk({ + ledger: input.ledger, + chunk: { + ref: input.ref, + kind: "read_chunk", + chunkId: input.chunk.chunkId, + content: input.chunk.content, + contentPreview: buildContentPreview(input.chunk.content), + chunkType: input.chunk.chunkType, + score: null, + sourceChunkPath: input.chunk.sourceChunkPath, + filePath: input.chunk.filePath, + metadata: input.chunk.metadata, + source: { + documentId: input.response.document.documentId, + sourceFileName: input.response.document.sourceFileName, + sectionPath: input.chunk.sectionPath, + }, + revisionKey: input.response.document.jobId, + ...(input.chunk.assetUrl ? { assetUrl: input.chunk.assetUrl } : {}), + }, + }) +} + +function addChunkFromGrepMatch(input: { + readonly ledger: MutableLedger + readonly response: KnowledgeGrepResponse + readonly match: KnowledgeGrepMatch + readonly ref: string +}): void { + addChunk({ + ledger: input.ledger, + chunk: { + ref: input.ref, + kind: "grep_match", + chunkId: input.match.chunkId, + content: input.match.snippet, + contentPreview: buildContentPreview(input.match.snippet), + chunkType: input.match.chunkType, + score: null, + sourceChunkPath: input.match.sourceChunkPath, + filePath: input.match.filePath, + metadata: { + position: input.match.position, + startOffset: input.match.startOffset, + endOffset: input.match.endOffset, + }, + source: { + documentId: input.response.document.documentId, + sourceFileName: input.response.document.sourceFileName, + sectionPath: input.match.sectionPath, + }, + revisionKey: input.response.document.jobId, + }, + }) +} + function addChunk(input: { readonly ledger: MutableLedger readonly chunk: Omit diff --git a/src/agent-harness/runtime.test.ts b/src/agent-harness/runtime.test.ts index 3e3c5fa..04614a7 100644 --- a/src/agent-harness/runtime.test.ts +++ b/src/agent-harness/runtime.test.ts @@ -14,8 +14,8 @@ import type { ContextPolicy, HarnessToolCallTrace, IntentFrame, + KnowhereToolRuntime, OutputManifest, - RetrievalCapability, } from "./types" describe("agent harness runtime", () => { @@ -28,59 +28,45 @@ describe("agent harness runtime", () => { expect(prompt).not.toContain("navigation action") }) - it("passes only outer retrieval parameters to KNOWHERE after intent and context policy are declared", async () => { - const query = vi.fn().mockResolvedValue( + it("tells the agent citation metadata is optional and ledger-resolved", () => { + const prompt = buildHarnessSystemPrompt(makeTurnInput()) + + expect(prompt).toContain( + "Citation label and source metadata are optional", + ) + expect(prompt).toContain( + "Notebook resolves citation metadata from evidence refs when possible", + ) + expect(prompt).not.toContain("must match the selected evidence ref exactly") + }) + + it("passes only outer retrieval parameters to KNOWHERE without planning-tool gating", async () => { + const query = vi.fn().mockResolvedValue( makeRetrievalResponse(), ) const state: { - intent?: IntentFrame - contextPolicy?: ContextPolicy toolCalls?: HarnessToolCallTrace[] } = {} const tools = createHarnessTools({ state, ledger: createEvidenceLedger(), - retrieval: { query }, + knowhereTools: makeKnowhereTools(query), recentTurns: [], }) - expect(await executeTool(tools.retrieve, { query: "q4 chart" })).toEqual({ - ok: false, - message: "declareIntent must be called before retrieve.", - }) - - await executeTool(tools.declareIntent, { - task: "show_media", - dependsOnPreviousTurn: false, - retrievalNeeded: "yes", - targetModalities: ["image"], - constraints: { desiredCount: 2, maxCount: 2 }, - groundingPolicy: "must_use_sources", - }) - expect(await executeTool(tools.retrieve, { query: "q4 chart" })).toEqual({ - ok: false, - message: "setContextPolicy must be called before retrieve.", - }) - - await executeTool(tools.setContextPolicy, { - carryHistory: "none", - reason: "The current request is unrelated to previous turns.", - activePriorTurnIds: [], - }) - const result = await executeTool(tools.retrieve, { + const result = await executeTool(tools.knowhere_search, { query: "q4 chart", - modalities: ["image"], + targetContent: "image", topK: 2, purpose: "Find the two requested charts.", }) - expect(result).toMatchObject({ - ok: true, - retrievalCount: 1, - }) + expect(result).toContain('') + expect(result).toContain('ref="r1:result:1"') + expect(result).toContain('ref="asset:r1:result:1"') expect(query).toHaveBeenCalledWith({ query: "q4 chart", - modalities: ["image"], + targetContent: "image", topK: 2, purpose: "Find the two requested charts.", signalPaths: undefined, @@ -91,17 +77,13 @@ describe("agent harness runtime", () => { "LegalAction", ) expect(state.toolCalls?.map((call) => [call.tool, call.ok])).toEqual([ - ["retrieve", false], - ["declareIntent", true], - ["retrieve", false], - ["setContextPolicy", true], - ["retrieve", true], + ["knowhere_search", true], ]) }) - it("returns only newly retrieved evidence in each retrieve tool result", async () => { + it("returns only newly searched evidence in each knowhere_search tool result", async () => { const query = vi - .fn() + .fn() .mockResolvedValueOnce(makeRetrievalResponse()) .mockResolvedValueOnce({ ...makeRetrievalResponse(), @@ -143,22 +125,17 @@ describe("agent harness runtime", () => { const tools = createHarnessTools({ state, ledger, - retrieval: { query }, + knowhereTools: makeKnowhereTools(query), recentTurns: [], }) - const firstResult = await executeTool(tools.retrieve, { query: "first" }) - const secondResult = await executeTool(tools.retrieve, { query: "second" }) + const firstResult = await executeTool(tools.knowhere_search, { query: "first" }) + const secondResult = await executeTool(tools.knowhere_search, { query: "second" }) - expect(firstResult).toMatchObject({ - retrievalCount: 1, - chunks: [{ ref: "r1:result:1" }], - }) - expect(secondResult).toMatchObject({ - retrievalCount: 2, - evidenceText: "Second evidence", - chunks: [{ ref: "r2:result:1" }], - }) + expect(firstResult).toContain('ref="r1:result:1"') + expect(secondResult).toContain('retrievalCount="2"') + expect(secondResult).toContain("Second evidence") + expect(secondResult).toContain('ref="r2:result:1"') expect(JSON.stringify(secondResult)).not.toContain("r1:result:1") expect(ledger.snapshot().chunks.map((chunk) => chunk.ref)).toEqual([ "r1:result:1", @@ -171,7 +148,7 @@ describe("agent harness runtime", () => { const tools = createHarnessTools({ state: {}, ledger: createEvidenceLedger(), - retrieval: { query: vi.fn() }, + knowhereTools: makeKnowhereTools(), inspectImages, recentTurns: [], }) @@ -183,12 +160,13 @@ describe("agent harness runtime", () => { expect(result).toEqual({ ok: false, - message: "retrieve must be called before inspectImage.", + message: + "Knowhere evidence tools must return image assets before inspectImage.", inspected: [], skipped: [ { ref: "asset:r1:result:1", - reason: "No retrieval evidence is available yet.", + reason: "No Knowhere evidence is available yet.", }, ], }) @@ -202,7 +180,7 @@ describe("agent harness runtime", () => { const tools = createHarnessTools({ state: {}, ledger, - retrieval: { query: vi.fn() }, + knowhereTools: makeKnowhereTools(), inspectImages, recentTurns: [], }) @@ -223,7 +201,7 @@ describe("agent harness runtime", () => { }, { ref: "missing", - reason: "Ref was not returned by retrieve as an asset.", + reason: "Ref was not returned by Knowhere as an asset.", }, ], }) @@ -236,7 +214,7 @@ describe("agent harness runtime", () => { const tools = createHarnessTools({ state: {}, ledger, - retrieval: { query: vi.fn() }, + knowhereTools: makeKnowhereTools(), inspectImages: vi.fn(), recentTurns: [], }) @@ -268,7 +246,7 @@ describe("agent harness runtime", () => { const tools = createHarnessTools({ state: {}, ledger, - retrieval: { query: vi.fn() }, + knowhereTools: makeKnowhereTools(), inspectImages, recentTurns: [], }) @@ -324,7 +302,7 @@ describe("agent harness runtime", () => { const tools = createHarnessTools({ state: {}, ledger, - retrieval: { query: vi.fn() }, + knowhereTools: makeKnowhereTools(), inspectImages, recentTurns: [], }) @@ -358,17 +336,15 @@ describe("agent harness runtime", () => { }) }) - it("blocks finalize until intent and context policy are declared", async () => { + it("accepts finalize output without planning-tool gating", async () => { const state: { - intent?: IntentFrame - contextPolicy?: ContextPolicy finalizedManifest?: OutputManifest finalized?: boolean } = {} const tools = createHarnessTools({ state, ledger: createEvidenceLedger(), - retrieval: { query: vi.fn() }, + knowhereTools: makeKnowhereTools(), recentTurns: [], }) @@ -379,31 +355,6 @@ describe("agent harness runtime", () => { unresolved: [], } - expect(await executeTool(tools.finalize, manifest)).toEqual({ - ok: false, - message: "declareIntent must be called before finalize.", - }) - expect(state.finalizedManifest).toBeUndefined() - - await executeTool(tools.declareIntent, { - task: "answer_question", - dependsOnPreviousTurn: false, - retrievalNeeded: "no", - targetModalities: ["text"], - constraints: {}, - groundingPolicy: "may_use_sources", - }) - expect(await executeTool(tools.finalize, manifest)).toEqual({ - ok: false, - message: "setContextPolicy must be called before finalize.", - }) - expect(state.finalizedManifest).toBeUndefined() - - await executeTool(tools.setContextPolicy, { - carryHistory: "none", - reason: "Self-contained request.", - activePriorTurnIds: [], - }) expect(await executeTool(tools.finalize, manifest)).toMatchObject({ ok: true, text: "Answer.", @@ -427,7 +378,7 @@ describe("agent harness runtime", () => { const tools = createHarnessTools({ state, ledger: createEvidenceLedger(), - retrieval: { query: vi.fn() }, + knowhereTools: makeKnowhereTools(), recentTurns: [ { id: "turn_1", @@ -459,7 +410,7 @@ describe("agent harness runtime", () => { const tools = createHarnessTools({ state, ledger: createEvidenceLedger(), - retrieval: { query: vi.fn() }, + knowhereTools: makeKnowhereTools(), recentTurns: [ { id: "turn_1", @@ -521,7 +472,7 @@ describe("agent harness runtime", () => { { type: "tool-call", toolCallId: "call_1", - toolName: "retrieve", + toolName: "knowhere_search", input: { query: "q4" }, providerOptions: { google: { @@ -537,12 +488,10 @@ describe("agent harness runtime", () => { { type: "tool-result", toolCallId: "call_1", - toolName: "retrieve", + toolName: "knowhere_search", output: { - type: "json", - value: { - ok: true, - }, + type: "text", + value: '', }, providerOptions: { google: { @@ -574,12 +523,10 @@ describe("agent harness runtime", () => { { type: "tool-result", toolCallId: "call_1", - toolName: "retrieve", + toolName: "knowhere_search", output: { - type: "json", - value: { - ok: true, - }, + type: "text", + value: '', }, }, ], @@ -619,13 +566,11 @@ describe("agent harness runtime", () => { { type: "tool-result", toolCallId: "call_1", - toolName: "retrieve", + toolName: "knowhere_search", output: { - type: "json", - value: { - ok: true, - assets: [{ ref: "asset:r1:referenced:1", type: "image" }], - }, + type: "text", + value: + '', }, }, ], @@ -645,13 +590,11 @@ describe("agent harness runtime", () => { { type: "tool-result", toolCallId: "call_1", - toolName: "retrieve", + toolName: "knowhere_search", output: { - type: "json", - value: { - ok: true, - assets: [{ ref: "asset:r1:referenced:1", type: "image" }], - }, + type: "text", + value: + '', }, }, ], @@ -673,13 +616,11 @@ describe("agent harness runtime", () => { { type: "tool-result", toolCallId: "call_1", - toolName: "retrieve", + toolName: "knowhere_search", output: { - type: "json", - value: { - ok: true, - chunks: [{ ref: "r1:result:1" }], - }, + type: "text", + value: + '', }, providerOptions: { google: { @@ -704,13 +645,11 @@ describe("agent harness runtime", () => { { type: "tool-result", toolCallId: "call_1", - toolName: "retrieve", + toolName: "knowhere_search", output: { - type: "json", - value: { - ok: true, - chunks: [{ ref: "r1:result:1" }], - }, + type: "text", + value: + '', }, }, ], @@ -729,6 +668,18 @@ function executeTool(tool: unknown, input: unknown): Promise { return (tool as { execute: (input: unknown) => Promise }).execute(input) } +function makeKnowhereTools( + search: KnowhereToolRuntime["search"] = vi.fn(), +): KnowhereToolRuntime { + return { + search, + listDocuments: vi.fn().mockResolvedValue({ documents: [] }), + getDocumentOutline: vi.fn().mockRejectedValue(new Error("Not configured.")), + readChunks: vi.fn().mockRejectedValue(new Error("Not configured.")), + grepChunks: vi.fn().mockRejectedValue(new Error("Not configured.")), + } +} + function makeTurnInput(overrides: Partial = {}): AgentTurnInput { return { surface: "notebook_chat", diff --git a/src/agent-harness/runtime.ts b/src/agent-harness/runtime.ts index d38028e..8963fdf 100644 --- a/src/agent-harness/runtime.ts +++ b/src/agent-harness/runtime.ts @@ -9,6 +9,7 @@ import { import { z } from "zod" import { createEvidenceLedger } from "./ledger" +import { knowhereToolText } from "./knowhere-text" import type { AgentTurn, AgentTurnInput, @@ -20,14 +21,12 @@ import type { ImageInspectionResponse, InspectImages, IntentFrame, + KnowhereSearchTargetContent, + KnowhereToolRuntime, OutputManifest, - RetrievalCapability, - TargetModality, } from "./types" -import { validateOutputManifest } from "./validator" const defaultMaxSteps = 14 -const defaultMaxRevisions = 1 const imageInspectionReminderStepNumber = 12 const forcedFinalizationStepNumber = 13 const imageInspectionRefLimit = 6 @@ -39,14 +38,9 @@ export type AgentHarnessModel = ToolLoopAgentSettings["model"] export type RunAgentHarnessInput = { readonly model: AgentHarnessModel readonly turn: AgentTurnInput - readonly retrieval: RetrievalCapability + readonly knowhereTools: KnowhereToolRuntime readonly inspectImages?: InspectImages readonly maxSteps?: number - /** - * How many times the agent may revise after a failed validation pass before - * the harness gives up and returns the last manifest with recorded errors. - */ - readonly maxRevisions?: number } type HarnessToolState = { @@ -71,6 +65,53 @@ type HarnessStepPreparation = { } const targetModalitySchema = z.enum(["text", "image", "table"]) +const knowhereSearchTargetContentSchema = z.enum([ + "all", + "text", + "image", + "table", + "text_image", + "text_table", +]) +const knowledgeChunkTypeSchema = z.enum(["text", "image", "table", "page"]) + +const knowhereDocumentReferenceSchema = z.object({ + localDocumentId: z.string().min(1).optional(), + documentId: z.string().min(1).optional(), + jobId: z.string().min(1).optional(), + revisionKey: z.string().min(1).optional(), +}) + +const knowhereSearchSchema = z.object({ + query: z.string().min(1), + targetContent: knowhereSearchTargetContentSchema.default("all"), + purpose: z.string().optional(), + topK: z.number().int().min(1).max(12).optional(), + signalPaths: z.array(z.string().min(1)).max(8).optional(), + filterMode: z.enum(["keep", "delete"]).optional(), + threshold: z.number().min(0).max(1).optional(), +}) + +const knowhereReadChunksSchema = knowhereDocumentReferenceSchema.extend({ + page: z.number().int().min(1).optional(), + pageSize: z.number().int().min(1).max(50).optional(), + sectionPath: z.string().min(1).optional(), + startChunk: z.number().int().min(0).optional(), + endChunk: z.number().int().min(0).optional(), + chunkId: z.string().min(1).optional(), + chunkType: knowledgeChunkTypeSchema.optional(), +}) + +const knowhereGrepChunksSchema = knowhereDocumentReferenceSchema.extend({ + pattern: z.string().min(1), + continuationCursor: z.string().min(1).optional(), + isRegex: z.boolean().optional(), + isCaseSensitive: z.boolean().optional(), + maxResults: z.number().int().min(1).max(50).optional(), + chunkType: knowledgeChunkTypeSchema.optional(), + sectionPathPrefix: z.string().min(1).optional(), + contextChars: z.number().int().min(0).max(2_000).optional(), +}) const intentFrameSchema = z.object({ task: z.enum([ @@ -116,12 +157,14 @@ const contextPolicySchema = z.object({ const outputCitationSchema = z.object({ ref: z.string().min(1), - label: z.string().min(1), - source: z.object({ - documentId: z.string().nullable().optional(), - sourceFileName: z.string().nullable().optional(), - sectionPath: z.string().nullable().optional(), - }), + label: z.string().min(1).optional(), + source: z + .object({ + documentId: z.string().nullable().optional(), + sourceFileName: z.string().nullable().optional(), + sectionPath: z.string().nullable().optional(), + }) + .optional(), }) const selectedOutputArtifactSchema = z.object({ @@ -166,7 +209,7 @@ export async function runAgentHarness( const tools = createHarnessTools({ state, ledger, - retrieval: input.retrieval, + knowhereTools: input.knowhereTools, inspectImages: input.inspectImages, recentTurns: input.turn.recentTurns, }) @@ -188,46 +231,11 @@ export async function runAgentHarness( ], }) - const maxRevisions = input.maxRevisions ?? defaultMaxRevisions - let messages = buildHarnessMessages(input.turn) - let manifest = buildFallbackManifest("") - let validationErrors: readonly string[] = [] - let revisionsUsed = 0 - - for (let attempt = 0; ; attempt += 1) { - const response = await agent.generate({ messages }) - manifest = - state.finalizedManifest ?? buildFallbackManifest(response.text.trim()) - - const validation = validateOutputManifest({ - manifest, - intent: state.intent, - contextPolicy: state.contextPolicy, - finalized: state.finalized === true, - ledger: ledger.snapshot(), - toolCalls: state.toolCalls, - surface: input.turn.surface, - }) - validationErrors = validation.errors - - if (validation.ok || attempt >= maxRevisions) break - - // Self-correction (reflexion): continue the same conversation with - // structured validator feedback and require a fresh finalize so the agent - // can repair its own contract violations instead of shipping them. - revisionsUsed += 1 - state.finalizedManifest = undefined - state.finalized = false - messages = [ - ...messages, - ...(response.response.messages as ModelMessage[]), - { - role: "user", - content: buildRevisionFeedback(validation.errors), - }, - ] - } - + const response = await agent.generate({ + messages: buildHarnessMessages(input.turn), + }) + const manifest = + state.finalizedManifest ?? buildFallbackManifest(response.text.trim()) const ledgerSnapshot = ledger.snapshot() return { manifest, @@ -238,8 +246,8 @@ export async function runAgentHarness( finalized: state.finalized === true, priorTurnReads: [...(state.priorTurnReads ?? [])], toolCalls: [...(state.toolCalls ?? [])], - validationErrors, - revisionsUsed, + validationErrors: [], + revisionsUsed: 0, }, } } @@ -347,24 +355,10 @@ type ModelMessageForRole = Extract< { readonly role: TRole } > -function buildRevisionFeedback(errors: readonly string[]): string { - return [ - "Your finalize output did not satisfy the output contract:", - ...errors.map((error) => `- ${error}`), - "", - "Fix every issue and call finalize again with a corrected manifest.", - "You must call finalize; freeform assistant text is not a valid final", - "answer contract.", - "Do not exceed the user's requested artifact count, only cite or display", - "evidence refs that exist in the evidence ledger, and do not fabricate", - "facts when evidence is missing.", - ].join("\n") -} - function buildForcedFinalizationFeedback(): string { return [ "The retrieval step budget has been reached.", - "Do not search again or call any evidence-reading tools.", + "Do not search again or call any Knowhere evidence-reading tools.", "Use only the evidence and tool results already available in this turn.", "Call finalize now with the best supported answer.", "If the existing evidence is insufficient, explain the gap in unresolved", @@ -396,14 +390,14 @@ function hasUninspectedImageAssets(input: { export function createHarnessTools(input: { readonly state: HarnessToolState readonly ledger: ReturnType - readonly retrieval: RetrievalCapability + readonly knowhereTools: KnowhereToolRuntime readonly inspectImages?: InspectImages readonly recentTurns: readonly AgentTurn[] }) { return { declareIntent: tool({ description: - "Declare the user's intent before any other action. This is working memory, not a final answer.", + "Declare the user's intent when it helps plan the response. This is working memory, not a final answer.", inputSchema: intentFrameSchema, execute: async (intent): Promise => traceToolCall(input.state, { @@ -433,82 +427,104 @@ export function createHarnessTools(input: { }), }), - retrieve: tool({ + knowhere_search: tool({ description: - "Ask KNOWHERE for evidence context. KNOWHERE handles internal navigation; this tool only submits a concise query and records returned evidence.", - inputSchema: z.object({ - query: z.string().min(1), - modalities: z.array(targetModalitySchema).default(["text"]), - purpose: z.string().optional(), - topK: z.number().int().min(1).max(12).optional(), - signalPaths: z.array(z.string().min(1)).max(8).optional(), - filterMode: z.enum(["keep", "delete"]).optional(), - threshold: z.number().min(0).max(1).optional(), - }), + "Search Knowhere for relevant Notebook evidence. Returns tagged text with evidence refs such as r1:result:1 and asset refs such as asset:r1:result:1.", + inputSchema: knowhereSearchSchema, execute: async (request) => traceToolCall(input.state, { - toolName: "retrieve", - inputSummary: summarizeRetrievalRequest(request), - execute: async () => { - if (!input.state.intent) { - return { - ok: false, - message: "declareIntent must be called before retrieve.", - } - } - if (!input.state.contextPolicy) { - return { - ok: false, - message: "setContextPolicy must be called before retrieve.", - } - } + toolName: "knowhere_search", + inputSummary: summarizeKnowhereSearchRequest(request), + execute: async () => + executeKnowhereSearch({ + ledger: input.ledger, + knowhereTools: input.knowhereTools, + request, + }), + summarizeOutput: summarizeKnowhereTextOutput, + }), + }), - const beforeSnapshot = input.ledger.snapshot() - const response = await input.retrieval.query({ - query: request.query, - modalities: request.modalities as TargetModality[], - purpose: request.purpose, - topK: request.topK, - signalPaths: request.signalPaths, - filterMode: request.filterMode, - threshold: request.threshold, - }) - const snapshot = input.ledger.addRetrievalResponse(response) - const currentChunks = snapshot.chunks.slice( - beforeSnapshot.chunks.length, - ) - const currentAssets = snapshot.assets.slice( - beforeSnapshot.assets.length, - ) - return { - ok: true, - retrievalCount: snapshot.retrievalCount, - evidenceText: response.evidenceText ?? "", - stopReason: response.stopReason ?? null, - failureReason: response.failureReason ?? null, - chunks: currentChunks.map((chunk) => ({ - ref: chunk.ref, - kind: chunk.kind, - type: chunk.chunkType, - preview: chunk.contentPreview, - source: chunk.source, - assetRef: chunk.assetRef, - })), - assets: currentAssets.map((asset) => ({ - ref: asset.ref, - type: asset.type, - label: asset.label, - source: asset.source, - })), - } - }, - summarizeOutput: summarizeRetrieveOutput, + knowhere_list_documents: tool({ + description: + "List ready visible Notebook/Knowhere documents available for this chat turn. Use this to discover documentId and revisionKey before outline/read/grep.", + inputSchema: z.object({}), + execute: async () => + traceToolCall(input.state, { + toolName: "knowhere_list_documents", + inputSummary: {}, + execute: async () => + executeKnowhereTextTool({ + operation: "list_documents", + execute: async () => + knowhereToolText.formatListDocuments( + await input.knowhereTools.listDocuments(), + ), + }), + summarizeOutput: summarizeKnowhereTextOutput, + }), + }), + + knowhere_get_document_outline: tool({ + description: + "Read a document outline from Knowhere parsed storage. Use documentId/revisionKey from knowhere_list_documents or search refs.", + inputSchema: knowhereDocumentReferenceSchema, + execute: async (request) => + traceToolCall(input.state, { + toolName: "knowhere_get_document_outline", + inputSummary: summarizeDocumentReference(request), + execute: async () => + executeKnowhereTextTool({ + operation: "get_document_outline", + validate: () => validateDocumentReference(request), + execute: async () => + knowhereToolText.formatOutline( + await input.knowhereTools.getDocumentOutline(request), + ), + }), + summarizeOutput: summarizeKnowhereTextOutput, + }), + }), + + knowhere_read_chunks: tool({ + description: + "Read complete chunk bodies from Knowhere parsed storage. This tool never slices individual chunk content; control read size with page/pageSize, sectionPath, startChunk/endChunk, chunkId, and chunkType.", + inputSchema: knowhereReadChunksSchema, + execute: async (request) => + traceToolCall(input.state, { + toolName: "knowhere_read_chunks", + inputSummary: summarizeReadChunksRequest(request), + execute: async () => + executeKnowhereReadChunks({ + ledger: input.ledger, + knowhereTools: input.knowhereTools, + request, + }), + summarizeOutput: summarizeKnowhereTextOutput, + }), + }), + + knowhere_grep_chunks: tool({ + description: + "Search chunk text with a literal or regex pattern. Returns bounded match snippets as grep refs such as grep1:match:1 and may include truncated=true with a continuationCursor.", + inputSchema: knowhereGrepChunksSchema, + execute: async (request) => + traceToolCall(input.state, { + toolName: "knowhere_grep_chunks", + inputSummary: summarizeGrepChunksRequest(request), + execute: async () => + executeKnowhereGrepChunks({ + ledger: input.ledger, + knowhereTools: input.knowhereTools, + request, + }), + summarizeOutput: summarizeKnowhereTextOutput, }), }), inspectImage: tool({ description: - "Inspect retrieved image asset refs visually for OCR, visual details, comparisons, or verification. Use only after retrieve has returned image assets.", + "Inspect Knowhere image asset refs visually for OCR, visual details, comparisons, or verification. Use only after a Knowhere tool has returned image assets.", inputSchema: z.object({ refs: z.array(z.string().min(1)).min(1).max(imageInspectionRefLimit), question: z.string().min(1), @@ -529,32 +545,6 @@ export function createHarnessTools(input: { }), }), - readEvidence: tool({ - description: - "Read more text from an evidence chunk already returned by KNOWHERE.", - inputSchema: z.object({ - ref: z.string().min(1), - offset: z.number().int().min(0).optional(), - limit: z.number().int().min(1).max(8_000).optional(), - }), - execute: async (request) => - traceToolCall(input.state, { - toolName: "readEvidence", - inputSummary: { - ref: request.ref, - offset: request.offset ?? 0, - limit: request.limit ?? 4_000, - }, - execute: async () => - input.ledger.read( - request.ref, - request.offset ?? 0, - request.limit ?? 4_000, - ), - summarizeOutput: summarizeReadEvidenceOutput, - }), - }), - readPriorTurn: tool({ description: "Read the full text and citation labels of a specific prior turn by id " + @@ -618,25 +608,13 @@ export function createHarnessTools(input: { description: "Finalize the user-facing output manifest. This is the only final answer " + "contract. Artifacts listed here with display=true are the exact set of " + - "images/tables shown to the user; cite only refs from the evidence ledger.", + "images/tables shown to the user; cite evidence refs when available.", inputSchema: outputManifestSchema, execute: async (manifest) => traceToolCall(input.state, { toolName: "finalize", inputSummary: summarizeManifest(manifest), execute: async () => { - if (!input.state.intent) { - return { - ok: false as const, - message: "declareIntent must be called before finalize.", - } - } - if (!input.state.contextPolicy) { - return { - ok: false as const, - message: "setContextPolicy must be called before finalize.", - } - } input.state.finalizedManifest = manifest input.state.finalized = true return { ok: true as const, ...manifest } @@ -689,14 +667,15 @@ async function inspectRetrievedImages(input: { } const snapshot = input.ledger.snapshot() - if (snapshot.retrievalCount === 0) { + if (snapshot.chunks.length === 0 && snapshot.assets.length === 0) { return { ok: false, - message: "retrieve must be called before inspectImage.", + message: + "Knowhere evidence tools must return image assets before inspectImage.", inspected: [], skipped: refs.map((ref) => ({ ref, - reason: "No retrieval evidence is available yet.", + reason: "No Knowhere evidence is available yet.", })), } } @@ -726,7 +705,7 @@ async function inspectRetrievedImages(input: { if (!asset) { skipped.push({ ref, - reason: "Ref was not returned by retrieve as an asset.", + reason: "Ref was not returned by Knowhere as an asset.", }) continue } @@ -804,6 +783,134 @@ async function inspectRetrievedImages(input: { } } +type KnowhereToolOperation = + | "search" + | "list_documents" + | "get_document_outline" + | "read_chunks" + | "grep_chunks" + +type KnowhereSearchToolRequest = z.infer +type KnowhereDocumentReferenceRequest = z.infer< + typeof knowhereDocumentReferenceSchema +> +type KnowhereReadChunksToolRequest = z.infer +type KnowhereGrepChunksToolRequest = z.infer +type DocumentReferenceSummary = { + readonly documentId?: string + readonly localDocumentId?: string + readonly hasJobId: boolean + readonly hasRevisionKey: boolean +} + +async function executeKnowhereSearch(input: { + readonly ledger: ReturnType + readonly knowhereTools: KnowhereToolRuntime + readonly request: KnowhereSearchToolRequest +}): Promise { + return executeKnowhereTextTool({ + operation: "search", + execute: async () => { + const beforeSnapshot = input.ledger.snapshot() + const response = await input.knowhereTools.search({ + query: input.request.query, + targetContent: input.request.targetContent, + purpose: input.request.purpose, + topK: input.request.topK, + signalPaths: input.request.signalPaths, + filterMode: input.request.filterMode, + threshold: input.request.threshold, + }) + const snapshot = input.ledger.addRetrievalResponse(response) + return knowhereToolText.formatSearch({ + response, + retrievalCount: snapshot.retrievalCount, + chunks: snapshot.chunks.slice(beforeSnapshot.chunks.length), + assets: snapshot.assets.slice(beforeSnapshot.assets.length), + }) + }, + }) +} + +async function executeKnowhereReadChunks(input: { + readonly ledger: ReturnType + readonly knowhereTools: KnowhereToolRuntime + readonly request: KnowhereReadChunksToolRequest +}): Promise { + return executeKnowhereTextTool({ + operation: "read_chunks", + validate: () => validateDocumentReference(input.request), + execute: async () => { + const beforeSnapshot = input.ledger.snapshot() + const response = await input.knowhereTools.readChunks(input.request) + const snapshot = input.ledger.addReadChunksResponse(response) + return knowhereToolText.formatReadChunks({ + response, + chunks: snapshot.chunks.slice(beforeSnapshot.chunks.length), + assets: snapshot.assets.slice(beforeSnapshot.assets.length), + }) + }, + }) +} + +async function executeKnowhereGrepChunks(input: { + readonly ledger: ReturnType + readonly knowhereTools: KnowhereToolRuntime + readonly request: KnowhereGrepChunksToolRequest +}): Promise { + return executeKnowhereTextTool({ + operation: "grep_chunks", + validate: () => validateDocumentReference(input.request), + execute: async () => { + const beforeSnapshot = input.ledger.snapshot() + const response = await input.knowhereTools.grepChunks(input.request) + const snapshot = input.ledger.addGrepChunksResponse(response) + return knowhereToolText.formatGrepChunks({ + response, + chunks: snapshot.chunks.slice(beforeSnapshot.chunks.length), + assets: snapshot.assets.slice(beforeSnapshot.assets.length), + }) + }, + }) +} + +async function executeKnowhereTextTool(input: { + readonly operation: KnowhereToolOperation + readonly validate?: () => string | null + readonly execute: () => Promise +}): Promise { + const validationError = input.validate?.() + if (validationError) { + return knowhereToolText.formatError({ + operation: input.operation, + message: validationError, + }) + } + + try { + return await input.execute() + } catch (error) { + return knowhereToolText.formatError({ + operation: input.operation, + message: formatUnknownError(error), + }) + } +} + +function validateDocumentReference( + request: KnowhereDocumentReferenceRequest, +): string | null { + if ( + request.documentId || + request.localDocumentId || + request.jobId + ) { + return null + } + + return "A documentId, localDocumentId, or jobId is required." +} + function getUniqueTrimmedRefs(refs: readonly string[]): string[] { const normalizedRefs: string[] = [] for (const ref of refs) { @@ -860,6 +967,7 @@ function recordToolCall( } function getToolTraceOk(output: unknown): boolean { + if (typeof output === "string") return !output.includes('status="error"') if (!isRecord(output)) return true if (typeof output.ok === "boolean") return output.ok if (typeof output.found === "boolean") return output.found @@ -884,9 +992,9 @@ function summarizeContextPolicy(policy: ContextPolicy): unknown { } } -function summarizeRetrievalRequest(request: { +function summarizeKnowhereSearchRequest(request: { readonly query: string - readonly modalities?: readonly TargetModality[] + readonly targetContent?: KnowhereSearchTargetContent readonly purpose?: string readonly topK?: number readonly signalPaths?: readonly string[] @@ -895,7 +1003,7 @@ function summarizeRetrievalRequest(request: { }): unknown { return { query: request.query, - modalities: request.modalities ?? ["text"], + targetContent: request.targetContent ?? "all", purpose: request.purpose, topK: request.topK, signalPathCount: request.signalPaths?.length ?? 0, @@ -904,27 +1012,67 @@ function summarizeRetrievalRequest(request: { } } -function summarizeRetrieveOutput(output: unknown): unknown { - if (!isRecord(output)) return output +function summarizeDocumentReference( + request: KnowhereDocumentReferenceRequest, +): DocumentReferenceSummary { return { - ok: output.ok, - retrievalCount: output.retrievalCount, - stopReason: output.stopReason, - failureReason: output.failureReason, - chunkCount: Array.isArray(output.chunks) ? output.chunks.length : 0, - assetCount: Array.isArray(output.assets) ? output.assets.length : 0, + documentId: request.documentId, + localDocumentId: request.localDocumentId, + hasJobId: typeof request.jobId === "string", + hasRevisionKey: typeof request.revisionKey === "string", } } -function summarizeReadEvidenceOutput(output: unknown): unknown { - if (!isRecord(output)) return output +function summarizeReadChunksRequest( + request: KnowhereReadChunksToolRequest, +): unknown { return { - found: output.found, - ref: output.ref, - contentLength: output.contentLength, - offset: output.offset, - limit: output.limit, - hasMoreContent: output.hasMoreContent, + ...summarizeDocumentReference(request), + page: request.page, + pageSize: request.pageSize, + sectionPath: request.sectionPath, + startChunk: request.startChunk, + endChunk: request.endChunk, + chunkId: request.chunkId, + chunkType: request.chunkType, + } +} + +function summarizeGrepChunksRequest( + request: KnowhereGrepChunksToolRequest, +): unknown { + return { + ...summarizeDocumentReference(request), + patternLength: request.pattern.trim().length, + continuationCursor: request.continuationCursor, + isRegex: request.isRegex, + isCaseSensitive: request.isCaseSensitive, + maxResults: request.maxResults, + chunkType: request.chunkType, + sectionPathPrefix: request.sectionPathPrefix, + contextChars: request.contextChars, + } +} + +function summarizeKnowhereTextOutput(output: unknown): unknown { + if (typeof output !== "string") return output + return { + ok: !output.includes('status="error"'), + textLength: output.length, + chunkCount: countOccurrences(output, " { return typeof value === "object" && value !== null } +function formatUnknownError(error: unknown): string { + if (error instanceof Error) return error.message + return String(error) +} + export function buildHarnessSystemPrompt(turn: AgentTurnInput): string { return [ "You are the outer Knowhere Agent Harness.", "KNOWHERE is only an evidence provider. Do not infer or control its internal navigation algorithm.", "Your job is to understand intent, decide context use, optionally retrieve evidence, select evidence/artifacts, create source-backed derived tables when useful, and finalize an output manifest.", "", - "Required workflow:", - "1. Call declareIntent first. Capture constraints like a requested image/table count in constraints.desiredCount.", - "2. Call setContextPolicy next, deciding how prior turns should influence this turn.", + "Recommended workflow:", + "1. Call declareIntent when it helps you plan the response. Capture constraints like a requested image/table count in constraints.desiredCount.", + "2. Call setContextPolicy when prior turns may influence this turn.", "3. When the policy needs prior-turn detail (references or corrections), call readPriorTurn for the relevant ids.", - "4. Call retrieve only when evidence is needed. The query must be concise and self-contained.", - "5. For pixel-level details, OCR, visual comparison, image verification, or when the likely answer is only visible on a retrieved page/image asset, call inspectImage only after retrieve returned image asset refs.", + "4. Call knowhere_search when relevance search is needed. Use knowhere_list_documents, knowhere_get_document_outline, knowhere_read_chunks, and knowhere_grep_chunks for focused document reads.", + "5. For pixel-level details, OCR, visual comparison, image verification, or when the likely answer is only visible on a returned page/image asset, call inspectImage only after a Knowhere tool returned image asset refs.", `6. inspectImage accepts at most ${imageInspectionRefLimit} image asset refs per call and per turn.`, - "7. Use readEvidence only for chunk refs already in the evidence ledger.", - "8. Call finalize with text, citations, artifacts, and unresolved issues. finalize requires declareIntent and setContextPolicy first.", + "7. knowhere_read_chunks returns complete chunk bodies; control size with page/pageSize, sectionPath, startChunk/endChunk, chunkId, and chunkType.", + "8. Call finalize with text, citations, artifacts, and unresolved issues when you are ready to answer.", "", "Context rules:", "- If the current user request is unrelated to prior turns, set carryHistory to none and do not reuse prior topics.", @@ -1022,11 +1175,12 @@ export function buildHarnessSystemPrompt(turn: AgentTurnInput): string { "- Final output is the OutputManifest passed to finalize, not freeform tool JSON or trailing text.", "- artifacts with display=true are the exact images/tables shown. Never display every candidate; honor constraints.desiredCount / maxCount.", "- Use type=derived_table only for tables you create from evidence; every derived_table.sourceRefs entry must reference evidence in the ledger.", - "- citations and selected image/table artifact refs may only reference refs returned by retrieve (in the evidence ledger).", + "- Prefer citation and selected image/table artifact refs returned by Knowhere tools in the evidence ledger.", + "- Citation label and source metadata are optional. Notebook resolves citation metadata from evidence refs when possible.", + "- If evidence is relevant but you cannot identify a supporting evidence ref, answer with unresolved issues instead of fabricating a ref.", "- inspectImage observations are inspection notes, not new source refs. Final citations and displayed image artifacts must use the original retrieved image asset refs.", "- If text evidence identifies a relevant page/image but does not include the exact fact, inspect the returned image asset for OCR/detail before saying the answer is unavailable.", "- If evidence is insufficient, list it in unresolved instead of fabricating facts.", - "- After a validation-feedback message, fix all listed issues and call finalize again.", `Surface: ${turn.surface}`, `Output capabilities: ${JSON.stringify(turn.outputCapabilities)}`, turn.sourceContext ? `Searchable source context:\n${turn.sourceContext}` : "", diff --git a/src/agent-harness/types.ts b/src/agent-harness/types.ts index 4601f6f..36b6a27 100644 --- a/src/agent-harness/types.ts +++ b/src/agent-harness/types.ts @@ -1,4 +1,10 @@ import type { + KnowledgeDocumentReference, + KnowledgeGrepParams, + KnowledgeGrepResponse, + KnowledgeOutline, + KnowledgeReadParams, + KnowledgeReadResponse, RetrievalQueryParams, RetrievalQueryResponse, } from "@ontos-ai/knowhere-sdk" @@ -77,23 +83,57 @@ export type AgentTurnInput = { } } -export type HarnessRetrievalRequest = Pick< +export type KnowhereSearchTargetContent = + | "all" + | "text" + | "image" + | "table" + | "text_image" + | "text_table" + +export type KnowhereSearchRequest = Pick< RetrievalQueryParams, "query" | "topK" | "signalPaths" | "filterMode" | "threshold" > & { - readonly modalities: readonly TargetModality[] + readonly targetContent?: KnowhereSearchTargetContent readonly purpose?: string } -export type RetrievalCapability = { - readonly query: ( - input: HarnessRetrievalRequest, +export type KnowhereDocumentSummary = { + readonly documentId?: string + readonly localDocumentId?: string + readonly revisionKey?: string + readonly namespace?: string + readonly sourceFileName: string + readonly title?: string + readonly status?: string + readonly chunkCount?: number + readonly typeCounts?: Readonly> +} + +export type KnowhereListDocumentsResponse = { + readonly documents: readonly KnowhereDocumentSummary[] +} + +export type KnowhereToolRuntime = { + readonly search: ( + input: KnowhereSearchRequest, ) => Promise + readonly listDocuments: () => Promise + readonly getDocumentOutline: ( + input: KnowledgeDocumentReference, + ) => Promise + readonly readChunks: ( + input: KnowledgeReadParams, + ) => Promise + readonly grepChunks: ( + input: KnowledgeGrepParams, + ) => Promise } export type EvidenceChunk = { readonly ref: string - readonly kind: "result" | "referenced_chunk" + readonly kind: "result" | "referenced_chunk" | "read_chunk" | "grep_match" readonly chunkId?: string readonly content: string readonly contentPreview: string @@ -169,8 +209,8 @@ export type EvidenceLedgerSnapshot = { export type OutputCitation = { readonly ref: string - readonly label: string - readonly source: EvidenceChunk["source"] + readonly label?: string + readonly source?: EvidenceChunk["source"] } export type OutputArtifact = { diff --git a/src/agent-harness/validator.test.ts b/src/agent-harness/validator.test.ts deleted file mode 100644 index be8de3e..0000000 --- a/src/agent-harness/validator.test.ts +++ /dev/null @@ -1,434 +0,0 @@ -import { describe, expect, it } from "vitest" - -import { validateOutputManifest } from "./validator" -import type { - ContextPolicy, - EvidenceLedgerSnapshot, - IntentFrame, - OutputManifest, -} from "./types" - -describe("validateOutputManifest", () => { - it("requires finalize to be the successful output path", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ text: "Freeform answer." }), - intent: makeIntent({ groundingPolicy: "no_retrieval" }), - contextPolicy: unrelatedContextPolicy, - finalized: false, - ledger: emptyLedger, - surface: "notebook_chat", - }) - - expect(validation.errors).toContain( - "Agent must call finalize to produce the output manifest.", - ) - }) - - it("requires the agent to declare intent and context policy before finalizing", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ text: "Answer." }), - ledger: emptyLedger, - surface: "notebook_chat", - }) - - expect(validation.errors).toContain( - "Agent must declare intent before finalizing.", - ) - expect(validation.errors).toContain( - "Agent must set context policy before finalizing.", - ) - }) - - it("limits displayed artifacts using the declared intent instead of hard-coded media rules", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ - artifacts: [ - { - type: "image", - ref: "asset:r1:result:1", - display: true, - reason: "front", - }, - { - type: "image", - ref: "asset:r1:result:2", - display: true, - reason: "back", - }, - { - type: "image", - ref: "asset:r1:result:3", - display: true, - reason: "extra candidate", - }, - ], - }), - intent: makeIntent({ desiredCount: 2, maxCount: 2 }), - contextPolicy: unrelatedContextPolicy, - ledger: { - ...emptyLedger, - assets: [ - makeAsset("asset:r1:result:1"), - makeAsset("asset:r1:result:2"), - makeAsset("asset:r1:result:3"), - ], - }, - surface: "notebook_chat", - }) - - expect(validation.errors).toContain( - "Displayed artifact count 3 exceeds desired count 2.", - ) - expect(validation.errors).toContain( - "Displayed artifact count 3 exceeds maximum count 2.", - ) - }) - - it("rejects grounded answers that use evidence without citations or selected artifacts", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ text: "Revenue increased." }), - intent: makeIntent({}), - contextPolicy: unrelatedContextPolicy, - ledger: { - ...emptyLedger, - chunks: [ - { - ref: "r1:result:1", - kind: "result", - content: "Revenue increased.", - contentPreview: "Revenue increased.", - chunkType: "text", - score: 0.9, - source: { - documentId: "doc_1", - sourceFileName: "report.pdf", - sectionPath: "Q4", - }, - }, - ], - }, - surface: "notebook_chat", - }) - - expect(validation.errors).toContain( - "Grounded output used evidence but did not cite or display any selected evidence.", - ) - }) - - it("accepts source-backed derived tables and rejects missing source refs", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ - artifacts: [ - { - type: "derived_table", - ref: "derived:table:1", - title: "Revenue comparison", - columns: ["Metric", "Value"], - rows: [["Revenue", "$10M"]], - sourceRefs: ["r1:result:1"], - display: true, - reason: "Structured comparison requested by the user.", - }, - { - type: "derived_table", - ref: "derived:table:2", - title: "Invalid table", - columns: ["Metric", "Value"], - rows: [["Revenue"]], - sourceRefs: ["missing"], - display: true, - reason: "Demonstrates validation.", - }, - ], - }), - intent: makeIntent({}), - contextPolicy: unrelatedContextPolicy, - ledger: { - ...emptyLedger, - chunks: [ - { - ref: "r1:result:1", - kind: "result", - content: "Revenue was $10M.", - contentPreview: "Revenue was $10M.", - chunkType: "text", - score: 0.9, - source: { - documentId: "doc_1", - sourceFileName: "report.pdf", - sectionPath: "Revenue", - }, - }, - ], - }, - surface: "notebook_chat", - }) - - expect(validation.errors).toContain( - "Derived table source ref 'missing' was not found in the evidence ledger.", - ) - expect(validation.errors).toContain( - "Derived table row 1 has 1 cells but expected 2.", - ) - }) - - it("requires compare outputs to cite at least two evidence refs", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ - text: "A is stronger than B.", - citations: [ - { - ref: "r1:result:1", - label: "report.pdf / A", - source: { - documentId: "doc_1", - sourceFileName: "report.pdf", - sectionPath: "A", - }, - }, - ], - }), - intent: { - ...makeIntent({}), - task: "compare", - }, - contextPolicy: unrelatedContextPolicy, - ledger: { - ...emptyLedger, - chunks: [ - { - ref: "r1:result:1", - kind: "result", - content: "A is strong.", - contentPreview: "A is strong.", - chunkType: "text", - score: 0.9, - source: { - documentId: "doc_1", - sourceFileName: "report.pdf", - sectionPath: "A", - }, - }, - ], - }, - surface: "notebook_chat", - }) - - expect(validation.errors).toContain( - "Compare outputs that must use sources require at least two evidence refs or an explicit unresolved reason.", - ) - }) - - it("keeps typing compose output insertion-ready", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ text: "- bullet\n- list" }), - intent: makeIntent({ groundingPolicy: "no_retrieval" }), - contextPolicy: unrelatedContextPolicy, - ledger: emptyLedger, - surface: "typing_compose", - }) - - expect(validation.errors).toContain( - "Typing compose output must be insertion-ready plain text.", - ) - }) - - it("rejects image inspection failure claims when inspectImage was not called", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ - text: "图像检查工具未能成功读取该页面的具体条款内容。", - citations: [ - { - ref: "r1:referenced:1", - label: "contract.pdf / page 8", - source: { - documentId: "doc_contract", - sourceFileName: "contract.pdf", - sectionPath: "Root / （6）现场工期进度管理方面的违约责任", - }, - }, - ], - }), - intent: makeIntent({}), - contextPolicy: unrelatedContextPolicy, - ledger: { - ...emptyLedger, - chunks: [ - { - ref: "r1:referenced:1", - kind: "referenced_chunk", - content: "", - contentPreview: "", - chunkType: "page", - score: null, - source: { - documentId: "doc_contract", - sourceFileName: "contract.pdf", - sectionPath: "Root / （6）现场工期进度管理方面的违约责任", - }, - assetRef: "asset:r1:referenced:1", - }, - ], - assets: [makeAsset("asset:r1:referenced:1")], - }, - toolCalls: [], - surface: "notebook_chat", - }) - - expect(validation.errors).toContain( - "Final output must not claim image/OCR inspection succeeded or failed unless inspectImage was called.", - ) - }) - - it("allows image inspection failure claims after inspectImage was called", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ - text: "对该页面的图像识别未成功提取具体数值。", - citations: [ - { - ref: "r1:referenced:1", - label: "contract.pdf / page 8", - source: { - documentId: "doc_contract", - sourceFileName: "contract.pdf", - sectionPath: "Root / （6）现场工期进度管理方面的违约责任", - }, - }, - ], - }), - intent: makeIntent({}), - contextPolicy: unrelatedContextPolicy, - ledger: { - ...emptyLedger, - chunks: [ - { - ref: "r1:referenced:1", - kind: "referenced_chunk", - content: "", - contentPreview: "", - chunkType: "page", - score: null, - source: { - documentId: "doc_contract", - sourceFileName: "contract.pdf", - sectionPath: "Root / （6）现场工期进度管理方面的违约责任", - }, - assetRef: "asset:r1:referenced:1", - }, - ], - assets: [makeAsset("asset:r1:referenced:1")], - }, - toolCalls: [ - { - tool: "inspectImage", - ok: false, - inputSummary: { refs: ["asset:r1:referenced:1"] }, - outputSummary: { inspectedCount: 0 }, - startedAt: "2026-07-06T00:00:00.000Z", - durationMs: 1, - }, - ], - surface: "notebook_chat", - }) - - expect(validation.errors).not.toContain( - "Final output must not claim image/OCR inspection succeeded or failed unless inspectImage was called.", - ) - }) - - it("rejects unreadable page claims when image assets are available but uninspected", () => { - const validation = validateOutputManifest({ - manifest: makeManifest({ - text: "相关的页面图像无法通过目前的检测工具获取。", - unresolved: ["由于无法直接读取该页面的详细条款内容，目前无法给出确切的赔偿金额。"], - }), - intent: makeIntent({}), - contextPolicy: unrelatedContextPolicy, - ledger: { - ...emptyLedger, - chunks: [ - { - ref: "r1:referenced:1", - kind: "referenced_chunk", - content: "", - contentPreview: "", - chunkType: "page", - score: null, - source: { - documentId: "doc_contract", - sourceFileName: "contract.pdf", - sectionPath: "Root / （6）现场工期进度管理方面的违约责任", - }, - assetRef: "asset:r1:referenced:1", - }, - ], - assets: [makeAsset("asset:r1:referenced:1")], - }, - toolCalls: [], - surface: "notebook_chat", - }) - - expect(validation.errors).toContain( - "Final output must inspect available image assets before claiming retrieved page/image content cannot be read.", - ) - expect(validation.errors).toContain( - "Final output must not claim image/OCR inspection succeeded or failed unless inspectImage was called.", - ) - }) -}) - -const emptyLedger: EvidenceLedgerSnapshot = { - retrievalCount: 0, - chunks: [], - assets: [], - evidenceText: [], - stopReasons: [], - failureReasons: [], - decisionTraces: [], -} - -const unrelatedContextPolicy: ContextPolicy = { - carryHistory: "none", - reason: "The current turn is self-contained.", - activePriorTurnIds: [], -} - -function makeIntent( - constraints: IntentFrame["constraints"] & { - readonly groundingPolicy?: IntentFrame["groundingPolicy"] - }, -): IntentFrame { - return { - task: "answer", - dependsOnPreviousTurn: false, - retrievalNeeded: constraints.groundingPolicy === "no_retrieval" ? "no" : "yes", - targetModalities: ["text"], - constraints, - groundingPolicy: constraints.groundingPolicy ?? "must_use_sources", - } -} - -function makeManifest(overrides: Partial): OutputManifest { - return { - text: "", - citations: [], - artifacts: [], - unresolved: [], - ...overrides, - } -} - -function makeAsset(ref: string): EvidenceLedgerSnapshot["assets"][number] { - return { - ref, - chunkRef: ref.replace("asset:", ""), - type: "image", - assetUrl: `https://assets.example/${ref}.png`, - label: ref, - source: { - documentId: "doc_1", - sourceFileName: "report.pdf", - sectionPath: ref, - }, - } -} diff --git a/src/agent-harness/validator.ts b/src/agent-harness/validator.ts deleted file mode 100644 index dd67343..0000000 --- a/src/agent-harness/validator.ts +++ /dev/null @@ -1,284 +0,0 @@ -import type { - ContextPolicy, - EvidenceLedgerSnapshot, - HarnessToolCallTrace, - IntentFrame, - OutputManifest, -} from "./types" - -export type ManifestValidationInput = { - readonly manifest: OutputManifest - readonly intent?: IntentFrame - readonly contextPolicy?: ContextPolicy - readonly finalized?: boolean - readonly ledger: EvidenceLedgerSnapshot - readonly toolCalls?: readonly HarnessToolCallTrace[] - readonly surface: "notebook_chat" | "typing_compose" | "typing_quick_ask" -} - -export type ManifestValidationResult = { - readonly ok: boolean - readonly errors: readonly string[] -} - -export function validateOutputManifest( - input: ManifestValidationInput, -): ManifestValidationResult { - const errors: string[] = [] - const text = input.manifest.text.trim() - - if (!text && input.manifest.artifacts.every((artifact) => !artifact.display)) { - errors.push("Final output must contain text or at least one displayed artifact.") - } - - validateWorkflow(input, errors) - validateArtifactRefs(input, errors) - validateArtifactCounts(input, errors) - validateGrounding(input, errors) - validateTaskEvidence(input, errors) - validateImageInspectionClaims(input, errors) - validateUnavailableImageContentClaims(input, errors) - validateTypingText(input, errors) - - return { - ok: errors.length === 0, - errors, - } -} - -function validateWorkflow( - input: ManifestValidationInput, - errors: string[], -): void { - if (input.finalized === false) { - errors.push("Agent must call finalize to produce the output manifest.") - } - - if (!input.intent) { - errors.push("Agent must declare intent before finalizing.") - } - - if (!input.contextPolicy) { - errors.push("Agent must set context policy before finalizing.") - } -} - -function validateArtifactRefs( - input: ManifestValidationInput, - errors: string[], -): void { - const knownRefs = new Set([ - ...input.ledger.chunks.map((chunk) => chunk.ref), - ...input.ledger.assets.map((asset) => asset.ref), - ]) - - for (const artifact of input.manifest.artifacts) { - if (artifact.type === "derived_table") { - artifact.rows.forEach((row, index) => { - if (row.length !== artifact.columns.length) { - errors.push( - `Derived table row ${index + 1} has ${row.length} cells but expected ${artifact.columns.length}.`, - ) - } - }) - - for (const ref of artifact.sourceRefs) { - if (!knownRefs.has(ref)) { - errors.push( - `Derived table source ref '${ref}' was not found in the evidence ledger.`, - ) - } - } - continue - } - - if (!knownRefs.has(artifact.ref)) { - errors.push( - `Artifact ref '${artifact.ref}' was not found in the evidence ledger.`, - ) - } - } - - for (const citation of input.manifest.citations) { - if (!knownRefs.has(citation.ref)) { - errors.push(`Citation ref '${citation.ref}' was not found in the evidence ledger.`) - } - } -} - -function validateArtifactCounts( - input: ManifestValidationInput, - errors: string[], -): void { - const displayedCount = input.manifest.artifacts.filter( - (artifact) => artifact.display, - ).length - const desiredCount = input.intent?.constraints.desiredCount - const maxCount = input.intent?.constraints.maxCount - - if (typeof desiredCount === "number" && displayedCount > desiredCount) { - errors.push( - `Displayed artifact count ${displayedCount} exceeds desired count ${desiredCount}.`, - ) - } - - if (typeof maxCount === "number" && displayedCount > maxCount) { - errors.push( - `Displayed artifact count ${displayedCount} exceeds maximum count ${maxCount}.`, - ) - } -} - -function validateGrounding( - input: ManifestValidationInput, - errors: string[], -): void { - if (input.intent?.groundingPolicy !== "must_use_sources") return - - const hasLedgerEvidence = - input.ledger.chunks.length > 0 || input.ledger.evidenceText.length > 0 - const hasOutputEvidence = - input.manifest.citations.length > 0 || - input.manifest.artifacts.some((artifact) => artifact.display) - const hasUnresolved = input.manifest.unresolved.length > 0 - - if (!hasLedgerEvidence && !hasUnresolved) { - errors.push( - "Grounded output requires evidence or an explicit unresolved reason.", - ) - } - - if (hasLedgerEvidence && !hasOutputEvidence && !hasUnresolved) { - errors.push( - "Grounded output used evidence but did not cite or display any selected evidence.", - ) - } -} - -function validateTaskEvidence( - input: ManifestValidationInput, - errors: string[], -): void { - if (input.intent?.groundingPolicy !== "must_use_sources") return - if (input.manifest.unresolved.length > 0) return - - const refs = getOutputEvidenceRefs(input.manifest) - - if (input.intent.task === "compare" && refs.size < 2) { - errors.push( - "Compare outputs that must use sources require at least two evidence refs or an explicit unresolved reason.", - ) - } - - if (input.intent.task === "summarize" && refs.size < 1) { - errors.push( - "Summaries that must use sources require at least one evidence ref or an explicit unresolved reason.", - ) - } -} - -function getOutputEvidenceRefs(manifest: OutputManifest): Set { - const refs = new Set() - - for (const citation of manifest.citations) refs.add(citation.ref) - - for (const artifact of manifest.artifacts) { - if (!artifact.display) continue - if (artifact.type === "derived_table") { - artifact.sourceRefs.forEach((ref) => refs.add(ref)) - } else { - refs.add(artifact.ref) - } - } - - return refs -} - -function validateTypingText( - input: ManifestValidationInput, - errors: string[], -): void { - if (input.surface !== "typing_compose") return - - const text = input.manifest.text - if (/```|^\s*#{1,6}\s|^\s*[-*]\s/m.test(text)) { - errors.push("Typing compose output must be insertion-ready plain text.") - } -} - -function validateImageInspectionClaims( - input: ManifestValidationInput, - errors: string[], -): void { - const claimText = [input.manifest.text, ...input.manifest.unresolved].join("\n") - if (!mentionsImageInspectionResult(claimText)) return - - if (hasImageInspectionToolCall(input)) return - - errors.push( - "Final output must not claim image/OCR inspection succeeded or failed unless inspectImage was called.", - ) -} - -function validateUnavailableImageContentClaims( - input: ManifestValidationInput, - errors: string[], -): void { - const hasImageAssets = input.ledger.assets.some((asset) => asset.type === "image") - if (!hasImageAssets) return - if (hasImageInspectionToolCall(input)) return - - const claimText = [input.manifest.text, ...input.manifest.unresolved].join("\n") - if (!mentionsUnavailableImageContent(claimText)) return - - errors.push( - "Final output must inspect available image assets before claiming retrieved page/image content cannot be read.", - ) -} - -function hasImageInspectionToolCall(input: ManifestValidationInput): boolean { - return input.toolCalls?.some((call) => call.tool === "inspectImage") === true -} - -function mentionsImageInspectionResult(value: string): boolean { - const normalized = value.replace(/\s+/g, " ").trim() - if (!normalized) return false - - return ( - /\b(?:image|visual)\s+(?:inspection|recognition|analysis)\s+(?:failed|did not|could not|was unable|found|showed|confirmed)/iu.test( - normalized, - ) || - /\bOCR\s+(?:failed|did not|could not|was unable|found|showed|confirmed|extracted)/iu.test( - normalized, - ) || - /(?:图像|图片|视觉).{0,8}(?:识别|检查|检视|查看|分析).{0,12}(?:未|没|无法|不能|不成功|失败|成功|显示|发现|提取|读取)/u.test( - normalized, - ) || - /(?:图像|图片|视觉).{0,12}(?:未|没|无法|不能|不成功|失败).{0,16}(?:识别|检查|检视|查看|分析|检测|提取|读取|获取)/u.test( - normalized, - ) || - /(?:未|没|无法|不能|不成功|失败).{0,12}(?:图像|图片|视觉|OCR).{0,12}(?:识别|检查|检视|查看|分析|检测|提取|读取|获取)/u.test( - normalized, - ) - ) -} - -function mentionsUnavailableImageContent(value: string): boolean { - const normalized = value.replace(/\s+/g, " ").trim() - if (!normalized) return false - - return ( - /\b(?:cannot|can't|could not|unable to|was unable to|did not)\s+(?:directly\s+)?(?:read|inspect|access|extract|see|view)\b.{0,48}\b(?:page|image|visual|OCR|content|clause|details|text)\b/iu.test( - normalized, - ) || - /\b(?:page|image|visual|OCR|content|clause|details|text)\b.{0,48}\b(?:cannot|can't|could not|unable to|was unable to|did not)\s+(?:directly\s+)?(?:read|inspect|access|extract|see|view)\b/iu.test( - normalized, - ) || - /(?:无法|不能|未能|没法|没有办法).{0,8}(?:直接)?(?:读取|查看|识别|检测|提取|看清|访问|获取).{0,16}(?:页面|页|图片|图像|条款|内容|细节|文字)/u.test( - normalized, - ) || - /(?:页面|页|图片|图像|条款|内容|细节|文字).{0,16}(?:无法|不能|未能|没法|没有办法).{0,16}(?:直接)?(?:读取|查看|识别|检测|提取|看清|访问|获取)/u.test( - normalized, - ) - ) -} diff --git a/src/app/api/demo-sources/[demoSourceId]/assets/[...assetPath]/route.test.ts b/src/app/api/demo-sources/[demoSourceId]/assets/[...assetPath]/route.test.ts new file mode 100644 index 0000000..aa1d09b --- /dev/null +++ b/src/app/api/demo-sources/[demoSourceId]/assets/[...assetPath]/route.test.ts @@ -0,0 +1,51 @@ +import { beforeEach, describe, expect, it, vi } from "vitest" + +const mocks = vi.hoisted(() => ({ + fetch: vi.fn(), +})) + +vi.mock("@/integrations/knowhere-demo", () => ({ + knowhereDemoApi: { + resolveApiURL: (pathname: string) => `https://demo.example${pathname}`, + }, +})) + +import { GET } from "./route" + +describe("GET /api/demo-sources/[demoSourceId]/assets/[...assetPath]", () => { + beforeEach(() => { + vi.clearAllMocks() + vi.stubGlobal("fetch", mocks.fetch) + }) + + it("serves demo image assets inline", async () => { + mocks.fetch.mockResolvedValue( + new Response("image bytes", { + status: 200, + headers: { + "content-type": "image/png", + }, + }), + ) + + const response = await GET( + new Request( + "http://localhost:3001/api/demo-sources/demo_1/assets/images/page-1.png", + ), + { + params: Promise.resolve({ + demoSourceId: "demo_1", + assetPath: ["images", "page-1.png"], + }), + }, + ) + + expect(response.status).toBe(200) + expect(response.headers.get("content-type")).toBe("image/png") + expect(response.headers.get("content-disposition")).toBe("inline") + expect(mocks.fetch).toHaveBeenCalledWith( + "https://demo.example/api/v1/demo/sources/demo_1/assets/images/page-1.png", + { cache: "no-store" }, + ) + }) +}) diff --git a/src/app/api/demo-sources/[demoSourceId]/assets/[...assetPath]/route.ts b/src/app/api/demo-sources/[demoSourceId]/assets/[...assetPath]/route.ts index 247f645..7dc89c3 100644 --- a/src/app/api/demo-sources/[demoSourceId]/assets/[...assetPath]/route.ts +++ b/src/app/api/demo-sources/[demoSourceId]/assets/[...assetPath]/route.ts @@ -33,6 +33,7 @@ export async function GET( status: 200, headers: { "content-type": response.headers.get("content-type") ?? "application/octet-stream", + "content-disposition": "inline", "cache-control": "public, max-age=3600", }, }) diff --git a/src/app/api/parsed-assets/inline/route.test.ts b/src/app/api/parsed-assets/inline/route.test.ts new file mode 100644 index 0000000..1bff9d3 --- /dev/null +++ b/src/app/api/parsed-assets/inline/route.test.ts @@ -0,0 +1,96 @@ +import { NextRequest } from "next/server" +import { beforeEach, describe, expect, it, vi } from "vitest" + +const mocks = vi.hoisted(() => ({ + getAuthenticated: vi.fn(), + getBlob: vi.fn(), +})) + +vi.mock("@/domains/workspace/request-context", () => ({ + notebookRequestContext: { + getAuthenticated: mocks.getAuthenticated, + }, +})) + +vi.mock("@vercel/blob", () => ({ + BlobNotFoundError: class BlobNotFoundError extends Error {}, + get: mocks.getBlob, +})) + +import { GET } from "./route" + +describe("GET /api/parsed-assets/inline", () => { + beforeEach(() => { + vi.clearAllMocks() + mocks.getAuthenticated.mockResolvedValue({ + workspace: { id: "workspace_1" }, + }) + }) + + it("serves Notebook parsed document image blobs inline", async () => { + const assetUrl = + "https://store.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_1/rev_1/page_citation_assets/page-1.png" + mocks.getBlob.mockResolvedValue({ + statusCode: 200, + stream: new Response("image bytes").body, + blob: { + contentType: "image/png", + }, + }) + + const response = await GET( + new NextRequest( + `http://localhost:3001/api/parsed-assets/inline?url=${encodeURIComponent( + assetUrl, + )}`, + ), + ) + + expect(response.status).toBe(200) + expect(response.headers.get("content-type")).toBe("image/png") + expect(response.headers.get("content-disposition")).toBe("inline") + await expect(response.text()).resolves.toBe("image bytes") + expect(mocks.getBlob).toHaveBeenCalledWith( + "workspaces/workspace_1/parsed-documents/doc_1/rev_1/page_citation_assets/page-1.png", + { access: "public" }, + ) + }) + + it("rejects blob URLs outside the authenticated workspace", async () => { + const assetUrl = + "https://store.public.blob.vercel-storage.com/workspaces/workspace_2/parsed-documents/doc_1/rev_1/page_citation_assets/page-1.png" + + const response = await GET( + new NextRequest( + `http://localhost:3001/api/parsed-assets/inline?url=${encodeURIComponent( + assetUrl, + )}`, + ), + ) + + expect(response.status).toBe(404) + expect(mocks.getBlob).not.toHaveBeenCalled() + }) + + it("does not serve non-image parsed assets through the inline image route", async () => { + const assetUrl = + "https://store.public.blob.vercel-storage.com/workspaces/workspace_1/sources/source_1/parsed-result/tables/table-1.html" + mocks.getBlob.mockResolvedValue({ + statusCode: 200, + stream: new Response("
").body, + blob: { + contentType: "text/html; charset=utf-8", + }, + }) + + const response = await GET( + new NextRequest( + `http://localhost:3001/api/parsed-assets/inline?url=${encodeURIComponent( + assetUrl, + )}`, + ), + ) + + expect(response.status).toBe(415) + }) +}) diff --git a/src/app/api/parsed-assets/inline/route.ts b/src/app/api/parsed-assets/inline/route.ts new file mode 100644 index 0000000..9e9beb3 --- /dev/null +++ b/src/app/api/parsed-assets/inline/route.ts @@ -0,0 +1,159 @@ +import { BlobNotFoundError, get } from "@vercel/blob" +import type { NextRequest } from "next/server" + +import { notebookRequestContext } from "@/domains/workspace/request-context" + +type ParsedAssetBlobPath = { + readonly pathname: string + readonly assetPath: string +} + +const notebookBlobHostSuffix = ".blob.vercel-storage.com" +const inlineImageContentTypes = new Set([ + "image/gif", + "image/jpeg", + "image/png", + "image/webp", +]) +const parsedDocumentsPathPattern = + /^\/?workspaces\/([^/]+)\/parsed-documents\/[^/]+\/[^/]+\/(.+)$/u +const parsedResultPathPattern = + /^\/?workspaces\/([^/]+)\/sources\/[^/]+\/parsed-result\/(.+)$/u + +export async function GET(request: NextRequest): Promise { + const { workspace } = await notebookRequestContext.getAuthenticated() + const parsedAssetBlobPath = getParsedAssetBlobPath( + request.nextUrl.searchParams.get("url"), + workspace.id, + ) + if (!parsedAssetBlobPath) { + return Response.json({ message: "Parsed asset not found." }, { status: 404 }) + } + + try { + const blob = await get(parsedAssetBlobPath.pathname, { access: "public" }) + if (!blob) { + return Response.json( + { message: "Parsed asset not found." }, + { status: 404 }, + ) + } + if (blob.statusCode !== 200) { + return new Response(null, { status: 304 }) + } + + const contentType = getInlineImageContentType( + blob.blob.contentType, + parsedAssetBlobPath.assetPath, + ) + if (!contentType) { + return Response.json( + { message: "Parsed asset is not an image." }, + { status: 415 }, + ) + } + + return new Response(blob.stream, { + status: 200, + headers: { + "content-type": contentType, + "content-disposition": "inline", + "cache-control": "public, max-age=3600", + }, + }) + } catch (error) { + if (error instanceof BlobNotFoundError) { + return Response.json( + { message: "Parsed asset not found." }, + { status: 404 }, + ) + } + throw error + } +} + +function getParsedAssetBlobPath( + assetUrl: string | null, + workspaceId: string, +): ParsedAssetBlobPath | null { + if (!assetUrl) return null + + const url = parseAbsoluteUrl(assetUrl) + if (!url) return null + if (url.protocol !== "https:") return null + if (!url.hostname.toLowerCase().endsWith(notebookBlobHostSuffix)) return null + + const pathname = decodePathname(url.pathname) + if (!pathname) return null + + const documentsMatch = parsedDocumentsPathPattern.exec(pathname) + if (documentsMatch) { + return getWorkspaceScopedAssetPath(documentsMatch, workspaceId, pathname) + } + + const parsedResultMatch = parsedResultPathPattern.exec(pathname) + if (parsedResultMatch) { + return getWorkspaceScopedAssetPath(parsedResultMatch, workspaceId, pathname) + } + + return null +} + +function getWorkspaceScopedAssetPath( + match: RegExpExecArray, + workspaceId: string, + pathname: string, +): ParsedAssetBlobPath | null { + const matchedWorkspaceId = match[1] + const assetPath = match[2] + if (matchedWorkspaceId !== workspaceId || !assetPath) return null + + return { + pathname: pathname.replace(/^\/+/u, ""), + assetPath, + } +} + +function getInlineImageContentType( + contentType: string | null, + assetPath: string, +): string | null { + const normalizedContentType = getBaseContentType(contentType) + if (normalizedContentType && inlineImageContentTypes.has(normalizedContentType)) { + return contentType + } + + return inferImageContentType(assetPath) +} + +function getBaseContentType(contentType: string | null): string | null { + const baseContentType = contentType?.split(";")[0]?.trim().toLowerCase() + return baseContentType && baseContentType.length > 0 ? baseContentType : null +} + +function inferImageContentType(assetPath: string): string | null { + const pathname = assetPath.toLowerCase().split("?")[0] ?? assetPath + if (pathname.endsWith(".png")) return "image/png" + if (pathname.endsWith(".jpg") || pathname.endsWith(".jpeg")) { + return "image/jpeg" + } + if (pathname.endsWith(".gif")) return "image/gif" + if (pathname.endsWith(".webp")) return "image/webp" + return null +} + +function parseAbsoluteUrl(value: string): URL | null { + try { + return new URL(value) + } catch { + return null + } +} + +function decodePathname(pathname: string): string | null { + try { + return decodeURIComponent(pathname) + } catch { + return null + } +} diff --git a/src/app/api/sources/[sourceId]/chunks/route.test.ts b/src/app/api/sources/[sourceId]/chunks/route.test.ts index 85fb1bd..f410285 100644 --- a/src/app/api/sources/[sourceId]/chunks/route.test.ts +++ b/src/app/api/sources/[sourceId]/chunks/route.test.ts @@ -8,6 +8,7 @@ const mocks = vi.hoisted(() => ({ fetchDemoChunkPage: vi.fn(), findSourceInWorkspace: vi.fn(), getCurrentUser: vi.fn(), + listChunks: vi.fn(), localizeRemoteDocument: vi.fn(), makeKnowhereClient: vi.fn(), makeKnowhereClientWithParsedStorage: vi.fn(), @@ -64,7 +65,7 @@ describe("GET /api/sources/[sourceId]/chunks", () => { beforeEach(() => { vi.clearAllMocks() mocks.makeKnowhereClientWithParsedStorage.mockReturnValue({ - client: { documents: { listChunks: vi.fn() } }, + client: { documents: { listChunks: mocks.listChunks } }, knowledge: { readChunks: mocks.readChunks }, }) }) @@ -421,7 +422,7 @@ describe("GET /api/sources/[sourceId]/chunks", () => { } }) - it("loads authenticated workspace chunks through the SDK durable read", async () => { + it("loads authenticated workspace chunks from parsed storage without durable hardening", async () => { mocks.getCurrentUser.mockResolvedValue({ id: "user_1", email: null, @@ -442,7 +443,10 @@ describe("GET /api/sources/[sourceId]/chunks", () => { ) mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") mocks.readChunks.mockResolvedValue({ - document: { localDocumentId: "doc_1" }, + document: { + localDocumentId: "doc_1", + resultDirectoryPath: "parsed-storage:doc_1", + }, chunks: [ { position: 1, @@ -498,7 +502,114 @@ describe("GET /api/sources/[sourceId]/chunks", () => { revisionKey: "job_1", page: 1, pageSize: 1, - assetUrlPolicy: "durable", + }) + expect(mocks.listChunks).not.toHaveBeenCalled() + }) + + it("falls back to raw Knowhere asset URLs for workspace chunk display when storage misses", async () => { + mocks.getCurrentUser.mockResolvedValue({ + id: "user_1", + email: null, + name: null, + }) + mocks.ensureWorkspace.mockResolvedValue({ + id: "workspace_1", + userId: "user_1", + namespace: "notebook-workspace_1", + createdAt: new Date("2026-05-10T00:00:00.000Z"), + }) + mocks.findSourceInWorkspace.mockResolvedValue( + makeReadySource({ + id: "00000000-0000-0000-0000-000000000002", + knowhereJobId: "job_1", + knowhereDocumentId: "doc_1", + }), + ) + mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") + mocks.readChunks.mockResolvedValue({ + document: { + localDocumentId: "doc_1", + resultDirectoryPath: "remote:doc_1", + }, + chunks: [ + { + position: 1, + chunkId: "parser_1", + chunkType: "image", + content: "Workspace chunk", + readableContent: "Workspace chunk", + sectionPath: "Summary", + sourceChunkPath: "Summary", + filePath: "images/chart.png", + metadata: {}, + }, + ], + page: 1, + pageSize: 1, + totalChunks: 1, + totalPages: 1, + }) + mocks.listChunks.mockResolvedValue({ + chunks: [ + { + id: "document_chunk_1", + chunkId: "parser_1", + chunkType: "image", + content: "Workspace chunk", + sectionId: null, + sectionPath: "Summary", + sourceChunkPath: "Summary", + filePath: "images/chart.png", + sortOrder: 0, + metadata: {}, + assetUrl: "https://knowhere.example/assets/chart.png", + }, + ], + pagination: { + page: 1, + pageSize: 1, + total: 1, + totalPages: 1, + }, + }) + + const response = await GET( + new NextRequest( + "http://localhost:3001/api/sources/00000000-0000-0000-0000-000000000002/chunks?page=1&pageSize=1", + ), + { + params: Promise.resolve({ + sourceId: "00000000-0000-0000-0000-000000000002", + }), + }, + ) + + await expect(response.json()).resolves.toMatchObject({ + chunks: [ + { + parserChunkId: "parser_1", + documentId: "doc_1", + assetUrl: "https://knowhere.example/assets/chart.png", + sourceTitle: "notes.pdf", + }, + ], + pagination: { + page: 1, + pageSize: 1, + total: 1, + }, + }) + expect(response.status).toBe(200) + expect(mocks.readChunks).toHaveBeenCalledWith({ + documentId: "doc_1", + revisionKey: "job_1", + page: 1, + pageSize: 1, + }) + expect(mocks.listChunks).toHaveBeenCalledWith("doc_1", { + page: 1, + pageSize: 1, + includeAssetUrls: true, }) }) @@ -548,6 +659,56 @@ describe("GET /api/sources/[sourceId]/chunks", () => { expect(mocks.readChunks).not.toHaveBeenCalled() }) + it("marks workspace chunks unavailable when the parsed document is missing remotely", async () => { + const notFoundError = new Error("Document not found") + notFoundError.name = "NotFoundError" + mocks.getCurrentUser.mockResolvedValue({ + id: "user_1", + email: null, + name: null, + }) + mocks.ensureWorkspace.mockResolvedValue({ + id: "workspace_1", + userId: "user_1", + namespace: "notebook-workspace_1", + createdAt: new Date("2026-05-10T00:00:00.000Z"), + }) + mocks.findSourceInWorkspace.mockResolvedValue( + makeReadySource({ + id: "00000000-0000-0000-0000-000000000002", + knowhereJobId: "job_1", + knowhereDocumentId: "doc_missing", + }), + ) + mocks.ensureApiKeyForWorkspace.mockResolvedValue("jwt_123") + mocks.readChunks.mockRejectedValue(notFoundError) + + const response = await GET( + new NextRequest( + "http://localhost:3001/api/sources/00000000-0000-0000-0000-000000000002/chunks?page=1&pageSize=1", + ), + { + params: Promise.resolve({ + sourceId: "00000000-0000-0000-0000-000000000002", + }), + }, + ) + + await expect(response.json()).resolves.toEqual({ + chunks: [], + pagination: { + page: 1, + pageSize: 1, + total: 0, + totalPages: 0, + }, + message: + "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", + isUnavailable: true, + }) + expect(response.status).toBe(200) + }) + it("materializes a remote source id and reads chunks through the SDK", async () => { const knowhereClient = { documents: { @@ -591,7 +752,10 @@ describe("GET /api/sources/[sourceId]/chunks", () => { }), ) mocks.readChunks.mockResolvedValue({ - document: { localDocumentId: "doc_remote" }, + document: { + localDocumentId: "doc_remote", + resultDirectoryPath: "parsed-storage:doc_remote", + }, chunks: [ { position: 1, @@ -660,7 +824,6 @@ describe("GET /api/sources/[sourceId]/chunks", () => { revisionKey: "job_result_1", page: 1, pageSize: 1, - assetUrlPolicy: "durable", }) }) }) diff --git a/src/app/home-content.tsx b/src/app/home-content.tsx new file mode 100644 index 0000000..c2f1a64 --- /dev/null +++ b/src/app/home-content.tsx @@ -0,0 +1,28 @@ +import { connection } from "next/server" +import { WorkspaceShell } from "@/components/workspace-shell" +import { loadWorkspaceShellInitialState } from "@/domains/workspace/initial-state" +import { effectOperation } from "@/lib/effect-operation" +import { summarizeUnknownError } from "@/lib/format-log-value" +import { logger } from "@/lib/logger" + +export async function HomeContent() { + await connection() + const initialState = await loadWorkspaceInitialState() + return +} + +async function loadWorkspaceInitialState(): ReturnType< + typeof loadWorkspaceShellInitialState +> { + try { + return await loadWorkspaceShellInitialState() + } catch (error) { + logger.error("workspace: initial state failed", { + error: summarizeUnknownError(error), + }) + throw effectOperation.createBoundaryError( + "Workspace initial state failed", + error, + ) + } +} diff --git a/src/app/page.test.ts b/src/app/page.test.ts index f7e3525..3e9f811 100644 --- a/src/app/page.test.ts +++ b/src/app/page.test.ts @@ -20,7 +20,7 @@ vi.mock("@/domains/workspace/initial-state", () => ({ loadWorkspaceShellInitialState: mocks.loadWorkspaceShellInitialState, })) -import { HomeContent } from "./page" +import { HomeContent } from "./home-content" import { makeWorkspaceInitialStateFailureFixture } from "@/test/workspace-initial-state-failure-fixture" describe("Home", () => { diff --git a/src/app/page.tsx b/src/app/page.tsx index 2ebc597..0a3c967 100644 --- a/src/app/page.tsx +++ b/src/app/page.tsx @@ -1,10 +1,5 @@ import { Suspense } from "react" -import { WorkspaceShell } from "@/components/workspace-shell" -import { loadWorkspaceShellInitialState } from "@/domains/workspace/initial-state" -import { effectOperation } from "@/lib/effect-operation" -import { summarizeUnknownError } from "@/lib/format-log-value" -import { logger } from "@/lib/logger" -import { connection } from "next/server" +import { HomeContent } from "./home-content" export default function Home() { return ( @@ -13,25 +8,3 @@ export default function Home() { ) } - -export async function HomeContent() { - await connection() - const initialState = await loadWorkspaceInitialState() - return -} - -async function loadWorkspaceInitialState(): ReturnType< - typeof loadWorkspaceShellInitialState -> { - try { - return await loadWorkspaceShellInitialState() - } catch (error) { - logger.error("workspace: initial state failed", { - error: summarizeUnknownError(error), - }) - throw effectOperation.createBoundaryError( - "Workspace initial state failed", - error, - ) - } -} diff --git a/src/components/chat-message-list.test.ts b/src/components/chat-message-list.test.ts index f94478e..9444128 100644 --- a/src/components/chat-message-list.test.ts +++ b/src/components/chat-message-list.test.ts @@ -146,6 +146,7 @@ describe("ChatMessageList", () => { score: 0.9, pageCitationAssetUrl: "https://blob.example/pages/page-000004.png", + pageCitationPageNumber: 4, source: { documentId: "doc_1", sourceFileName: "report.pdf", @@ -162,18 +163,17 @@ describe("ChatMessageList", () => { const citationButton = screen.getByRole("button", { name: "Open source report.pdf", }); - const pageImageLink = screen.getByRole("link", { - name: "Open page image for report.pdf", - }); - - expect(pageImageLink.getAttribute("href")).toBe( - "https://blob.example/pages/page-000004.png", - ); + expect( + screen.queryByRole("link", { + name: "Open page image for report.pdf", + }), + ).toBeNull(); await user.click(citationButton); expect(onCitationClick).toHaveBeenCalledWith( expect.objectContaining({ pageCitationAssetUrl: "https://blob.example/pages/page-000004.png", + pageCitationPageNumber: 4, }), "assistant_1:0", ); diff --git a/src/components/chat-message-list.tsx b/src/components/chat-message-list.tsx index 531c71f..38c5059 100644 --- a/src/components/chat-message-list.tsx +++ b/src/components/chat-message-list.tsx @@ -2,7 +2,7 @@ import { type CSSProperties, type ReactElement } from "react"; import { type VirtualItem } from "@tanstack/react-virtual"; -import { ExternalLink, ImageIcon, MessageCircle } from "lucide-react"; +import { ImageIcon, MessageCircle } from "lucide-react"; import ReactMarkdown, { defaultUrlTransform, type Components, @@ -633,10 +633,6 @@ function CitationChip({ citationId: string, ) => void; }): ReactElement { - const pageCitationAssetUrl = getTrimmedCitationField( - citation.pageCitationAssetUrl, - ); - return ( @@ -646,11 +642,7 @@ function CitationChip({ disabled={!onCitationClick || isPending} onClick={() => onCitationClick?.(citation, citationId)} aria-busy={isPending} - className={`inline-flex h-8 max-w-[250px] cursor-pointer items-center border border-primary/20 bg-primary/10 px-3 text-left font-mono text-xs font-semibold leading-none text-primary shadow-[0_1px_0_rgba(15,23,42,0.06)] transition-[background-color,border-color,color,box-shadow,transform] hover:border-primary/35 hover:bg-primary/15 hover:text-primary hover:shadow-[0_0_0_2px_rgba(37,99,235,0.12)] active:translate-y-px active:bg-primary/20 focus:outline-none focus:ring-4 focus:ring-ring/15 focus:ring-offset-2 focus:ring-offset-background disabled:cursor-wait disabled:opacity-75 disabled:hover:border-primary/20 disabled:hover:bg-primary/10 disabled:hover:text-primary disabled:hover:shadow-[0_1px_0_rgba(15,23,42,0.06)] dark:border-transparent dark:bg-[#5c606b] dark:text-[#cfd3dc] dark:shadow-none dark:hover:border-[#8f96a8] dark:hover:bg-[#4f535e] dark:hover:text-white dark:hover:shadow-[0_0_0_2px_rgba(143,150,168,0.22)] dark:active:bg-[#454955] dark:disabled:hover:border-transparent dark:disabled:hover:bg-[#5c606b] dark:disabled:hover:text-[#cfd3dc] dark:disabled:hover:shadow-none ${ - pageCitationAssetUrl - ? "rounded-l-md rounded-r-none" - : "rounded-md" - }`} + className="inline-flex h-8 max-w-[250px] cursor-pointer items-center rounded-md border border-primary/20 bg-primary/10 px-3 text-left font-mono text-xs font-semibold leading-none text-primary shadow-[0_1px_0_rgba(15,23,42,0.06)] transition-[background-color,border-color,color,box-shadow,transform] hover:border-primary/35 hover:bg-primary/15 hover:text-primary hover:shadow-[0_0_0_2px_rgba(37,99,235,0.12)] active:translate-y-px active:bg-primary/20 focus:outline-none focus:ring-4 focus:ring-ring/15 focus:ring-offset-2 focus:ring-offset-background disabled:cursor-wait disabled:opacity-75 disabled:hover:border-primary/20 disabled:hover:bg-primary/10 disabled:hover:text-primary disabled:hover:shadow-[0_1px_0_rgba(15,23,42,0.06)] dark:border-transparent dark:bg-[#5c606b] dark:text-[#cfd3dc] dark:shadow-none dark:hover:border-[#8f96a8] dark:hover:bg-[#4f535e] dark:hover:text-white dark:hover:shadow-[0_0_0_2px_rgba(143,150,168,0.22)] dark:active:bg-[#454955] dark:disabled:hover:border-transparent dark:disabled:hover:bg-[#5c606b] dark:disabled:hover:text-[#cfd3dc] dark:disabled:hover:shadow-none" aria-label={`Open source ${label}`} > {label} @@ -664,28 +656,6 @@ function CitationChip({ {tooltipLabel} - {pageCitationAssetUrl ? ( - - - - - - - - Open page image - - - ) : null} ); } diff --git a/src/components/chunks-panel-state.test.ts b/src/components/chunks-panel-state.test.ts index bcee0f0..53cb15a 100644 --- a/src/components/chunks-panel-state.test.ts +++ b/src/components/chunks-panel-state.test.ts @@ -139,6 +139,65 @@ describe("chunksPanelState", () => { ]) }) + it("deduplicates page-asset chunks by page number", () => { + const chunks: ParsedChunkView[] = [ + { + chunkId: "page_4_first", + type: "page", + content: "First page 4 summary.", + sourceTitle: "manual.pdf", + pageNums: [4], + pageAssets: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-4-a.png", + contentType: "image/png", + }, + ], + }, + { + chunkId: "page_4_duplicate", + type: "page", + content: "Duplicate page 4 summary.", + sourceTitle: "manual.pdf", + pageNums: [4], + pageAssets: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-4-b.png", + contentType: "image/png", + }, + ], + }, + { + chunkId: "page_5", + type: "page", + content: "Page 5 summary.", + sourceTitle: "manual.pdf", + pageAssets: [ + { + pageNumber: 5, + assetUrl: "https://assets.example/page-5.png", + contentType: "image/png", + }, + ], + }, + { + chunkId: "table_1", + type: "table", + content: "", + sourceTitle: "manual.pdf", + pageNums: [4], + }, + ] + + expect( + chunksPanelState + .getPageAssetChunksWithoutDuplicatePages(chunks) + .map((chunk) => chunk.chunkId), + ).toEqual(["page_4_first", "page_5", "table_1"]) + }) + it("formats Knowhere section paths and reference labels for display", () => { expect( chunksPanelState.formatChunkSectionPath( diff --git a/src/components/chunks-panel-state.ts b/src/components/chunks-panel-state.ts index dfd2ac7..0c64515 100644 --- a/src/components/chunks-panel-state.ts +++ b/src/components/chunks-panel-state.ts @@ -42,6 +42,9 @@ type ChunksPanelStateModule = { chunks: readonly ParsedChunkView[], focusedChunkId: string | null, ) => readonly ParsedChunkView[] + readonly getPageAssetChunksWithoutDuplicatePages: ( + chunks: readonly ParsedChunkView[], + ) => readonly ParsedChunkView[] readonly getReferenceLabel: (connection: ParsedChunkConnection) => string readonly getRenderableReferences: ( chunk: ParsedChunkView, @@ -157,6 +160,32 @@ function dedupeChunksById( return uniqueChunks } +function getPageAssetChunksWithoutDuplicatePages( + chunks: readonly ParsedChunkView[], +): readonly ParsedChunkView[] { + const seenPageNumbers = new Set() + + return chunks.filter((chunk) => { + if (chunk.type !== "page") return true + + const pageNumber = getPageAssetChunkPageNumber(chunk) + if (pageNumber === null) return true + if (seenPageNumbers.has(pageNumber)) return false + + seenPageNumbers.add(pageNumber) + return true + }) +} + +function getPageAssetChunkPageNumber(chunk: ParsedChunkView): number | null { + const pageAssetNumbers = (chunk.pageAssets ?? []) + .map((pageAsset) => pageAsset.pageNumber) + .filter(isPositivePageNumber) + if (pageAssetNumbers.length > 0) return Math.min(...pageAssetNumbers) + + return getFirstPageNumber(chunk) +} + function createMutableSectionTreeNode(input: { readonly id: string readonly kind: ChunkSectionTreeNodeKind @@ -336,6 +365,10 @@ function getFirstPageNumber(chunk: ParsedChunkView): number | null { return Math.min(...finitePageNumbers) } +function isPositivePageNumber(pageNumber: number): boolean { + return Number.isFinite(pageNumber) && pageNumber > 0 +} + function formatChunkSectionPath( sectionPath: ParsedChunkView["sectionPath"], ): string | null { @@ -454,6 +487,7 @@ export const chunksPanelState: ChunksPanelStateModule = { formatChunkSectionPath, formatReferenceLabel, getChunksWithFocusedFirst, + getPageAssetChunksWithoutDuplicatePages, getReferenceLabel, getRenderableReferences, } diff --git a/src/components/chunks-panel.test.ts b/src/components/chunks-panel.test.ts index f99a68b..3031cb8 100644 --- a/src/components/chunks-panel.test.ts +++ b/src/components/chunks-panel.test.ts @@ -37,6 +37,7 @@ describe("ChunksPanel", () => { unobserve() {} disconnect() {} }; + Element.prototype.scrollIntoView = vi.fn(); }); afterEach(async () => { @@ -71,6 +72,201 @@ describe("ChunksPanel", () => { expect(screen.getByText(/Showing all parsed chunks from/)).toBeTruthy(); }); + it("renders page chunk assets inside the normal chunk list", async () => { + mockVisibleVirtualViewport(); + + render( + React.createElement(C, { + chunks: [ + { + chunkId: "page_4", + type: "page", + content: "Page 4 summary", + sourceTitle: "report.pdf", + pageNums: [4], + pageAssets: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-000004.png", + contentType: "image/png", + width: 1200, + height: 1600, + }, + ], + }, + { + chunkId: "table_1", + type: "table", + content: "
Budget
", + sourceTitle: "report.pdf", + }, + ], + selectedSource: "report.pdf", + selectedSourceView: { + id: "source_1", + title: "report.pdf", + mimeType: "application/pdf", + status: "ready", + documentPresentation: { kind: "page-assets", pageCount: 4 }, + }, + }), + ); + + expect(screen.getByRole("heading", { name: "Parsed Chunks" })).toBeTruthy(); + expect( + await screen.findByRole("img", { name: "Page 4" }), + ).toBeTruthy(); + expect(screen.getByText("image/png")).toBeTruthy(); + expect(screen.getByText("Budget")).toBeTruthy(); + expect( + screen.queryByRole("button", { name: "Original" }), + ).toBeNull(); + expect(screen.queryByRole("button", { name: "Tree" })).toBeNull(); + expect( + screen.queryByRole("button", { name: /original file/i }), + ).toBeNull(); + }); + + it("renders only the first page-asset chunk for each page number", async () => { + mockVisibleVirtualViewport(); + + render( + React.createElement(C, { + chunks: [ + { + chunkId: "page_4_first", + type: "page", + content: "First page 4 summary", + sourceTitle: "report.pdf", + pageNums: [4], + pageAssets: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-000004-a.png", + contentType: "image/png", + }, + ], + }, + { + chunkId: "page_4_duplicate", + type: "page", + content: "Duplicate page 4 summary", + sourceTitle: "report.pdf", + pageNums: [4], + pageAssets: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-000004-b.png", + contentType: "image/png", + }, + ], + }, + { + chunkId: "page_5", + type: "page", + content: "Page 5 summary", + sourceTitle: "report.pdf", + pageNums: [5], + pageAssets: [ + { + pageNumber: 5, + assetUrl: "https://assets.example/page-000005.png", + contentType: "image/png", + }, + ], + }, + ], + selectedSource: "report.pdf", + selectedSourceView: { + id: "source_1", + title: "report.pdf", + mimeType: "application/pdf", + status: "ready", + documentPresentation: { kind: "page-assets", pageCount: 5 }, + }, + }), + ); + + expect(await screen.findAllByRole("img", { name: "Page 4" })) + .toHaveLength(1); + expect(screen.getByRole("img", { name: "Page 5" })).toBeTruthy(); + expect(screen.queryByTestId("chunk-card-shell-page_4_duplicate")) + .toBeNull(); + }); + + it("renders page chunks normally when no page assets exist", async () => { + mockVisibleVirtualViewport(); + + render( + React.createElement(C, { + chunks: [ + { + chunkId: "page_4", + type: "page", + content: "Page 4 summary", + readableContent: "Page 4 summary", + sourceTitle: "report.pdf", + pageNums: [4], + }, + ], + selectedSource: "report.pdf", + selectedSourceView: { + id: "source_1", + title: "report.pdf", + mimeType: "application/pdf", + status: "ready", + documentPresentation: { kind: "page-assets", pageCount: 4 }, + }, + }), + ); + + expect(await screen.findByText("Page 4 summary")).toBeTruthy(); + expect(screen.queryByRole("img", { name: "Page 4" })).toBeNull(); + }); + + it("shows a page-level placeholder when a page asset image fails to load", async () => { + mockVisibleVirtualViewport(); + + render( + React.createElement(C, { + chunks: [ + { + chunkId: "page_4", + type: "page", + content: "Page 4 summary", + sourceTitle: "report.pdf", + pageNums: [4], + pageAssets: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-000004.png", + contentType: "image/png", + width: 1200, + height: 1600, + }, + ], + }, + ], + selectedSource: "report.pdf", + selectedSourceView: { + id: "source_1", + title: "report.pdf", + mimeType: "application/pdf", + status: "ready", + documentPresentation: { kind: "page-assets", pageCount: 4 }, + }, + }), + ); + + const pageImage = await screen.findByRole("img", { name: "Page 4" }); + fireEvent.error(pageImage); + + expect( + screen.getByTestId("page-asset-image-unavailable-4"), + ).toBeTruthy(); + expect(screen.getByText("Page image unavailable.")).toBeTruthy(); + }); + it("defaults parsed chunks into a section tree view", () => { render( React.createElement(C, { diff --git a/src/components/chunks-panel.tsx b/src/components/chunks-panel.tsx index 78cd170..f49df7e 100644 --- a/src/components/chunks-panel.tsx +++ b/src/components/chunks-panel.tsx @@ -42,16 +42,22 @@ import { MAX_UPLOAD_MB } from "@/domains/sources/validation"; import { useSourceOriginalPreviewWarmup } from "@/components/source-original-preview-warmup"; import { sourceOriginalPreviewModel } from "@/components/source-original-preview-model"; import type { ParsedChunkView } from "@/domains/chunks/types"; -import type { SourceOriginalFileView, SourceView } from "@/domains/sources/types"; +import type { + SourceOriginalFileView, + SourceView, +} from "@/domains/sources/types"; import type { AnalyticsContext } from "@/lib/posthog"; import { cn } from "@/lib/utils"; export type ChunksPanelProps = { chunks: ParsedChunkView[]; selectedSource?: string | null; + selectedSourceView?: SourceView | null; selectedSourceFile?: SourceOriginalFileView | null; focusedChunkId?: string | null; focusedChunkRequestId?: number; + focusedPageNumber?: number | null; + focusedPageRequestId?: number; citationListViewRequestId?: number; isLoading?: boolean; isLoadingMore?: boolean; @@ -77,6 +83,7 @@ type ChunkDisplayModeState = { export function ChunksPanel({ chunks = [], selectedSource = null, + selectedSourceView = null, selectedSourceFile = null, focusedChunkId = null, focusedChunkRequestId = 0, @@ -93,8 +100,23 @@ export function ChunksPanel({ analyticsContext, sourceCountSnapshot = 0, }: Partial = {}) { + const hasPageAssetChunks = chunks.some( + (chunk) => (chunk.pageAssets?.length ?? 0) > 0, + ); + const isPageAssetSource = + selectedSourceView?.documentPresentation?.kind === "page-assets" || + hasPageAssetChunks; + const displayChunks = useMemo( + () => + isPageAssetSource + ? chunksPanelState.getPageAssetChunksWithoutDuplicatePages(chunks) + : chunks, + [chunks, isPageAssetSource], + ); + const effectiveVisibleView = isPageAssetSource ? "parsed" : undefined; const originalPreviewCacheKey = selectedSourceFile?.url ?? null; const isOriginalPreviewAvailable = + !isPageAssetSource && sourceOriginalPreviewModel.canPreviewOriginalFile( selectedSource, selectedSourceFile, @@ -129,7 +151,7 @@ export function ChunksPanel({ visibleChunks, visibleView, } = useChunksPanelWorkflow({ - chunks, + chunks: displayChunks, selectedSource, selectedSourceFile, focusedChunkId, @@ -139,6 +161,7 @@ export function ChunksPanel({ isLoadingMore, onLoadMore, }); + const activeVisibleView = effectiveVisibleView ?? visibleView; useSourceOriginalPreviewWarmup({ sourceTitle: selectedSource, @@ -240,14 +263,21 @@ export function ChunksPanel({ citationListViewRequestId ? "list" : chunkDisplayModeState.mode; - const headerTitle = focusedChunkId ? "Referenced Chunks" : "Parsed Chunks"; + const headerTitle = activeVisibleView === "original" + ? "Original File" + : focusedChunkId + ? "Referenced Chunks" + : "Parsed Chunks"; const shouldMountOriginalPreview = - visibleView === "original" || - (originalPreviewCacheKey !== null && - mountedOriginalPreviewKey === originalPreviewCacheKey); + !isPageAssetSource && + (activeVisibleView === "original" || + (originalPreviewCacheKey !== null && + mountedOriginalPreviewKey === originalPreviewCacheKey)); const isTreeModeVisible = - visibleView === "parsed" && chunkDisplayMode === "tree"; - const headerSubtitle = visibleView === "original" ? ( + !isPageAssetSource && + activeVisibleView === "parsed" && + chunkDisplayMode === "tree"; + const headerSubtitle = activeVisibleView === "original" ? ( selectedSource ? ( <> Showing the original file for{" "} @@ -299,14 +329,16 @@ export function ChunksPanel({

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

{headerSubtitle}

- {visibleView === "parsed" && chunks.length > 0 ? ( + {!isPageAssetSource && + activeVisibleView === "parsed" && + displayChunks.length > 0 ? (
) : null} - {hasOriginalView ? ( + {!isPageAssetSource && hasOriginalView ? (
@@ -350,7 +382,7 @@ export function ChunksPanel({
- + {isLoading ? ( - ) : chunks.length === 0 ? ( + ) : displayChunks.length === 0 && processingMessage ? ( + + ) : displayChunks.length === 0 ? ( selectedSource ? ( ) : ( @@ -377,7 +411,7 @@ export function ChunksPanel({ ) : isTreeModeVisible ? ( - {isTreeModeVisible ? ( + {!isPageAssetSource && isTreeModeVisible ? (
{shouldMountOriginalPreview ? ( - + + +

+ {message} +

+
+ ); +} + function truncateTreeLabel(value: string): string { const maxLength = 42; if (value.length <= maxLength) return value; diff --git a/src/components/parsed-chunk-card.test.ts b/src/components/parsed-chunk-card.test.ts index 946cffb..8bf5a42 100644 --- a/src/components/parsed-chunk-card.test.ts +++ b/src/components/parsed-chunk-card.test.ts @@ -1,6 +1,6 @@ // @vitest-environment jsdom import React from "react"; -import { cleanup, render, screen } from "@testing-library/react"; +import { cleanup, fireEvent, render, screen } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { afterEach, describe, expect, it, vi } from "vitest"; @@ -81,6 +81,105 @@ describe("ParsedChunkCard", () => { ).toBeTruthy(); }); + it("renders page citation assets instead of page summary content", () => { + render( + React.createElement(ParsedChunkCard, { + chunk: { + chunkId: "page_4", + type: "page", + content: "The summary should not be primary when an image exists.", + readableContent: "The summary should not be primary when an image exists.", + sourceTitle: "manual.pdf", + pageNums: [4], + pageAssets: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-4.png", + contentType: "image/png", + width: 1200, + height: 1600, + }, + ], + }, + isFocused: false, + isOriginalPreviewAvailable: true, + onChunkClick: vi.fn(), + onReferenceClick: vi.fn(), + }), + ); + + const pageImage = screen.getByRole("img", { name: "Page 4" }); + + expect(pageImage.getAttribute("src")).toBe( + "https://assets.example/page-4.png", + ); + expect(screen.getByText("image/png")).toBeTruthy(); + expect(screen.queryByText(/summary should not be primary/i)).toBeNull(); + expect( + screen.queryByRole("button", { name: /original file/i }), + ).toBeNull(); + + fireEvent.error(pageImage); + expect(screen.getByTestId("page-asset-image-unavailable-4")).toBeTruthy(); + }); + + it("routes Notebook Blob page assets through the inline image endpoint", () => { + const assetUrl = + "https://store.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_1/rev_1/page_citation_assets/page-4.png"; + + render( + React.createElement(ParsedChunkCard, { + chunk: { + chunkId: "page_4", + type: "page", + content: "The summary should not be primary when an image exists.", + readableContent: "The summary should not be primary when an image exists.", + sourceTitle: "manual.pdf", + pageNums: [4], + pageAssets: [ + { + pageNumber: 4, + assetUrl, + contentType: "image/png", + }, + ], + }, + isFocused: false, + onReferenceClick: vi.fn(), + }), + ); + + const pageImage = screen.getByRole("img", { name: "Page 4" }); + expect(pageImage.getAttribute("src")).toBe( + `/api/parsed-assets/inline?url=${encodeURIComponent(assetUrl)}`, + ); + }); + + it("routes Notebook Blob image chunks through the inline image endpoint", () => { + const assetUrl = + "https://store.public.blob.vercel-storage.com/workspaces/workspace_1/sources/source_1/parsed-result/images/image-1.jpg"; + + render( + React.createElement(ParsedChunkCard, { + chunk: { + chunkId: "image_1", + type: "image", + content: "", + sourceTitle: "manual.pdf", + assetUrl, + summary: "A scanned diagram", + }, + isFocused: false, + onReferenceClick: vi.fn(), + }), + ); + + const image = screen.getByRole("img", { name: "A scanned diagram" }); + expect(image.getAttribute("src")).toBe( + `/api/parsed-assets/inline?url=${encodeURIComponent(assetUrl)}`, + ); + }); + it("routes resolved artifact reference clicks to the target chunk", async () => { const user = userEvent.setup(); const onReferenceClick = vi.fn(); diff --git a/src/components/parsed-chunk-card.tsx b/src/components/parsed-chunk-card.tsx index 1ba233d..a8ddd6d 100644 --- a/src/components/parsed-chunk-card.tsx +++ b/src/components/parsed-chunk-card.tsx @@ -1,7 +1,15 @@ "use client"; -import { useMemo, type MouseEvent, type ReactNode } from "react"; -import { FileSearch, FileText, ImageIcon, Table2, Tags, TextQuote } from "lucide-react"; +import { useMemo, useState, type MouseEvent, type ReactNode } from "react"; +import { + FileSearch, + FileText, + ImageIcon, + ImageOff, + Table2, + Tags, + TextQuote, +} from "lucide-react"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; @@ -185,7 +193,9 @@ function ChunkSourcePanel({ ) : null}
- {onChunkClick && firstPageNumber !== null ? ( + {onChunkClick && + firstPageNumber !== null && + !hasPageCitationAssets(chunk) ? ( void; }): ReactNode { + const pageAssets = chunk.pageAssets ?? []; + return ( - -

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

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

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

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

+ Page image unavailable. +

+
+ ); +} + function ImageChunkCard({ chunk, isFocused, @@ -417,6 +510,9 @@ function ImageChunkCard({ readonly sourceOriginalFile: SourceOriginalFileView | null; }): ReactNode { const imageAssetUrl = getImageChunkAssetUrl(chunk, sourceOriginalFile); + const inlineImageAssetUrl = imageAssetUrl + ? getInlineImageAssetUrl(imageAssetUrl) + : null; return ( - {imageAssetUrl ? ( + {inlineImageAssetUrl ? (
{/* eslint-disable-next-line @next/next/no-img-element -- Parsed artifact dimensions are not known before render. */} {chunk.summary @@ -467,6 +563,33 @@ function getImageChunkAssetUrl( return sourceOriginalFile.url; } +function getInlineImageAssetUrl(assetUrl: string): string { + if (!isNotebookBlobAssetUrl(assetUrl)) return assetUrl; + + return `/api/parsed-assets/inline?url=${encodeURIComponent(assetUrl)}`; +} + +function isNotebookBlobAssetUrl(assetUrl: string): boolean { + try { + const url = new URL(assetUrl); + if (!url.hostname.toLowerCase().endsWith(".blob.vercel-storage.com")) { + return false; + } + + const pathname = decodeURIComponent(url.pathname).toLowerCase(); + return ( + pathname.includes("/parsed-documents/") || + pathname.includes("/parsed-result/") + ); + } catch { + return false; + } +} + +function hasPageCitationAssets(chunk: ParsedChunkView): boolean { + return (chunk.pageAssets?.length ?? 0) > 0; +} + function renderTextChunkContent( chunk: ParsedChunkView, onReferenceClick: (chunkId: string) => void, diff --git a/src/components/source-row.test.ts b/src/components/source-row.test.ts index 427e09a..97c4bf9 100644 --- a/src/components/source-row.test.ts +++ b/src/components/source-row.test.ts @@ -87,6 +87,27 @@ describe("SourceRow", () => { ); }); + it("labels page-asset documents with page counts", () => { + render( + React.createElement(SourceRow, { + isArchiving: false, + isSelected: false, + onSelect: vi.fn(), + source: { + id: "source_1", + mimeType: "application/pdf", + title: "scan.pdf", + status: "ready", + chunkCount: 4, + documentPresentation: { kind: "page-assets", pageCount: 4 }, + }, + }), + ); + + expect(screen.getByText("Processed · 4 pages")).toBeTruthy(); + expect(screen.queryByText("Processed · 4 chunks")).toBeNull(); + }); + it("shows source archive loading locally", () => { render( React.createElement(SourceRow, { diff --git a/src/components/source-row.tsx b/src/components/source-row.tsx index 9fb87d3..473453f 100644 --- a/src/components/source-row.tsx +++ b/src/components/source-row.tsx @@ -197,6 +197,9 @@ function getReadySourceLabel(source: SourceView): string { function getReadySourceStatusText(source: SourceView): string { if (source.kind === "remote") return getReadySourceLabel(source); + if (source.documentPresentation?.kind === "page-assets") { + return `${getReadySourceLabel(source)} · ${source.documentPresentation.pageCount} pages`; + } if (typeof source.chunkCount !== "number") return getReadySourceLabel(source); return `${getReadySourceLabel(source)} · ${source.chunkCount} chunks`; } diff --git a/src/components/sources-panel.tsx b/src/components/sources-panel.tsx index 70a0285..d1950ef 100644 --- a/src/components/sources-panel.tsx +++ b/src/components/sources-panel.tsx @@ -351,6 +351,7 @@ function getSelectedSourcePage( } function getChunkTreeHref(source: SourceView): string | undefined { + if (source.documentPresentation?.kind === "page-assets") return undefined; return source.documentId ? `/inspect/${encodeURIComponent(source.documentId)}/chunks` : undefined; diff --git a/src/components/workspace-citation-focus.test.ts b/src/components/workspace-citation-focus.test.ts index baa5a1a..21fce5a 100644 --- a/src/components/workspace-citation-focus.test.ts +++ b/src/components/workspace-citation-focus.test.ts @@ -169,6 +169,63 @@ describe("useWorkspaceCitationFocus", () => { expect(result.current.pendingCitationId).toBeNull(); }); + it("focuses page-asset citations through the loaded page chunk", async () => { + const pageChunk: ParsedChunkView = { + chunkId: "page_4", + documentId: "document_1", + type: "page", + content: "Page 4 summary", + sourceTitle: "Contract.pdf", + pageAssets: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-000004.png", + contentType: "image/png", + }, + ], + }; + const fetchChunks = vi.fn(async () => [pageChunk]); + const selectSource = vi.fn(); + const pageAssetSource: SourceView = { + ...readySource, + documentPresentation: { kind: "page-assets", pageCount: 8 }, + }; + const pageCitation: ChatCitationView = { + chunkType: "page", + score: 0.9, + pageCitationAssetUrl: "https://assets.example/page-000004.png", + pageCitationPageNumber: 4, + source: { + documentId: "document_1", + sourceFileName: "Contract.pdf", + sectionPath: "Page 4", + }, + }; + + const { result } = renderHook(() => + useWorkspaceCitationFocus({ + fetchChunks, + onSelectSource: selectSource, + selectedSourceId: null, + sources: [pageAssetSource], + }), + { wrapper: createSWRWrapper }, + ); + + await act(async () => { + await result.current.handleCitationClick(pageCitation, "message_1:0"); + }); + + expect(fetchChunks).toHaveBeenCalledWith("source_1"); + expect(selectSource).toHaveBeenCalledWith("source_1"); + expect(result.current.focusedChunk.chunkId).toBe("page_4"); + expect(result.current.focusedPage).toEqual({ + pageNumber: null, + requestId: 0, + }); + expect(result.current.citationListViewRequestId).toBe(1); + }); + it("reuses cached chunks for a different source without refetching", async () => { const fetchChunks = vi.fn(async () => [prefetchedChunk]); const selectSource = vi.fn(); diff --git a/src/components/workspace-citation-focus.ts b/src/components/workspace-citation-focus.ts index 357cd4e..70b0acf 100644 --- a/src/components/workspace-citation-focus.ts +++ b/src/components/workspace-citation-focus.ts @@ -13,6 +13,11 @@ type FocusedChunkState = { readonly requestId: number } +type FocusedPageState = { + readonly pageNumber: number | null + readonly requestId: number +} + type PrefetchedChunksBySourceId = Readonly> type PrefetchedChunksUpdater = ( current: PrefetchedChunksBySourceId, @@ -21,7 +26,6 @@ type PrefetchedChunksUpdater = ( type WorkspaceCitationFocusInput = { readonly fetchChunks: (sourceId: string) => Promise readonly initialPrefetchedChunksBySourceId?: PrefetchedChunksBySourceId - readonly onRemoteSourceChunksLoaded?: (sourceId: string) => void readonly onSelectSource: (sourceId: string | null) => void readonly selectedSourceId: string | null readonly sources: readonly SourceView[] @@ -30,6 +34,7 @@ type WorkspaceCitationFocusInput = { type WorkspaceCitationFocus = { readonly citationListViewRequestId: number readonly focusedChunk: FocusedChunkState + readonly focusedPage: FocusedPageState readonly handleCitationClick: ( citation: ChatCitationView, citationId: string, @@ -52,7 +57,6 @@ type WorkspaceCitationFocus = { export function useWorkspaceCitationFocus({ fetchChunks, initialPrefetchedChunksBySourceId = {}, - onRemoteSourceChunksLoaded, onSelectSource, selectedSourceId, sources, @@ -61,6 +65,10 @@ export function useWorkspaceCitationFocus({ chunkId: null, requestId: 0, }) + const [focusedPage, setFocusedPage] = useState({ + pageNumber: null, + requestId: 0, + }) const [pendingCitationId, setPendingCitationId] = useState( null, ) @@ -90,7 +98,6 @@ export function useWorkspaceCitationFocus({ selectedSourceId, sources, prefetchedChunksBySourceId, - onRemoteSourceChunksLoaded, }) const requestChunkFocus = useCallback( @@ -102,6 +109,12 @@ export function useWorkspaceCitationFocus({ }, [], ) + const requestPageFocus = useCallback((pageNumber: number | null): void => { + setFocusedPage((current) => ({ + pageNumber, + requestId: current.requestId + 1, + })) + }, []) const updatePrefetchedChunksBySourceId = useCallback( (updater: PrefetchedChunksUpdater): void => { @@ -122,10 +135,12 @@ export function useWorkspaceCitationFocus({ ) } requestChunkFocus(null) + requestPageFocus(null) }, [ onSelectSource, requestChunkFocus, + requestPageFocus, selectedSourceId, updatePrefetchedChunksBySourceId, ], @@ -192,6 +207,7 @@ export function useWorkspaceCitationFocus({ citation, ) if (!source) return + setCitationListViewRequestId((current) => current + 1) const loadedChunkId = workspaceCitationState.getLoadedCitationChunkId({ @@ -271,6 +287,7 @@ export function useWorkspaceCitationFocus({ return { citationListViewRequestId, focusedChunk, + focusedPage, handleCitationClick, handleLoadAllChunks, handleLoadMoreChunks, diff --git a/src/components/workspace-citation-state.test.ts b/src/components/workspace-citation-state.test.ts index ce5851e..cce0e6b 100644 --- a/src/components/workspace-citation-state.test.ts +++ b/src/components/workspace-citation-state.test.ts @@ -91,6 +91,44 @@ describe("workspaceCitationState", () => { ).toBeNull() }) + it("focuses loaded page chunks from page citation metadata", () => { + const citation: ChatCitationView = { + chunkType: "page", + score: 0.93, + pageCitationPageNumber: 4, + pageCitationAssetUrl: "https://assets.example/page-4.png", + source: { + documentId: "document_1", + sectionPath: "Page 4", + }, + } + + expect( + workspaceCitationState.getLoadedCitationChunkId({ + citation, + selectedSourceId: "source_1", + sourceId: "source_1", + selectedChunks: [ + { + chunkId: "page_4", + documentId: "document_1", + type: "page", + content: "Page 4 summary", + sourceTitle: "Contract.pdf", + pageAssets: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-4.png", + contentType: "image/png", + }, + ], + }, + ], + hasMoreSelectedChunks: true, + }), + ).toBe("page_4") + }) + describe("hasExactCitationTargetHint", () => { it("returns true when the citation has content text", () => { const citation: ChatCitationView = { @@ -120,6 +158,22 @@ describe("workspaceCitationState", () => { ).toBe(true) }) + it("returns true when the citation has page asset metadata", () => { + const citation: ChatCitationView = { + chunkType: "page", + score: 0.5, + pageCitationPageNumber: 4, + source: { + documentId: "document_1", + sectionPath: "Root", + }, + } + + expect( + workspaceCitationState.hasExactCitationTargetHint(citation), + ).toBe(true) + }) + it("returns false for source-only citations with no useful target hint", () => { const citation: ChatCitationView = { chunkType: "text", diff --git a/src/components/workspace-citation-state.ts b/src/components/workspace-citation-state.ts index d8331af..0353409 100644 --- a/src/components/workspace-citation-state.ts +++ b/src/components/workspace-citation-state.ts @@ -27,6 +27,7 @@ type WorkspaceCitationStateModule = { readonly hasExactCitationTargetHint: ( citation: ChatCitationView, ) => boolean + readonly getCitationPageNumber: (citation: ChatCitationView) => number | null readonly upsertPrefetchedChunks: ( current: PrefetchedChunksBySourceId, sourceId: string, @@ -61,7 +62,10 @@ function getLoadedCitationChunkId( ? resolveCitationChunkByContent(input.citation, input.selectedChunks) : resolveCitationChunk(input.citation, input.selectedChunks) - return focusedChunk?.chunkId ?? null + return ( + focusedChunk?.chunkId ?? + resolveCitationPageChunkId(input.citation, input.selectedChunks) + ) } function hasExactCitationTargetHint(citation: ChatCitationView): boolean { @@ -69,6 +73,13 @@ function hasExactCitationTargetHint(citation: ChatCitationView): boolean { return true } + if ( + getCitationPageNumber(citation) !== null || + typeof citation.pageCitationAssetUrl === "string" + ) { + return true + } + const sectionPath = citation.source.sectionPath if (typeof sectionPath !== "string") return false @@ -79,6 +90,62 @@ function hasExactCitationTargetHint(citation: ChatCitationView): boolean { return true } +function getCitationPageNumber(citation: ChatCitationView): number | null { + if ( + typeof citation.pageCitationPageNumber === "number" && + Number.isSafeInteger(citation.pageCitationPageNumber) && + citation.pageCitationPageNumber > 0 + ) { + return citation.pageCitationPageNumber + } + + const sectionPath = citation.source.sectionPath + if (typeof sectionPath !== "string") return null + + const match = /\bpage\s+(\d+)\b/i.exec(sectionPath) + if (!match) return null + + const pageNumber = Number.parseInt(match[1]!, 10) + return Number.isSafeInteger(pageNumber) && pageNumber > 0 ? pageNumber : null +} + +function resolveCitationPageChunkId( + citation: ChatCitationView, + chunks: readonly ParsedChunkView[], +): string | null { + const pageNumber = getCitationPageNumber(citation) + if (pageNumber === null) return null + + const documentChunks = citation.source.documentId + ? chunks.filter((chunk) => chunk.documentId === citation.source.documentId) + : chunks + const matches = documentChunks.filter((chunk) => + isChunkForPageNumber(chunk, pageNumber), + ) + if (matches.length === 0) return null + + const pageAssetChunk = matches.find( + (chunk) => + chunk.type === "page" && + (chunk.pageAssets ?? []).some( + (pageAsset) => pageAsset.pageNumber === pageNumber, + ), + ) + + return pageAssetChunk?.chunkId ?? matches[0]?.chunkId ?? null +} + +function isChunkForPageNumber( + chunk: ParsedChunkView, + pageNumber: number, +): boolean { + return ( + (chunk.pageAssets ?? []).some( + (pageAsset) => pageAsset.pageNumber === pageNumber, + ) || (chunk.pageNums ?? []).includes(pageNumber) + ) +} + function upsertPrefetchedChunks( current: PrefetchedChunksBySourceId, sourceId: string, @@ -120,6 +187,7 @@ function removePrefetchedChunks( export const workspaceCitationState: WorkspaceCitationStateModule = { findCitationSource, getLoadedCitationChunkId, + getCitationPageNumber, hasExactCitationTargetHint, upsertPrefetchedChunks, removePrefetchedChunks, diff --git a/src/components/workspace-selected-chunks.test.ts b/src/components/workspace-selected-chunks.test.ts index 45db25a..59cf03a 100644 --- a/src/components/workspace-selected-chunks.test.ts +++ b/src/components/workspace-selected-chunks.test.ts @@ -39,7 +39,7 @@ describe("useWorkspaceSelectedChunks", () => { }); }); - it("returns prefetched chunks while checking the visible page for media", async () => { + it("returns prefetched chunks while loading the visible chunk page", async () => { const { result } = renderHook( () => useWorkspaceSelectedChunks({ @@ -161,8 +161,140 @@ describe("useWorkspaceSelectedChunks", () => { expect(result.current.selectedChunks).toEqual([]); }); - it("requests a source refresh after loading an unlocalized remote source", async () => { - const onRemoteSourceChunksLoaded = vi.fn(); + it("surfaces unavailable chunk messages without a loading state", async () => { + fetchChunkPageMock.mockResolvedValue({ + chunks: [], + isUnavailable: true, + message: + "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", + pagination: { + page: 1, + pageSize: 50, + total: 0, + totalPages: 0, + }, + }); + + const { result } = renderHook( + () => + useWorkspaceSelectedChunks({ + selectedSourceId: "source_1", + sources: [readySource], + prefetchedChunksBySourceId: {}, + }), + { wrapper: createSWRWrapper }, + ); + + await waitFor(() => + expect(result.current.selectedChunksMessage).toBe( + "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", + ), + ); + expect(result.current.isSelectedChunksLoading).toBe(false); + expect(result.current.selectedChunks).toEqual([]); + }); + + it("loads chunk pages for page-asset sources", async () => { + const pageAssetSource: SourceView = { + ...readySource, + documentPresentation: { kind: "page-assets", pageCount: 4 }, + }; + fetchChunkPageMock.mockResolvedValue({ + chunks: [ + { + chunkId: "page_1", + type: "page", + content: "Page summary", + sourceTitle: "lecture.pdf", + pageAssets: [ + { + pageNumber: 1, + assetUrl: "https://blob.example/page-1.png", + contentType: "image/png", + }, + ], + }, + ], + pagination: { + page: 1, + pageSize: 50, + total: 1, + totalPages: 1, + }, + }); + + const { result } = renderHook( + () => + useWorkspaceSelectedChunks({ + selectedSourceId: "source_1", + sources: [pageAssetSource], + prefetchedChunksBySourceId: {}, + }), + { wrapper: createSWRWrapper }, + ); + + await waitFor(() => + expect(result.current.selectedChunks.map((chunk) => chunk.chunkId)).toEqual([ + "page_1", + ]), + ); + expect(result.current.isSelectedChunksLoading).toBe(false); + expect(fetchChunkPageMock).toHaveBeenCalledWith("source_1", 1); + }); + + it("detects page assets from selected source chunks without a second request", async () => { + fetchChunkPageMock.mockResolvedValue({ + chunks: [ + { + chunkId: "page_1", + type: "page", + content: "Page summary", + sourceTitle: "lecture.pdf", + pageAssets: [ + { + pageNumber: 1, + assetUrl: "https://blob.example/page-1.png", + contentType: "image/png", + }, + { + pageNumber: 20, + assetUrl: "https://blob.example/page-20.png", + contentType: "image/png", + }, + ], + }, + ], + pagination: { + page: 1, + pageSize: 50, + total: 1, + totalPages: 1, + }, + }); + + const { result } = renderHook( + () => + useWorkspaceSelectedChunks({ + selectedSourceId: "source_1", + sources: [readySource], + prefetchedChunksBySourceId: {}, + }), + { wrapper: createSWRWrapper }, + ); + + await waitFor(() => + expect(result.current.selectedSource?.documentPresentation).toEqual({ + kind: "page-assets", + pageCount: 20, + }), + ); + expect(result.current.selectedChunks.map((chunk) => chunk.chunkId)).toEqual([ + "page_1", + ]); + expect(fetchChunkPageMock).toHaveBeenCalledWith("source_1", 1); + }); + + it("loads an unlocalized remote source without requesting a source refresh", async () => { const remoteSource: SourceView = { ...readySource, id: "knowhere-doc:default:doc_remote", @@ -187,30 +319,24 @@ describe("useWorkspaceSelectedChunks", () => { }, }); - const { rerender } = renderHook( - (input: { - readonly onRemoteSourceChunksLoaded: (sourceId: string) => void; - }) => + const { result } = renderHook( + () => useWorkspaceSelectedChunks({ selectedSourceId: "knowhere-doc:default:doc_remote", sources: [remoteSource], prefetchedChunksBySourceId: {}, - onRemoteSourceChunksLoaded: input.onRemoteSourceChunksLoaded, }), - { - initialProps: { onRemoteSourceChunksLoaded }, - wrapper: createSWRWrapper, - }, + { wrapper: createSWRWrapper }, ); await waitFor(() => - expect(onRemoteSourceChunksLoaded).toHaveBeenCalledWith( - "knowhere-doc:default:doc_remote", - ), + expect(result.current.selectedChunks.map((chunk) => chunk.chunkId)).toEqual([ + "chunk_1", + ]), + ); + expect(result.current.selectedSource?.id).toBe( + "knowhere-doc:default:doc_remote", ); - - rerender({ onRemoteSourceChunksLoaded }); - expect(onRemoteSourceChunksLoaded).toHaveBeenCalledTimes(1); }); it("returns an empty chunk list when no source is selected", () => { diff --git a/src/components/workspace-selected-chunks.ts b/src/components/workspace-selected-chunks.ts index f00ac68..260a6f5 100644 --- a/src/components/workspace-selected-chunks.ts +++ b/src/components/workspace-selected-chunks.ts @@ -1,6 +1,6 @@ "use client" -import { useEffect, useMemo, useRef } from "react" +import { useMemo } from "react" import useSWRInfinite from "swr/infinite" import { workspaceClient } from "@/domains/workspace/client" @@ -17,7 +17,6 @@ type WorkspaceSelectedChunksInput = { readonly selectedSourceId: string | null readonly sources: readonly SourceView[] readonly prefetchedChunksBySourceId: Readonly> - readonly onRemoteSourceChunksLoaded?: (sourceId: string) => void } type WorkspaceSelectedChunks = { @@ -34,10 +33,11 @@ export function useWorkspaceSelectedChunks({ selectedSourceId, sources, prefetchedChunksBySourceId, - onRemoteSourceChunksLoaded, }: WorkspaceSelectedChunksInput): WorkspaceSelectedChunks { - const selectedSource = sources.find((source) => source.id === selectedSourceId) - const remoteSourceRefreshRequestedIdsRef = useRef>(new Set()) + const rawSelectedSource = sources.find( + (source) => source.id === selectedSourceId, + ) + const selectedSource = rawSelectedSource const prefetchedSelectedChunks = selectedSourceId ? prefetchedChunksBySourceId[selectedSourceId] : undefined @@ -66,6 +66,8 @@ export function useWorkspaceSelectedChunks({ }, ) const selectedChunksMessage = getSelectedChunksMessage(selectedChunkPages) + const hasProcessingSelectedChunkPage = + hasProcessingChunkPage(selectedChunkPages) const pagedSelectedChunks = useMemo( () => resolveChunkConnectionTargets( @@ -83,9 +85,17 @@ export function useWorkspaceSelectedChunks({ : undefined, [pagedSelectedChunks, prefetchedSelectedChunks], ) - const selectedChunks = selectedSourceId - ? (resolvedPrefetchedChunks ?? pagedSelectedChunks) - : [] + const selectedChunks = useMemo( + () => + selectedSourceId + ? (resolvedPrefetchedChunks ?? pagedSelectedChunks) + : [], + [pagedSelectedChunks, resolvedPrefetchedChunks, selectedSourceId], + ) + const resolvedSelectedSource = useMemo( + () => getResolvedSelectedSource(selectedSource, selectedChunks), + [selectedChunks, selectedSource], + ) const hasMoreSelectedChunks = !prefetchedSelectedChunks && workspaceClientCache.hasMoreChunkPages(selectedChunkPages) @@ -97,22 +107,12 @@ export function useWorkspaceSelectedChunks({ typeof selectedChunkPages[selectedChunkPageCount - 1] === "undefined", ) const isSelectedChunksLoading = - Boolean(selectedChunksMessage) || + hasProcessingSelectedChunkPage || (selectedChunkSourceId !== null && !prefetchedSelectedChunks && !selectedChunkPages && isChunksLoading) - useEffect(() => { - const sourceId = selectedSource?.id - if (!sourceId || selectedSource.kind !== "remote") return - if (!selectedChunkPages || selectedChunkPages.length === 0) return - if (remoteSourceRefreshRequestedIdsRef.current.has(sourceId)) return - - remoteSourceRefreshRequestedIdsRef.current.add(sourceId) - onRemoteSourceChunksLoaded?.(sourceId) - }, [onRemoteSourceChunksLoaded, selectedChunkPages, selectedSource]) - function handleLoadMoreChunks(): void { if (!hasMoreSelectedChunks || isSelectedChunksLoadingMore) return void setSelectedChunkPageCount(selectedChunkPageCount + 1) @@ -125,7 +125,7 @@ export function useWorkspaceSelectedChunks({ isSelectedChunksLoadingMore, selectedChunksMessage, selectedChunks, - selectedSource, + selectedSource: resolvedSelectedSource, } } @@ -137,6 +137,36 @@ function fetchChunksByKey([ return workspaceClient.fetchChunkPage(sourceId, page) } +function getResolvedSelectedSource( + source: SourceView | undefined, + chunks: readonly ParsedChunkView[], +): SourceView | undefined { + if (!source) return undefined + if (source.documentPresentation?.kind === "page-assets") return source + + const pageCount = getLoadedPageAssetCount(chunks) + if (!pageCount) return source + + return { + ...source, + documentPresentation: { + kind: "page-assets", + pageCount, + }, + } +} + +function getLoadedPageAssetCount( + chunks: readonly ParsedChunkView[], +): number | null { + const pageNumbers = chunks.flatMap((chunk) => + (chunk.pageAssets ?? []).map((pageAsset) => pageAsset.pageNumber), + ) + if (pageNumbers.length === 0) return null + + return Math.max(...pageNumbers) +} + function hasProcessingChunkPage( pages: readonly SourceChunksResponse[] | undefined, ): boolean { @@ -148,7 +178,8 @@ function getSelectedChunksMessage( ): string | null { const page = pages?.find( (candidate) => - candidate.isProcessing && typeof candidate.message === "string", + (candidate.isProcessing || candidate.isUnavailable) && + typeof candidate.message === "string", ) return page?.message ?? null } diff --git a/src/components/workspace-shell-layout.tsx b/src/components/workspace-shell-layout.tsx index 3b9e997..0a71d90 100644 --- a/src/components/workspace-shell-layout.tsx +++ b/src/components/workspace-shell-layout.tsx @@ -40,6 +40,11 @@ type FocusedChunkState = { readonly requestId: number } +type FocusedPageState = { + readonly pageNumber: number | null + readonly requestId: number +} + type WorkspaceShellUser = { readonly id: string readonly name: string | null @@ -66,6 +71,7 @@ export type WorkspaceShellLayoutProps = { readonly dashboardUrl?: string readonly desktopPanelWidths: Readonly readonly focusedChunk: FocusedChunkState + readonly focusedPage?: FocusedPageState readonly hasMessages: boolean readonly hasMoreSelectedChunks: boolean readonly contentView: ContentView @@ -84,6 +90,7 @@ export type WorkspaceShellLayoutProps = { readonly selectedSourceFile: SourceOriginalFileView | null readonly selectedSourceId: string | null readonly selectedSourceTitle: string | null + readonly selectedSourceView?: SourceView readonly sourceTitlesByDocumentId: Readonly> readonly sources: readonly SourceView[] readonly officialLibrarySources: readonly OfficialLibrarySourceView[] @@ -146,6 +153,10 @@ export function WorkspaceShellLayout( props.desktopPanelWidths.chat <= workspaceShellState.desktopSidePanelCompactThreshold const isSourcesPanelNarrow = props.desktopPanelWidths.sources < 220 + const selectedSource = + props.selectedSourceView ?? + props.sources.find((source) => source.id === props.selectedSourceId) + const focusedPage = props.focusedPage ?? { pageNumber: null, requestId: 0 } const handleDesktopLayoutRef = useCallback( (element: HTMLDivElement | null): void => { onDesktopLayoutElementChange(element) @@ -265,10 +276,13 @@ export function WorkspaceShellLayout( { expect(countFetches(fetch, "/api/sources/source_2/chunks")).toBe(0); }); + it("keeps a remote document open without refreshing sources after chunks load", async () => { + const remoteSourceId = "knowhere-doc:default:doc_remote"; + const encodedRemoteSourceId = encodeURIComponent(remoteSourceId); + const fetch = vi.fn(async (input) => { + const url = getRequestURL(input); + + if ( + url.pathname === `/api/sources/${encodedRemoteSourceId}/chunks` + ) { + return Response.json({ + chunks: [ + { + chunkId: "remote_page_1", + documentId: "doc_remote", + sectionPath: "Page 1", + type: "page", + content: "Remote page summary.", + sourceTitle: "remote.pdf", + pageNums: [1], + pageAssets: [ + { + pageNumber: 1, + assetUrl: "https://assets.example/page-1.png", + contentType: "image/png", + }, + ], + }, + ], + pagination: { + page: Number(url.searchParams.get("page") ?? "1"), + pageSize: 50, + total: 1, + totalPages: 1, + }, + }); + } + + if (url.pathname === "/api/sources") { + return Response.json({ + sources: [ + { + id: "source_localized", + kind: "workspace", + title: "remote.pdf", + status: "ready", + documentId: "doc_remote", + }, + ], + }); + } + + return Response.json({ message: "Unexpected request" }, { status: 404 }); + }); + vi.stubGlobal("fetch", fetch); + + render( + React.createElement(C, { + sources: [ + { + id: remoteSourceId, + kind: "remote", + title: "remote.pdf", + status: "ready", + documentId: "doc_remote", + excludedFromQuery: false, + }, + ], + }), + ); + + const desktopChunksPanel = within(screen.getByTestId("desktop-chunks-panel")); + await waitFor(() => { + expect(desktopChunksPanel.getByRole("img", { name: "Page 1" })) + .toBeTruthy(); + }); + + expect(countFetches(fetch, "/api/sources")).toBe(0); + expect( + countFetches( + fetch, + `/api/sources/${encodedRemoteSourceId}/chunks`, + ), + ).toBe(1); + }); + it("focuses guest citations on desktop using loaded demo chunks", async () => { const fetch = vi.fn(async (input) => { const url = getRequestURL(input); diff --git a/src/components/workspace-shell.tsx b/src/components/workspace-shell.tsx index 64663a0..7dd43dc 100644 --- a/src/components/workspace-shell.tsx +++ b/src/components/workspace-shell.tsx @@ -114,7 +114,6 @@ function WorkspaceShellContent({ fetchChunks: workspaceClient.fetchChunks, initialPrefetchedChunksBySourceId: initialPrefetchedChunksBySourceId ?? undefined, - onRemoteSourceChunksLoaded: sourceWorkflow.handleSourcesRefresh, onSelectSource: handleCitationSourceSelected, selectedSourceId: sourceWorkflow.selectedSourceId, sources: sourceWorkflow.sources, @@ -223,6 +222,7 @@ function WorkspaceShellContent({ dashboardUrl={dashboardUrl} citationListViewRequestId={citationFocus.citationListViewRequestId} focusedChunk={citationFocus.focusedChunk} + focusedPage={citationFocus.focusedPage} hasMessages={hasMessages} hasMoreSelectedChunks={citationFocus.hasMoreSelectedChunks} contentView={contentView} @@ -241,6 +241,7 @@ function WorkspaceShellContent({ selectedSourceFile={citationFocus.selectedSource?.originalFile ?? null} selectedSourceId={sourceWorkflow.selectedSourceId} selectedSourceTitle={selectedSourceTitle} + selectedSourceView={citationFocus.selectedSource} sourceTitlesByDocumentId={sourceWorkflow.sourceTitlesByDocumentId} sources={sourceWorkflow.sources} officialLibrarySources={officialLibrarySources ?? []} diff --git a/src/domains/chat/chat-citation-persistence.ts b/src/domains/chat/chat-citation-persistence.ts index 54e89de..5b0b26e 100644 --- a/src/domains/chat/chat-citation-persistence.ts +++ b/src/domains/chat/chat-citation-persistence.ts @@ -87,6 +87,7 @@ function toCitationView( score: citation.score, assetUrl: citation.assetUrl, pageCitationAssetUrl: citation.pageCitationAssetUrl, + pageCitationPageNumber: citation.pageCitationPageNumber, description: "description" in citation ? citation.description : undefined, source: { documentId: citation.source.documentId, diff --git a/src/domains/chat/citations.ts b/src/domains/chat/citations.ts index b6b4038..cdc9dda 100644 --- a/src/domains/chat/citations.ts +++ b/src/domains/chat/citations.ts @@ -20,6 +20,9 @@ export function toChatCitationViews( ...(result.pageCitationAssetUrl ? { pageCitationAssetUrl: result.pageCitationAssetUrl } : {}), + ...(result.pageCitationPageNumber + ? { pageCitationPageNumber: result.pageCitationPageNumber } + : {}), ...(description ? { description } : {}), source: { documentId: result.source.documentId, diff --git a/src/domains/chat/contracts.ts b/src/domains/chat/contracts.ts index 3400af5..e43b232 100644 --- a/src/domains/chat/contracts.ts +++ b/src/domains/chat/contracts.ts @@ -1,16 +1,22 @@ import type { + Knowledge, RetrievalQueryParams, RetrievalQueryResponse, } from "@ontos-ai/knowhere-sdk" import type { Source } from "@/infrastructure/db/schema" -import type { HarnessRunResult, InspectImages } from "@/agent-harness" +import type { + HarnessRunResult, + InspectImages, + KnowhereToolRuntime, +} from "@/agent-harness" import type { ChatArtifactView, ChatCitationView, } from "@/domains/chat/types" import type { HardenMediaAssetUrls } from "./media-asset-hardening" import type { HardenChatAssetUrl } from "./media-assets" +import type { NotebookKnowhereRemoteDocumentClient } from "./knowhere-tools" export type RetrievalClient = { query(params: RetrievalQueryParams): Promise @@ -57,6 +63,7 @@ export type GenerateAnswer = (input: { sources: readonly Source[] excludedSourceIds: readonly string[] searchSources: SearchSources + knowhereTools?: KnowhereToolRuntime inspectImages?: InspectImages }) => Promise @@ -67,6 +74,8 @@ export type AnswerQuestionInput = { sources: readonly Source[] excludedSourceIds: readonly string[] retrieval: RetrievalClient + knowledge?: Knowledge + remoteDocumentClient?: NotebookKnowhereRemoteDocumentClient generateAnswer: GenerateAnswer hardenChatAssetUrl?: HardenChatAssetUrl hardenMediaAssetUrls?: HardenMediaAssetUrls diff --git a/src/domains/chat/index.test.ts b/src/domains/chat/index.test.ts index 8c239a9..697fe77 100644 --- a/src/domains/chat/index.test.ts +++ b/src/domains/chat/index.test.ts @@ -1,5 +1,11 @@ import { afterEach, describe, expect, it, vi } from "vitest" -import type { RetrievalResult } from "@ontos-ai/knowhere-sdk" +import type { + Knowledge, + KnowledgeGrepResponse, + KnowledgeOutline, + KnowledgeReadResponse, + RetrievalResult, +} from "@ontos-ai/knowhere-sdk" import { Effect } from "effect" import { ToolLoopAgent } from "ai" import type { HarnessRunResult } from "@/agent-harness" @@ -8,6 +14,7 @@ import { answerQuestionWithRetrieval, generateAgenticOutputManifest, parseChatRequestBody, + type GenerateAnswer, type SearchSources, } from "." import type { @@ -78,7 +85,7 @@ describe("answerQuestionWithRetrieval", () => { namespace: "notebook-workspace", query: "What does the document say?", topK: 8, - useAgentic: true, + useAgentic: false, dataType: 1, excludeDocumentIds: ["doc_excluded", "doc_remote"], }); @@ -88,14 +95,185 @@ describe("answerQuestionWithRetrieval", () => { sources, excludedSourceIds: ["source_2", "knowhere-doc:default:doc_remote"], searchSources: expect.any(Function), + knowhereTools: expect.any(Object), }); expect(answer).toEqual({ answer: "The answer is grounded.", - citations: [result], + citations: [], + artifacts: [], + }); + }); + + it("does not create source chips from retrieval results when the manifest has no citations", async () => { + const unrelatedResult = makeRetrievalResult({ + content: "Information hiding is unrelated to the requested source.", + source: { + documentId: "doc_information_hiding", + sourceFileName: "information_hiding.pdf", + sectionPath: "Root", + }, + }); + const retrieval = { + query: vi.fn().mockResolvedValue({ + results: [unrelatedResult], + evidenceText: "Information hiding evidence.", + referencedChunks: [], + namespace: "notebook-workspace", + query: "requested fact", + routerUsed: "workflow_single_step", + answerText: null, + }), + }; + const generateAnswer = vi.fn(async ({ searchSources }) => { + await searchSources({ query: "requested fact" }); + return makeHarnessRunResult("The answer omits citations."); + }); + + const answer = await Effect.runPromise( + answerQuestionWithRetrieval({ + question: "What does the selected source say?", + namespace: "notebook-workspace", + sources: [makeSource()], + excludedSourceIds: [], + retrieval, + generateAnswer, + messages: [], + }), + ); + + expect(answer).toEqual({ + answer: "The answer omits citations.", + citations: [], artifacts: [], }); }); + it("exposes search, list, outline, read, and grep through the Knowhere tool runtime", async () => { + const result = makeRetrievalResult({ + chunkType: "image", + source: { + documentId: "doc_included", + sourceFileName: "notes.txt", + sectionPath: "images/diagram.png", + }, + }); + const retrieval = { + query: vi.fn().mockResolvedValue({ + results: [result], + evidenceText: "Diagram evidence.", + referencedChunks: [], + namespace: "notebook-workspace", + query: "diagram", + routerUsed: "workflow_single_step", + answerText: null, + }), + }; + const getDocumentOutline = vi.fn().mockResolvedValue(makeKnowledgeOutline()); + const readChunks = vi.fn().mockResolvedValue( + makeKnowledgeReadResponse("Full diagram chunk body."), + ); + const grepChunks = vi.fn().mockResolvedValue(makeKnowledgeGrepResponse()); + const knowledge = { + getDocumentOutline, + readChunks, + grepChunks, + } as unknown as Knowledge; + const listDocuments = vi.fn().mockResolvedValue({ + documents: [ + { + documentId: "doc_remote", + namespace: "default", + status: "ready", + currentJobResultId: "job_remote", + sourceFileName: "remote.pdf", + }, + ], + }); + const generateAnswer = vi.fn( + async ({ knowhereTools }: Parameters[0]) => { + if (!knowhereTools) throw new Error("Knowhere tools were not provided."); + + const searchResponse = await knowhereTools.search({ + query: "diagram", + targetContent: "image", + topK: 2, + }); + const documents = await knowhereTools.listDocuments(); + await knowhereTools.getDocumentOutline({ + documentId: "doc_included", + revisionKey: "job_123", + }); + await knowhereTools.readChunks({ + documentId: "doc_included", + revisionKey: "job_123", + page: 1, + pageSize: 2, + }); + await knowhereTools.grepChunks({ + documentId: "doc_included", + revisionKey: "job_123", + pattern: "diagram", + maxResults: 3, + }); + + expect(searchResponse.results).toEqual([result]); + expect( + documents.documents.map((document) => document.documentId), + ).toEqual(["doc_included", "doc_remote"]); + return makeHarnessRunResult("Runtime answer."); + }, + ); + const sources = [ + makeSource({ id: "source_included", knowhereDocumentId: "doc_included" }), + makeSource({ id: "source_excluded", knowhereDocumentId: "doc_excluded" }), + ]; + + const answer = await Effect.runPromise( + answerQuestionWithRetrieval({ + question: "Show the diagram.", + namespace: "notebook-workspace", + sources, + excludedSourceIds: ["source_excluded"], + retrieval, + knowledge, + remoteDocumentClient: { documents: { list: listDocuments } }, + generateAnswer, + messages: [], + }), + ); + + expect(retrieval.query).toHaveBeenCalledWith({ + namespace: "notebook-workspace", + query: "diagram", + topK: 2, + useAgentic: false, + dataType: 3, + excludeDocumentIds: ["doc_excluded"], + }); + expect(listDocuments).toHaveBeenCalledWith({ + namespace: "default", + page: 1, + pageSize: 200, + }); + expect(getDocumentOutline).toHaveBeenCalledWith({ + documentId: "doc_included", + revisionKey: "job_123", + }); + expect(readChunks).toHaveBeenCalledWith({ + documentId: "doc_included", + revisionKey: "job_123", + page: 1, + pageSize: 2, + }); + expect(grepChunks).toHaveBeenCalledWith({ + documentId: "doc_included", + revisionKey: "job_123", + pattern: "diagram", + maxResults: 3, + }); + expect(answer.answer).toBe("Runtime answer."); + }); + it("does not carry no-evidence metadata from default into a successful legacy namespace result", async () => { const legacyResult = makeRetrievalResult({ source: { @@ -170,7 +348,7 @@ describe("answerQuestionWithRetrieval", () => { ); expect(answer).toEqual({ answer: "The legacy answer is grounded.", - citations: [legacyResult], + citations: [], artifacts: [], }); }); @@ -428,8 +606,18 @@ describe("answerQuestionWithRetrieval", () => { }; const generateAnswer = vi.fn(async ({ searchSources }) => { await searchSources({ query: "What improved?" }); - return makeHarnessRunResult( + return makeHarnessRunResultWithLedger( "Revenue improved [Source 1: revenue growth]. Margins expanded [Source 2: margin expansion].", + { + citations: [ + makeOutputCitation("r1:result:1", firstResult), + makeOutputCitation("r1:result:2", secondResult), + ], + chunks: [ + makeEvidenceChunkFromRetrievalResult("r1:result:1", firstResult), + makeEvidenceChunkFromRetrievalResult("r1:result:2", secondResult), + ], + }, ); }); @@ -472,8 +660,9 @@ describe("answerQuestionWithRetrieval", () => { }; const generateAnswer = vi.fn(async ({ searchSources }) => { await searchSources({ query: "Tesla xAI investment" }); - return makeHarnessRunResult( + return makeCitedHarnessRunResult( "Tesla invested in xAI [Source 1: xAI investment].", + result, ); }); const sources = [ @@ -501,6 +690,7 @@ describe("answerQuestionWithRetrieval", () => { sources, excludedSourceIds: [], searchSources: expect.any(Function), + knowhereTools: expect.any(Object), }); const expectedResult = { ...result, @@ -543,7 +733,10 @@ describe("answerQuestionWithRetrieval", () => { targetContent: "image", purpose: "Find visual rocket launch chunks.", }); - return makeHarnessRunResult(`Use this launch photo. ${upstreamAssetUrl}`); + return makeCitedHarnessRunResult( + `Use this launch photo. ${upstreamAssetUrl}`, + result, + ); }); const hardenChatAssetUrl = vi .fn() @@ -579,7 +772,7 @@ describe("answerQuestionWithRetrieval", () => { namespace: "notebook-workspace", query: "SpaceX rocket photos", topK: 8, - useAgentic: true, + useAgentic: false, dataType: 3, }); expect(answer.answer).toBe("Use this launch photo."); @@ -803,7 +996,10 @@ describe("answerQuestionWithRetrieval", () => { }; const generateAnswer = vi.fn(async ({ searchSources }) => { await searchSources({ query: "page four evidence" }); - return makeHarnessRunResult(`This page has the answer. ${storedPageAssetUrl}`); + return makeCitedHarnessRunResult( + `This page has the answer. ${storedPageAssetUrl}`, + result, + ); }); const hardenMediaAssetUrls = vi.fn( async ({ @@ -887,6 +1083,35 @@ describe("answerQuestionWithRetrieval", () => { "https://blob.example/workspaces/workspace_1/sources/source_pages/parsed-result/page_citation_assets/page-6.png"; const hardenedPageAssetUrl = "https://blob.example/workspaces/workspace_1/chat-assets/source-source_pages/page-6.png"; + const pageMetadata = { + pageNums: [6], + pageAssets: [ + { + pageNum: 6, + artifactRef: "page_citation_assets/page-6.png", + assetUrl: rawPageAssetUrl, + contentType: "image/png", + }, + ], + }; + const referencedPageChunk: HarnessRunResult["trace"]["ledger"]["chunks"][number] = { + ref: "r1:referenced:1", + kind: "referenced_chunk", + chunkId: "chunk_page_6", + content: "", + contentPreview: "", + chunkType: "page", + score: null, + filePath: null, + metadata: pageMetadata, + source: { + documentId: "doc_pages", + sourceFileName: null, + sectionPath: "Page 6", + }, + revisionKey: "job_1", + assetUrl: rawPageAssetUrl, + }; const retrieval = { query: vi.fn().mockResolvedValue({ results: [], @@ -900,17 +1125,7 @@ describe("answerQuestionWithRetrieval", () => { filePath: null, jobId: "job_1", assetUrl: rawPageAssetUrl, - metadata: { - pageNums: [6], - pageAssets: [ - { - pageNum: 6, - artifactRef: "page_citation_assets/page-6.png", - assetUrl: rawPageAssetUrl, - contentType: "image/png", - }, - ], - }, + metadata: pageMetadata, }, ], namespace: "notebook-workspace", @@ -921,7 +1136,16 @@ describe("answerQuestionWithRetrieval", () => { }; const generateAnswer = vi.fn(async ({ searchSources }) => { await searchSources({ query: "page six evidence" }); - return makeHarnessRunResult("This page has referenced evidence."); + return makeHarnessRunResultWithLedger("This page has referenced evidence.", { + citations: [ + { + ref: "r1:referenced:1", + label: "Page 6", + source: referencedPageChunk.source, + }, + ], + chunks: [referencedPageChunk], + }); }); const hardenMediaAssetUrls = vi.fn( async ({ @@ -1222,7 +1446,7 @@ describe("answerQuestionWithRetrieval", () => { ]); }); - it("returns a safe fallback when the harness still has validation errors", async () => { + it("returns agent output when a legacy harness trace has validation errors", async () => { const retrieval = { query: vi.fn().mockResolvedValue({ results: [makeRetrievalResult()], @@ -1261,13 +1485,112 @@ describe("answerQuestionWithRetrieval", () => { ); expect(answer).toEqual({ - answer: - "I couldn't safely finish that response because the agent output did not pass Notebook's validation checks. Please try again.", + answer: "This invalid answer should not ship.", citations: [], artifacts: [], }); }); + it("renders answer and resolves citations when manifest source metadata is wrong", async () => { + process.env.AI_GATEWAY_API_KEY = "test_gateway_key"; + const result = makeRetrievalResult({ + content: "Information hiding is a module design principle.", + source: { + documentId: "doc_information_hiding", + sourceFileName: "information_hiding.pdf", + sectionPath: "Root / Module Design", + }, + }); + const retrieval = { + query: vi.fn().mockResolvedValue({ + results: [result], + evidenceText: "Information hiding evidence.", + referencedChunks: [], + namespace: "notebook-workspace", + query: "information hiding", + routerUsed: "workflow_single_step", + answerText: null, + }), + }; + let generateCallCount = 0; + vi.spyOn(ToolLoopAgent.prototype, "generate").mockImplementation( + async function mockGenerate( + this: ToolLoopAgent, + ): ReturnType { + generateCallCount += 1; + const tools = this.tools as unknown as Record< + string, + { execute: (input: unknown) => Promise } + >; + + if (generateCallCount === 1) { + await tools.declareIntent?.execute({ + task: "answer", + dependsOnPreviousTurn: false, + retrievalNeeded: "yes", + targetModalities: ["text"], + constraints: { citationRequired: true }, + groundingPolicy: "must_use_sources", + }); + await tools.setContextPolicy?.execute({ + carryHistory: "none", + reason: "Self-contained request.", + activePriorTurnIds: [], + }); + await tools.knowhere_search?.execute({ + query: "information hiding", + targetContent: "text", + }); + } + + await tools.finalize?.execute({ + text: "Information hiding is a module design principle.", + citations: [ + { + ref: "r1:result:1", + label: "claimed-source.pdf / Claimed", + source: { + documentId: "doc_claimed", + sourceFileName: "claimed-source.pdf", + sectionPath: "Claimed", + }, + }, + ], + artifacts: [], + unresolved: [], + }); + + return { + text: "ignored", + response: { messages: [] }, + } as unknown as Awaited>; + }, + ); + + const answer = await Effect.runPromise( + answerQuestionWithRetrieval({ + question: "What is information hiding?", + namespace: "notebook-workspace", + sources: [makeSource()], + excludedSourceIds: [], + retrieval, + generateAnswer: generateAgenticOutputManifest, + messages: [], + }), + ); + + expect(generateCallCount).toBe(1); + expect(answer.answer).toBe("Information hiding is a module design principle."); + expect(answer.artifacts).toEqual([]); + expect(answer.citations.map((citation) => citation.source)).toEqual([ + { + documentId: "doc_information_hiding", + sourceFileName: "information_hiding.pdf", + sectionPath: "Root / Module Design", + }, + ]); + }); + it("keeps image-only harness output instead of treating it as no results", async () => { const assetUrl = "https://blob.example/images/diagram.png"; const retrieval = { @@ -1568,12 +1891,13 @@ describe("answerQuestionWithRetrieval", () => { sources, excludedSourceIds: [], searchSources: expect.any(Function), + knowhereTools: expect.any(Object), }); expect(retrieval.query).toHaveBeenCalledWith({ namespace: "notebook-workspace", query: "公民身份证明 图片", topK: 8, - useAgentic: true, + useAgentic: false, dataType: 3, }); const imageCitations = answer.citations.filter( @@ -1664,7 +1988,7 @@ describe("answerQuestionWithRetrieval", () => { namespace: "notebook-workspace", query: "Tesla Q4 2025 Update energy generation and storage deployments", topK: 8, - useAgentic: true, + useAgentic: false, dataType: 1, }); expect(generateAnswer).toHaveBeenCalledWith({ @@ -1673,6 +1997,7 @@ describe("answerQuestionWithRetrieval", () => { sources: [makeSource({ title: "TSLA-Q4-2025-Update.pdf" })], excludedSourceIds: [], searchSources: expect.any(Function), + knowhereTools: expect.any(Object), }); }); @@ -1720,7 +2045,7 @@ describe("answerQuestionWithRetrieval", () => { namespace: "notebook-workspace", query: "Tesla energy storage deployments", topK: 8, - useAgentic: true, + useAgentic: false, dataType: 1, }); expect(JSON.stringify(queryInput)).not.toContain( @@ -1729,6 +2054,23 @@ describe("answerQuestionWithRetrieval", () => { }); it("uses structured referenced chunks from RetrievalQueryResponse as citations", async () => { + const referencedImageChunk: HarnessRunResult["trace"]["ledger"]["chunks"][number] = { + ref: "r1:referenced:1", + kind: "referenced_chunk", + chunkId: "chunk_1", + content: "", + contentPreview: "", + chunkType: "image", + score: null, + filePath: "images/launch.jpg", + source: { + documentId: "doc_spacex", + sourceFileName: null, + sectionPath: "Assets / images / launch.jpg", + }, + revisionKey: "job_1", + assetUrl: "https://blob.example/images/launch.jpg", + }; const retrieval = { query: vi.fn().mockResolvedValue({ results: [], @@ -1755,7 +2097,16 @@ describe("answerQuestionWithRetrieval", () => { query: "SpaceX launch image", targetContent: "image", }); - return makeHarnessRunResult("Here is the launch image."); + return makeHarnessRunResultWithLedger("Here is the launch image.", { + citations: [ + { + ref: "r1:referenced:1", + label: "launch image", + source: referencedImageChunk.source, + }, + ], + chunks: [referencedImageChunk], + }); }); const answer = await Effect.runPromise( @@ -1821,9 +2172,9 @@ describe("generateAgenticOutputManifest", () => { reason: "The current request is self-contained.", activePriorTurnIds: [], }); - await tools.retrieve?.execute({ + await tools.knowhere_search?.execute({ query: "冯荣洲 身份证 图片", - modalities: ["text", "image"], + targetContent: "text_image", topK: 2, purpose: "Find exactly the requested identity-card images.", }); @@ -1835,7 +2186,7 @@ describe("generateAgenticOutputManifest", () => { label: "商务标文件.pdf / 身份证正面", source: { documentId: "doc_identity", - sourceFileName: "商务标文件.pdf", + sourceFileName: "document-generated.pdf", sectionPath: "身份证正面", }, }, @@ -1954,9 +2305,9 @@ describe("generateAgenticOutputManifest", () => { reason: "The current request is self-contained.", activePriorTurnIds: [], }); - await tools.retrieve?.execute({ + await tools.knowhere_search?.execute({ query: "identity card front image", - modalities: ["image"], + targetContent: "image", topK: 1, purpose: "Find the ID card image to inspect.", }); @@ -1972,7 +2323,7 @@ describe("generateAgenticOutputManifest", () => { label: "identity.pdf / images/id-front.png", source: { documentId: "doc_identity", - sourceFileName: "identity.pdf", + sourceFileName: "generated.pdf", sectionPath: "images/id-front.png", }, }, @@ -2093,9 +2444,9 @@ describe("generateAgenticOutputManifest", () => { reason: "The current request is self-contained.", activePriorTurnIds: [], }); - await tools.retrieve?.execute({ + await tools.knowhere_search?.execute({ query: "进度计划 违约金 承包人", - modalities: ["text"], + targetContent: "text", topK: 6, purpose: "Find the contract clause and page for the liquidated damages amount.", }); @@ -2112,7 +2463,7 @@ describe("generateAgenticOutputManifest", () => { label: "投标书 / （6）现场工期进度管理方面的违约责任", source: { documentId: "doc_contract", - sourceFileName: "投标书.pdf", + sourceFileName: null, sectionPath: "Root / （6）现场工期进度管理方面的违约责任", }, }, @@ -2221,7 +2572,7 @@ describe("generateAgenticOutputManifest", () => { expect(result.trace.validationErrors).toEqual([]); }); - it("self-corrects an over-budget manifest via a validation-feedback revision", async () => { + it("keeps an over-budget manifest after the first generation", async () => { process.env.AI_GATEWAY_API_KEY = "test_gateway_key"; let generateCallCount = 0; vi.spyOn(ToolLoopAgent.prototype, "generate").mockImplementation( @@ -2248,15 +2599,25 @@ describe("generateAgenticOutputManifest", () => { reason: "Self-contained request.", activePriorTurnIds: [], }); - await tools.retrieve?.execute({ + await tools.knowhere_search?.execute({ query: "身份证 图片", - modalities: ["image"], + targetContent: "image", topK: 3, purpose: "Find requested identity images.", }); await tools.finalize?.execute({ text: "见下方图片。", - citations: [{ ref: "r1:result:1", label: "id" }], + citations: [ + { + ref: "r1:result:1", + label: "ids.pdf / 身份证 1", + source: { + documentId: "doc_identity", + sourceFileName: "ids.pdf", + sectionPath: "身份证 1", + }, + }, + ], artifacts: [1, 2, 3].map((index) => ({ type: "image", ref: `asset:r1:result:${index}`, @@ -2268,7 +2629,17 @@ describe("generateAgenticOutputManifest", () => { } else { await tools.finalize?.execute({ text: "见下方图片。", - citations: [{ ref: "r1:result:1", label: "id" }], + citations: [ + { + ref: "r1:result:1", + label: "ids.pdf / 身份证 1", + source: { + documentId: "doc_identity", + sourceFileName: "ids.pdf", + sectionPath: "身份证 1", + }, + }, + ], artifacts: [1, 2].map((index) => ({ type: "image", ref: `asset:r1:result:${index}`, @@ -2318,12 +2689,12 @@ describe("generateAgenticOutputManifest", () => { searchSources, }); - expect(generateCallCount).toBe(2); - expect(result.trace.revisionsUsed).toBe(1); + expect(generateCallCount).toBe(1); + expect(result.trace.revisionsUsed).toBe(0); expect(result.trace.validationErrors).toEqual([]); expect( result.manifest.artifacts.filter((artifact) => artifact.display).length, - ).toBe(2); + ).toBe(3); }); }); @@ -2421,6 +2792,171 @@ function makeHarnessRunResult(text: string): HarnessRunResult { }; } +function makeCitedHarnessRunResult( + text: string, + result: RetrievalResult, + ref = "r1:result:1", +): HarnessRunResult { + return makeHarnessRunResultWithLedger(text, { + citations: [makeOutputCitation(ref, result)], + chunks: [makeEvidenceChunkFromRetrievalResult(ref, result)], + }); +} + +function makeHarnessRunResultWithLedger( + text: string, + input: { + readonly citations?: HarnessRunResult["manifest"]["citations"] + readonly chunks?: HarnessRunResult["trace"]["ledger"]["chunks"] + readonly assets?: HarnessRunResult["trace"]["ledger"]["assets"] + readonly artifacts?: HarnessRunResult["manifest"]["artifacts"] + }, +): HarnessRunResult { + const chunks = input.chunks ?? []; + return { + manifest: { + text, + citations: input.citations ?? [], + artifacts: input.artifacts ?? [], + unresolved: [], + }, + trace: { + ...makeHarnessRunResult("").trace, + ledger: { + retrievalCount: chunks.length > 0 ? 1 : 0, + chunks, + assets: input.assets ?? [], + evidenceText: [], + stopReasons: [], + failureReasons: [], + decisionTraces: [], + }, + }, + }; +} + +function makeOutputCitation( + ref: string, + result: RetrievalResult, +): HarnessRunResult["manifest"]["citations"][number] { + return { + ref, + label: [result.source.sourceFileName, result.source.sectionPath] + .filter(Boolean) + .join(" / "), + source: { + documentId: result.source.documentId, + sourceFileName: result.source.sourceFileName, + sectionPath: result.source.sectionPath, + }, + }; +} + +function makeEvidenceChunkFromRetrievalResult( + ref: string, + result: RetrievalResult, +): HarnessRunResult["trace"]["ledger"]["chunks"][number] { + return { + ref, + kind: "result", + ...(result.chunkId ? { chunkId: result.chunkId } : {}), + content: result.content, + contentPreview: result.content, + chunkType: result.chunkType, + score: result.score, + ...(result.sourceChunkPath ? { sourceChunkPath: result.sourceChunkPath } : {}), + ...(result.filePath ? { filePath: result.filePath } : {}), + ...(result.metadata ? { metadata: result.metadata } : {}), + source: { + documentId: result.source.documentId, + sourceFileName: result.source.sourceFileName, + sectionPath: result.source.sectionPath, + }, + ...(result.assetUrl ? { assetUrl: result.assetUrl } : {}), + }; +} + +function makeKnowledgeOutline(): KnowledgeOutline { + return { + document: makeLocalKnowledgeDocument(), + totalChunks: 1, + typeCounts: { text: 1, image: 0, table: 0, page: 0 }, + sections: [ + { + sectionPath: "Root / Diagram", + sectionTitle: "Diagram", + sectionLevel: 2, + summary: "Diagram section.", + startChunk: 1, + endChunk: 1, + chunkCount: 1, + typeCounts: { text: 1, image: 0, table: 0, page: 0 }, + children: [], + }, + ], + sectionTree: [], + }; +} + +function makeKnowledgeReadResponse(content: string): KnowledgeReadResponse { + return { + document: makeLocalKnowledgeDocument(), + chunks: [ + { + position: 1, + chunkId: "chunk_1", + chunkType: "text", + content, + readableContent: content, + sectionPath: "Root / Diagram", + sourceChunkPath: "chunks/chunk-1.md", + filePath: "notes.txt", + metadata: {}, + }, + ], + page: 1, + pageSize: 1, + totalChunks: 1, + totalPages: 1, + }; +} + +function makeKnowledgeGrepResponse(): KnowledgeGrepResponse { + return { + document: makeLocalKnowledgeDocument(), + matches: [ + { + position: 1, + chunkId: "chunk_1", + chunkType: "text", + sectionPath: "Root / Diagram", + sourceChunkPath: "chunks/chunk-1.md", + filePath: "notes.txt", + startOffset: 0, + endOffset: 7, + snippet: "diagram", + }, + ], + scannedChunks: 1, + truncated: false, + }; +} + +function makeLocalKnowledgeDocument() { + return { + localDocumentId: "doc_included", + documentId: "doc_included", + jobId: "job_123", + namespace: "notebook-workspace", + sourceFileName: "notes.txt", + chunkCount: 1, + typeCounts: { text: 1, image: 0, table: 0, page: 0 }, + resultDirectoryPath: "parsed-storage:doc_included/job_123", + createdAt: new Date("2026-01-01T00:00:00Z"), + updatedAt: new Date("2026-01-01T00:00:00Z"), + }; +} + type KnowhereQueryResponseLogMeta = { readonly query: string readonly resultCount: number diff --git a/src/domains/chat/index.ts b/src/domains/chat/index.ts index fa07579..d4dbcc5 100644 --- a/src/domains/chat/index.ts +++ b/src/domains/chat/index.ts @@ -16,7 +16,6 @@ import type { EvidenceChunk, HarnessRunResult, OutputArtifact, - OutputCitation, } from "@/agent-harness" import { toChatCitationViews, @@ -40,8 +39,12 @@ import { } from "./media-assets" import { enrichRetrievalResultsWithPageCitationAssetUrls } from "./page-citation-assets" import type { HardenableRetrievalResult } from "./media-asset-hardening" +import { notebookKnowhereTools } from "./knowhere-tools" const DEFAULT_TOP_K = 8 +const NOTEBOOK_USE_AGENTIC_RETRIEVAL: NonNullable< + RetrievalQueryParams["useAgentic"] +> = false const MAX_AGENTIC_TOP_K = 12 const MAX_AGENTIC_MERGED_RESULT_COUNT = 24 const MAX_AGENTIC_MERGED_REFERENCED_CHUNK_COUNT = 24 @@ -51,8 +54,6 @@ const KNOWHERE_RESPONSE_TEXT_LOG_LIMIT = 200 const KNOWHERE_CHUNK_LOG_LIMIT = 100 const KNOWHERE_RESPONSE_LOG_ITEM_LIMIT = 20 const NO_RESULTS_ANSWER = "I couldn't find that in your sources." -const HARNESS_VALIDATION_FAILURE_ANSWER = - "I couldn't safely finish that response because the agent output did not pass Notebook's validation checks. Please try again." const RAW_URL_PATTERN = /https?:\/\/[^\s)\]}>"']+/g const REDACTED_MEDIA_URL = "[media asset URL hidden]" const RETRIEVAL_TARGET_CONTENT_DATA_TYPES: Readonly< @@ -214,6 +215,14 @@ export const answerQuestionWithRetrieval = ( sources: input.sources, excludedSourceIds: input.excludedSourceIds, searchSources, + knowhereTools: notebookKnowhereTools.createRuntime({ + namespace: input.namespace, + sources: input.sources, + excludedSourceIds: input.excludedSourceIds, + searchSources, + knowledge: input.knowledge, + remoteDocumentClient: input.remoteDocumentClient, + }), ...(input.inspectImages ? { inspectImages: input.inspectImages } : {}), }), ) @@ -222,28 +231,11 @@ export const answerQuestionWithRetrieval = ( answerLength: generatedAnswer.manifest.text.length, retrievalCallCount: retrievalResponses.length, citationCount: generatedAnswer.manifest.citations.length, - harnessValidationErrorCount: generatedAnswer.trace.validationErrors.length, - revisionsUsed: generatedAnswer.trace.revisionsUsed, + finalized: generatedAnswer.trace.finalized, }) - if (generatedAnswer.trace.validationErrors.length > 0) { - logger.warn("chat-agent: validation failed; returning safe fallback", { - validationErrors: generatedAnswer.trace.validationErrors, - revisionsUsed: generatedAnswer.trace.revisionsUsed, - finalized: generatedAnswer.trace.finalized, - intentTask: generatedAnswer.trace.intent?.task ?? null, - retrievalCallCount: retrievalResponses.length, - }) - return { - answer: HARNESS_VALIDATION_FAILURE_ANSWER, - citations: [] as ChatCitationView[], - artifacts: [] as ChatArtifactView[], - } - } const rawResults = selectCitationRawResults({ generatedAnswer, - retrievalResponses, - sources: input.sources, }) if ( rawResults.length === 0 && @@ -427,7 +419,7 @@ function toChatArtifactView(input: { } function normalizeHarnessSource( - source: OutputCitation["source"], + source: EvidenceChunk["source"], sources: readonly AnswerQuestionInput["sources"][number][], ): ChatCitationView["source"] { const sourceTitle = source.documentId @@ -795,7 +787,7 @@ function buildRetrievalQueryParams(input: { namespace: input.namespace, query, topK: normalizeTopK(input.input.topK), - useAgentic: true, + useAgentic: NOTEBOOK_USE_AGENTIC_RETRIEVAL, dataType, ...(input.input.signalPaths && input.input.signalPaths.length > 0 ? { signalPaths: input.input.signalPaths } @@ -846,14 +838,12 @@ function normalizeTopK(value: number | undefined): number { /** * Display citations come from the agent-curated manifest (the refs it chose to - * cite), resolved against the evidence ledger. Only when the agent cited - * nothing do we fall back to the full set of retrieved results, so a grounded - * answer still shows its sources instead of appearing unsupported. + * cite), resolved against the evidence ledger. If the manifest has no citations, + * only displayed artifacts produce citation chips; arbitrary retrieval results + * stay hidden so Notebook does not imply support from an unselected source. */ function selectCitationRawResults(input: { readonly generatedAnswer: HarnessRunResult - readonly retrievalResponses: readonly RetrievalQueryResponse[] - readonly sources: readonly AnswerQuestionInput["sources"][number][] }): RetrievalResult[] { const curated = mapManifestCitationsToResults(input.generatedAnswer) if (curated.length > 0) return curated @@ -861,7 +851,7 @@ function selectCitationRawResults(input: { input.generatedAnswer, ) if (displayedArtifacts.length > 0) return displayedArtifacts - return collectRetrievalResults(input.retrievalResponses, input.sources) + return [] } function mapManifestCitationsToResults( @@ -996,49 +986,6 @@ function hasDisplayedManifestArtifacts(result: HarnessRunResult): boolean { return result.manifest.artifacts.some((artifact) => artifact.display) } -function collectRetrievalResults( - responses: readonly RetrievalQueryResponse[], - sources: readonly AnswerQuestionInput["sources"][number][], -): RetrievalResult[] { - const results: RetrievalResult[] = [] - const seenKeys = new Set() - const sourceTitlesByDocumentId = new Map( - sources.flatMap((source): readonly [string, string][] => - source.knowhereDocumentId ? [[source.knowhereDocumentId, source.title]] : [], - ), - ) - - for (const response of responses) { - for (const result of [ - ...response.results, - ...response.referencedChunks.map((chunk): RetrievalResult => ({ - chunkId: chunk.chunkId, - content: "", - chunkType: chunk.chunkType, - score: null, - ...(chunk.assetUrl ? { assetUrl: chunk.assetUrl } : {}), - ...(chunk.sourceChunkPath ? { sourceChunkPath: chunk.sourceChunkPath } : {}), - ...(chunk.filePath ? { filePath: chunk.filePath } : {}), - ...(chunk.metadata ? { metadata: chunk.metadata } : {}), - source: { - documentId: chunk.documentId, - sourceFileName: sourceTitlesByDocumentId.get(chunk.documentId), - sectionPath: chunk.sectionPath, - }, - })), - ]) { - const key = getRetrievalResultKey(result) - if (seenKeys.has(key)) continue - - seenKeys.add(key) - results.push(result) - if (results.length >= MAX_CITATION_RESULTS) return results - } - } - - return results -} - function formatRetrievalEvidenceText( responses: readonly RetrievalQueryResponse[], ): string | undefined { diff --git a/src/domains/chat/knowhere-tools.ts b/src/domains/chat/knowhere-tools.ts new file mode 100644 index 0000000..210b487 --- /dev/null +++ b/src/domains/chat/knowhere-tools.ts @@ -0,0 +1,155 @@ +import { Effect } from "effect" +import type { Knowledge } from "@ontos-ai/knowhere-sdk" + +import type { + KnowhereDocumentSummary, + KnowhereToolRuntime, +} from "@/agent-harness" +import type { Source } from "@/infrastructure/db/schema" +import { + listRemoteLibraryDocuments, + type RemoteLibraryDocument, +} from "@/domains/sources/remote-library" +import type { SearchSources } from "./contracts" +import { excludeDocuments } from "./retrieval" + +type RemoteDocumentClient = Parameters< + typeof listRemoteLibraryDocuments +>[0]["client"] + +export type NotebookKnowhereRemoteDocumentClient = RemoteDocumentClient + +type NotebookKnowhereToolsInput = { + readonly namespace: string + readonly sources: readonly Source[] + readonly excludedSourceIds: readonly string[] + readonly searchSources: SearchSources + readonly knowledge?: Knowledge + readonly remoteDocumentClient?: RemoteDocumentClient +} + +type SearchOnlyRuntimeInput = { + readonly searchSources: SearchSources +} + +export const notebookKnowhereTools = { + createRuntime(input: NotebookKnowhereToolsInput): KnowhereToolRuntime { + return { + search: (request) => input.searchSources(request), + listDocuments: async () => ({ + documents: await listVisibleDocuments(input), + }), + getDocumentOutline: async (request) => { + const knowledge = requireKnowledge(input.knowledge) + return knowledge.getDocumentOutline(request) + }, + readChunks: async (request) => { + const knowledge = requireKnowledge(input.knowledge) + return knowledge.readChunks(request) + }, + grepChunks: async (request) => { + const knowledge = requireKnowledge(input.knowledge) + return knowledge.grepChunks(request) + }, + } + }, + + createSearchOnlyRuntime(input: SearchOnlyRuntimeInput): KnowhereToolRuntime { + return { + search: (request) => input.searchSources(request), + listDocuments: async () => ({ documents: [] }), + getDocumentOutline: async () => { + throw new Error("Knowhere document outline is not configured.") + }, + readChunks: async () => { + throw new Error("Knowhere chunk reads are not configured.") + }, + grepChunks: async () => { + throw new Error("Knowhere grep is not configured.") + }, + } + }, +} as const + +async function listVisibleDocuments( + input: NotebookKnowhereToolsInput, +): Promise { + const excludedSourceIds = new Set(input.excludedSourceIds) + const excludedDocumentIds = new Set( + excludeDocuments(input.sources, input.excludedSourceIds) + .excludeDocumentIds ?? [], + ) + const localDocuments = input.sources + .filter( + (source): source is Source & { readonly knowhereDocumentId: string } => + source.status === "ready" && + Boolean(source.knowhereDocumentId) && + !excludedSourceIds.has(source.id) && + !excludedDocumentIds.has(source.knowhereDocumentId ?? ""), + ) + .map((source): KnowhereDocumentSummary => ({ + documentId: source.knowhereDocumentId, + revisionKey: source.knowhereJobId ?? undefined, + namespace: input.namespace, + sourceFileName: source.title, + title: source.title, + status: source.status, + })) + + const remoteDocuments = await listVisibleRemoteDocuments({ + input, + localDocuments, + excludedDocumentIds, + }) + + return [...localDocuments, ...remoteDocuments] +} + +async function listVisibleRemoteDocuments(input: { + readonly input: NotebookKnowhereToolsInput + readonly localDocuments: readonly KnowhereDocumentSummary[] + readonly excludedDocumentIds: ReadonlySet +}): Promise { + if (!input.input.remoteDocumentClient) return [] + + const localDocumentIds = new Set( + input.localDocuments.flatMap((document): string[] => + document.documentId ? [document.documentId] : [], + ), + ) + const documents = await Effect.runPromise( + listRemoteLibraryDocuments({ + workspace: { namespace: input.input.namespace }, + client: input.input.remoteDocumentClient, + localSources: input.input.sources, + }), + ) + + return documents + .filter( + (document) => + document.status === "ready" && + !localDocumentIds.has(document.documentId) && + !input.excludedDocumentIds.has(document.documentId), + ) + .map(toRemoteDocumentSummary) +} + +function toRemoteDocumentSummary( + document: RemoteLibraryDocument, +): KnowhereDocumentSummary { + return { + documentId: document.documentId, + revisionKey: document.revisionKey, + namespace: document.namespace, + sourceFileName: + document.sourceFileName ?? document.title ?? document.documentId, + title: document.title, + status: document.status, + } +} + +function requireKnowledge(knowledge: Knowledge | undefined): Knowledge { + if (knowledge) return knowledge + throw new Error("Knowhere parsed-document reads are not configured.") +} diff --git a/src/domains/chat/media-asset-hardening.test.ts b/src/domains/chat/media-asset-hardening.test.ts index 7c106c9..2afa6d9 100644 --- a/src/domains/chat/media-asset-hardening.test.ts +++ b/src/domains/chat/media-asset-hardening.test.ts @@ -24,7 +24,7 @@ afterEach(() => { describe("hardenChatMediaAssetUrls", () => { it("keeps an already Notebook-owned asset URL without calling the hardener", async () => { const ownedUrl = - "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_1/rev_1/assets/images/a.png" + "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_1/rev_1/images/a.png" const hardenChatAssetUrl = vi.fn(async () => null) const result = await hardenChatMediaAssetUrls({ @@ -52,7 +52,7 @@ describe("hardenChatMediaAssetUrls", () => { const rawAssetUrl = "https://knowhere-storage.example/results/job_1/images/id-front.jpg?AWSAccessKeyId=test" const durableUrl = - "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_identity/rev_1/assets/images/id-front.jpg" + "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_identity/rev_1/images/id-front.jpg" const hardenChatAssetUrl = vi.fn().mockResolvedValue(durableUrl) const result = await hardenChatMediaAssetUrls({ @@ -116,7 +116,7 @@ describe("hardenChatMediaAssetUrls", () => { const rawAssetUrl = "https://knowhere-storage.example/results/job_1/images/front.jpg?AWSAccessKeyId=test" const durableUrl = - "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_identity/rev_1/assets/images/front.jpg" + "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_identity/rev_1/images/front.jpg" const hardenChatAssetUrl = vi.fn().mockResolvedValue(durableUrl) const result = await hardenChatMediaAssetUrls({ @@ -164,7 +164,7 @@ describe("isNotebookOwnedAssetUrl", () => { ).toBe(true) expect( isNotebookOwnedAssetUrl( - "https://cdn.example/workspaces/w/parsed-documents/d/r/assets/a.png", + "https://cdn.example/workspaces/w/parsed-documents/d/r/images/a.png", ), ).toBe(true) expect( diff --git a/src/domains/chat/page-citation-assets.ts b/src/domains/chat/page-citation-assets.ts index 566cead..223498b 100644 --- a/src/domains/chat/page-citation-assets.ts +++ b/src/domains/chat/page-citation-assets.ts @@ -7,6 +7,7 @@ import type { HardenChatAssetUrl } from "./media-assets" export type PageCitationAssetRetrievalResult = RetrievalResult & { readonly pageCitationAssetUrl?: string + readonly pageCitationPageNumber?: number } type EnrichRetrievalResultsWithPageCitationAssetUrlsInput = { @@ -62,10 +63,13 @@ async function enrichRetrievalResultWithPageCitationAssetUrl(input: { sourcesByDocumentId: input.sourcesByDocumentId, hardenChatAssetUrl: input.hardenChatAssetUrl, }) - if (sourceAssetUrl) { + if (sourceAssetUrl || directAsset?.pageNum) { return { ...input.result, - pageCitationAssetUrl: sourceAssetUrl, + ...(sourceAssetUrl ? { pageCitationAssetUrl: sourceAssetUrl } : {}), + ...(directAsset?.pageNum + ? { pageCitationPageNumber: directAsset.pageNum } + : {}), } } diff --git a/src/domains/chat/prompt.ts b/src/domains/chat/prompt.ts index efee008..a4ae326 100644 --- a/src/domains/chat/prompt.ts +++ b/src/domains/chat/prompt.ts @@ -8,17 +8,15 @@ import { runAgentHarness, type AgentTurn, type AgentTurnInput, - type HarnessRetrievalRequest, type HarnessRunResult, type InspectImages, - type TargetModality, + type KnowhereToolRuntime, } from "@/agent-harness" import type { - AgenticRetrievalQuery, - AgenticRetrievalTargetContent, ChatHistoryMessage, SearchSources, } from "./contracts" +import { notebookKnowhereTools } from "./knowhere-tools" const RECENT_CONTEXT_MESSAGE_LIMIT = 8 const CONTEXT_CONTENT_CHAR_LIMIT = 900 @@ -30,6 +28,7 @@ type GenerateAgenticOutputManifestInput = { sources: readonly Source[] excludedSourceIds: readonly string[] searchSources: SearchSources + knowhereTools?: KnowhereToolRuntime inspectImages?: InspectImages } @@ -59,10 +58,11 @@ export const generateAgenticOutputManifestEffect = ( runAgentHarness({ model: CHAT_MODEL, turn, - retrieval: { - query: (request) => - input.searchSources(toAgenticRetrievalQuery(request)), - }, + knowhereTools: + input.knowhereTools ?? + notebookKnowhereTools.createSearchOnlyRuntime({ + searchSources: input.searchSources, + }), ...(input.inspectImages ? { inspectImages: input.inspectImages } : {}), }), ) @@ -74,7 +74,7 @@ export const generateAgenticOutputManifestEffect = ( citationCount: result.manifest.citations.length, artifactCount: result.manifest.artifacts.length, unresolvedCount: result.manifest.unresolved.length, - validationErrorCount: result.trace.validationErrors.length, + finalized: result.trace.finalized, intentTask: result.trace.intent?.task ?? null, carryHistory: result.trace.contextPolicy?.carryHistory ?? null, }) @@ -124,36 +124,6 @@ function getCitationLabels( .filter((label) => label.length > 0) } -function toAgenticRetrievalQuery( - request: HarnessRetrievalRequest, -): AgenticRetrievalQuery { - return { - query: request.query, - targetContent: toAgenticRetrievalTargetContent(request.modalities), - purpose: request.purpose, - topK: request.topK, - signalPaths: request.signalPaths, - filterMode: request.filterMode, - threshold: request.threshold, - } -} - -function toAgenticRetrievalTargetContent( - modalities: readonly TargetModality[], -): AgenticRetrievalTargetContent { - const requestedModalities = new Set(modalities) - if (requestedModalities.has("image") && requestedModalities.has("text")) { - return "text_image" - } - if (requestedModalities.has("table") && requestedModalities.has("text")) { - return "text_table" - } - if (requestedModalities.has("image")) return "image" - if (requestedModalities.has("table")) return "table" - if (requestedModalities.has("text")) return "text" - return "all" -} - function formatSourceContext( sources: readonly Source[], excludedSourceIds: readonly string[], diff --git a/src/domains/chat/route-answer.ts b/src/domains/chat/route-answer.ts index 2f8ae6b..3f9fc9d 100644 --- a/src/domains/chat/route-answer.ts +++ b/src/domains/chat/route-answer.ts @@ -27,6 +27,7 @@ import type { Source } from "@/infrastructure/db/schema" import { CHAT_MODEL } from "@/lib/ai" import type { HardenChatAssetUrl } from "./media-assets" import { isAuthError } from "@/integrations/dashboard/api-key-service" +import { makeKnowhereClientWithParsedStorage } from "@/integrations/knowhere" import { summarizeUnknownError } from "@/lib/format-log-value" import { logger } from "@/lib/logger" import { routeResult, type RouteResult } from "@/lib/route-result" @@ -98,6 +99,9 @@ const answerChatEffect = (input: AnswerChatInput) => const parsedStorage = new BlobParsedDocumentStorage({ workspaceId: workspace.id, }) + const knowhereResources = makeKnowhereClientWithParsedStorage(apiKey, { + workspaceId: workspace.id, + }) const hardenChatAssetUrl: HardenChatAssetUrl = async ({ source, sourcePath, @@ -140,6 +144,8 @@ const answerChatEffect = (input: AnswerChatInput) => threadId: body.value.threadId, excludedSourceIds: body.value.excludedSourceIds, retrieval: client.retrieval, + knowledge: knowhereResources.knowledge, + remoteDocumentClient: client, generateAnswer: generateAgenticOutputManifest, hardenChatAssetUrl, hardenMediaAssetUrls: ({ results, artifacts }) => diff --git a/src/domains/chat/route-service.test.ts b/src/domains/chat/route-service.test.ts index c771f12..d66ae68 100644 --- a/src/domains/chat/route-service.test.ts +++ b/src/domains/chat/route-service.test.ts @@ -183,7 +183,7 @@ describe("chat route services", () => { }) const rawUrl = "https://knowhere-storage.example/results/job_1/pages/page-1.png" const durableUrl = - "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_legacy/job_1/assets/pages/page-1.png" + "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_legacy/job_1/pages/page-1.png" mocks.parsedStorageWriteAsset.mockResolvedValue({ url: durableUrl }) mocks.getAuthenticatedWithClient.mockResolvedValue({ user: { id: "user_1" }, @@ -249,7 +249,7 @@ describe("chat route services", () => { const rawUrl = "https://knowhere-storage.example/results/job_1/images/id-front.png?AWSAccessKeyId=test" const durableUrl = - "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_identity/job_1/assets/images/id-front.png" + "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_identity/job_1/images/id-front.png" mocks.parsedStorageWriteAsset.mockResolvedValue({ url: durableUrl }) mocks.generateText.mockResolvedValue({ text: `The card number is visible. ${durableUrl} ${rawUrl}`, @@ -364,7 +364,7 @@ describe("chat route services", () => { knowhereJobId: "job_1", }) const durableUrl = - "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_contract/job_1/assets/page_citation_assets/page-8.png" + "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_contract/job_1/page_citation_assets/page-8.png" mocks.parsedStorageGetAssetUrl.mockResolvedValue(durableUrl) mocks.generateText.mockResolvedValue({ text: "The page states 5000 yuan per occurrence.", @@ -467,7 +467,7 @@ describe("chat route services", () => { const rawUrl = "https://knowhere-storage.example/results/job_remote/page_citation_assets/page-8.png?AWSAccessKeyId=test" const durableUrl = - "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_remote/job_remote/assets/page_citation_assets/page-8.png" + "https://fake.public.blob.vercel-storage.com/workspaces/workspace_1/parsed-documents/doc_remote/job_remote/page_citation_assets/page-8.png" mocks.parsedStorageWriteAsset.mockResolvedValue({ url: durableUrl }) mocks.generateText.mockResolvedValue({ text: "The page states 5000 yuan per occurrence.", diff --git a/src/domains/chat/service.test.ts b/src/domains/chat/service.test.ts index 7c0ce92..3a1ab82 100644 --- a/src/domains/chat/service.test.ts +++ b/src/domains/chat/service.test.ts @@ -48,7 +48,7 @@ describe("handleChatTurn", () => { { role: "assistant", content: "Grounded answer.", - citations: [makeRetrievalResult()], + citations: undefined, }, ], }); @@ -57,7 +57,7 @@ describe("handleChatTurn", () => { namespace: "notebook-namespace", query: "What does the document say?", topK: 8, - useAgentic: true, + useAgentic: false, dataType: 1, excludeDocumentIds: ["doc_excluded"], }); @@ -67,6 +67,7 @@ describe("handleChatTurn", () => { sources, excludedSourceIds: ["source_excluded"], searchSources: expect.any(Function), + knowhereTools: expect.any(Object), }); expect(repository.appendMessageToThread).toHaveBeenNthCalledWith(1, "workspace_1", { threadId: "thread_1", @@ -77,7 +78,7 @@ describe("handleChatTurn", () => { threadId: "thread_1", role: "assistant", content: "Grounded answer.", - citations: [makeRetrievalResult()], + citations: [], artifacts: [], }); }); @@ -216,12 +217,13 @@ describe("handleChatTurn", () => { sources, excludedSourceIds: [], searchSources: expect.any(Function), + knowhereTools: expect.any(Object), }); expect(retrieval.query).toHaveBeenCalledWith({ namespace: "notebook-namespace", query: "Tesla Q4 2025 Update energy generation and storage deployments", topK: 8, - useAgentic: true, + useAgentic: false, dataType: 1, }); }); diff --git a/src/domains/chat/service.ts b/src/domains/chat/service.ts index bfa7e42..7a7b2b5 100644 --- a/src/domains/chat/service.ts +++ b/src/domains/chat/service.ts @@ -66,6 +66,8 @@ type ChatTurnInput = { threadId?: string excludedSourceIds: readonly string[] retrieval: RetrievalClient + knowledge?: AnswerQuestionInput["knowledge"] + remoteDocumentClient?: AnswerQuestionInput["remoteDocumentClient"] generateAnswer: GenerateAnswer hardenChatAssetUrl?: AnswerQuestionInput["hardenChatAssetUrl"] hardenMediaAssetUrls?: AnswerQuestionInput["hardenMediaAssetUrls"] @@ -126,6 +128,8 @@ export const handleChatTurnEffect = (input: ChatTurnInput) => sources: readySources, excludedSourceIds: input.excludedSourceIds, retrieval: input.retrieval, + knowledge: input.knowledge, + remoteDocumentClient: input.remoteDocumentClient, generateAnswer: input.generateAnswer, hardenChatAssetUrl: input.hardenChatAssetUrl, hardenMediaAssetUrls: input.hardenMediaAssetUrls, diff --git a/src/domains/chat/types.ts b/src/domains/chat/types.ts index a0469e0..e430e4c 100644 --- a/src/domains/chat/types.ts +++ b/src/domains/chat/types.ts @@ -7,6 +7,7 @@ export type RetrievalResultView = { readonly score: number | null readonly assetUrl?: string readonly pageCitationAssetUrl?: string + readonly pageCitationPageNumber?: number readonly source: { readonly documentId?: string | null readonly sourceFileName?: string | null diff --git a/src/domains/chat/view.ts b/src/domains/chat/view.ts index ed39295..f5c7f3a 100644 --- a/src/domains/chat/view.ts +++ b/src/domains/chat/view.ts @@ -52,6 +52,7 @@ function toPersistedCitationViews(value: unknown): ChatCitationView[] | undefine score: getNumber(item.score) ?? 0, assetUrl: getString(item.assetUrl), pageCitationAssetUrl: getString(item.pageCitationAssetUrl), + pageCitationPageNumber: getNumber(item.pageCitationPageNumber), description: getString(item.description), source: { documentId: getString(item.source.documentId), @@ -85,6 +86,9 @@ function toPersistedArtifactViews(value: unknown): ChatArtifactView[] | undefine pageCitationAssetUrl: getString( item.citation.pageCitationAssetUrl, ), + pageCitationPageNumber: getNumber( + item.citation.pageCitationPageNumber, + ), description: getString(item.citation.description), source: { documentId: getString(item.citation.source.documentId), diff --git a/src/domains/chunks/index.test.ts b/src/domains/chunks/index.test.ts index 316bfd6..586d32b 100644 --- a/src/domains/chunks/index.test.ts +++ b/src/domains/chunks/index.test.ts @@ -169,6 +169,42 @@ describe("toParsedChunkView", () => { }); }); + it("maps usable page citation assets on page chunks", () => { + const chunk = makeDocumentChunk({ + id: "document_page_1", + chunkId: "parser_page_1", + chunkType: "page" as DocumentChunk["chunkType"], + metadata: { + pageAssets: [ + { + pageNum: 4, + assetUrl: "https://assets.example/page-4.png", + contentType: "image/png", + width: 1200, + height: 1600, + }, + { + pageNum: 5, + assetUrl: " ", + contentType: "image/png", + }, + ], + }, + }); + + expect(toParsedChunkView(chunk, "manual.pdf", "doc_123")).toMatchObject({ + pageAssets: [ + { + pageNumber: 4, + assetUrl: "https://assets.example/page-4.png", + contentType: "image/png", + width: 1200, + height: 1600, + }, + ], + }); + }); + it("maps SDK-normalized page number metadata", () => { const chunk = makeDocumentChunk({ metadata: { diff --git a/src/domains/chunks/index.ts b/src/domains/chunks/index.ts index eca1830..e78293e 100644 --- a/src/domains/chunks/index.ts +++ b/src/domains/chunks/index.ts @@ -17,6 +17,7 @@ export type ChunkKnowhereClient = { params: { page: number pageSize: number + chunkType?: "text" | "image" | "table" | "page" includeAssetUrls: boolean }, ): Promise<{ diff --git a/src/domains/chunks/normalization.ts b/src/domains/chunks/normalization.ts index c6590b4..3ab6bae 100644 --- a/src/domains/chunks/normalization.ts +++ b/src/domains/chunks/normalization.ts @@ -35,6 +35,10 @@ function createParsedChunkView( metadata: input.metadata, }) const summary = getStringMetadata(input.metadata, "summary") + const pageAssets = + type === "page" + ? getPageAssetViews(input.metadata["pageAssets"], assetUrl) + : undefined return { chunkId: input.chunkId, @@ -47,6 +51,7 @@ function createParsedChunkView( readableContent: getReadableContent({ type, content, summary }), filePath, assetUrl, + pageAssets, summary, keywords: getStringArrayMetadata(input.metadata, "keywords"), pageNums: getPageNumbers( @@ -58,6 +63,38 @@ function createParsedChunkView( } } +function getPageAssetViews( + value: unknown, + fallbackAssetUrl?: string, +): ParsedChunkView["pageAssets"] | undefined { + if (!Array.isArray(value)) return undefined + + const pageAssets = value.flatMap((item): NonNullable => { + if (!isRecord(item)) return [] + const pageNumber = + getPositiveInteger(item["pageNum"]) ?? getPositiveInteger(item["pageNumber"]) + const assetUrl = getString(item["assetUrl"]) ?? fallbackAssetUrl + const contentType = getString(item["contentType"]) + if (!pageNumber || !assetUrl || !contentType) return [] + + return [ + { + pageNumber, + assetUrl, + contentType, + ...(getPositiveNumber(item["width"]) !== undefined + ? { width: getPositiveNumber(item["width"]) } + : {}), + ...(getPositiveNumber(item["height"]) !== undefined + ? { height: getPositiveNumber(item["height"]) } + : {}), + }, + ] + }) + + return pageAssets.length > 0 ? pageAssets : undefined +} + function resolveConnectionTargets( chunks: readonly ParsedChunkView[], ): ParsedChunkView[] { @@ -281,6 +318,20 @@ function getPageNumbers(value: unknown): number[] | undefined { return uniquePageNumbers.length > 0 ? uniquePageNumbers : undefined } +function getPositiveInteger(value: unknown): number | undefined { + return typeof value === "number" && + Number.isSafeInteger(value) && + value > 0 + ? value + : undefined +} + +function getPositiveNumber(value: unknown): number | undefined { + return typeof value === "number" && Number.isFinite(value) && value > 0 + ? value + : undefined +} + function isRecord(value: unknown): value is Readonly> { return typeof value === "object" && value !== null } diff --git a/src/domains/chunks/read.test.ts b/src/domains/chunks/read.test.ts index e3c1000..f63c897 100644 --- a/src/domains/chunks/read.test.ts +++ b/src/domains/chunks/read.test.ts @@ -1,5 +1,5 @@ import { describe, expect, it, vi } from "vitest" -import type { Knowledge } from "@ontos-ai/knowhere-sdk" +import type { DocumentChunk, Knowledge } from "@ontos-ai/knowhere-sdk" import { readAllSourceChunks, readSourceChunkPage } from "./read" @@ -20,9 +20,13 @@ function makeReadChunk(overrides: Record = {}) { } describe("readSourceChunkPage", () => { - it("reads a page with durable asset URLs and maps chunks to the view model", async () => { + it("uses parsed-storage chunks without durable asset hardening", async () => { + const listChunks = vi.fn() const readChunks = vi.fn(async () => ({ - document: { localDocumentId: "doc_1" }, + document: { + localDocumentId: "doc_1", + resultDirectoryPath: "parsed-storage:doc_1", + }, chunks: [ makeReadChunk({ chunkType: "image", @@ -38,6 +42,7 @@ describe("readSourceChunkPage", () => { const knowledge = { readChunks } as unknown as Knowledge const result = await readSourceChunkPage({ + client: { documents: { listChunks } }, knowledge, source: { documentId: "doc_1", title: "notes.pdf", revisionKey: "rev_1" }, params: { page: 2, pageSize: 50 }, @@ -48,8 +53,8 @@ describe("readSourceChunkPage", () => { revisionKey: "rev_1", page: 2, pageSize: 50, - assetUrlPolicy: "durable", }) + expect(listChunks).not.toHaveBeenCalled() expect(result.pagination).toEqual({ page: 2, pageSize: 50, @@ -64,9 +69,100 @@ describe("readSourceChunkPage", () => { }) }) + it("falls back to Knowhere asset URLs when the parsed-storage probe reads remote chunks", async () => { + const readChunks = vi.fn(async () => ({ + document: { + localDocumentId: "doc_1", + resultDirectoryPath: "remote:doc_1", + }, + chunks: [makeReadChunk({ chunkType: "image", filePath: "images/a.png" })], + page: 1, + pageSize: 50, + totalChunks: 1, + totalPages: 1, + })) + const listChunks = vi.fn(async () => ({ + chunks: [ + makeRemoteDocumentChunk({ + id: "document_chunk_1", + chunkId: "parser_1", + chunkType: "image", + content: "Body", + sectionId: null, + sectionPath: "Summary", + sourceChunkPath: "Summary", + filePath: "images/a.png", + sortOrder: 0, + metadata: {}, + assetUrl: "https://knowhere.example/assets/a.png", + }), + ], + pagination: { page: 1, pageSize: 50, total: 1, totalPages: 1 }, + })) + const knowledge = { readChunks } as unknown as Knowledge + + const result = await readSourceChunkPage({ + client: { documents: { listChunks } }, + knowledge, + source: { documentId: "doc_1", title: "notes.pdf", revisionKey: "rev_1" }, + params: { page: 1, pageSize: 50 }, + }) + + expect(readChunks).toHaveBeenCalledWith({ + documentId: "doc_1", + revisionKey: "rev_1", + page: 1, + pageSize: 50, + }) + expect(listChunks).toHaveBeenCalledWith("doc_1", { + page: 1, + pageSize: 50, + includeAssetUrls: true, + }) + expect(result.chunks[0]?.assetUrl).toBe( + "https://knowhere.example/assets/a.png", + ) + }) + + it("uses SDK remote chunks directly when the SDK returns usable asset URLs", async () => { + const listChunks = vi.fn() + const readChunks = vi.fn(async () => ({ + document: { + localDocumentId: "doc_1", + resultDirectoryPath: "remote:doc_1", + }, + chunks: [ + makeReadChunk({ + chunkType: "image", + filePath: "images/a.png", + assetUrl: "https://sdk.example/assets/a.png", + }), + ], + page: 1, + pageSize: 50, + totalChunks: 1, + totalPages: 1, + })) + const knowledge = { readChunks } as unknown as Knowledge + + const result = await readSourceChunkPage({ + client: { documents: { listChunks } }, + knowledge, + source: { documentId: "doc_1", title: "notes.pdf", revisionKey: "rev_1" }, + params: { page: 1, pageSize: 50 }, + }) + + expect(listChunks).not.toHaveBeenCalled() + expect(result.chunks[0]?.assetUrl).toBe("https://sdk.example/assets/a.png") + }) + it("omits revisionKey when the source has none", async () => { + const listChunks = vi.fn() const readChunks = vi.fn(async () => ({ - document: { localDocumentId: "doc_1" }, + document: { + localDocumentId: "doc_1", + resultDirectoryPath: "parsed-storage:doc_1", + }, chunks: [], page: 1, pageSize: 50, @@ -76,6 +172,7 @@ describe("readSourceChunkPage", () => { const knowledge = { readChunks } as unknown as Knowledge await readSourceChunkPage({ + client: { documents: { listChunks } }, knowledge, source: { documentId: "doc_1", title: "notes.pdf", revisionKey: null }, params: { page: 1, pageSize: 50 }, @@ -85,17 +182,39 @@ describe("readSourceChunkPage", () => { documentId: "doc_1", page: 1, pageSize: 50, - assetUrlPolicy: "durable", }) }) }) +function makeRemoteDocumentChunk( + overrides: Partial = {}, +): DocumentChunk { + return { + id: "document_chunk_1", + chunkId: "parser_1", + chunkType: "text", + content: "Body", + sectionId: null, + sectionPath: "Summary", + sourceChunkPath: "Summary", + filePath: null, + sortOrder: 0, + metadata: {}, + assetUrl: null, + ...overrides, + } +} + describe("readAllSourceChunks", () => { it("pages the SDK to exhaustion", async () => { + const listChunks = vi.fn() const readChunks = vi .fn() .mockResolvedValueOnce({ - document: { localDocumentId: "doc_1" }, + document: { + localDocumentId: "doc_1", + resultDirectoryPath: "parsed-storage:doc_1", + }, chunks: [makeReadChunk({ chunkId: "c1" })], page: 1, pageSize: 200, @@ -103,7 +222,10 @@ describe("readAllSourceChunks", () => { totalPages: 2, }) .mockResolvedValueOnce({ - document: { localDocumentId: "doc_1" }, + document: { + localDocumentId: "doc_1", + resultDirectoryPath: "parsed-storage:doc_1", + }, chunks: [makeReadChunk({ chunkId: "c2" })], page: 2, pageSize: 200, @@ -113,6 +235,7 @@ describe("readAllSourceChunks", () => { const knowledge = { readChunks } as unknown as Knowledge const chunks = await readAllSourceChunks({ + client: { documents: { listChunks } }, knowledge, source: { documentId: "doc_1", title: "notes.pdf", revisionKey: "rev_1" }, }) @@ -123,8 +246,8 @@ describe("readAllSourceChunks", () => { revisionKey: "rev_1", page: 1, pageSize: 200, - assetUrlPolicy: "durable", }) + expect(listChunks).not.toHaveBeenCalled() expect(chunks.map((chunk) => chunk.parserChunkId)).toEqual(["c1", "c2"]) }) }) diff --git a/src/domains/chunks/read.ts b/src/domains/chunks/read.ts index 1bf4790..ae474cf 100644 --- a/src/domains/chunks/read.ts +++ b/src/domains/chunks/read.ts @@ -1,25 +1,61 @@ import "server-only" -import type { Knowledge } from "@ontos-ai/knowhere-sdk" +import type { + DocumentChunk, + Knowledge, + KnowledgeReadChunk, + KnowledgeReadResponse, +} from "@ontos-ai/knowhere-sdk" -import { toParsedChunkViewFromReadChunk, type ChunkPage, type ChunkPageParams } from "@/domains/chunks" +import { + toParsedChunkView, + toParsedChunkViewFromReadChunk, + type ChunkPage, + type ChunkPageParams, +} from "@/domains/chunks" import type { ParsedChunkView } from "@/domains/chunks/types" const loadAllPageSize = 200 +type DisplayReadChunkType = "text" | "image" | "table" | "page" + type ReadableSource = { readonly documentId: string readonly title: string readonly revisionKey?: string | null } +type DisplayReadClient = { + readonly documents: { + listChunks( + documentId: string, + params: { + readonly page: number + readonly pageSize: number + readonly chunkType?: DisplayReadChunkType + readonly includeAssetUrls: true + }, + ): Promise<{ + readonly chunks: readonly DocumentChunk[] + readonly pagination?: { + readonly page?: number + readonly pageSize?: number + readonly total?: number + readonly totalPages?: number + } + }> + } +} + /** * Read a single display page of parsed chunks through the SDK. The SDK serves - * from configured Blob storage when fresh and falls back to Knowhere remote - * otherwise. Display reads request durable asset URLs so media chunks remain - * usable whether the page came from Blob or the remote Knowhere fallback. + * from configured Blob storage when fresh and falls back to Knowhere remote. + * If the probe falls through to remote without asset URLs, fetch the same page + * from Knowhere with short-lived asset URLs instead of hardening assets into + * Notebook storage. */ export async function readSourceChunkPage(input: { + readonly client: DisplayReadClient readonly knowledge: Knowledge readonly source: ReadableSource readonly params: ChunkPageParams @@ -29,22 +65,59 @@ export async function readSourceChunkPage(input: { ...(input.source.revisionKey ? { revisionKey: input.source.revisionKey } : {}), page: input.params.page, pageSize: input.params.pageSize, - assetUrlPolicy: "durable", }) + if (shouldUseKnowledgeChunkResponse(response)) { + return toChunkPageFromKnowledgeResponse(response, input.source, input.params) + } + + const remoteResponse = await input.client.documents.listChunks( + input.source.documentId, + { + page: input.params.page, + pageSize: input.params.pageSize, + includeAssetUrls: true, + }, + ) + const chunks = remoteResponse.chunks.map((chunk) => + toParsedChunkView( + chunk, + input.source.title, + input.source.documentId, + ), + ) + + return { + chunks, + pagination: { + page: remoteResponse.pagination?.page ?? input.params.page, + pageSize: remoteResponse.pagination?.pageSize ?? input.params.pageSize, + total: remoteResponse.pagination?.total ?? chunks.length, + totalPages: + remoteResponse.pagination?.totalPages ?? + Math.max(1, Math.ceil(chunks.length / input.params.pageSize)), + }, + } +} + +function toChunkPageFromKnowledgeResponse( + response: KnowledgeReadResponse, + source: ReadableSource, + params: ChunkPageParams, +): ChunkPage { const chunks = response.chunks.map((chunk) => - toParsedChunkViewFromReadChunk(chunk, input.source.title, input.source.documentId), + toParsedChunkViewFromReadChunk(chunk, source.title, source.documentId), ) return { chunks, pagination: { - page: response.page ?? input.params.page, - pageSize: response.pageSize ?? input.params.pageSize, + page: response.page ?? params.page, + pageSize: response.pageSize ?? params.pageSize, total: response.totalChunks ?? chunks.length, totalPages: response.totalPages ?? - Math.max(1, Math.ceil(chunks.length / input.params.pageSize)), + Math.max(1, Math.ceil(chunks.length / params.pageSize)), }, } } @@ -54,6 +127,7 @@ export async function readSourceChunkPage(input: { * the tree view and load-all display mode. */ export async function readAllSourceChunks(input: { + readonly client: DisplayReadClient readonly knowledge: Knowledge readonly source: ReadableSource }): Promise { @@ -69,20 +143,85 @@ export async function readAllSourceChunks(input: { : {}), page, pageSize: loadAllPageSize, - assetUrlPolicy: "durable", }) - for (const chunk of response.chunks) { + if (shouldUseKnowledgeChunkResponse(response)) { + for (const chunk of response.chunks) { + chunks.push( + toParsedChunkViewFromReadChunk( + chunk, + input.source.title, + input.source.documentId, + ), + ) + } + totalPages = Math.max(1, response.totalPages ?? 1) + page += 1 + continue + } + + const remoteResponse = await input.client.documents.listChunks( + input.source.documentId, + { + page, + pageSize: loadAllPageSize, + includeAssetUrls: true, + }, + ) + for (const chunk of remoteResponse.chunks) { chunks.push( - toParsedChunkViewFromReadChunk( + toParsedChunkView( chunk, input.source.title, input.source.documentId, ), ) } - totalPages = Math.max(1, response.totalPages ?? 1) + totalPages = Math.max(1, remoteResponse.pagination?.totalPages ?? 1) page += 1 } while (page <= totalPages) return chunks } + +function shouldUseKnowledgeChunkResponse( + response: KnowledgeReadResponse, +): boolean { + return isParsedStorageReadResponse(response) || hasUsableReadChunkAssetUrl(response.chunks) +} + +function isParsedStorageReadResponse(response: KnowledgeReadResponse): boolean { + const resultDirectoryPath = response.document.resultDirectoryPath + return ( + typeof resultDirectoryPath === "string" && + resultDirectoryPath.startsWith("parsed-storage:") + ) +} + +function hasUsableReadChunkAssetUrl( + chunks: readonly KnowledgeReadChunk[], +): boolean { + return chunks.some( + (chunk) => + hasNonEmptyString(chunk.assetUrl) || + hasMetadataPageAssetUrl(chunk.metadata), + ) +} + +function hasMetadataPageAssetUrl( + metadata: Readonly>, +): boolean { + const value = metadata.pageAssets + if (!Array.isArray(value)) return false + + return value.some( + (item) => isRecord(item) && hasNonEmptyString(item.assetUrl), + ) +} + +function hasNonEmptyString(value: unknown): value is string { + return typeof value === "string" && value.trim().length > 0 +} + +function isRecord(value: unknown): value is Readonly> { + return typeof value === "object" && value !== null +} diff --git a/src/domains/chunks/types.ts b/src/domains/chunks/types.ts index 2e84189..fedb17a 100644 --- a/src/domains/chunks/types.ts +++ b/src/domains/chunks/types.ts @@ -1,3 +1,5 @@ +import type { SourcePageAssetView } from "@/domains/sources/types" + export type ChunkType = "text" | "image" | "table" | "page" export type ParsedChunkConnection = { @@ -30,6 +32,7 @@ export type ParsedChunkView = { readonly filePath?: string /** Public Blob URL for parsed media/table artifacts when Notebook stored it. */ readonly assetUrl?: string + readonly pageAssets?: readonly SourcePageAssetView[] readonly summary?: string readonly keywords?: readonly string[] readonly pageNums?: readonly number[] diff --git a/src/domains/sources/counts.test.ts b/src/domains/sources/counts.test.ts index 22f0b8c..72d132f 100644 --- a/src/domains/sources/counts.test.ts +++ b/src/domains/sources/counts.test.ts @@ -2,6 +2,7 @@ import { describe, expect, it, vi } from "vitest" import { Effect } from "effect" import type Knowhere from "@ontos-ai/knowhere-sdk" +import type { Knowledge } from "@ontos-ai/knowhere-sdk" import type { Source } from "@/infrastructure/db/schema" @@ -32,8 +33,13 @@ function makeSource(overrides: Partial = {}): Source { describe("countChunksBySourceId", () => { it("counts ready source chunks from the document total", async () => { const listChunks = vi.fn(async () => ({ pagination: { total: 12 } })) + const readChunks = vi.fn(async () => ({ + chunks: [], + totalChunks: 0, + })) const mockClient = { documents: { listChunks }, + knowledge: { readChunks }, } as unknown as Knowhere const { countChunksBySourceId } = await import("./counts") @@ -54,6 +60,14 @@ describe("countChunksBySourceId", () => { ) expect(listChunks).toHaveBeenCalledTimes(1) + expect(readChunks).toHaveBeenCalledWith({ + documentId: "doc_ready", + revisionKey: "job_1", + chunkType: "page", + page: 1, + pageSize: 1, + assetUrlPolicy: "durable", + }) expect(listChunks).toHaveBeenCalledWith("doc_ready", { page: 1, pageSize: 1, @@ -65,8 +79,10 @@ describe("countChunksBySourceId", () => { const listChunks = vi.fn(async () => { throw new Error("temporary outage") }) + const readChunks = vi.fn(async () => ({ chunks: [], totalChunks: 0 })) const mockClient = { documents: { listChunks }, + knowledge: { readChunks }, } as unknown as Knowhere const { countChunksBySourceId } = await import("./counts") @@ -86,8 +102,10 @@ describe("countChunksBySourceId", () => { const listChunks = vi.fn().mockResolvedValue({ pagination: { total: 70 }, }) + const readChunks = vi.fn(async () => ({ chunks: [], totalChunks: 0 })) const mockClient = { documents: { listChunks }, + knowledge: { readChunks }, } as unknown as Knowhere const { countChunksBySourceId } = await import("./counts") @@ -106,6 +124,212 @@ describe("countChunksBySourceId", () => { ) expect(listChunks).not.toHaveBeenCalled() + expect(readChunks).not.toHaveBeenCalled() expect(counts.size).toBe(0) }) }) + +describe("sourceViewOptionsBySourceId", () => { + it("detects page count from many page assets in a single SDK page chunk", async () => { + const listChunks = vi.fn(async () => ({ pagination: { total: 12 } })) + const readChunks = vi.fn(async () => ({ + chunks: [ + { + chunkId: "page_bundle", + chunkType: "page", + metadata: { + pageAssets: Array.from({ length: 20 }, (_, index) => ({ + pageNum: index + 1, + artifactRef: `pages/page-${String(index + 1).padStart(6, "0")}.png`, + assetUrl: `https://assets.example/page-${index + 1}.png`, + })), + }, + }, + ], + totalChunks: 1, + })) + const mockClient = { + documents: { listChunks }, + knowledge: { readChunks }, + } as unknown as Knowhere + + const { sourceViewOptionsBySourceId } = await import("./counts") + + const options = await Effect.runPromise( + sourceViewOptionsBySourceId( + [makeSource({ id: "ready", knowhereDocumentId: "doc_ready" })], + mockClient, + ), + ) + + expect(options.get("ready")).toEqual({ + chunkCount: 20, + documentPresentation: { kind: "page-assets", pageCount: 20 }, + }) + expect(listChunks).not.toHaveBeenCalled() + }) + + it("detects page-asset documents from SDK page chunks", async () => { + const listChunks = vi.fn(async () => ({ pagination: { total: 12 } })) + const readChunks = vi.fn(async () => ({ + chunks: [ + { + chunkId: "page_1", + chunkType: "page", + metadata: { + pageAssets: [ + { + pageNum: 1, + artifactRef: "pages/page-000001.png", + assetUrl: "https://assets.example/page-000001.png", + }, + ], + }, + }, + ], + totalChunks: 4, + })) + const mockClient = { + documents: { listChunks }, + knowledge: { readChunks }, + } as unknown as Knowhere + + const { sourceViewOptionsBySourceId } = await import("./counts") + + const options = await Effect.runPromise( + sourceViewOptionsBySourceId( + [makeSource({ id: "ready", knowhereDocumentId: "doc_ready" })], + mockClient, + ), + ) + + expect(options.get("ready")).toEqual({ + chunkCount: 4, + documentPresentation: { kind: "page-assets", pageCount: 4 }, + }) + expect(listChunks).not.toHaveBeenCalled() + }) + + it("uses a source-specific knowledge reader for presentation detection", async () => { + const defaultReadChunks = vi.fn(async () => ({ chunks: [], totalChunks: 0 })) + const parsedStorageReadChunks = vi.fn(async () => ({ + chunks: [ + { + chunkId: "page_1", + chunkType: "page", + metadata: { + pageAssets: [ + { + pageNum: 1, + artifactRef: "pages/page-000001.png", + assetUrl: "https://assets.example/page-000001.png", + }, + ], + }, + }, + ], + totalChunks: 1, + })) + const listChunks = vi.fn(async () => ({ pagination: { total: 12 } })) + const mockClient = { + documents: { listChunks }, + knowledge: { readChunks: defaultReadChunks }, + } as unknown as Knowhere + + const { sourceViewOptionsBySourceId } = await import("./counts") + + const options = await Effect.runPromise( + sourceViewOptionsBySourceId( + [makeSource({ id: "ready", knowhereDocumentId: "doc_ready" })], + mockClient, + { + getKnowledgeForSource: () => + ({ readChunks: parsedStorageReadChunks }) as unknown as Knowledge, + }, + ), + ) + + expect(options.get("ready")).toEqual({ + chunkCount: 1, + documentPresentation: { kind: "page-assets", pageCount: 1 }, + }) + expect(parsedStorageReadChunks).toHaveBeenCalledWith({ + documentId: "doc_ready", + revisionKey: "job_1", + chunkType: "page", + page: 1, + pageSize: 1, + assetUrlPolicy: "durable", + }) + expect(defaultReadChunks).not.toHaveBeenCalled() + }) + + it("falls back to chunk counts when page presentation detection fails", async () => { + const listChunks = vi.fn(async () => ({ pagination: { total: 12 } })) + const readChunks = vi.fn(async () => { + throw new Error("parsed storage and remote unavailable") + }) + const mockClient = { + documents: { listChunks }, + knowledge: { readChunks }, + } as unknown as Knowhere + + const { sourceViewOptionsBySourceId } = await import("./counts") + + const options = await Effect.runPromise( + sourceViewOptionsBySourceId( + [makeSource({ id: "ready", knowhereDocumentId: "doc_ready" })], + mockClient, + ), + ) + + expect(options.get("ready")).toEqual({ chunkCount: 12 }) + expect(listChunks).toHaveBeenCalledWith("doc_ready", { + page: 1, + pageSize: 1, + }) + }) + + it("skips page presentation reads when detection is disabled", async () => { + const listChunks = vi.fn(async () => ({ pagination: { total: 12 } })) + const readChunks = vi.fn(async () => ({ + chunks: [ + { + chunkId: "page_1", + chunkType: "page", + metadata: { + pageAssets: [ + { + pageNum: 1, + artifactRef: "pages/page-000001.png", + assetUrl: "https://assets.example/page-000001.png", + }, + ], + }, + }, + ], + totalChunks: 1, + })) + const mockClient = { + documents: { listChunks }, + knowledge: { readChunks }, + } as unknown as Knowhere + + const { sourceViewOptionsBySourceId } = await import("./counts") + + const options = await Effect.runPromise( + sourceViewOptionsBySourceId( + [makeSource({ id: "ready", knowhereDocumentId: "doc_ready" })], + mockClient, + { documentPresentationDetection: "disabled" }, + ), + ) + + expect(readChunks).not.toHaveBeenCalled() + expect(listChunks).toHaveBeenCalledWith("doc_ready", { + page: 1, + pageSize: 1, + }) + expect(options.get("ready")).toEqual({ chunkCount: 12 }) + }) +}) diff --git a/src/domains/sources/counts.ts b/src/domains/sources/counts.ts index e6e0c9b..d161adf 100644 --- a/src/domains/sources/counts.ts +++ b/src/domains/sources/counts.ts @@ -2,8 +2,15 @@ import "server-only" import { Effect } from "effect" import type Knowhere from "@ontos-ai/knowhere-sdk" +import type { Knowledge, KnowledgeReadChunk } from "@ontos-ai/knowhere-sdk" import type { Source } from "@/infrastructure/db/schema" +import type { SourceDocumentPresentation } from "./types" + +type PageAssetDocumentPresentation = Extract< + SourceDocumentPresentation, + { readonly kind: "page-assets" } +> type CountChunksClient = { readonly documents: { @@ -14,11 +21,35 @@ type CountChunksClient = { readonly pagination?: { readonly total?: number } }> } + readonly knowledge: { + readChunks(params: { + readonly documentId: string + readonly revisionKey?: string + readonly chunkType: "page" + readonly page: number + readonly pageSize: number + readonly assetUrlPolicy: "durable" + }): Promise<{ + readonly chunks: readonly KnowledgeReadChunk[] + readonly totalChunks?: number + }> + } } -export const countChunksBySourceId = ( +export type SourceViewOptionsLoadOptions = { + readonly documentPresentationDetection?: "enabled" | "disabled" + readonly getKnowledgeForSource?: (source: Source) => Knowledge +} + +export type SourceViewOptions = { + readonly chunkCount?: number + readonly documentPresentation?: SourceDocumentPresentation +} + +export const sourceViewOptionsBySourceId = ( sources: readonly Source[], client: Knowhere, + options: SourceViewOptionsLoadOptions = {}, ) => Effect.gen(function* () { const countClient = client as unknown as CountChunksClient @@ -28,15 +59,19 @@ export const countChunksBySourceId = ( source.status === "ready" && source.knowhereDocumentId, ) - if (readySources.length === 0) return new Map() + if (readySources.length === 0) return new Map() const entries = yield* Effect.all( readySources.map((source) => Effect.gen(function* () { - const total = yield* Effect.tryPromise(() => - loadSourceChunkCount(countClient, source.knowhereDocumentId!), - ).pipe(Effect.catchAll(() => Effect.succeed(undefined))) - return [source.id, total] as const + const loadedOptions = yield* Effect.tryPromise(() => + loadSourceViewOptions(countClient, source, options), + ).pipe( + Effect.catchAll(() => + Effect.sync((): SourceViewOptions | undefined => undefined), + ), + ) + return [source.id, loadedOptions] as const }), ), { concurrency: "unbounded" }, @@ -44,26 +79,87 @@ export const countChunksBySourceId = ( return new Map( entries.filter( - (entry): entry is readonly [string, number] => - typeof entry[1] === "number", + (entry): entry is readonly [string, SourceViewOptions] => + entry[1] !== undefined, ), ) }) -export const sourceViewOptionsBySourceId = ( +export const countChunksBySourceId = ( sources: readonly Source[], client: Knowhere, ) => Effect.gen(function* () { - const counts = yield* countChunksBySourceId(sources, client) - return new Map( - sources.map((source) => [ - source.id, - { chunkCount: counts.get(source.id) }, - ]), - ) + const sourceOptions = yield* sourceViewOptionsBySourceId(sources, client) + const countEntries: [string, number][] = [] + for (const [sourceId, options] of sourceOptions.entries()) { + if (typeof options.chunkCount === "number") { + countEntries.push([sourceId, options.chunkCount]) + } + } + return new Map(countEntries) }) +async function loadSourceViewOptions( + client: CountChunksClient, + source: Source, + options: SourceViewOptionsLoadOptions, +): Promise { + const documentId = source.knowhereDocumentId + if (!documentId) return undefined + + if (options.documentPresentationDetection !== "disabled") { + const pagePresentation = await loadPageAssetPresentation( + client, + source, + options, + ) + if (pagePresentation) { + return { + chunkCount: pagePresentation.pageCount, + documentPresentation: pagePresentation, + } + } + } + + const chunkCount = await loadSourceChunkCount(client, documentId) + return typeof chunkCount === "number" ? { chunkCount } : undefined +} + +async function loadPageAssetPresentation( + client: CountChunksClient, + source: Source, + options: SourceViewOptionsLoadOptions, +): Promise { + const documentId = source.knowhereDocumentId + if (!documentId) return undefined + + try { + const knowledge = options.getKnowledgeForSource?.(source) ?? client.knowledge + const response = await knowledge.readChunks({ + documentId, + ...(source.knowhereJobId ? { revisionKey: source.knowhereJobId } : {}), + chunkType: "page", + page: 1, + pageSize: 1, + assetUrlPolicy: "durable", + }) + const firstChunk = response.chunks[0] + if (!firstChunk || firstChunk.chunkType !== "page") return undefined + const maxPageAssetNumber = getMaxUsablePageAssetNumber( + firstChunk.metadata.pageAssets, + ) + if (!maxPageAssetNumber) return undefined + + const totalChunks = getPositiveFiniteNumber(response.totalChunks) ?? 0 + const pageCount = Math.max(totalChunks, maxPageAssetNumber) + + return { kind: "page-assets", pageCount } + } catch { + return undefined + } +} + async function loadSourceChunkCount( client: CountChunksClient, documentId: string, @@ -75,3 +171,33 @@ async function loadSourceChunkCount( const total = response.pagination?.total return typeof total === "number" && Number.isFinite(total) ? total : undefined } + +function getMaxUsablePageAssetNumber(value: unknown): number | undefined { + if (!Array.isArray(value)) return undefined + + const pageNumbers = value.flatMap((item): number[] => { + if (!isRecord(item)) return [] + const pageNum = item.pageNum + const artifactRef = item.artifactRef + const assetUrl = item.assetUrl + const isUsable = + typeof pageNum === "number" && + Number.isSafeInteger(pageNum) && + pageNum > 0 && + ((typeof artifactRef === "string" && artifactRef.trim().length > 0) || + (typeof assetUrl === "string" && assetUrl.trim().length > 0)) + return isUsable ? [pageNum] : [] + }) + if (pageNumbers.length === 0) return undefined + return Math.max(...pageNumbers) +} + +function getPositiveFiniteNumber(value: unknown): number | undefined { + return typeof value === "number" && Number.isFinite(value) && value > 0 + ? value + : undefined +} + +function isRecord(value: unknown): value is Readonly> { + return typeof value === "object" && value !== null +} diff --git a/src/domains/sources/display-read-unavailable.ts b/src/domains/sources/display-read-unavailable.ts new file mode 100644 index 0000000..302cd85 --- /dev/null +++ b/src/domains/sources/display-read-unavailable.ts @@ -0,0 +1,47 @@ +const sourceUnavailableMessage: string = + "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere." + +function isDisplayReadUnavailableError(error: unknown): boolean { + const details = readErrorDetails(error) + if ( + details.name === "NotFoundError" && + details.message.toLowerCase().includes("document not found") + ) { + return true + } + + const causeDetails = readErrorDetails(details.cause) + return ( + causeDetails.name === "NotFoundError" && + causeDetails.message.toLowerCase().includes("document not found") + ) +} + +function readErrorDetails(error: unknown): { + readonly name: string + readonly message: string + readonly cause: unknown +} { + if (!isRecord(error)) { + return { + name: "", + message: String(error), + cause: undefined, + } + } + + return { + name: typeof error.name === "string" ? error.name : "", + message: typeof error.message === "string" ? error.message : "", + cause: error.cause, + } +} + +function isRecord(value: unknown): value is Readonly> { + return typeof value === "object" && value !== null +} + +export const displayReadUnavailable = { + isError: isDisplayReadUnavailableError, + message: sourceUnavailableMessage, +} as const diff --git a/src/domains/sources/parsed-document-blob-storage.test.ts b/src/domains/sources/parsed-document-blob-storage.test.ts index 8a3c427..b5fb005 100644 --- a/src/domains/sources/parsed-document-blob-storage.test.ts +++ b/src/domains/sources/parsed-document-blob-storage.test.ts @@ -1,9 +1,5 @@ import { describe, expect, it } from "vitest" -import type { - KnowhereParsedSnapshotChunkPage, - KnowhereParsedSnapshotManifest, - ParsedDocumentSyncProgress, -} from "@ontos-ai/knowhere-sdk" +import type { ParsedDocumentSyncProgress } from "@ontos-ai/knowhere-sdk" import { BlobParsedDocumentStorage, @@ -34,6 +30,8 @@ function createFakeBlobStore(): { return { statusCode: 200, stream: new Response(body).body as ReadableStream, + url: toUrl(pathname), + contentType: object.contentType, } }, put: async (pathname, body, options) => { @@ -55,70 +53,85 @@ function createFakeBlobStore(): { const documentId = "doc_123" const revisionKey = "job_result_1" -const manifest: KnowhereParsedSnapshotManifest = { - version: 1, - kind: "knowhere-parsed-result-snapshot", +const manifest = { + version: "2.0", jobId: "job_1", - revisionKey, - documentId, sourceFileName: "example.pdf", - totalChunks: 2, - chunkPageSize: 200, - chunkPages: [{ page: 1, pageSize: 200, chunkCount: 2, key: "chunks/page-1.json" }], - assetUrlsByFilePath: {}, - createdAt: "2026-07-04T00:00:00.000Z", -} + statistics: { + totalChunks: 2, + textChunks: 1, + imageChunks: 1, + tableChunks: 0, + pageChunks: 0, + }, +} satisfies Record -const chunkPage: KnowhereParsedSnapshotChunkPage = { - version: 1, - jobId: "job_1", - revisionKey, - documentId, - sourceFileName: "example.pdf", - page: 1, - pageSize: 200, - total: 2, - totalPages: 1, +const chunks = { chunks: [ { - id: "c1", - chunkId: "c1", - chunkType: "text", + chunk_id: "c1", + type: "text", content: "hello", - sourceChunkPath: "example.pdf", - sortOrder: 0, + path: "example.pdf", metadata: {}, }, ], -} +} satisfies Record describe("BlobParsedDocumentStorage", () => { - it("round-trips a manifest keyed by workspace/document/revision", async () => { + it("round-trips a manifest at the result-relative manifest path", async () => { const { store, objects } = createFakeBlobStore() const storage = new BlobParsedDocumentStorage({ workspaceId: "ws_1", blobStore: store, }) - await storage.writeManifest({ documentId, revisionKey, manifest }) + await storage.writeObject({ + documentId, + revisionKey, + path: "manifest.json", + body: Buffer.from(JSON.stringify(manifest), "utf8"), + contentType: "application/json; charset=utf-8", + }) expect([...objects.keys()]).toContain( - "workspaces/ws_1/parsed-documents/doc_123/job_result_1/manifest/current.json", + "workspaces/ws_1/parsed-documents/doc_123/job_result_1/manifest.json", + ) + const read = await storage.readObject({ + documentId, + revisionKey, + path: "manifest.json", + }) + expect(read ? Buffer.from(read.body).toString("utf8") : null).toBe( + JSON.stringify(manifest), ) - const read = await storage.readManifest({ documentId, revisionKey }) - expect(read).toEqual(manifest) }) - it("round-trips a chunk page", async () => { - const { store } = createFakeBlobStore() + it("round-trips chunks at the result-relative chunks path", async () => { + const { store, objects } = createFakeBlobStore() const storage = new BlobParsedDocumentStorage({ workspaceId: "ws_1", blobStore: store, }) - await storage.writeChunkPage({ documentId, revisionKey, page: chunkPage }) - const read = await storage.readChunkPage({ documentId, revisionKey, page: 1 }) - expect(read).toEqual(chunkPage) + await storage.writeObject({ + documentId, + revisionKey, + path: "chunks.json", + body: Buffer.from(JSON.stringify(chunks), "utf8"), + contentType: "application/json; charset=utf-8", + }) + expect([...objects.keys()]).toContain( + "workspaces/ws_1/parsed-documents/doc_123/job_result_1/chunks.json", + ) + const read = await storage.readObject({ + documentId, + revisionKey, + path: "chunks.json", + }) + expect(read ? Buffer.from(read.body).toString("utf8") : null).toBe( + JSON.stringify(chunks), + ) }) it("round-trips sync progress", async () => { @@ -131,7 +144,6 @@ describe("BlobParsedDocumentStorage", () => { documentId, revisionKey, nextChunkPage: 3, - nextAssetIndex: 0, status: "running", updatedAt: "2026-07-04T00:00:00.000Z", } @@ -141,7 +153,7 @@ describe("BlobParsedDocumentStorage", () => { expect(read).toEqual(progress) }) - it("writes an asset and resolves a durable URL", async () => { + it("writes an asset without adding a Notebook logical assets prefix", async () => { const { store } = createFakeBlobStore() const storage = new BlobParsedDocumentStorage({ workspaceId: "ws_1", @@ -156,7 +168,7 @@ describe("BlobParsedDocumentStorage", () => { contentType: "image/png", }) expect(written.url).toContain( - "workspaces/ws_1/parsed-documents/doc_123/job_result_1/assets/images/fig-1.png", + "workspaces/ws_1/parsed-documents/doc_123/job_result_1/images/fig-1.png", ) const url = await storage.getAssetUrl({ @@ -167,16 +179,75 @@ describe("BlobParsedDocumentStorage", () => { expect(url).toBe(written.url) }) - it("returns null for a missing manifest, chunk page, progress, and asset", async () => { + it("round-trips arbitrary result-relative objects", async () => { + const { store, objects } = createFakeBlobStore() + const storage = new BlobParsedDocumentStorage({ + workspaceId: "ws_1", + blobStore: store, + }) + + await storage.writeObject({ + documentId, + revisionKey, + path: "page_citation_assets/page-1.png", + body: new Uint8Array([4, 5, 6]), + contentType: "image/png", + }) + await storage.writeObject({ + documentId, + revisionKey, + path: "images/fig-1.png", + body: new Uint8Array([7, 8, 9]), + contentType: "image/png", + }) + + expect([...objects.keys()]).toContain( + "workspaces/ws_1/parsed-documents/doc_123/job_result_1/page_citation_assets/page-1.png", + ) + expect([...objects.keys()]).toContain( + "workspaces/ws_1/parsed-documents/doc_123/job_result_1/images/fig-1.png", + ) + const read = await storage.readObject({ + documentId, + revisionKey, + path: "page_citation_assets/page-1.png", + }) + expect(read?.body).toEqual(Buffer.from([4, 5, 6])) + expect(read?.contentType).toBe("image/png") + }) + + it("resolves legacy asset-prefixed URLs without writing new assets there", async () => { + const { store, objects } = createFakeBlobStore() + const storage = new BlobParsedDocumentStorage({ + workspaceId: "ws_1", + blobStore: store, + }) + const legacyKey = + "workspaces/ws_1/parsed-documents/doc_123/job_result_1/assets/images/fig-1.png" + objects.set(legacyKey, { + body: Buffer.from([1]), + contentType: "image/png", + }) + + const url = await storage.getAssetUrl({ + documentId, + revisionKey, + sourcePath: "images/fig-1.png", + }) + expect(url).toBe( + "https://fake.public.blob.vercel-storage.com/workspaces/ws_1/parsed-documents/doc_123/job_result_1/assets/images/fig-1.png", + ) + }) + + it("returns null for a missing object, progress, and asset", async () => { const { store } = createFakeBlobStore() const storage = new BlobParsedDocumentStorage({ workspaceId: "ws_1", blobStore: store, }) - expect(await storage.readManifest({ documentId, revisionKey })).toBeNull() expect( - await storage.readChunkPage({ documentId, revisionKey, page: 9 }), + await storage.readObject({ documentId, revisionKey, path: "manifest.json" }), ).toBeNull() expect(await storage.readSyncProgress({ documentId, revisionKey })).toBeNull() expect( @@ -191,10 +262,17 @@ describe("BlobParsedDocumentStorage", () => { blobStore: store, }) - await storage.writeManifest({ documentId, revisionKey, manifest }) - const staleRead = await storage.readManifest({ + await storage.writeObject({ + documentId, + revisionKey, + path: "manifest.json", + body: Buffer.from(JSON.stringify(manifest), "utf8"), + contentType: "application/json; charset=utf-8", + }) + const staleRead = await storage.readObject({ documentId, revisionKey: "job_result_2", + path: "manifest.json", }) expect(staleRead).toBeNull() }) @@ -207,10 +285,18 @@ describe("BlobParsedDocumentStorage", () => { }) await expect( - storage.readManifest({ documentId: "../escape", revisionKey }), + storage.readObject({ + documentId: "../escape", + revisionKey, + path: "manifest.json", + }), ).rejects.toThrow(/Invalid parsed storage segment/) await expect( - storage.readManifest({ documentId, revisionKey: "a/b" }), + storage.readObject({ + documentId, + revisionKey: "a/b", + path: "manifest.json", + }), ).rejects.toThrow(/Invalid parsed storage segment/) await expect( storage.getAssetUrl({ @@ -219,5 +305,12 @@ describe("BlobParsedDocumentStorage", () => { sourcePath: "../../etc/passwd", }), ).rejects.toThrow(/Invalid parsed storage path/) + await expect( + storage.readObject({ + documentId, + revisionKey, + path: "images/../escape.png", + }), + ).rejects.toThrow(/Invalid parsed storage path/) }) }) diff --git a/src/domains/sources/parsed-document-blob-storage.ts b/src/domains/sources/parsed-document-blob-storage.ts index 23a93a0..5a05e5c 100644 --- a/src/domains/sources/parsed-document-blob-storage.ts +++ b/src/domains/sources/parsed-document-blob-storage.ts @@ -2,15 +2,15 @@ import "server-only" import { del, get, head, put, BlobNotFoundError } from "@vercel/blob" import type { - KnowhereParsedSnapshotChunkPage, - KnowhereParsedSnapshotManifest, + ParsedDocumentObject, + ParsedDocumentObjectHead, + ParsedDocumentObjectParams, + ParsedDocumentRevisionParams, ParsedDocumentStorage, - ParsedDocumentStorageAsset, - ParsedDocumentStorageAssetParams, - ParsedDocumentStorageChunkPageParams, ParsedDocumentStorageDocument, - ParsedDocumentStorageManifestParams, ParsedDocumentSyncProgress, + ParsedDocumentWriteObjectParams, + ParsedDocumentWriteObjectResult, } from "@ontos-ai/knowhere-sdk" /** @@ -26,14 +26,16 @@ import type { */ const parsedDocumentsDirectoryName = "parsed-documents" -const manifestStoragePath = "manifest/current.json" const syncProgressStoragePath = "sync-progress.json" const jsonContentType = "application/json; charset=utf-8" +const binaryContentType = "application/octet-stream" type BlobGetResult = | { readonly statusCode: 200 readonly stream: ReadableStream + readonly url: string + readonly contentType?: string } | { readonly statusCode: 304 @@ -77,8 +79,35 @@ export type BlobParsedDocumentStorageInput = { readonly blobStore?: ParsedDocumentBlobStore } +type ParsedDocumentStorageObject = ParsedDocumentStorageDocument & { + readonly path: string +} + +type LegacyParsedDocumentStorageAsset = { + readonly sourcePath: string + readonly body: string | Buffer | Uint8Array + readonly contentType: string + readonly metadata?: Readonly> +} + +type LegacyParsedDocumentStorageAssetParams = ParsedDocumentStorageDocument & { + readonly sourcePath: string +} + const vercelBlobStore: ParsedDocumentBlobStore = { - get: (pathname, options) => get(pathname, options), + get: async (pathname, options) => { + const blob = await get(pathname, options) + if (!blob) return null + if (blob.statusCode === 304) { + return { statusCode: 304, stream: null } + } + return { + statusCode: 200, + stream: blob.stream, + url: blob.blob.url, + contentType: blob.blob.contentType, + } + }, put: async (pathname, body, options) => { const blob = await put(pathname, body, { access: options.access, @@ -110,54 +139,15 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { this.blobStore = input.blobStore ?? vercelBlobStore } - async readManifest( - params: ParsedDocumentStorageManifestParams, - ): Promise { - return this.readJson( - this.getManifestKey(params.documentId, params.revisionKey), - ) - } - - async writeManifest(params: { - readonly documentId: string - readonly revisionKey: string - readonly manifest: KnowhereParsedSnapshotManifest - }): Promise { - await this.writeJson( - this.getManifestKey(params.documentId, params.revisionKey), - params.manifest, - ) - } - - async readChunkPage( - params: ParsedDocumentStorageChunkPageParams, - ): Promise { - // chunkType filtering happens SDK-side after read; storage returns the full page. - return this.readJson( - this.getChunkPageKey(params.documentId, params.revisionKey, params.page), - ) - } - - async writeChunkPage(params: { - readonly documentId: string - readonly revisionKey: string - readonly page: KnowhereParsedSnapshotChunkPage - }): Promise { - await this.writeJson( - this.getChunkPageKey( - params.documentId, - params.revisionKey, - params.page.page, - ), - params.page, - ) - } - async writeAsset( - params: ParsedDocumentStorageDocument & ParsedDocumentStorageAsset, + params: ParsedDocumentStorageDocument & LegacyParsedDocumentStorageAsset, ): Promise<{ readonly sourcePath: string; readonly url?: string }> { const blob = await this.blobStore.put( - this.getAssetKey(params.documentId, params.revisionKey, params.sourcePath), + this.getObjectKey({ + documentId: params.documentId, + revisionKey: params.revisionKey, + path: params.sourcePath, + }), Buffer.from(params.body), { access: "public", @@ -170,29 +160,109 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { } async getAssetUrl( - params: ParsedDocumentStorageAssetParams, + params: LegacyParsedDocumentStorageAssetParams, ): Promise { const result = await this.blobStore.head( - this.getAssetKey(params.documentId, params.revisionKey, params.sourcePath), + this.getObjectKey({ + documentId: params.documentId, + revisionKey: params.revisionKey, + path: params.sourcePath, + }), ) - return result?.url ?? null + if (result) return result.url + + const legacyResult = await this.blobStore.head( + this.getObjectKey({ + documentId: params.documentId, + revisionKey: params.revisionKey, + path: `assets/${params.sourcePath}`, + }), + ) + return legacyResult?.url ?? null } async readSyncProgress( - params: ParsedDocumentStorageDocument, + params: ParsedDocumentRevisionParams, ): Promise { return this.readJson( - this.getSyncProgressKey(params.documentId, params.revisionKey), + this.getObjectKey({ + documentId: params.documentId, + revisionKey: params.revisionKey, + path: syncProgressStoragePath, + }), ) } async writeSyncProgress(params: ParsedDocumentSyncProgress): Promise { await this.writeJson( - this.getSyncProgressKey(params.documentId, params.revisionKey), + this.getObjectKey({ + documentId: params.documentId, + revisionKey: params.revisionKey, + path: syncProgressStoragePath, + }), params, ) } + async readObject( + params: ParsedDocumentObjectParams, + ): Promise { + const object = await this.readBlobObject(this.getObjectKey(params)) + if (!object) return null + + return { + documentId: params.documentId, + revisionKey: params.revisionKey, + path: params.path, + body: object.body, + ...(object.contentType ? { contentType: object.contentType } : {}), + contentLength: object.body.byteLength, + url: object.url, + } + } + + async writeObject( + params: ParsedDocumentWriteObjectParams, + ): Promise { + const blob = await this.blobStore.put( + this.getObjectKey(params), + Buffer.from(params.body), + { + access: "public", + allowOverwrite: true, + contentType: params.contentType ?? binaryContentType, + multipart: true, + }, + ) + return { + documentId: params.documentId, + revisionKey: params.revisionKey, + path: params.path, + url: blob.url, + } + } + + async headObject( + params: ParsedDocumentObjectParams, + ): Promise { + const result = await this.blobStore.head(this.getObjectKey(params)) + if (!result) return null + + return { + documentId: params.documentId, + revisionKey: params.revisionKey, + path: params.path, + url: result.url, + } + } + + async getObjectUrl( + params: ParsedDocumentObjectParams, + ): Promise { + const result = await this.blobStore.head(this.getObjectKey(params)) + return result?.url ?? null + } + private getRevisionPrefix(documentId: string, revisionKey: string): string { return [ "workspaces", @@ -203,28 +273,8 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { ].join("/") } - private getManifestKey(documentId: string, revisionKey: string): string { - return `${this.getRevisionPrefix(documentId, revisionKey)}/${manifestStoragePath}` - } - - private getChunkPageKey( - documentId: string, - revisionKey: string, - page: number, - ): string { - return `${this.getRevisionPrefix(documentId, revisionKey)}/chunks/page-${page}.json` - } - - private getSyncProgressKey(documentId: string, revisionKey: string): string { - return `${this.getRevisionPrefix(documentId, revisionKey)}/${syncProgressStoragePath}` - } - - private getAssetKey( - documentId: string, - revisionKey: string, - sourcePath: string, - ): string { - return `${this.getRevisionPrefix(documentId, revisionKey)}/assets/${normalizeRelativeStoragePath(sourcePath)}` + private getObjectKey(params: ParsedDocumentStorageObject): string { + return `${this.getRevisionPrefix(params.documentId, params.revisionKey)}/${normalizeRelativeStoragePath(params.path)}` } private async readJson(key: string): Promise { @@ -246,10 +296,24 @@ export class BlobParsedDocumentStorage implements ParsedDocumentStorage { } private async readBlobText(key: string): Promise { + const object = await this.readBlobObject(key) + return object ? object.body.toString("utf8") : null + } + + private async readBlobObject(key: string): Promise<{ + readonly body: Buffer + readonly contentType?: string + readonly url?: string + } | null> { try { const result = await this.blobStore.get(key, { access: "public" }) if (!result || result.statusCode !== 200) return null - return await new Response(result.stream).text() + const body = Buffer.from(await new Response(result.stream).arrayBuffer()) + return { + body, + ...(result.contentType ? { contentType: result.contentType } : {}), + url: result.url, + } } catch (error) { if (error instanceof BlobNotFoundError) return null throw error diff --git a/src/domains/sources/route-chunks.ts b/src/domains/sources/route-chunks.ts index 7d97920..c636cb6 100644 --- a/src/domains/sources/route-chunks.ts +++ b/src/domains/sources/route-chunks.ts @@ -6,13 +6,14 @@ import { resolveChunkConnectionTargets } from "@/domains/chunks" import type { DemoChunkPage } from "@/integrations/knowhere-demo" import { logger } from "@/lib/logger" import { routeResult } from "@/lib/route-result" +import { displayReadUnavailable } from "./display-read-unavailable" import { decodeRemoteSourceId, findRemoteLibraryDocumentBySourceId, } from "./remote-library" import { getClientForWorkspace, - getKnowledgeForSource, + getKnowledgeResourcesForSource, } from "./route-dependencies" import { sourceRowRepository } from "./source-row-repository" import type { @@ -98,7 +99,7 @@ const loadSourceChunksEffect = ( const apiKey = yield* Effect.tryPromise(() => deps.ensureApiKeyForWorkspace(workspace.id, input.cookieHeader), ) - const knowledge = getKnowledgeForSource({ + const readResources = getKnowledgeResourcesForSource({ apiKey, workspaceId: workspace.id, sourceId: source.id, @@ -112,20 +113,31 @@ const loadSourceChunksEffect = ( } if (input.shouldLoadAll) { - const chunks = yield* Effect.tryPromise(() => - readAllSourceChunks({ knowledge, source: readableSource }), + return yield* Effect.tryPromise(() => + readAllSourceChunks({ + client: readResources.client, + knowledge: readResources.knowledge, + source: readableSource, + }), + ).pipe( + Effect.map((chunks) => + routeResult.ok({ chunks: resolveChunkConnectionTargets(chunks) }), + ), + Effect.catchAll((error) => recoverUnavailableChunks(input, error)), ) - return routeResult.ok({ chunks: resolveChunkConnectionTargets(chunks) }) } - const chunkPage = yield* Effect.tryPromise(() => + return yield* Effect.tryPromise(() => readSourceChunkPage({ - knowledge, + client: readResources.client, + knowledge: readResources.knowledge, source: readableSource, params: input.pageParams, }), + ).pipe( + Effect.map((chunkPage) => routeResult.ok(chunkPage)), + Effect.catchAll((error) => recoverUnavailableChunks(input, error)), ) - return routeResult.ok(chunkPage) }) const loadRemoteChunkPageEffect = ( @@ -168,7 +180,7 @@ const loadRemoteChunkPageEffect = ( ) const documentId = source.knowhereDocumentId ?? remoteDocument.documentId - const knowledge = getKnowledgeForSource({ + const readResources = getKnowledgeResourcesForSource({ apiKey, workspaceId: workspace.id, sourceId: source.id, @@ -182,20 +194,31 @@ const loadRemoteChunkPageEffect = ( } if (input.shouldLoadAll) { - const chunks = yield* Effect.tryPromise(() => - readAllSourceChunks({ knowledge, source: readableSource }), + return yield* Effect.tryPromise(() => + readAllSourceChunks({ + client: readResources.client, + knowledge: readResources.knowledge, + source: readableSource, + }), + ).pipe( + Effect.map((chunks) => + routeResult.ok({ chunks: resolveChunkConnectionTargets(chunks) }), + ), + Effect.catchAll((error) => recoverUnavailableChunks(input, error)), ) - return routeResult.ok({ chunks: resolveChunkConnectionTargets(chunks) }) } - const chunkPage = yield* Effect.tryPromise(() => + return yield* Effect.tryPromise(() => readSourceChunkPage({ - knowledge, + client: readResources.client, + knowledge: readResources.knowledge, source: readableSource, params: input.pageParams, }), + ).pipe( + Effect.map((chunkPage) => routeResult.ok(chunkPage)), + Effect.catchAll((error) => recoverUnavailableChunks(input, error)), ) - return routeResult.ok(chunkPage) }) const loadDemoChunkPageEffect = ( @@ -337,4 +360,47 @@ function sourceSnapshotProcessing( ) } +function sourceChunksUnavailable( + input: LoadSourceChunksInput, +): JsonRouteResult<{ + readonly chunks: [] + readonly pagination?: { + readonly page: number + readonly pageSize: number + readonly total: 0 + readonly totalPages: 0 + } + readonly message: string + readonly isUnavailable: true +}> { + if (input.shouldLoadAll) { + return routeResult.ok({ + chunks: [], + message: displayReadUnavailable.message, + isUnavailable: true, + }) + } + + return routeResult.ok({ + chunks: [], + pagination: { + page: input.pageParams.page, + pageSize: input.pageParams.pageSize, + total: 0, + totalPages: 0, + }, + message: displayReadUnavailable.message, + isUnavailable: true, + }) +} + +function recoverUnavailableChunks( + input: LoadSourceChunksInput, + error: unknown, +): Effect.Effect, unknown> { + return displayReadUnavailable.isError(error) + ? Effect.succeed(sourceChunksUnavailable(input)) + : Effect.fail(error) +} + export { createRouteChunks } diff --git a/src/domains/sources/route-dependencies.ts b/src/domains/sources/route-dependencies.ts index a39635f..b0cc887 100644 --- a/src/domains/sources/route-dependencies.ts +++ b/src/domains/sources/route-dependencies.ts @@ -28,10 +28,11 @@ const defaultDependencies: SourceRouteServiceDependencies = { ensureApiKeyForWorkspace, ensureWorkspace: workspaceService.ensureWorkspace, getCurrentUser, - getSourceViewOptionsBySourceId: (sources, client) => + getSourceViewOptionsBySourceId: (sources, client, options) => getDefaultSourceViewOptionsBySourceId( sources, client as ReturnType, + options, ), makeKnowhereClient: (apiKey: string) => makeDefaultKnowhereClient(apiKey) as SourceRouteKnowhereClient, @@ -98,6 +99,16 @@ function getKnowledgeForSource(input: { readonly documentId: string readonly revisionKey?: string | null }): Knowledge { + return getKnowledgeResourcesForSource(input).knowledge +} + +function getKnowledgeResourcesForSource(input: { + readonly apiKey: string + readonly workspaceId: string + readonly sourceId: string + readonly documentId: string + readonly revisionKey?: string | null +}): { readonly client: SourceRouteKnowhereClient; readonly knowledge: Knowledge } { const scheduler = createParsedDocumentSyncScheduler({ workspaceId: input.workspaceId, sourceId: input.sourceId, @@ -105,15 +116,16 @@ function getKnowledgeForSource(input: { apiKey: input.apiKey, revisionKey: input.revisionKey ?? undefined, }) - const { knowledge } = makeKnowhereClientWithParsedStorage(input.apiKey, { + const resources = makeKnowhereClientWithParsedStorage(input.apiKey, { workspaceId: input.workspaceId, scheduler, }) - return knowledge + return { client: resources.client, knowledge: resources.knowledge } } export { createSourceRouteDependencies, getClientForWorkspace, getKnowledgeForSource, + getKnowledgeResourcesForSource, } diff --git a/src/domains/sources/route-listing.ts b/src/domains/sources/route-listing.ts index 35c0c72..08c17b2 100644 --- a/src/domains/sources/route-listing.ts +++ b/src/domains/sources/route-listing.ts @@ -113,6 +113,9 @@ const listSourcesEffect = ( const sourceOptions = yield* deps.getSourceViewOptionsBySourceId( sourcesNeedingChunkCount, client, + { + documentPresentationDetection: "disabled", + }, ) const hiddenDemoSourceIds = new Set( yield* Effect.tryPromise(() => diff --git a/src/domains/sources/route-service.test.ts b/src/domains/sources/route-service.test.ts index a056427..cb9c7da 100644 --- a/src/domains/sources/route-service.test.ts +++ b/src/domains/sources/route-service.test.ts @@ -472,6 +472,9 @@ describe("source route service", () => { expect(getSourceViewOptionsBySourceId).toHaveBeenCalledWith( [source], knowhereClient, + expect.objectContaining({ + documentPresentationDetection: "disabled", + }), ); expect(result).toEqual({ status: 200, @@ -546,6 +549,9 @@ describe("source route service", () => { expect(getSourceViewOptionsBySourceId).toHaveBeenCalledWith( [], knowhereClient, + expect.objectContaining({ + documentPresentationDetection: "disabled", + }), ); expect(result).toEqual({ status: 200, @@ -628,6 +634,9 @@ describe("source route service", () => { expect(getSourceViewOptionsBySourceId).toHaveBeenCalledWith( [], knowhereClient, + expect.objectContaining({ + documentPresentationDetection: "disabled", + }), ); expect(result).toEqual({ status: 200, @@ -700,6 +709,9 @@ describe("source route service", () => { expect(getSourceViewOptionsBySourceId).toHaveBeenCalledWith( [], knowhereClient, + expect.objectContaining({ + documentPresentationDetection: "disabled", + }), ); expect(knowhereClient.documents.listChunks).not.toHaveBeenCalled(); expect(result).toEqual({ diff --git a/src/domains/sources/route-types.ts b/src/domains/sources/route-types.ts index 96bd80c..01d1883 100644 --- a/src/domains/sources/route-types.ts +++ b/src/domains/sources/route-types.ts @@ -4,7 +4,10 @@ import type { ChunkPageParams, } from "@/domains/chunks" import type { ParsedChunkView } from "@/domains/chunks/types" -import type { SourceStatus, SourceView } from "@/domains/sources/types" +import type { + SourceStatus, + SourceView, +} from "@/domains/sources/types" import type { AuthUser } from "@/infrastructure/auth" import type { Source, Workspace } from "@/infrastructure/db/schema" import type { @@ -13,7 +16,10 @@ import type { } from "@/integrations/knowhere-demo" import type { RouteResult } from "@/lib/route-result" import type { SourceBlobUploadInput } from "./blob-upload" -import type { sourceViewOptionsBySourceId } from "./counts" +import type { + sourceViewOptionsBySourceId, + SourceViewOptionsLoadOptions, +} from "./counts" import type { UploadKnowhereClient } from "./upload" type SourceRouteKnowhereClient = UploadKnowhereClient & @@ -102,6 +108,17 @@ type SourceChunksBody = readonly chunks: readonly ParsedChunkView[] } | ChunkPage + | { + readonly chunks: readonly [] + readonly pagination?: { + readonly page: number + readonly pageSize: number + readonly total: 0 + readonly totalPages: 0 + } + readonly message: string + readonly isUnavailable: true + } | { readonly message: string } @@ -226,6 +243,7 @@ type SourceRouteServiceDependencies = { readonly getSourceViewOptionsBySourceId: ( sources: readonly Source[], client: SourceRouteKnowhereClient, + options?: SourceViewOptionsLoadOptions, ) => ReturnType readonly makeKnowhereClient: (apiKey: string) => SourceRouteKnowhereClient readonly listSourcesForWorkspace: (workspaceId: string) => Promise diff --git a/src/domains/sources/types.ts b/src/domains/sources/types.ts index 8479f87..e0ff9be 100644 --- a/src/domains/sources/types.ts +++ b/src/domains/sources/types.ts @@ -16,6 +16,18 @@ export type SourceOfficialLibraryView = { readonly sourceUrl: string } +export type SourceDocumentPresentation = + | { readonly kind: "parsed-chunks" } + | { readonly kind: "page-assets"; readonly pageCount: number } + +export type SourcePageAssetView = { + readonly pageNumber: number + readonly assetUrl: string + readonly contentType: string + readonly width?: number + readonly height?: number +} + export type OfficialLibrarySourceView = { readonly librarySourceId: string readonly categoryId: string @@ -50,6 +62,11 @@ export type SourceView = { readonly officialLibrary?: SourceOfficialLibraryView /** Count from the Notebook parsed snapshot manifest when available. */ readonly chunkCount?: number + /** + * Preferred source content presentation. Missing values are treated as + * parsed chunks so older views and cached responses remain compatible. + */ + readonly documentPresentation?: SourceDocumentPresentation /** User opt-out for this query session. Drives excludeDocumentIds. */ readonly excludedFromQuery?: boolean } diff --git a/src/domains/sources/view.test.ts b/src/domains/sources/view.test.ts index cc5b223..0069f6a 100644 --- a/src/domains/sources/view.test.ts +++ b/src/domains/sources/view.test.ts @@ -121,6 +121,18 @@ describe("toSourceView", () => { }); }); + it("includes document presentation options when provided", () => { + expect( + toSourceView(makeSource(), { + chunkCount: 4, + documentPresentation: { kind: "page-assets", pageCount: 4 }, + }), + ).toMatchObject({ + chunkCount: 4, + documentPresentation: { kind: "page-assets", pageCount: 4 }, + }); + }); + it("does not expose legacy demo original proxy routes", () => { const view = toSourceView( makeSource({ diff --git a/src/domains/sources/view.ts b/src/domains/sources/view.ts index 4096a81..5c6fe50 100644 --- a/src/domains/sources/view.ts +++ b/src/domains/sources/view.ts @@ -3,6 +3,7 @@ import { Schema } from "effect"; import type { Source } from "@/infrastructure/db/schema"; import type { SourceView } from "@/domains/sources/types"; import { sourceFailureMessage } from "./failure-message"; +import type { SourceDocumentPresentation } from "./types"; const SourceStatus = Schema.Literal( "uploading", @@ -13,7 +14,10 @@ const SourceStatus = Schema.Literal( export function toSourceView( source: Source, - options: { chunkCount?: number } = {}, + options: { + readonly chunkCount?: number + readonly documentPresentation?: SourceDocumentPresentation + } = {}, ): SourceView { const originalFile = getSourceOriginalFile(source) const status = toSourceStatus(source.status) @@ -35,6 +39,9 @@ export function toSourceView( ...(options.chunkCount !== undefined ? { chunkCount: options.chunkCount } : {}), + ...(options.documentPresentation !== undefined + ? { documentPresentation: options.documentPresentation } + : {}), }; } diff --git a/src/domains/workspace/client.test.ts b/src/domains/workspace/client.test.ts index bc19bba..47878bf 100644 --- a/src/domains/workspace/client.test.ts +++ b/src/domains/workspace/client.test.ts @@ -90,6 +90,36 @@ describe("workspaceClient", () => { }) }) + it("preserves source chunk unavailable messages without processing state", async () => { + mockRouteClient.getJson.mockResolvedValue({ + chunks: [], + message: + "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", + isUnavailable: true, + pagination: { + page: 1, + pageSize: 50, + total: 0, + totalPages: 0, + }, + }) + + const page = await workspaceClient.fetchChunkPage("source_1", 1) + + expect(page).toEqual({ + chunks: [], + isUnavailable: true, + message: + "Source is unavailable. The parsed document is not available locally and could not be loaded from Knowhere.", + pagination: { + page: 1, + pageSize: 50, + total: 0, + totalPages: 0, + }, + }) + }) + it("throws materialization route errors instead of treating them as empty sources", async () => { mockRouteClient.postJsonWithStatus.mockResolvedValue({ status: 502, diff --git a/src/domains/workspace/client.ts b/src/domains/workspace/client.ts index a650b6f..a415d6a 100644 --- a/src/domains/workspace/client.ts +++ b/src/domains/workspace/client.ts @@ -4,7 +4,9 @@ import type { ChatThreadView, } from "@/domains/chat/types" import type { ParsedChunkView } from "@/domains/chunks/types" -import type { SourceView } from "@/domains/sources/types" +import type { + SourceView, +} from "@/domains/sources/types" import { workspaceRouteClient } from "./route-client" const workspaceClientKeys = { @@ -25,6 +27,7 @@ const workspaceClientConfig = { type SourceChunksResponse = { chunks?: ParsedChunkView[] isProcessing?: boolean + isUnavailable?: boolean message?: string pagination?: { page: number @@ -131,7 +134,10 @@ async function fetchChunkPage( return { chunks: Array.isArray(body.chunks) ? body.chunks : [], ...(typeof body.message === "string" ? { message: body.message } : {}), - ...(typeof body.message === "string" ? { isProcessing: true } : {}), + ...(body.isUnavailable === true ? { isUnavailable: true } : {}), + ...(typeof body.message === "string" && body.isUnavailable !== true + ? { isProcessing: true } + : {}), pagination: body.pagination, } } diff --git a/src/domains/workspace/initial-state.test.ts b/src/domains/workspace/initial-state.test.ts index 23bd50f..9fc1450 100644 --- a/src/domains/workspace/initial-state.test.ts +++ b/src/domains/workspace/initial-state.test.ts @@ -172,6 +172,9 @@ describe("loadWorkspaceShellInitialState", () => { expect(sourceViewOptionsBySourceId).toHaveBeenCalledWith( [source], expect.any(Object), + expect.objectContaining({ + documentPresentationDetection: "disabled", + }), ) expect(state.sources).toEqual([ { @@ -203,7 +206,13 @@ describe("loadWorkspaceShellInitialState", () => { const state = await loadWorkspaceShellInitialState(deps) - expect(sourceViewOptionsBySourceId).toHaveBeenCalledWith([], expect.any(Object)) + expect(sourceViewOptionsBySourceId).toHaveBeenCalledWith( + [], + expect.any(Object), + expect.objectContaining({ + documentPresentationDetection: "disabled", + }), + ) expect(state.sources).toEqual([ expect.objectContaining({ id: "source_demo", @@ -231,7 +240,13 @@ describe("loadWorkspaceShellInitialState", () => { const state = await loadWorkspaceShellInitialState(deps) - expect(sourceViewOptionsBySourceId).toHaveBeenCalledWith([], expect.any(Object)) + expect(sourceViewOptionsBySourceId).toHaveBeenCalledWith( + [], + expect.any(Object), + expect.objectContaining({ + documentPresentationDetection: "disabled", + }), + ) expect(state.sources).toEqual([ expect.objectContaining({ id: "demo-tsla-q4-2025", @@ -259,7 +274,13 @@ describe("loadWorkspaceShellInitialState", () => { const state = await loadWorkspaceShellInitialState(deps) - expect(sourceViewOptionsBySourceId).toHaveBeenCalledWith([], expect.any(Object)) + expect(sourceViewOptionsBySourceId).toHaveBeenCalledWith( + [], + expect.any(Object), + expect.objectContaining({ + documentPresentationDetection: "disabled", + }), + ) expect(state.sources).toEqual([ expect.objectContaining({ id: "demo-tsla-q4-2025", diff --git a/src/domains/workspace/initial-state.ts b/src/domains/workspace/initial-state.ts index 851a34b..f959ad9 100644 --- a/src/domains/workspace/initial-state.ts +++ b/src/domains/workspace/initial-state.ts @@ -149,6 +149,7 @@ type WorkspaceShellInitialStateDependencies = { readonly sourceViewOptionsBySourceId: ( sources: readonly Source[], client: WorkspaceShellInitialStateClient, + options?: Parameters[2], ) => ReturnType } @@ -356,6 +357,9 @@ export const loadWorkspaceShellInitialStateEffect = ( deps.sourceViewOptionsBySourceId( sourcesNeedingChunkCount, client, + { + documentPresentationDetection: "disabled", + }, ), ) diff --git a/src/integrations/knowhere.ts b/src/integrations/knowhere.ts index c0d074d..e55fc0b 100644 --- a/src/integrations/knowhere.ts +++ b/src/integrations/knowhere.ts @@ -14,10 +14,8 @@ import { BlobParsedDocumentStorage } from "@/domains/sources/parsed-document-blo * SDK returns `completed:false` and the caller re-enqueues to continue. */ const defaultParsedStorageLimits: ParsedDocumentStorageLimits = { - chunkPageSize: 200, remotePageSize: 100, maxPagesPerSync: 10, - maxAssetsPerSync: 20, syncDeadlineMs: 8000, grepMaxPages: 50, grepDeadlineMs: 8000,