From 99cfe1df6fb4a142459eea6907280d2d81c81c28 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Fri, 11 Sep 2026 13:00:49 +0800 Subject: [PATCH 1/2] refactor(knowhere): streamline response handling and remove unused types - Removed unused types and functions related to KnowledgeReadResponse and KnowledgeGrepResponse from the knowhere-text and ledger modules. - Updated the knowhereToolText to handle chunk formatting with a new chunkPickStart parameter for better citation management. - Adjusted tests to reflect changes in response handling and ensure proper functionality with the updated code structure. This refactor enhances code clarity and maintains focus on relevant response types, improving overall maintainability. --- src/agent-harness/knowhere-text.test.ts | 190 +-------- src/agent-harness/knowhere-text.ts | 162 +------- src/agent-harness/ledger.test.ts | 234 +++-------- src/agent-harness/ledger.ts | 164 +------- src/agent-harness/runtime.test.ts | 151 +++++-- src/agent-harness/runtime.ts | 377 +++++------------- src/agent-harness/types.ts | 34 +- src/components/chat-message-list.tsx | 12 +- src/domains/chat/contracts.ts | 2 - src/domains/chat/index.test.ts | 243 +---------- src/domains/chat/index.ts | 19 +- src/domains/chat/knowhere-tools.ts | 145 +------ src/domains/chat/media-assets.test.ts | 25 ++ src/domains/chat/media-assets.ts | 27 +- src/domains/chat/page-citation-assets.test.ts | 25 ++ src/domains/chat/page-citation-assets.ts | 7 +- src/domains/chat/prompt.ts | 2 +- src/domains/chat/route-answer.ts | 1 - src/domains/chat/service.ts | 2 - 19 files changed, 448 insertions(+), 1374 deletions(-) diff --git a/src/agent-harness/knowhere-text.test.ts b/src/agent-harness/knowhere-text.test.ts index 5ec8abbb..af1094e9 100644 --- a/src/agent-harness/knowhere-text.test.ts +++ b/src/agent-harness/knowhere-text.test.ts @@ -1,10 +1,5 @@ import { describe, expect, it } from "vitest" -import type { - KnowledgeGrepResponse, - KnowledgeOutline, - KnowledgeReadResponse, - RetrievalQueryResponse, -} from "@ontos-ai/knowhere-sdk" +import type { RetrievalQueryResponse } from "@ontos-ai/knowhere-sdk" import { createEvidenceLedger } from "./ledger" import { knowhereToolText } from "./knowhere-text" @@ -18,11 +13,13 @@ describe("knowhereToolText", () => { const text = knowhereToolText.formatSearch({ response, retrievalCount: snapshot.retrievalCount, + chunkPickStart: 0, chunks: snapshot.chunks, assets: snapshot.assets, }) expect(text).toContain('') + expect(text).toContain('pick="1"') expect(text).toContain('ref="r1:result:1"') expect(text).toContain('ref="asset:r1:result:1"') expect(text).toContain("Page one summary.") @@ -33,77 +30,17 @@ describe("knowhereToolText", () => { 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.", + operation: "search", + message: "Knowhere search failed.", }), ).toBe( [ - '', + '', "", - "A documentId is required.", + "Knowhere search failed.", "", "", ].join("\n"), @@ -147,116 +84,3 @@ function makeSearchResponse(): RetrievalQueryResponse { 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 index 2140e7ec..5f55ff78 100644 --- a/src/agent-harness/knowhere-text.ts +++ b/src/agent-harness/knowhere-text.ts @@ -1,15 +1,6 @@ -import type { - KnowledgeGrepResponse, - KnowledgeOutline, - KnowledgeReadResponse, - RetrievalQueryResponse, -} from "@ontos-ai/knowhere-sdk" +import type { RetrievalQueryResponse } from "@ontos-ai/knowhere-sdk" -import type { - EvidenceAsset, - EvidenceChunk, - KnowhereListDocumentsResponse, -} from "./types" +import type { EvidenceAsset, EvidenceChunk } from "./types" type EvidenceDelta = { readonly chunks: readonly EvidenceChunk[] @@ -19,14 +10,7 @@ type EvidenceDelta = { type SearchTextInput = EvidenceDelta & { readonly response: RetrievalQueryResponse readonly retrievalCount: number -} - -type ReadChunksTextInput = EvidenceDelta & { - readonly response: KnowledgeReadResponse -} - -type GrepChunksTextInput = EvidenceDelta & { - readonly response: KnowledgeGrepResponse + readonly chunkPickStart: number } type ErrorTextInput = { @@ -34,12 +18,7 @@ type ErrorTextInput = { readonly message: string } -type KnowhereOperation = - | "search" - | "list_documents" - | "get_document_outline" - | "read_chunks" - | "grep_chunks" +type KnowhereOperation = "search" const assetInstruction = "Notebook returned image/page asset refs. Call inspectImage with the asset refs you will cite before finalize so OCR/visual context and provenance boxes exist. Do not expose raw asset URLs." @@ -58,99 +37,12 @@ export const knowhereToolText = { }), // Model grounding comes from (results). Do not also inject // evidenceText — same bodies, no citeable refs, doubles context. - formatEvidenceChunks(input.chunks), + formatEvidenceChunks(input.chunks, input.chunkPickStart), 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", { @@ -174,14 +66,18 @@ function wrapKnowhereBlock( ].join("\n") } -function formatEvidenceChunks(chunks: readonly EvidenceChunk[]): string { +function formatEvidenceChunks( + chunks: readonly EvidenceChunk[], + chunkPickStart: number, +): string { if (chunks.length === 0) return "" return [ "", - ...chunks.map((chunk) => + ...chunks.map((chunk, index) => [ formatOpenTag("chunk", { + pick: String(chunkPickStart + index + 1), ref: chunk.ref, kind: chunk.kind, chunkId: chunk.chunkId, @@ -228,42 +124,6 @@ function formatAssetInstruction(assets: readonly EvidenceAsset[]): string { 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") } diff --git a/src/agent-harness/ledger.test.ts b/src/agent-harness/ledger.test.ts index 3af12759..6d3bdb32 100644 --- a/src/agent-harness/ledger.test.ts +++ b/src/agent-harness/ledger.test.ts @@ -1,9 +1,5 @@ import { describe, expect, it } from "vitest" -import type { - KnowledgeGrepResponse, - KnowledgeReadResponse, - RetrievalQueryResponse, -} from "@ontos-ai/knowhere-sdk" +import type { RetrievalQueryResponse } from "@ontos-ai/knowhere-sdk" import { createEvidenceLedger } from "./ledger" @@ -117,118 +113,80 @@ describe("createEvidenceLedger", () => { ) }) - it("adds read chunk refs and page image assets", () => { + it("does not throw when a result is missing chunkType", () => { + // Knowhere's chunkType is declared as a required string in the SDK + // type, but real retrieval results can omit it at runtime, same as the + // referencedChunks case above. Unlike referencedChunks (which carry no + // real content), results carry real content, so the chunk must still be + // kept in the ledger -- only asset-type detection should be guarded. 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", - }), - }), - ]) - }) - - it("copies page metadata onto grep matches from the same chunk id", () => { - const ledger = createEvidenceLedger() - ledger.addReadChunksResponse(makeReadResponse()) - - const snapshot = ledger.addGrepChunksResponse({ - ...makeGrepResponse(), - matches: [ + const snapshot = ledger.addRetrievalResponse({ + namespace: "default", + query: "hypertension target", + routerUsed: "agent_explore", + answerText: "", + evidenceText: "[E1] some evidence", + stopReason: "finished", + failureReason: null, + results: [ { - position: 1, - chunkId: "chunk_page_1", - chunkType: "page", - sectionPath: "Root / Page 1", - sourceChunkPath: "pages/page-1.md", - startOffset: 0, - endOffset: 12, - snippet: "Full page snip", + content: "Target BP is <130/80 mmHg.", + chunkType: undefined as unknown as string, + score: 0.9, + assetUrl: "https://assets.example/images/chart.png", + source: { + documentId: "doc_1", + sourceFileName: "guideline.pdf", + sectionPath: "BP targets", + }, }, ], + referencedChunks: [], }) - expect(snapshot.chunks[1]).toEqual( - expect.objectContaining({ - ref: "grep1:match:1", - kind: "grep_match", - chunkId: "chunk_page_1", - metadata: expect.objectContaining({ - pageNums: [1], - position: 1, - startOffset: 0, - endOffset: 12, - }), - }), - ) + expect(snapshot.chunks.map((chunk) => chunk.ref)).toEqual(["r1:result:1"]) + expect(snapshot.chunks[0]?.content).toBe("Target BP is <130/80 mmHg.") + // chunkType is unknown, but the assetUrl itself has an image extension, + // so asset detection still recognizes it as an image via the URL check. + expect(snapshot.assets).toEqual([ + expect.objectContaining({ ref: "asset:r1:result:1", type: "image" }), + ]) }) - it("copies pageNumbers from the grep match when the SDK provides them", () => { + it("skips agent_explore referencedChunks that lack chunkType instead of throwing", () => { const ledger = createEvidenceLedger() - const snapshot = ledger.addGrepChunksResponse({ - ...makeGrepResponse(), - matches: [ + const snapshot = ledger.addRetrievalResponse({ + namespace: "default", + query: "hypertension CAD blood pressure target", + routerUsed: "agent_explore", + answerText: "", + evidenceText: "[E1] some evidence", + stopReason: "finished", + failureReason: null, + results: [ { - position: 1, - chunkId: "chunk_page_4", - chunkType: "page", - sectionPath: "FINANCIAL SUMMARY", - sourceChunkPath: "pages/page-4.md", - startOffset: 0, - endOffset: 12, - snippet: "automotive revenues", - pageNumbers: [4], + content: "Target BP is <130/80 mmHg.", + chunkType: "text", + score: 0.9, + source: { + documentId: "doc_1", + sourceFileName: "guideline.pdf", + sectionPath: "BP targets", + }, }, ], + // Real agent_explore responses can return referencedChunks entries + // that only carry a summary id, with no chunkType/chunkId/documentId + // even though the SDK type declares those as required strings. + referencedChunks: [ + { summary: "8da0776b-c52b-5602-8579-25c421706f5f" }, + ] as unknown as RetrievalQueryResponse["referencedChunks"], }) - expect(snapshot.chunks[0]).toEqual( - expect.objectContaining({ - ref: "grep1:match:1", - kind: "grep_match", - chunkId: "chunk_page_4", - metadata: expect.objectContaining({ - pageNums: [4], - position: 1, - startOffset: 0, - endOffset: 12, - }), - }), - ) + expect(snapshot.chunks.map((chunk) => chunk.ref)).toEqual(["r1:result:1"]) + expect(snapshot.chunks[0]?.content).toBe("Target BP is <130/80 mmHg.") }) }) @@ -302,77 +260,3 @@ 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 42f31207..9c80448d 100644 --- a/src/agent-harness/ledger.ts +++ b/src/agent-harness/ledger.ts @@ -1,8 +1,4 @@ import type { - KnowledgeGrepMatch, - KnowledgeGrepResponse, - KnowledgeReadChunk, - KnowledgeReadResponse, RetrievalQueryResponse, RetrievalResult, } from "@ontos-ai/knowhere-sdk" @@ -18,8 +14,6 @@ const imageExtensions = [".jpg", ".jpeg", ".png", ".gif", ".webp", ".svg"] as co type MutableLedger = { retrievalCount: number - readCount: number - grepCount: number chunks: EvidenceChunk[] assets: EvidenceAsset[] evidenceText: string[] @@ -47,8 +41,6 @@ export type EvidenceLedger = ReturnType export function createEvidenceLedger() { const ledger: MutableLedger = { retrievalCount: 0, - readCount: 0, - grepCount: 0, chunks: [], assets: [], evidenceText: [], @@ -84,6 +76,15 @@ export function createEvidenceLedger() { }) response.referencedChunks.forEach((chunk, index) => { + // Knowhere's agent_explore router returns referencedChunks entries + // that may carry only a summary id with no chunkType/content + // (despite the SDK type declaring chunkType as required). Skip + // entries missing a usable chunkType: they have no real content + // (content is always "" here) and chunkType is required downstream + // (asset-type detection calls chunkType.toLowerCase()). + if (typeof chunk.chunkType !== "string" || chunk.chunkType.trim().length === 0) { + return + } const content = "" addChunk({ ledger, @@ -112,38 +113,6 @@ 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) { @@ -177,13 +146,6 @@ export function createEvidenceLedger() { return ledger.chunks.length > 0 || ledger.evidenceText.length > 0 }, - hasRef(ref: string): boolean { - return ( - ledger.chunks.some((chunk) => chunk.ref === ref) || - ledger.assets.some((asset) => asset.ref === ref) - ) - }, - snapshot(): EvidenceLedgerSnapshot { return snapshot(ledger) }, @@ -219,100 +181,6 @@ 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 { - const donor = input.ledger.chunks.find( - (chunk) => - chunk.chunkId === input.match.chunkId && hasPageMetadata(chunk.metadata), - ) - const pageNums = - input.match.pageNumbers && input.match.pageNumbers.length > 0 - ? [...input.match.pageNumbers] - : undefined - - 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: { - ...(donor?.metadata ?? {}), - ...(pageNums ? { pageNums } : {}), - 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 hasPageMetadata( - metadata: Readonly> | undefined, -): boolean { - if (!metadata) return false - const values = [ - metadata.pageNums, - metadata.page_nums, - metadata.pageNum, - metadata.page_num, - metadata.pageAssets, - metadata.page_assets, - ] - return values.some((value) => { - if (Array.isArray(value)) return value.length > 0 - if (typeof value === "number") return Number.isSafeInteger(value) && value > 0 - if (typeof value === "string") return value.trim().length > 0 - return false - }) -} - function addChunk(input: { readonly ledger: MutableLedger readonly chunk: Omit @@ -347,8 +215,16 @@ function buildContentPreview(content: string): string { return `${normalized.slice(0, contentPreviewLimit)}...` } +// Knowhere's chunkType is declared as a required string in the SDK type, +// but real API responses (seen on referencedChunks; results are the same +// contract) can omit it. Normalize defensively instead of calling +// .toLowerCase() on a value that may be undefined at runtime. +function normalizeChunkType(chunkType: string): string { + return typeof chunkType === "string" ? chunkType.toLowerCase() : "" +} + function isRenderableAsset(chunkType: string, assetUrl: string): boolean { - const normalizedChunkType = chunkType.toLowerCase() + const normalizedChunkType = normalizeChunkType(chunkType) return ( normalizedChunkType === "image" || normalizedChunkType === "table" || @@ -379,7 +255,7 @@ function getEvidenceAssetCandidate( function getPageCitationAssetCandidate( chunk: Omit, ): EvidenceAssetCandidate | null { - if (chunk.chunkType.toLowerCase() !== "page") return null + if (normalizeChunkType(chunk.chunkType) !== "page") return null const candidates = [ ...parsePageCitationAssetCandidates(chunk.metadata?.pageAssets), @@ -409,7 +285,7 @@ function getPageCitationAssetCandidate( } function getAssetType(chunkType: string, assetUrl: string): "image" | "table" { - return chunkType.toLowerCase() === "table" && !isImageAssetUrl(assetUrl) + return normalizeChunkType(chunkType) === "table" && !isImageAssetUrl(assetUrl) ? "table" : "image" } diff --git a/src/agent-harness/runtime.test.ts b/src/agent-harness/runtime.test.ts index b459d5b8..11ac3d94 100644 --- a/src/agent-harness/runtime.test.ts +++ b/src/agent-harness/runtime.test.ts @@ -28,6 +28,14 @@ describe("agent harness runtime", () => { expect(prompt).toContain( "Call knowhere_search only when memory is insufficient and groundingPolicy requires citing source documents", ) + expect(prompt).toContain( + "call knowhere_search again with a refined query", + ) + expect(prompt).toContain("Refine knowhere_search at most twice") + expect(prompt).not.toContain("knowhere_list_documents") + expect(prompt).not.toContain("knowhere_get_document_outline") + expect(prompt).not.toContain("knowhere_read_chunks") + expect(prompt).not.toContain("knowhere_grep_chunks") expect(prompt).toContain( "Do not treat every question as a document-retrieval task", ) @@ -42,16 +50,14 @@ describe("agent harness runtime", () => { expect(prompt).not.toContain("navigation action") }) - it("tells the agent citation metadata is optional and ledger-resolved", () => { + it("tells the agent to cite by search-chunk pick numbers", () => { 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") + expect(prompt).toContain("citations is a list of { pick }") + expect(prompt).toContain("Notebook writes the citation list from those picks") + expect(prompt).toContain("Do not pass documentId or evidence refs as citations") + expect(prompt).not.toContain("Citation refs must be evidence refs") + expect(prompt).not.toContain("Citation label and source metadata are optional") }) it("tells the agent to emit [[cite:n]] markers instead of title/pN or [1]", () => { @@ -98,6 +104,7 @@ describe("agent harness runtime", () => { }) expect(result).toContain('') + expect(result).toContain('pick="1"') expect(result).toContain('ref="r1:result:1"') expect(result).toContain('ref="asset:r1:result:1"') expect(query).toHaveBeenCalledWith({ @@ -169,10 +176,12 @@ describe("agent harness runtime", () => { const firstResult = await executeTool(tools.knowhere_search, { query: "first" }) const secondResult = await executeTool(tools.knowhere_search, { query: "second" }) + expect(firstResult).toContain('pick="1"') expect(firstResult).toContain('ref="r1:result:1"') expect(secondResult).toContain('retrievalCount="2"') expect(secondResult).toContain("Second retrieval evidence.") expect(secondResult).not.toContain("") + expect(secondResult).toContain('pick="2"') 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([ @@ -471,7 +480,7 @@ describe("agent harness runtime", () => { expect( await executeTool(tools.finalize, { text: "Revenue was $24.9B [[cite:1]] [[cite:2]].", - citations: [{ ref: "r1:result:1" }, { ref: "r1:result:2" }], + citations: [{ pick: 1 }, { pick: 2 }], memoryCitations: [], artifacts: [], unresolved: [], @@ -598,7 +607,7 @@ describe("agent harness runtime", () => { const finalize = await executeTool(tools.finalize, { text: "The amount is 5000 yuan [[cite:1]].", - citations: [{ ref: "r1:referenced:1" }], + citations: [{ pick: 1 }], memoryCitations: [], artifacts: [], unresolved: [], @@ -610,6 +619,98 @@ describe("agent harness runtime", () => { expect(state.finalized).not.toBe(true) }) + it("writes citation refs from ledger picks and rejects picks outside the ledger", async () => { + const ledger = createEvidenceLedger() + ledger.addRetrievalResponse(makeRetrievalResponse()) + const state: { + finalized?: boolean + finalizedManifest?: OutputManifest + } = {} + const tools = createHarnessTools({ + state, + ledger, + memoryTools: makeMemoryTools(), + knowhereTools: makeKnowhereTools(), + recentTurns: [], + }) + + const rejected = await executeTool(tools.finalize, { + text: "Target is <130/80 mmHg [[cite:1]].", + citations: [{ pick: 99 }], + memoryCitations: [], + artifacts: [], + unresolved: [], + }) + + expect(rejected).toMatchObject({ + ok: false, + unknownPicks: [99], + }) + expect(String((rejected as { message: string }).message)).toContain( + "Available picks: 1-1", + ) + expect(state.finalized).not.toBe(true) + + const accepted = await executeTool(tools.finalize, { + text: "Target is <130/80 mmHg [[cite:1]].", + citations: [{ pick: 1 }], + memoryCitations: [], + artifacts: [], + unresolved: [], + }) + expect(accepted).toMatchObject({ + ok: true, + citations: [{ ref: "r1:result:1" }], + }) + expect(state.finalized).toBe(true) + expect(state.finalizedManifest?.citations).toEqual([{ ref: "r1:result:1" }]) + }) + + it("maps finalize picks across successive searches", async () => { + const ledger = createEvidenceLedger() + ledger.addRetrievalResponse(makeRetrievalResponse()) + ledger.addRetrievalResponse({ + ...makeRetrievalResponse(), + query: "second query", + results: [ + { + content: "Second retrieval evidence.", + chunkType: "text", + score: 0.8, + source: { + documentId: "doc_2", + sourceFileName: "second.pdf", + sectionPath: "Second", + }, + }, + ], + }) + const state: { + finalizedManifest?: OutputManifest + } = {} + const tools = createHarnessTools({ + state, + ledger, + memoryTools: makeMemoryTools(), + knowhereTools: makeKnowhereTools(), + recentTurns: [], + }) + + const accepted = await executeTool(tools.finalize, { + text: "Second source [[cite:1]].", + citations: [{ pick: 2 }], + memoryCitations: [], + artifacts: [], + unresolved: [], + }) + + expect(accepted).toMatchObject({ + ok: true, + citations: [{ ref: "r2:result:1" }], + }) + expect(state.finalizedManifest?.citations).toEqual([{ ref: "r2:result:1" }]) + }) + it("accepts finalize output without planning-tool gating", async () => { const state: { finalizedManifest?: OutputManifest @@ -719,7 +820,7 @@ describe("agent harness runtime", () => { }) const manifest = { text: "The contractor pays 5000 yuan per occurrence [[cite:1]].", - citations: [{ ref: "r1:referenced:1" }], + citations: [{ pick: 1 }], memoryCitations: [], artifacts: [], unresolved: [], @@ -781,7 +882,7 @@ describe("agent harness runtime", () => { const result = await executeTool(tools.finalize, { text: "The contractor pays 5000 yuan [[cite:1]].", - citations: [{ ref: "r1:result:1" }], + citations: [{ pick: 1 }], memoryCitations: [], artifacts: [], unresolved: [], @@ -1024,6 +1125,15 @@ describe("agent harness runtime", () => { expect(result.activeTools).not.toContain("knowhere_search") }) + it("forces a tool call on ordinary steps so the model cannot skip finalize with bare text", () => { + const result = prepareHarnessStep({ + stepNumber: 3, + messages: [], + }) + + expect(result.toolChoice).toBe("required") + }) + it("opens memory_search and knowhere_search together as peers when sources are required", () => { const result = prepareHarnessStep({ stepNumber: 3, @@ -1039,15 +1149,12 @@ describe("agent harness runtime", () => { }) expect(result.activeTools).toEqual( - expect.arrayContaining([ - "memory_search", - "knowhere_search", - "knowhere_list_documents", - "knowhere_get_document_outline", - "knowhere_read_chunks", - "knowhere_grep_chunks", - ]), + expect.arrayContaining(["memory_search", "knowhere_search"]), ) + expect(result.activeTools).not.toContain("knowhere_list_documents") + expect(result.activeTools).not.toContain("knowhere_get_document_outline") + expect(result.activeTools).not.toContain("knowhere_read_chunks") + expect(result.activeTools).not.toContain("knowhere_grep_chunks") }) it("forces image inspection before forced finalization when image assets are available", () => { @@ -1195,10 +1302,6 @@ function makeKnowhereTools( ): 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.")), } } diff --git a/src/agent-harness/runtime.ts b/src/agent-harness/runtime.ts index 0f4aab0a..43d9329a 100644 --- a/src/agent-harness/runtime.ts +++ b/src/agent-harness/runtime.ts @@ -18,6 +18,7 @@ import type { ContextPolicy, EvidenceAsset, EvidenceChunk, + EvidenceLedgerSnapshot, HarnessRunResult, HarnessToolCallTrace, HarnessTrace, @@ -30,6 +31,8 @@ import type { KnowhereToolRuntime, MemorySearchKind, MemoryToolRuntime, + OutputArtifactView, + OutputCitation, OutputManifest, } from "./types" import { memorySearchKinds } from "./types" @@ -67,10 +70,12 @@ type HarnessTools = ReturnType type HarnessStepPreparation = { messages: ModelMessage[] activeTools?: Array> - toolChoice?: { - type: "tool" - toolName: Extract - } + toolChoice?: + | "required" + | { + type: "tool" + toolName: Extract + } } const targetModalitySchema = z.enum(["text", "image", "table"]) @@ -82,14 +87,6 @@ const knowhereSearchTargetContentSchema = z.enum([ "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), @@ -101,27 +98,6 @@ const knowhereSearchSchema = z.object({ 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([ "answer", @@ -164,16 +140,8 @@ const contextPolicySchema = z.object({ activePriorTurnIds: z.array(z.string()).default([]), }) -const outputCitationSchema = z.object({ - ref: z.string().min(1), - 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 citationPickSchema = z.object({ + pick: z.number().int().positive(), }) const memoryCitationSchema = z.object({ @@ -210,9 +178,9 @@ const outputArtifactSchema = z.union([ derivedTableArtifactSchema, ]) -const outputManifestSchema = z.object({ +const finalizeManifestSchema = z.object({ text: z.string(), - citations: z.array(outputCitationSchema).default([]), + citations: z.array(citationPickSchema).default([]), memoryCitations: z.array(memoryCitationSchema).default([]), artifacts: z.array(outputArtifactSchema).default([]), unresolved: z.array(z.string()).default([]), @@ -286,13 +254,7 @@ const alwaysAvailableTools = [ const fluidRetrievalTools = ["memory_search"] as const -const crystalRetrievalTools = [ - "knowhere_search", - "knowhere_list_documents", - "knowhere_get_document_outline", - "knowhere_read_chunks", - "knowhere_grep_chunks", -] as const +const crystalRetrievalTools = ["knowhere_search"] as const /** Reserved third retrieval slot (cognition). Not registered this round. */ const cognitionRetrievalTools = [] as const @@ -358,6 +320,10 @@ export function prepareHarnessStep(input: { activeTools: selectHarnessActiveTools({ intent: input.intent, }), + // finalize is the only output contract (see its tool description). Force + // a tool call every step so the model cannot end the turn with a bare + // text response that skips finalize's citation/artifact validation. + toolChoice: "required", } } @@ -448,7 +414,7 @@ type ModelMessageForRole = Extract< function buildForcedFinalizationFeedback(): string { return [ "The retrieval step budget has been reached.", - "Do not search again or call any Knowhere evidence-reading tools.", + "Do not search again.", "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", @@ -551,7 +517,7 @@ export function createHarnessTools(input: { knowhere_search: tool({ description: - "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.", + "Search Knowhere for relevant Notebook evidence. Returns tagged text with pick numbers for finalize citations, evidence refs such as r1:result:1, and asset refs such as asset:r1:result:1.", inputSchema: knowhereSearchSchema, execute: async (request) => traceToolCall(input.state, { @@ -567,83 +533,6 @@ export function createHarnessTools(input: { }), }), - 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 cited Knowhere page/image asset refs for OCR, visual details, and provenance boxes. Call this after retrieval and before finalize whenever the answer cites page or image assets.", @@ -730,17 +619,42 @@ 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 evidence refs when available. " + - "Use citations for Knowhere evidence and memoryCitations for fluid memory refs. " + + "images/tables shown to the user. citations is the list of evidence picks " + + "you used; each pick is the pick number on a Knowhere search chunk. " + + "Notebook writes citation refs from the evidence ledger. " + + "Use memoryCitations for fluid memory refs. " + "Cited page/image assets must be inspected with inspectImage first.", - inputSchema: outputManifestSchema, + inputSchema: finalizeManifestSchema, execute: async (manifest) => traceToolCall(input.state, { toolName: "finalize", inputSummary: summarizeManifest(manifest), execute: async () => { + const resolvedCitations = resolveCitationPicks({ + citations: manifest.citations, + ledger: input.ledger, + }) + if (!resolvedCitations.ok) { + return { + ok: false as const, + message: buildFinalizeRequiresPicksMessage({ + unknownPicks: resolvedCitations.unknownPicks, + ledger: input.ledger.snapshot(), + }), + unknownPicks: resolvedCitations.unknownPicks, + } + } + + const outputManifest: OutputManifest = { + text: manifest.text, + citations: resolvedCitations.citations, + memoryCitations: manifest.memoryCitations, + artifacts: manifest.artifacts, + unresolved: manifest.unresolved, + } + const inspectRefs = getUninspectedCitedImageRefs({ - manifest, + manifest: outputManifest, ledger: input.ledger, inspectedImageRefs: input.state.inspectedImageRefs ?? [], inspectImagesAvailable: input.inspectImages !== undefined, @@ -753,9 +667,9 @@ export function createHarnessTools(input: { } } - input.state.finalizedManifest = manifest + input.state.finalizedManifest = outputManifest input.state.finalized = true - return { ok: true as const, ...manifest } + return { ok: true as const, ...outputManifest } }, summarizeOutput: summarizeFinalizeOutput, }), @@ -924,6 +838,50 @@ async function inspectRetrievedImages(input: { } } +function resolveCitationPicks(input: { + readonly citations: readonly { pick: number }[] + readonly ledger: ReturnType +}): + | { ok: true; citations: OutputCitation[] } + | { ok: false; unknownPicks: number[] } { + const chunks = input.ledger.snapshot().chunks + const unknownPicks: number[] = [] + const citations: OutputCitation[] = [] + + for (const citation of input.citations) { + const chunk = Number.isInteger(citation.pick) + ? chunks[citation.pick - 1] + : undefined + if (!chunk) { + if (!unknownPicks.includes(citation.pick)) { + unknownPicks.push(citation.pick) + } + continue + } + citations.push({ ref: chunk.ref }) + } + + if (unknownPicks.length > 0) { + return { ok: false, unknownPicks } + } + return { ok: true, citations } +} + +function buildFinalizeRequiresPicksMessage(input: { + readonly unknownPicks: readonly number[] + readonly ledger: EvidenceLedgerSnapshot +}): string { + const availableCount = input.ledger.chunks.length + return [ + "Citations must use pick numbers from Knowhere search chunks.", + `Unknown citation picks: ${input.unknownPicks.join(" ")}.`, + availableCount > 0 + ? `Available picks: 1-${availableCount}.` + : "No evidence picks are available; list the gap in unresolved and omit citations.", + "Call finalize again using only available picks.", + ].join(" ") +} + function getUninspectedCitedImageRefs(input: { readonly manifest: OutputManifest readonly ledger: ReturnType @@ -1009,26 +967,10 @@ function buildFinalizeRequiresInspectionMessage( ].join(" ") } -type KnowhereToolOperation = - | "search" - | "list_documents" - | "get_document_outline" - | "read_chunks" - | "grep_chunks" +type KnowhereToolOperation = "search" type MemorySearchToolRequest = z.infer 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 executeMemorySearch(input: { readonly memoryTools: MemoryToolRuntime @@ -1070,48 +1012,7 @@ async function executeKnowhereSearch(input: { 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, + chunkPickStart: beforeSnapshot.chunks.length, chunks: snapshot.chunks.slice(beforeSnapshot.chunks.length), assets: snapshot.assets.slice(beforeSnapshot.assets.length), }) @@ -1121,17 +1022,8 @@ async function executeKnowhereGrepChunks(input: { 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) { @@ -1142,20 +1034,6 @@ async function executeKnowhereTextTool(input: { } } -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) { @@ -1276,48 +1154,6 @@ function summarizeKnowhereSearchRequest(request: { } } -function summarizeDocumentReference( - request: KnowhereDocumentReferenceRequest, -): DocumentReferenceSummary { - return { - documentId: request.documentId, - localDocumentId: request.localDocumentId, - hasJobId: typeof request.jobId === "string", - hasRevisionKey: typeof request.revisionKey === "string", - } -} - -function summarizeReadChunksRequest( - request: KnowhereReadChunksToolRequest, -): unknown { - return { - ...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 { @@ -1377,7 +1213,13 @@ function summarizeReadPriorTurnOutput(output: unknown): unknown { } } -function summarizeManifest(manifest: OutputManifest): unknown { +function summarizeManifest(manifest: { + readonly text: string + readonly citations: readonly unknown[] + readonly memoryCitations: readonly unknown[] + readonly artifacts: readonly OutputArtifactView[] + readonly unresolved: readonly string[] +}): unknown { return { textLength: manifest.text.length, citationCount: manifest.citations.length, @@ -1428,15 +1270,15 @@ export function buildHarnessSystemPrompt(turn: AgentTurnInput): string { "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 memory_search first when known fluid memory may answer the request. Call knowhere_search only when memory is insufficient and groundingPolicy requires citing source documents. Use knowhere_list_documents, knowhere_get_document_outline, knowhere_read_chunks, and knowhere_grep_chunks for focused document reads.", + "4. Call memory_search first when known fluid memory may answer the request. Call knowhere_search when memory is insufficient and groundingPolicy requires citing source documents. If the returned evidence is not enough to answer, or the query needs to focus differently, call knowhere_search again with a refined query (different keywords / topK / targetContent) instead of trying to browse documents directly. Knowhere's own retrieval agent already navigates the corpus internally.", "5. After Knowhere returns image/page asset refs, call inspectImage on the page/image assets you will cite before finalize. This supplies OCR/visual context and provenance boxes.", "6. Inspect each unique cited page once; retrieval already bounds the available evidence set.", - "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.", + "7. Call finalize with text, citations, artifacts, and unresolved issues when you are ready to answer.", "", "Retrieval rules:", "- First use memory_search to see whether known fluid memory can answer directly.", "- Call knowhere_search only when memory is insufficient and groundingPolicy requires citing source documents.", + "- Refine knowhere_search at most twice. If two refined searches still do not add new relevant evidence, call finalize and list the gap in unresolved.", "- Do not treat every question as a document-retrieval task.", "", "Context rules:", @@ -1448,13 +1290,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.", - "- Prefer citation and selected image/table artifact refs returned by Knowhere tools in the evidence ledger.", + "- citations is a list of { pick }. pick is the 1-based number on the search chunk you are using. Notebook writes the citation list from those picks. Do not pass documentId or evidence refs as citations.", "- Place [[cite:n]] immediately after the supported claim. n is the 1-based index into the citations array passed to finalize.", "- Write one marker per index: [[cite:1]] [[cite:3]] [[cite:5]]. Never group indices as [[cite:1, 3, 5]].", "- Do not write title/pN, [1], Markdown footnotes, or [Source N: ...] in the answer text. Notebook renders chips from [[cite:n]] and citation metadata.", "- Repeat [[cite:n]] when another claim uses the same page. Do not collapse same-page citations to one row.", - "- Citation label and source metadata are optional. Notebook resolves citation metadata from evidence refs when possible.", - "- If evidence is relevant but you cannot identify a supporting evidence ref, answer with unresolved issues instead of fabricating a ref.", + "- If you have no supporting evidence pick, omit citations and list the gap in unresolved.", "- inspectImage observations are inspection notes, not new source refs. Final citations and displayed image artifacts must use the original retrieved image asset refs.", "- Do not finalize cited page/image assets from chunk text alone when inspectImage is available. Inspect those asset refs first, then write the answer using the inspection notes.", "- 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.", diff --git a/src/agent-harness/types.ts b/src/agent-harness/types.ts index 82f38ab5..ad686ad6 100644 --- a/src/agent-harness/types.ts +++ b/src/agent-harness/types.ts @@ -1,10 +1,4 @@ import type { - KnowledgeDocumentReference, - KnowledgeGrepParams, - KnowledgeGrepResponse, - KnowledgeOutline, - KnowledgeReadParams, - KnowledgeReadResponse, RetrievalQueryParams, RetrievalQueryResponse, } from "@ontos-ai/knowhere-sdk" @@ -99,36 +93,10 @@ export type KnowhereSearchRequest = Pick< readonly purpose?: string } -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 const memorySearchKinds = [ @@ -170,7 +138,7 @@ export type MemoryCitation = { export type EvidenceChunk = { readonly ref: string - readonly kind: "result" | "referenced_chunk" | "read_chunk" | "grep_match" + readonly kind: "result" | "referenced_chunk" readonly chunkId?: string readonly content: string readonly contentPreview: string diff --git a/src/components/chat-message-list.tsx b/src/components/chat-message-list.tsx index 796f2eab..2b7a5d13 100644 --- a/src/components/chat-message-list.tsx +++ b/src/components/chat-message-list.tsx @@ -931,10 +931,14 @@ function isImageCitation( citation: ChatCitationView, assetUrl: string, ): boolean { - return ( - citation.chunkType.toLowerCase() === "image" || - hasImageFileExtension(assetUrl) - ); + // chunkType is sourced from Knowhere retrieval results, which can omit it + // at runtime even though the type says it's a required string. Normalize + // defensively instead of calling .toLowerCase() on a possibly-missing value. + const chunkType = + typeof citation.chunkType === "string" + ? citation.chunkType.toLowerCase() + : ""; + return chunkType === "image" || hasImageFileExtension(assetUrl); } function hasImageFileExtension(assetUrl: string): boolean { diff --git a/src/domains/chat/contracts.ts b/src/domains/chat/contracts.ts index 527b52c4..9accd573 100644 --- a/src/domains/chat/contracts.ts +++ b/src/domains/chat/contracts.ts @@ -16,7 +16,6 @@ import type { } 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 @@ -76,7 +75,6 @@ export type AnswerQuestionInput = { useAgentic?: boolean 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 021ebb22..b65be0f7 100644 --- a/src/domains/chat/index.test.ts +++ b/src/domains/chat/index.test.ts @@ -1,9 +1,6 @@ import { afterEach, describe, expect, it, vi } from "vitest" import type { Knowledge, - KnowledgeGrepResponse, - KnowledgeOutline, - KnowledgeReadResponse, RetrievalQueryParams, RetrievalQueryResponse, RetrievalResult, @@ -150,7 +147,7 @@ describe("answerQuestionWithRetrieval", () => { }); }); - it("exposes search, list, outline, read, and grep through the Knowhere tool runtime", async () => { + it("exposes search through the Knowhere tool runtime", async () => { const result = makeRetrievalResult({ chunkType: "image", source: { @@ -170,36 +167,6 @@ describe("answerQuestionWithRetrieval", () => { 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", - documentMetadata: { - createdByClient: "cli", - }, - }, - { - documentId: "doc_untagged", - namespace: "default", - status: "ready", - sourceFileName: "dummy.pdf", - }, - ], - }); const generateAnswer = vi.fn( async ({ knowhereTools }: Parameters[0]) => { if (!knowhereTools) throw new Error("Knowhere tools were not provided."); @@ -209,28 +176,8 @@ describe("answerQuestionWithRetrieval", () => { 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."); }, ); @@ -246,8 +193,6 @@ describe("answerQuestionWithRetrieval", () => { sources, excludedSourceIds: ["source_excluded"], retrieval, - knowledge, - remoteDocumentClient: { documents: { list: listDocuments } }, generateAnswer, messages: [], }), @@ -261,27 +206,6 @@ describe("answerQuestionWithRetrieval", () => { 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."); }); @@ -1431,10 +1355,10 @@ describe("answerQuestionWithRetrieval", () => { expect(answer.citations[0]?.pageCitationAssetUrl).toBe(hardenedPageAssetUrl); }); - it("hydrates page numbers for grep citations from the matching parsed chunk", async () => { + it("hydrates page numbers for citations missing page metadata from the matching parsed chunk", async () => { const grepChunk = { - ref: "grep1:match:1", - kind: "grep_match" as const, + ref: "r1:result:1", + kind: "result" as const, chunkId: "chunk_financial_summary", content: "ept percentages and per share data)\nTotal automotive revenues\n17,693", contentPreview: "ept percentages and per share data)", @@ -1479,7 +1403,7 @@ describe("answerQuestionWithRetrieval", () => { makeHarnessRunResultWithLedger( "Automotive revenue was $17,693 million [[cite:1]].", { - citations: [{ ref: "grep1:match:1" }], + citations: [{ ref: "r1:result:1" }], chunks: [grepChunk], }, ), @@ -2054,17 +1978,7 @@ describe("answerQuestionWithRetrieval", () => { 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", - }, - }, - ], + citations: [{ pick: 1 }], memoryCitations: [], artifacts: [], unresolved: [], @@ -2698,17 +2612,7 @@ describe("generateAgenticOutputManifest", () => { }); await tools.finalize?.execute({ text: "已找到相关身份证图片，见下方图片。", - citations: [ - { - ref: "r1:result:1", - label: "商务标文件.pdf / 身份证正面", - source: { - documentId: "doc_identity", - sourceFileName: "document-generated.pdf", - sectionPath: "身份证正面", - }, - }, - ], + citations: [{ pick: 1 }], memoryCitations: [], artifacts: [ { @@ -2837,17 +2741,7 @@ describe("generateAgenticOutputManifest", () => { }); await tools.finalize?.execute({ text: "The inspected image appears to show the requested ID card.", - citations: [ - { - ref: "asset:r1:result:1", - label: "identity.pdf / images/id-front.png", - source: { - documentId: "doc_identity", - sourceFileName: "generated.pdf", - sectionPath: "images/id-front.png", - }, - }, - ], + citations: [{ pick: 1 }], memoryCitations: [], artifacts: [ { @@ -2926,7 +2820,7 @@ describe("generateAgenticOutputManifest", () => { ], }); expect(result.manifest.citations.map((citation) => citation.ref)).toEqual([ - "asset:r1:result:1", + "r1:result:1", ]); expect(result.manifest.artifacts).toEqual([ { @@ -2979,17 +2873,7 @@ describe("generateAgenticOutputManifest", () => { }); await tools.finalize?.execute({ text: "承包人自行修改发包人审批的进度计划，应按每次 5000 元赔偿违约金。", - citations: [ - { - ref: "asset:r1:referenced:1", - label: "投标书 / （6）现场工期进度管理方面的违约责任", - source: { - documentId: "doc_contract", - sourceFileName: null, - sectionPath: "Root / （6）现场工期进度管理方面的违约责任", - }, - }, - ], + citations: [{ pick: 1 }], memoryCitations: [], artifacts: [], unresolved: [], @@ -3088,7 +2972,7 @@ describe("generateAgenticOutputManifest", () => { }); expect(result.manifest.text).toContain("5000 元"); expect(result.manifest.citations.map((citation) => citation.ref)).toEqual([ - "asset:r1:referenced:1", + "r1:referenced:1", ]); expect(result.trace.toolCalls.map((call) => call.tool)).toContain( "inspectImage", @@ -3131,17 +3015,7 @@ describe("generateAgenticOutputManifest", () => { }); await tools.finalize?.execute({ text: "见下方图片。", - citations: [ - { - ref: "r1:result:1", - label: "ids.pdf / 身份证 1", - source: { - documentId: "doc_identity", - sourceFileName: "ids.pdf", - sectionPath: "身份证 1", - }, - }, - ], + citations: [{ pick: 1 }], memoryCitations: [], artifacts: [1, 2, 3].map((index) => ({ type: "image", @@ -3154,17 +3028,7 @@ describe("generateAgenticOutputManifest", () => { } else { await tools.finalize?.execute({ text: "见下方图片。", - citations: [ - { - ref: "r1:result:1", - label: "ids.pdf / 身份证 1", - source: { - documentId: "doc_identity", - sourceFileName: "ids.pdf", - sectionPath: "身份证 1", - }, - }, - ], + citations: [{ pick: 1 }], memoryCitations: [], artifacts: [1, 2].map((index) => ({ type: "image", @@ -3452,87 +3316,6 @@ function makeEvidenceChunkFromRetrievalResult( }; } -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 11a7b85c..75f1dafa 100644 --- a/src/domains/chat/index.ts +++ b/src/domains/chat/index.ts @@ -247,12 +247,7 @@ export const answerQuestionWithRetrieval = ( 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 } : {}), }), @@ -1056,12 +1051,16 @@ function mapManifestCitationsToResults( ) const results: RetrievalResult[] = [] + const droppedRefs: string[] = [] for (const citation of result.manifest.citations) { const chunk = chunksByRef.get(citation.ref) ?? resolveChunkForAssetRef(citation.ref, assetsByRef, chunksByRef) - if (!chunk) continue + if (!chunk) { + droppedRefs.push(citation.ref) + continue + } const retrievalResult = toRetrievalResultFromEvidenceChunk( mergeChunkPageMetadata(chunk, result.trace.ledger.chunks), @@ -1080,6 +1079,14 @@ function mapManifestCitationsToResults( if (results.length >= MAX_CITATION_RESULTS) break } + if (droppedRefs.length > 0) { + logger.warn("chat-agent: dropped unresolved citation refs", { + droppedRefs, + ledgerChunkRefs: result.trace.ledger.chunks.map((chunk) => chunk.ref), + ledgerAssetRefs: result.trace.ledger.assets.map((asset) => asset.ref), + }) + } + return results } diff --git a/src/domains/chat/knowhere-tools.ts b/src/domains/chat/knowhere-tools.ts index 5c70fd7b..bac01f33 100644 --- a/src/domains/chat/knowhere-tools.ts +++ b/src/domains/chat/knowhere-tools.ts @@ -1,35 +1,7 @@ -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, - isNotebookVisibleRemoteDocument, - type RemoteLibraryDocument, -} from "@/domains/sources/remote-library" +import type { KnowhereToolRuntime } from "@/agent-harness" 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 } @@ -37,121 +9,6 @@ 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) => - isNotebookVisibleRemoteDocument(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-assets.test.ts b/src/domains/chat/media-assets.test.ts index 1d744826..c4ee72f5 100644 --- a/src/domains/chat/media-assets.test.ts +++ b/src/domains/chat/media-assets.test.ts @@ -2,6 +2,7 @@ import { describe, expect, it, vi } from "vitest" import type { RetrievalResult } from "@ontos-ai/knowhere-sdk" import { + dedupeMediaCitationResults, enrichRetrievalResultsWithAssetUrls, formatRetrievedMediaAssetContext, isImageAssetUrl, @@ -9,6 +10,14 @@ import { } from "./media-assets" import type { Source } from "@/infrastructure/db/schema" +vi.mock("@/lib/logger", () => ({ + logger: { + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + }, +})) + describe("chat media assets", () => { it("enriches retrieved image chunks from Notebook parsed asset URLs", async () => { const hardenChatAssetUrl = vi @@ -259,6 +268,22 @@ describe("chat media assets", () => { expect(answer).toBe("{\"name\":\"冯荣洲\",\"status\":\"matched\"}") }) + + it("dedupes results with a missing chunkType instead of throwing", () => { + // Knowhere's chunkType is declared as a required string in the SDK + // type, but real retrieval results can omit it at runtime. + const resultWithoutChunkType = makeRetrievalResult({ + chunkType: undefined as unknown as string, + assetUrl: "https://blob.example/images/launch.jpg", + }) + + expect(() => + dedupeMediaCitationResults([resultWithoutChunkType]), + ).not.toThrow() + + const deduped = dedupeMediaCitationResults([resultWithoutChunkType]) + expect(deduped).toEqual([resultWithoutChunkType]) + }) }) function makeRetrievalResult( diff --git a/src/domains/chat/media-assets.ts b/src/domains/chat/media-assets.ts index 423a8788..09c9cd07 100644 --- a/src/domains/chat/media-assets.ts +++ b/src/domains/chat/media-assets.ts @@ -1,7 +1,15 @@ import type { RetrievalResult } from "@ontos-ai/knowhere-sdk" +import { logger } from "@/lib/logger" import type { Source } from "@/infrastructure/db/schema" +// Knowhere's chunkType is declared as a required string in the SDK type, +// but real retrieval results can omit it. Normalize defensively instead of +// calling .toLowerCase() on a value that may be undefined at runtime. +function normalizeChunkType(chunkType: string): string { + return typeof chunkType === "string" ? chunkType.toLowerCase() : "" +} + const retrievedMediaAssetLimit = 6 const imageExtensions = [".jpg", ".jpeg", ".png", ".gif", ".webp", ".svg"] as const const internalMetadataKeys = new Set([ @@ -79,7 +87,7 @@ export function dedupeMediaCitationResults( const existingResult = dedupedResults[existingIndex] const existingAssetUrl = getTrimmedString(existingResult?.assetUrl) - if ( + const keepCurrent = existingResult && existingAssetUrl && compareMediaCitationResult( @@ -88,7 +96,16 @@ export function dedupeMediaCitationResults( assetUrl, existingAssetUrl, ) > 0 - ) { + logger.info("chat-agent: merged duplicate media citation", { + assetKey, + keptChunkId: keepCurrent + ? (result.chunkId ?? null) + : (existingResult?.chunkId ?? null), + droppedChunkId: keepCurrent + ? (existingResult?.chunkId ?? null) + : (result.chunkId ?? null), + }) + if (keepCurrent) { dedupedResults[existingIndex] = result } } @@ -161,7 +178,7 @@ function getMediaCitationResultScore( result: RetrievalResult, assetUrl: string, ): number { - const chunkType = result.chunkType.toLowerCase() + const chunkType = normalizeChunkType(result.chunkType) const isImageAsset = isImageAssetUrl(assetUrl) const isTableAsset = chunkType === "table" const source = result.source @@ -234,7 +251,7 @@ async function addAssetCitationResults( source: Source, hardenChatAssetUrl: HardenChatAssetUrl, ): Promise { - if (result.chunkType.toLowerCase() === "page") return [result] + if (normalizeChunkType(result.chunkType) === "page") return [result] const existingAssetUrl = getTrimmedString(result.assetUrl) if (existingAssetUrl && isNotebookOwnedAssetUrl(existingAssetUrl)) return [result] @@ -433,7 +450,7 @@ function isRenderableMediaAsset( result: RetrievalResult, assetUrl: string, ): boolean { - const chunkType = result.chunkType.toLowerCase() + const chunkType = normalizeChunkType(result.chunkType) return chunkType === "image" || chunkType === "table" || isImageAssetUrl(assetUrl) } diff --git a/src/domains/chat/page-citation-assets.test.ts b/src/domains/chat/page-citation-assets.test.ts index 5b162641..06f679ed 100644 --- a/src/domains/chat/page-citation-assets.test.ts +++ b/src/domains/chat/page-citation-assets.test.ts @@ -203,6 +203,31 @@ describe("enrichRetrievalResultsWithPageCitationAssetUrls", () => { expect(result?.pageCitationAssetUrl).toBeUndefined() expect(result?.pageCitationPageNumber).toBe(4) }) + + it("does not throw when a result is missing chunkType", async () => { + // Knowhere's chunkType is declared as a required string in the SDK + // type, but real retrieval results can omit it at runtime. + const resultWithoutChunkType = makeRetrievalResult({ + chunkType: undefined as unknown as string, + metadata: { + pageNums: [4], + pageAssets: [ + { + pageNum: 4, + artifactRef: "page_citation_assets/page-4.png", + assetUrl: "https://assets.example/pages/page-4.png", + }, + ], + }, + }) + + const [result] = await enrichRetrievalResultsWithPageCitationAssetUrls({ + results: [resultWithoutChunkType], + sources: [makeSource()], + }) + + expect(result?.pageCitationAssetUrl).toBeUndefined() + }) }) function makeRetrievalResult( diff --git a/src/domains/chat/page-citation-assets.ts b/src/domains/chat/page-citation-assets.ts index 3007f60d..07e93197 100644 --- a/src/domains/chat/page-citation-assets.ts +++ b/src/domains/chat/page-citation-assets.ts @@ -122,7 +122,12 @@ async function getStoredPageCitationAssetUrl(input: { } function isPageResult(result: RetrievalResult): boolean { - return result.chunkType.toLowerCase() === "page" + // Knowhere's chunkType is declared as a required string in the SDK type, + // but real retrieval results can omit it. Normalize defensively instead + // of calling .toLowerCase() on a value that may be undefined at runtime. + return typeof result.chunkType === "string" + ? result.chunkType.toLowerCase() === "page" + : false } function getDirectPageCitationAsset( diff --git a/src/domains/chat/prompt.ts b/src/domains/chat/prompt.ts index 4f721019..00f3eadf 100644 --- a/src/domains/chat/prompt.ts +++ b/src/domains/chat/prompt.ts @@ -62,7 +62,7 @@ export const generateAgenticOutputManifestEffect = ( turn, knowhereTools: input.knowhereTools ?? - notebookKnowhereTools.createSearchOnlyRuntime({ + notebookKnowhereTools.createRuntime({ searchSources: input.searchSources, }), memoryTools: mementoMemoryTools.createRuntime({ diff --git a/src/domains/chat/route-answer.ts b/src/domains/chat/route-answer.ts index 0a7cbb4d..34cc4d7d 100644 --- a/src/domains/chat/route-answer.ts +++ b/src/domains/chat/route-answer.ts @@ -145,7 +145,6 @@ const answerChatEffect = (input: AnswerChatInput) => excludedSourceIds: body.value.excludedSourceIds, retrieval: client.retrieval, knowledge: knowhereResources.knowledge, - remoteDocumentClient: client, hardenChatAssetUrl, hardenMediaAssetUrls: ({ results, artifacts }) => hardenChatMediaAssetUrls({ diff --git a/src/domains/chat/service.ts b/src/domains/chat/service.ts index 7839ab02..8139a42f 100644 --- a/src/domains/chat/service.ts +++ b/src/domains/chat/service.ts @@ -68,7 +68,6 @@ type ChatTurnInput = { excludedSourceIds: readonly string[] retrieval: RetrievalClient knowledge?: AnswerQuestionInput["knowledge"] - remoteDocumentClient?: AnswerQuestionInput["remoteDocumentClient"] generateAnswer: GenerateAnswer hardenChatAssetUrl?: AnswerQuestionInput["hardenChatAssetUrl"] hardenMediaAssetUrls?: AnswerQuestionInput["hardenMediaAssetUrls"] @@ -131,7 +130,6 @@ export const handleChatTurnEffect = (input: ChatTurnInput) => excludedSourceIds: input.excludedSourceIds, retrieval: input.retrieval, knowledge: input.knowledge, - remoteDocumentClient: input.remoteDocumentClient, generateAnswer: input.generateAnswer, hardenChatAssetUrl: input.hardenChatAssetUrl, hardenMediaAssetUrls: input.hardenMediaAssetUrls, From 568587f4cb4f53edbbd348999cdc55b1938f18f3 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Fri, 11 Sep 2026 16:05:24 +0800 Subject: [PATCH 2/2] fix(chat): honor document scope and block source-required finalize before search Pass include/exclude document IDs through to Knowhere and keep finalize closed until a required document search has run. Co-authored-by: Cursor --- src/agent-harness/runtime.test.ts | 53 ++++++++++- src/agent-harness/runtime.ts | 38 +++++++- src/agent-harness/types.ts | 3 + src/domains/chat/contracts.ts | 9 +- src/domains/chat/index.test.ts | 91 +++++++++++++++++++ src/domains/chat/index.ts | 5 +- src/domains/chat/knowhere-tools.ts | 27 +++++- src/domains/chat/prompt.ts | 1 + src/domains/chat/retrieval.ts | 20 +++- .../knowhere-retrieval-wire.test.ts | 90 ++++++++++++++++++ 10 files changed, 320 insertions(+), 17 deletions(-) create mode 100644 src/integrations/knowhere-retrieval-wire.test.ts diff --git a/src/agent-harness/runtime.test.ts b/src/agent-harness/runtime.test.ts index 11ac3d94..d36ca1be 100644 --- a/src/agent-harness/runtime.test.ts +++ b/src/agent-harness/runtime.test.ts @@ -124,6 +124,19 @@ describe("agent harness runtime", () => { ]) }) + it.each([ + { includeDocumentIds: ["doc_1"], excludeDocumentIds: ["doc_2"] }, + { includeDocumentIds: [], excludeDocumentIds: [] }, + ])("preserves document scopes in the actual search tool: %j", async (scope) => { + const search = vi.fn().mockResolvedValue(makeRetrievalResponse()) + const tools = createHarnessTools({ + state: {}, ledger: createEvidenceLedger(), recentTurns: [], + memoryTools: makeMemoryTools(), knowhereTools: makeKnowhereTools(search), + }) + await executeTool(tools.knowhere_search, { query: "question", ...scope }) + expect(search).toHaveBeenCalledWith(expect.objectContaining(scope)) + }) + it("returns only newly searched evidence in each knowhere_search tool result", async () => { const query = vi .fn() @@ -1083,8 +1096,8 @@ describe("agent harness runtime", () => { "setContextPolicy", "inspectImage", "readPriorTurn", - "finalize", ]) + expect(result.activeTools).not.toContain("finalize") expect(result.activeTools).not.toContain("memory_search") expect(result.activeTools).not.toContain("knowhere_search") }) @@ -1104,6 +1117,7 @@ describe("agent harness runtime", () => { }) expect(result.activeTools).toContain("memory_search") + expect(result.activeTools).toContain("finalize") expect(result.activeTools).not.toContain("knowhere_search") }) @@ -1121,6 +1135,7 @@ describe("agent harness runtime", () => { messages: [], }) + expect(result.activeTools).toContain("finalize") expect(result.activeTools).not.toContain("memory_search") expect(result.activeTools).not.toContain("knowhere_search") }) @@ -1132,6 +1147,22 @@ describe("agent harness runtime", () => { }) expect(result.toolChoice).toBe("required") + expect(result.activeTools).not.toContain("finalize") + }) + + it("blocks finalize on the first step so a document question cannot skip search", () => { + const result = prepareHarnessStep({ + stepNumber: 1, + messages: [ + { + role: "user", + content: "高血压合并冠心病，血压目标一般怎么定？", + }, + ], + }) + + expect(result.activeTools).not.toContain("finalize") + expect(result.toolChoice).toBe("required") }) it("opens memory_search and knowhere_search together as peers when sources are required", () => { @@ -1151,12 +1182,32 @@ describe("agent harness runtime", () => { expect(result.activeTools).toEqual( expect.arrayContaining(["memory_search", "knowhere_search"]), ) + expect(result.activeTools).not.toContain("finalize") expect(result.activeTools).not.toContain("knowhere_list_documents") expect(result.activeTools).not.toContain("knowhere_get_document_outline") expect(result.activeTools).not.toContain("knowhere_read_chunks") expect(result.activeTools).not.toContain("knowhere_grep_chunks") }) + it("reopens finalize after must_use_sources has called knowhere_search, including empty results", () => { + const result = prepareHarnessStep({ + stepNumber: 4, + hasKnowhereSearch: true, + intent: { + task: "answer", + dependsOnPreviousTurn: false, + retrievalNeeded: "yes", + targetModalities: ["text"], + constraints: {}, + groundingPolicy: "must_use_sources", + }, + messages: [], + }) + + expect(result.activeTools).toContain("finalize") + expect(result.activeTools).toContain("knowhere_search") + }) + it("forces image inspection before forced finalization when image assets are available", () => { const result = prepareHarnessStep({ stepNumber: 12, diff --git a/src/agent-harness/runtime.ts b/src/agent-harness/runtime.ts index 43d9329a..dbe91a58 100644 --- a/src/agent-harness/runtime.ts +++ b/src/agent-harness/runtime.ts @@ -90,6 +90,12 @@ const knowhereSearchTargetContentSchema = z.enum([ const knowhereSearchSchema = z.object({ query: z.string().min(1), + includeDocumentIds: z.array(z.string().trim().min(1)).optional().describe( + "Only search these verified document IDs. Omit for unrestricted search; [] searches no documents. Use IDs from source context or previous search results, never filenames or guessed IDs.", + ), + excludeDocumentIds: z.array(z.string().trim().min(1)).optional().describe( + "Exclude these verified document IDs. Exclusions take precedence over includeDocumentIds. If the ID is unknown, describe the document constraint in query instead.", + ), targetContent: knowhereSearchTargetContentSchema.default("all"), purpose: z.string().optional(), topK: z.number().int().min(1).max(12).optional(), @@ -215,6 +221,9 @@ export async function runAgentHarness( hasUninspectedImageAssets: input.inspectImages !== undefined && hasUninspectedImageAssets({ state, ledger }), + hasKnowhereSearch: (state.toolCalls ?? []).some( + (call) => call.tool === "knowhere_search", + ), }), stopWhen: [ () => state.finalized === true, @@ -249,7 +258,6 @@ const alwaysAvailableTools = [ "setContextPolicy", "inspectImage", "readPriorTurn", - "finalize", ] as const const fluidRetrievalTools = ["memory_search"] as const @@ -273,6 +281,7 @@ export function prepareHarnessStep(input: { readonly stepNumber: number readonly messages: readonly ModelMessage[] readonly hasUninspectedImageAssets?: boolean + readonly hasKnowhereSearch?: boolean readonly intent?: IntentFrame }): HarnessStepPreparation { const messages = sanitizeHarnessModelMessagesForStep(input.messages) @@ -319,6 +328,7 @@ export function prepareHarnessStep(input: { messages, activeTools: selectHarnessActiveTools({ intent: input.intent, + hasKnowhereSearch: input.hasKnowhereSearch === true, }), // finalize is the only output contract (see its tool description). Force // a tool call every step so the model cannot end the turn with a bare @@ -329,10 +339,14 @@ export function prepareHarnessStep(input: { function selectHarnessActiveTools(input: { readonly intent?: IntentFrame + readonly hasKnowhereSearch: boolean }): Array> { const tools: Array> = [ ...alwaysAvailableTools, ] + if (allowsFinalize(input)) { + tools.push("finalize") + } if (!allowsRetrieval(input.intent)) { return tools } @@ -355,6 +369,21 @@ function allowsRetrieval(intent?: IntentFrame): boolean { ) } +function allowsFinalize(input: { + readonly intent?: IntentFrame + readonly hasKnowhereSearch: boolean +}): boolean { + if (!input.intent) return false + if ( + input.intent.groundingPolicy === "must_use_sources" && + allowsRetrieval(input.intent) && + !input.hasKnowhereSearch + ) { + return false + } + return true +} + export function sanitizeHarnessModelMessagesForStep( messages: readonly ModelMessage[], ): ModelMessage[] { @@ -1001,6 +1030,12 @@ async function executeKnowhereSearch(input: { const beforeSnapshot = input.ledger.snapshot() const response = await input.knowhereTools.search({ query: input.request.query, + ...(input.request.includeDocumentIds !== undefined + ? { includeDocumentIds: input.request.includeDocumentIds } + : {}), + ...(input.request.excludeDocumentIds !== undefined + ? { excludeDocumentIds: input.request.excludeDocumentIds } + : {}), targetContent: input.request.targetContent, purpose: input.request.purpose, topK: input.request.topK, @@ -1280,6 +1315,7 @@ export function buildHarnessSystemPrompt(turn: AgentTurnInput): string { "- Call knowhere_search only when memory is insufficient and groundingPolicy requires citing source documents.", "- Refine knowhere_search at most twice. If two refined searches still do not add new relevant evidence, call finalize and list the gap in unresolved.", "- Do not treat every question as a document-retrieval task.", + "- For document-scoped searches, use includeDocumentIds/excludeDocumentIds only with verified IDs from source context or prior search results. If IDs are unknown, preserve the document requirement in query so Knowhere can locate it. Never invent IDs or substitute filenames. Exclusions win; an empty includeDocumentIds means no documents.", "", "Context rules:", "- If the current user request is unrelated to prior turns, set carryHistory to none and do not reuse prior topics.", diff --git a/src/agent-harness/types.ts b/src/agent-harness/types.ts index ad686ad6..b5a94168 100644 --- a/src/agent-harness/types.ts +++ b/src/agent-harness/types.ts @@ -89,6 +89,9 @@ export type KnowhereSearchRequest = Pick< RetrievalQueryParams, "query" | "topK" | "signalPaths" | "filterMode" | "threshold" > & { + /** Omitted means all documents; [] means none. Exclusions take precedence. */ + readonly includeDocumentIds?: string[] + readonly excludeDocumentIds?: string[] readonly targetContent?: KnowhereSearchTargetContent readonly purpose?: string } diff --git a/src/domains/chat/contracts.ts b/src/domains/chat/contracts.ts index 9accd573..a38a6f98 100644 --- a/src/domains/chat/contracts.ts +++ b/src/domains/chat/contracts.ts @@ -9,6 +9,7 @@ import type { HarnessRunResult, InspectImages, KnowhereToolRuntime, + KnowhereSearchRequest, } from "@/agent-harness" import type { ChatArtifactView, @@ -40,13 +41,7 @@ export type AgenticRetrievalPlan = { purpose: string | null } -export type AgenticRetrievalQuery = Pick< - RetrievalQueryParams, - "query" | "topK" | "signalPaths" | "filterMode" | "threshold" -> & { - readonly targetContent?: AgenticRetrievalTargetContent - readonly purpose?: string -} +export type AgenticRetrievalQuery = KnowhereSearchRequest export type AgenticRetrievalResponse = RetrievalQueryResponse & { retrievalPlan?: AgenticRetrievalPlan diff --git a/src/domains/chat/index.test.ts b/src/domains/chat/index.test.ts index b65be0f7..26fd0385 100644 --- a/src/domains/chat/index.test.ts +++ b/src/domains/chat/index.test.ts @@ -1,4 +1,7 @@ import { afterEach, describe, expect, it, vi } from "vitest" +import { createServer } from "node:http" +import { once } from "node:events" +import Knowhere from "@ontos-ai/knowhere-sdk" import type { Knowledge, RetrievalQueryParams, @@ -103,6 +106,94 @@ describe("answerQuestionWithRetrieval", () => { }); }); + it.each([true, false])("sends document scope through the installed SDK (agentic=%s)", async (useAgentic) => { + const received: Record[] = [] + const result = makeRetrievalResult({ + content: "Evidence from the allowed document.", + source: { documentId: "doc_included", sourceFileName: "notes.txt", sectionPath: "Overview" }, + }) + const server = createServer((request, response) => { + let body = "" + request.setEncoding("utf8") + request.on("data", (chunk: string) => { body += chunk }) + request.on("end", () => { + received.push(JSON.parse(body)) + response.writeHead(200, { "Content-Type": "application/json" }) + response.end(JSON.stringify({ + namespace: "default", query: "document question", router_used: "test_fixture", + results: [{ content: result.content, chunk_type: "text", score: 1, + source: { document_id: "doc_included", source_file_name: "notes.txt", section_path: "Overview" } }], + referenced_chunks: [], evidence_text: result.content, answer_text: "", + })) + }) + }) + try { + server.listen(0, "127.0.0.1") + await once(server, "listening") + const address = server.address() + if (!address || typeof address === "string") throw new Error("Expected TCP address") + const client = new Knowhere({ apiKey: "scope-test", baseURL: `http://127.0.0.1:${address.port}`, maxRetries: 0 }) + const scopes = [ + {}, + { includeDocumentIds: ["doc_included", "doc_other"] }, + { excludeDocumentIds: ["doc_other"] }, + { includeDocumentIds: ["doc_included", "doc_other"], excludeDocumentIds: ["doc_other"] }, + { includeDocumentIds: [] }, + ] + const answer = await Effect.runPromise(answerQuestionWithRetrieval({ + question: "document question", namespace: "default", useAgentic, + sources: [makeSource(), makeSource({ id: "source_other", knowhereDocumentId: "doc_other" })], + excludedSourceIds: ["knowhere-doc:default:doc_user_excluded"], + retrieval: client.retrieval, + generateAnswer: async ({ knowhereTools }) => { + if (!knowhereTools) throw new Error("Missing Knowhere tools") + for (const scope of scopes) await knowhereTools.search({ query: "document question", ...scope }) + return makeCitedHarnessRunResult("Evidence [[cite:1]].", result) + }, + messages: [], + })) + expect(received).toEqual(scopes.map((scope) => ({ + namespace: "default", query: "document question", top_k: 8, use_agentic: useAgentic, data_type: 1, + ...(scope.includeDocumentIds !== undefined ? { include_document_ids: scope.includeDocumentIds } : {}), + exclude_document_ids: ["doc_user_excluded", ...(scope.excludeDocumentIds ?? [])], + }))) + expect(answer.answer).toBe("Evidence [[cite:1]].") + expect(answer.citations).toHaveLength(1) + expect(answer.citations[0]?.source.documentId).toBe("doc_included") + } finally { + if (server.listening) { + await new Promise((resolve, reject) => { + server.close((error) => error ? reject(error) : resolve()) + server.closeAllConnections() + }) + } + } + }) + + it("rejects invented document IDs but accepts IDs discovered by a prior search", async () => { + const result = makeRetrievalResult({ source: { + documentId: "doc_discovered", sourceFileName: "named-document.pdf", sectionPath: "Overview", + } }) + const retrieval = { query: vi.fn().mockResolvedValue({ + namespace: "default", query: "named document", routerUsed: "agent_explore", + results: [result], referencedChunks: [], evidenceText: result.content, answerText: "", + }) } + await Effect.runPromise(answerQuestionWithRetrieval({ + question: "Only search the named document", namespace: "default", + sources: [makeSource()], excludedSourceIds: [], retrieval, messages: [], + generateAnswer: async ({ knowhereTools }) => { + if (!knowhereTools) throw new Error("Missing Knowhere tools") + await expect(knowhereTools.search({ query: "question", includeDocumentIds: ["named-document.pdf"] })).rejects.toThrow("unverified ID") + await expect(knowhereTools.search({ query: "question", excludeDocumentIds: ["doc_invented"] })).rejects.toThrow("unverified ID") + expect(retrieval.query).not.toHaveBeenCalled() + await knowhereTools.search({ query: "Only search named-document.pdf" }) + await knowhereTools.search({ query: "question", includeDocumentIds: ["doc_discovered"] }) + expect(retrieval.query).toHaveBeenLastCalledWith(expect.objectContaining({ includeDocumentIds: ["doc_discovered"] })) + return makeCitedHarnessRunResult("Evidence [[cite:1]].", result) + }, + })) + }) + 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.", diff --git a/src/domains/chat/index.ts b/src/domains/chat/index.ts index 75f1dafa..9f73b20b 100644 --- a/src/domains/chat/index.ts +++ b/src/domains/chat/index.ts @@ -32,7 +32,7 @@ import type { AnswerQuestionResult, } from "./contracts" import { - excludeDocuments, + getRetrievalDocumentScope, normalizeRetrievalQuery, } from "./retrieval" import { @@ -248,6 +248,7 @@ export const answerQuestionWithRetrieval = ( searchSources, knowhereTools: notebookKnowhereTools.createRuntime({ searchSources, + sources: input.sources, }), ...(input.inspectImages ? { inspectImages: input.inspectImages } : {}), }), @@ -976,7 +977,7 @@ function buildRetrievalQueryParams(input: { ...(typeof input.input.threshold === "number" ? { threshold: input.input.threshold } : {}), - ...excludeDocuments(input.sources, input.excludedSourceIds), + ...getRetrievalDocumentScope(input.sources, input.excludedSourceIds, input.input), } } diff --git a/src/domains/chat/knowhere-tools.ts b/src/domains/chat/knowhere-tools.ts index bac01f33..49bdbb92 100644 --- a/src/domains/chat/knowhere-tools.ts +++ b/src/domains/chat/knowhere-tools.ts @@ -1,14 +1,39 @@ import type { KnowhereToolRuntime } from "@/agent-harness" import type { SearchSources } from "./contracts" +import type { Source } from "@/infrastructure/db/schema" type NotebookKnowhereToolsInput = { readonly searchSources: SearchSources + readonly sources: readonly Source[] } export const notebookKnowhereTools = { createRuntime(input: NotebookKnowhereToolsInput): KnowhereToolRuntime { + const knownDocumentIds = new Set( + input.sources.flatMap((source) => + source.knowhereDocumentId ? [source.knowhereDocumentId] : [], + ), + ) return { - search: (request) => input.searchSources(request), + search: async (request) => { + const requestedIds = [ + ...(request.includeDocumentIds ?? []), + ...(request.excludeDocumentIds ?? []), + ] + if (requestedIds.some((id) => !knownDocumentIds.has(id))) { + throw new Error( + "Document scope contains an unverified ID. Use document IDs from source context or prior search results; otherwise keep the document requirement in query so Knowhere can locate it.", + ) + } + const response = await input.searchSources(request) + for (const result of response.results) { + if (result.source.documentId) knownDocumentIds.add(result.source.documentId) + } + for (const ref of response.referencedChunks) { + if (ref.documentId) knownDocumentIds.add(ref.documentId) + } + return response + }, } }, } as const diff --git a/src/domains/chat/prompt.ts b/src/domains/chat/prompt.ts index 00f3eadf..4d24c99a 100644 --- a/src/domains/chat/prompt.ts +++ b/src/domains/chat/prompt.ts @@ -64,6 +64,7 @@ export const generateAgenticOutputManifestEffect = ( input.knowhereTools ?? notebookKnowhereTools.createRuntime({ searchSources: input.searchSources, + sources: input.sources, }), memoryTools: mementoMemoryTools.createRuntime({ workspaceId: input.workspaceId, diff --git a/src/domains/chat/retrieval.ts b/src/domains/chat/retrieval.ts index 4d1103fa..e5391136 100644 --- a/src/domains/chat/retrieval.ts +++ b/src/domains/chat/retrieval.ts @@ -1,4 +1,4 @@ -import type { RetrievalQueryParams } from "@ontos-ai/knowhere-sdk" +import type { KnowhereSearchRequest } from "@/agent-harness/types" import type { Source } from "@/infrastructure/db/schema" import { decodeRemoteSourceId } from "@/domains/sources/remote-library" @@ -21,10 +21,11 @@ export function normalizeRetrievalQuery(value: string, fallback: string): string return normalized.slice(0, RETRIEVAL_QUERY_CHAR_LIMIT) } -export function excludeDocuments( +export function getRetrievalDocumentScope( sources: readonly Source[], excludedSourceIds: readonly string[], -): Pick { + request: Pick, +): Pick { const excluded = new Set(excludedSourceIds) const localDocumentIds = sources .filter((source) => excluded.has(source.id)) @@ -34,10 +35,19 @@ export function excludeDocuments( .map((sourceId) => decodeRemoteSourceId(sourceId)?.documentId) .filter((documentId): documentId is string => Boolean(documentId)) const documentIds = Array.from( - new Set([...localDocumentIds, ...remoteDocumentIds]), + new Set([ + ...localDocumentIds, + ...remoteDocumentIds, + ...(request.excludeDocumentIds ?? []), + ]), ) - return documentIds.length > 0 ? { excludeDocumentIds: documentIds } : {} + return { + ...(request.includeDocumentIds !== undefined + ? { includeDocumentIds: [...new Set(request.includeDocumentIds)] } + : {}), + ...(documentIds.length > 0 ? { excludeDocumentIds: documentIds } : {}), + } } function stripWrappingQuotes(value: string): string { diff --git a/src/integrations/knowhere-retrieval-wire.test.ts b/src/integrations/knowhere-retrieval-wire.test.ts new file mode 100644 index 00000000..cf698a2e --- /dev/null +++ b/src/integrations/knowhere-retrieval-wire.test.ts @@ -0,0 +1,90 @@ +import { createServer } from "node:http" +import { once } from "node:events" +import { afterEach, describe, expect, it } from "vitest" +import type { RetrievalQueryParams } from "@ontos-ai/knowhere-sdk" + +import { makeKnowhereClient } from "./knowhere" + +describe("makeKnowhereClient retrieval wire payload", () => { + const originalBaseURL = process.env.KNOWHERE_BASE_URL + + afterEach(() => { + restoreEnv("KNOWHERE_BASE_URL", originalBaseURL) + }) + + it.each<{ + name: string + params: RetrievalQueryParams + body: string + }>([ + { + name: "serializes overlapping include and exclude IDs", + params: { + query: "battery charging", + includeDocumentIds: ["doc_123", "doc_old"], + excludeDocumentIds: ["doc_old"], + } as RetrievalQueryParams, + body: '{"query":"battery charging","include_document_ids":["doc_123","doc_old"],"exclude_document_ids":["doc_old"]}', + }, + { + name: "preserves empty inclusion and exclusion arrays", + params: { + query: "battery charging", + includeDocumentIds: [], + excludeDocumentIds: [], + } as RetrievalQueryParams, + body: '{"query":"battery charging","include_document_ids":[],"exclude_document_ids":[]}', + }, + { + name: "omits document filters when not supplied", + params: { query: "battery charging" }, + body: '{"query":"battery charging"}', + }, + ])("$name", async ({ params, body }) => { + let receivedBody = "" + let receivedMethod: string | undefined + let receivedUrl: string | undefined + const server = createServer((request, response) => { + receivedMethod = request.method + receivedUrl = request.url + request.setEncoding("utf8") + request.on("data", (chunk: string) => { + receivedBody += chunk + }) + request.on("end", () => { + response.writeHead(200, { "Content-Type": "application/json" }) + response.end('{"results":[]}') + }) + }) + + try { + server.listen(0, "127.0.0.1") + await once(server, "listening") + const address = server.address() + if (!address || typeof address === "string") { + throw new Error("Expected TCP server address") + } + process.env.KNOWHERE_BASE_URL = `http://127.0.0.1:${address.port}` + const client = makeKnowhereClient("scope-test") + + await client.retrieval.query(params) + + expect(receivedMethod).toBe("POST") + expect(receivedUrl).toBe("/v2/retrieval/query") + expect(receivedBody).toBe(body) + } finally { + await new Promise((resolve, reject) => { + server.close((error) => (error ? reject(error) : resolve())) + server.closeAllConnections() + }) + } + }) +}) + +function restoreEnv(key: string, value: string | undefined): void { + if (value === undefined) { + delete process.env[key] + return + } + process.env[key] = value +}