diff --git a/src/agent-harness/runtime.test.ts b/src/agent-harness/runtime.test.ts index 5470690..ef7e5e0 100644 --- a/src/agent-harness/runtime.test.ts +++ b/src/agent-harness/runtime.test.ts @@ -98,6 +98,73 @@ describe("agent harness runtime", () => { ]) }) + it("returns only newly retrieved evidence in each retrieve tool result", async () => { + const query = vi + .fn() + .mockResolvedValueOnce(makeRetrievalResponse()) + .mockResolvedValueOnce({ + ...makeRetrievalResponse(), + query: "second query", + evidenceText: "Second evidence", + results: [ + { + content: "Second retrieval evidence.", + chunkType: "text", + score: 0.8, + source: { + documentId: "doc_2", + sourceFileName: "second.pdf", + sectionPath: "Second", + }, + }, + ], + }) + const state: { + intent?: IntentFrame + contextPolicy?: ContextPolicy + toolCalls?: HarnessToolCallTrace[] + } = { + intent: { + task: "answer", + dependsOnPreviousTurn: false, + retrievalNeeded: "yes", + targetModalities: ["text"], + constraints: {}, + groundingPolicy: "must_use_sources", + }, + contextPolicy: { + carryHistory: "none", + reason: "Self-contained request.", + activePriorTurnIds: [], + }, + } + const ledger = createEvidenceLedger() + const tools = createHarnessTools({ + state, + ledger, + retrieval: { query }, + recentTurns: [], + }) + + const firstResult = await executeTool(tools.retrieve, { query: "first" }) + const secondResult = await executeTool(tools.retrieve, { query: "second" }) + + expect(firstResult).toMatchObject({ + retrievalCount: 1, + chunks: [{ ref: "r1:result:1" }], + }) + expect(secondResult).toMatchObject({ + retrievalCount: 2, + evidenceText: "Second evidence", + chunks: [{ ref: "r2:result:1" }], + }) + expect(JSON.stringify(secondResult)).not.toContain("r1:result:1") + expect(ledger.snapshot().chunks.map((chunk) => chunk.ref)).toEqual([ + "r1:result:1", + "r2:result:1", + ]) + }) + it("blocks finalize until intent and context policy are declared", async () => { const state: { intent?: IntentFrame diff --git a/src/agent-harness/runtime.ts b/src/agent-harness/runtime.ts index 895dfc5..fabc627 100644 --- a/src/agent-harness/runtime.ts +++ b/src/agent-harness/runtime.ts @@ -353,6 +353,7 @@ export function createHarnessTools(input: { } } + const beforeSnapshot = input.ledger.snapshot() const response = await input.retrieval.query({ query: request.query, modalities: request.modalities as TargetModality[], @@ -363,13 +364,19 @@ export function createHarnessTools(input: { threshold: request.threshold, }) const snapshot = input.ledger.addRetrievalResponse(response) + const currentChunks = snapshot.chunks.slice( + beforeSnapshot.chunks.length, + ) + const currentAssets = snapshot.assets.slice( + beforeSnapshot.assets.length, + ) return { ok: true, retrievalCount: snapshot.retrievalCount, evidenceText: response.evidenceText ?? "", stopReason: response.stopReason ?? null, failureReason: response.failureReason ?? null, - chunks: snapshot.chunks.map((chunk) => ({ + chunks: currentChunks.map((chunk) => ({ ref: chunk.ref, kind: chunk.kind, type: chunk.chunkType, @@ -377,7 +384,7 @@ export function createHarnessTools(input: { source: chunk.source, assetRef: chunk.assetRef, })), - assets: snapshot.assets.map((asset) => ({ + assets: currentAssets.map((asset) => ({ ref: asset.ref, type: asset.type, label: asset.label, diff --git a/src/domains/chat/index.test.ts b/src/domains/chat/index.test.ts index 0be6b78..d16f26e 100644 --- a/src/domains/chat/index.test.ts +++ b/src/domains/chat/index.test.ts @@ -8,6 +8,7 @@ import { answerQuestionWithRetrieval, generateAgenticOutputManifest, parseChatRequestBody, + type SearchSources, } from "." import type { HardenableRetrievalResult, @@ -174,6 +175,91 @@ describe("answerQuestionWithRetrieval", () => { }); }); + it("bounds merged retrieval evidence before passing it to the answer agent", async () => { + const defaultResults = Array.from({ length: 40 }, (_, index) => + makeRetrievalResult({ + content: `Default namespace result ${index + 1}`, + source: { + documentId: `doc_default_${index + 1}`, + sourceFileName: "default.pdf", + sectionPath: `Default ${index + 1}`, + }, + }), + ); + const workspaceResults = Array.from({ length: 40 }, (_, index) => + makeRetrievalResult({ + content: `Workspace result ${index + 1}`, + source: { + documentId: `doc_workspace_${index + 1}`, + sourceFileName: "workspace.pdf", + sectionPath: `Workspace ${index + 1}`, + }, + }), + ); + const referencedChunks = Array.from({ length: 40 }, (_, index) => ({ + chunkId: `chunk_${index + 1}`, + documentId: `doc_reference_${index + 1}`, + chunkType: "text" as const, + sectionPath: `Reference ${index + 1}`, + })); + const retrieval = { + query: vi + .fn() + .mockResolvedValueOnce({ + results: defaultResults, + evidenceText: "Default evidence", + referencedChunks, + namespace: "default", + query: "large response", + routerUsed: "workflow_single_step", + answerText: null, + stopReason: "answer_done", + failureReason: null, + }) + .mockResolvedValueOnce({ + results: workspaceResults, + evidenceText: "Workspace evidence", + referencedChunks, + namespace: "notebook-workspace", + query: "large response", + routerUsed: "workflow_single_step", + answerText: null, + stopReason: "answer_done", + failureReason: null, + }), + }; + const generateAnswer = vi.fn( + async ({ searchSources }: { searchSources: SearchSources }) => { + const response = await searchSources({ + query: "large response", + topK: 3, + }); + expect(response.results).toHaveLength(6); + expect(response.referencedChunks).toHaveLength(6); + expect(response.results.map((result) => result.content)).toEqual( + [ + ...defaultResults.slice(0, 3), + ...workspaceResults.slice(0, 3), + ].map((result) => result.content), + ); + return makeHarnessRunResult("The answer is grounded."); + }, + ); + + await Effect.runPromise( + answerQuestionWithRetrieval({ + question: "What does the document say?", + namespace: "notebook-workspace", + namespaces: ["default", "notebook-workspace"], + sources: [makeSource()], + excludedSourceIds: [], + retrieval, + generateAnswer, + messages: [], + }), + ); + }); + it("does not hide a failed namespace query behind an empty namespace result", async () => { const retrievalError = new Error("Legacy namespace query failed."); const retrieval = { @@ -235,7 +321,14 @@ describe("answerQuestionWithRetrieval", () => { }); const retrieval = { query: vi.fn().mockResolvedValue({ - results: [result], + results: [ + result, + ...Array.from({ length: 30 }, (_, index) => + makeRetrievalResult({ + content: `extra result ${index + 1}`, + }), + ), + ], evidenceText: `Evidence https://blob.example/evidence.jpg ${"evidence ".repeat( 80, )}`, @@ -251,6 +344,12 @@ describe("answerQuestionWithRetrieval", () => { jobId: "job_1", assetUrl: "https://blob.example/id.jpg", }, + ...Array.from({ length: 30 }, (_, index) => ({ + chunkId: `chunk_extra_${index + 1}`, + documentId: "doc_identity", + chunkType: "text" as const, + sectionPath: `Extra referenced chunk ${index + 1}`, + })), ], namespace: "notebook-workspace", query: "冯荣洲 身份证 ID card", @@ -284,19 +383,13 @@ describe("answerQuestionWithRetrieval", () => { const response = meta.response as KnowhereQueryResponseLogMeta; expect(response).toMatchObject({ query: "冯荣洲 身份证 ID card", - resultCount: 1, - referencedChunkCount: 1, - results: [ - { - chunkType: "image", - }, - ], - referencedChunks: [ - { - chunkType: "image", - }, - ], + resultCount: 31, + referencedChunkCount: 31, }); + expect(response.results).toHaveLength(20); + expect(response.referencedChunks).toHaveLength(20); + expect(response.results[0]).toMatchObject({ chunkType: "image" }); + expect(response.referencedChunks[0]).toMatchObject({ chunkType: "image" }); expect(response.answerText.length).toBeLessThanOrEqual(203); expect(response.evidenceText.length).toBeLessThanOrEqual(203); expect(response.results[0]?.content.length).toBeLessThanOrEqual(103); diff --git a/src/domains/chat/index.ts b/src/domains/chat/index.ts index dc10453..21e6ef8 100644 --- a/src/domains/chat/index.ts +++ b/src/domains/chat/index.ts @@ -43,9 +43,13 @@ import type { HardenableRetrievalResult } from "./media-asset-hardening" const DEFAULT_TOP_K = 8 const MAX_AGENTIC_TOP_K = 12 +const MAX_AGENTIC_MERGED_RESULT_COUNT = 24 +const MAX_AGENTIC_MERGED_REFERENCED_CHUNK_COUNT = 24 +const MAX_AGENTIC_MERGED_TEXT_CHARS = 12_000 const MAX_CITATION_RESULTS = 20 const KNOWHERE_RESPONSE_TEXT_LOG_LIMIT = 200 const KNOWHERE_CHUNK_LOG_LIMIT = 100 +const KNOWHERE_RESPONSE_LOG_ITEM_LIMIT = 20 const NO_RESULTS_ANSWER = "I couldn't find that in your sources." const HARNESS_VALIDATION_FAILURE_ANSWER = "I couldn't safely finish that response because the agent output did not pass Notebook's validation checks. Please try again." @@ -88,6 +92,13 @@ type KnowhereReferencedChunkLog = { readonly summary: string } +type AgenticMergedEvidenceLimits = { + readonly resultCountPerResponse: number + readonly referencedChunkCountPerResponse: number + readonly resultCount: number + readonly referencedChunkCount: number +} + export type { AnswerQuestionInput, AnswerQuestionResult, @@ -186,7 +197,14 @@ export const answerQuestionWithRetrieval = ( ) { throw queryFailures[0] } - return mergeRetrievalResponses(queryResponses, retrievalPlan) + return mergeRetrievalResponses( + queryResponses, + retrievalPlan, + getAgenticMergedEvidenceLimits({ + namespaceCount: queryResponses.length, + topK: queryInput.topK, + }), + ) } const generatedAnswer = yield* Effect.tryPromise(() => @@ -593,10 +611,12 @@ function formatKnowhereQueryResponseForLog( response.evidenceText ?? "", KNOWHERE_RESPONSE_TEXT_LOG_LIMIT, ), - results: response.results.map(formatKnowhereResultChunkForLog), - referencedChunks: response.referencedChunks.map( - formatKnowhereReferencedChunkForLog, - ), + results: response.results + .slice(0, KNOWHERE_RESPONSE_LOG_ITEM_LIMIT) + .map(formatKnowhereResultChunkForLog), + referencedChunks: response.referencedChunks + .slice(0, KNOWHERE_RESPONSE_LOG_ITEM_LIMIT) + .map(formatKnowhereReferencedChunkForLog), } } @@ -654,6 +674,7 @@ function getRetrievalNamespaces(input: AnswerQuestionInput): readonly string[] { function mergeRetrievalResponses( responses: readonly RetrievalQueryResponse[], retrievalPlan: AgenticRetrievalPlan, + evidenceLimits: AgenticMergedEvidenceLimits, ): AgenticRetrievalResponse { const [first] = responses if (!first) { @@ -661,16 +682,27 @@ function mergeRetrievalResponses( } const statusResponses = getRetrievalStatusResponses(responses) - const results = responses.flatMap((response) => response.results) - const referencedChunks = responses.flatMap( - (response) => response.referencedChunks, - ) + const results = responses + .flatMap((response) => + response.results.slice(0, evidenceLimits.resultCountPerResponse), + ) + .slice(0, evidenceLimits.resultCount) + const referencedChunks = responses + .flatMap((response) => + response.referencedChunks.slice( + 0, + evidenceLimits.referencedChunkCountPerResponse, + ), + ) + .slice(0, evidenceLimits.referencedChunkCount) const evidenceTexts = responses .map((response) => response.evidenceText) .filter((value): value is string => Boolean(value)) + .map(truncateAgenticModelText) const answerTexts = responses .map((response) => response.answerText) .filter((value): value is string => Boolean(value)) + .map(truncateAgenticModelText) return { ...first, @@ -693,6 +725,30 @@ function mergeRetrievalResponses( } } +function getAgenticMergedEvidenceLimits(input: { + readonly namespaceCount: number + readonly topK: number | undefined +}): AgenticMergedEvidenceLimits { + const namespaceCount = Math.max(input.namespaceCount, 1) + const perResponseCount = normalizeTopK(input.topK) + const requestedCount = perResponseCount * namespaceCount + return { + resultCountPerResponse: perResponseCount, + referencedChunkCountPerResponse: perResponseCount, + resultCount: Math.min(requestedCount, MAX_AGENTIC_MERGED_RESULT_COUNT), + referencedChunkCount: Math.min( + requestedCount, + MAX_AGENTIC_MERGED_REFERENCED_CHUNK_COUNT, + ), + } +} + +function truncateAgenticModelText(value: string): string { + const trimmed = value.trim() + if (trimmed.length <= MAX_AGENTIC_MERGED_TEXT_CHARS) return trimmed + return `${trimmed.slice(0, MAX_AGENTIC_MERGED_TEXT_CHARS)}\n...[truncated]` +} + function getRetrievalStatusResponses( responses: readonly RetrievalQueryResponse[], ): readonly RetrievalQueryResponse[] {