diff --git a/src/core/think/gather.ts b/src/core/think/gather.ts index 881df4c39..a13c806e0 100644 --- a/src/core/think/gather.ts +++ b/src/core/think/gather.ts @@ -19,6 +19,8 @@ import type { BrainEngine, TakeHit, Take } from '../engine.ts'; import { hybridSearch } from '../search/hybrid.ts'; import type { SearchResult } from '../types.ts'; import { sanitizeQueryForPrompt } from '../search/expansion.ts'; +import { ensureWellFormed } from '../text-safe.ts'; +import { CJK_SLUG_CHARS } from '../cjk.ts'; export interface ThinkGatherOpts { question: string; @@ -187,20 +189,256 @@ export async function runGather( }; } +const EXCERPT_STOP_WORDS = new Set([ + 'a', 'about', 'an', 'and', 'are', 'as', 'at', 'be', 'been', 'being', 'by', + 'can', 'did', 'do', 'does', 'for', 'from', 'had', 'has', 'have', 'how', 'i', + 'if', 'in', 'including', 'into', 'is', 'it', 'its', 'me', 'my', 'of', 'on', + 'or', 'our', 'so', 'than', 'that', 'the', 'their', 'them', 'then', 'these', + 'they', 'this', 'those', 'to', 'was', 'were', 'what', 'when', 'where', + 'which', 'who', 'why', 'will', 'with', 'would', 'you', 'your', +]); + +const MAX_EXCERPT_QUERY_TERMS = 24; +const EXCERPT_TOKEN_PATTERN = + `[${CJK_SLUG_CHARS}]+|(?:(?![${CJK_SLUG_CHARS}])[\\p{L}\\p{N}])+`; +const CJK_TOKEN_PATTERN = new RegExp(`^[${CJK_SLUG_CHARS}]+$`, 'u'); + +function normalizeExcerptToken(value: string): string { + return value.normalize('NFKD').replace(/\p{M}/gu, '').toLocaleLowerCase('en'); +} + +interface ExcerptToken { + normalized: string; + start: number; + end: number; +} + +interface ExcerptQueryTerm { + normalized: string; + keys: string[]; + weight: number; +} + +interface MatchedExcerptToken extends ExcerptToken { + term: ExcerptQueryTerm; +} + +/** Tokenize while preserving offsets in the original, un-normalized string. */ +function excerptTokens(value: string): ExcerptToken[] { + const tokens: ExcerptToken[] = []; + for (const match of value.matchAll(new RegExp(EXCERPT_TOKEN_PATTERN, 'gu'))) { + const raw = match[0]; + const start = match.index; + if (CJK_TOKEN_PATTERN.test(raw)) { + if (raw.length === 1) { + tokens.push({ normalized: raw, start, end: start + 1 }); + continue; + } + for (let offset = 0; offset < raw.length - 1; offset++) { + tokens.push({ + normalized: raw.slice(offset, offset + 2), + start: start + offset, + end: start + offset + 2, + }); + } + continue; + } + tokens.push({ + normalized: normalizeExcerptToken(raw), + start, + end: start + raw.length, + }); + } + return tokens; +} + +/** Small, deterministic inflection set for lexical matches already accepted by search. */ +function excerptMatchKeys(term: string): string[] { + const keys = new Set([term]); + const addRoot = (root: string): void => { + if (root.length >= 4) keys.add(root); + }; + if (term.length >= 6 && term.endsWith('ies')) addRoot(`${term.slice(0, -3)}y`); + if (term.length >= 7 && term.endsWith('ing')) addRoot(term.slice(0, -3)); + if (term.length >= 6 && term.endsWith('ed')) addRoot(term.slice(0, -2)); + if (term.length >= 6 && term.endsWith('es')) addRoot(term.slice(0, -2)); + if (term.length >= 5 && term.endsWith('s') && !/(?:ss|us|is)$/.test(term)) { + addRoot(term.slice(0, -1)); + } + if (term.length >= 5 && term.endsWith('e')) addRoot(term.slice(0, -1)); + return Array.from(keys); +} + +function boundedExcerptTerms(terms: ExcerptQueryTerm[]): ExcerptQueryTerm[] { + if (terms.length <= MAX_EXCERPT_QUERY_TERMS) return terms; + const edgeSize = MAX_EXCERPT_QUERY_TERMS / 2; + return [...terms.slice(0, edgeSize), ...terms.slice(-edgeSize)]; +} + +function isHighSurrogate(code: number): boolean { + return code >= 0xd800 && code <= 0xdbff; +} + +function isLowSurrogate(code: number): boolean { + return code >= 0xdc00 && code <= 0xdfff; +} + +function surrogateSafeWindowStart(content: string, requested: number): number { + const start = Math.max(0, Math.min(requested, content.length)); + if (start <= 0 || start >= content.length) return start; + const startsAtLow = isLowSurrogate(content.charCodeAt(start)); + const followsHigh = isHighSurrogate(content.charCodeAt(start - 1)); + return startsAtLow && followsHigh ? start + 1 : start; +} + +function surrogateSafeWindowEnd(content: string, requested: number): number { + const end = Math.max(0, Math.min(requested, content.length)); + if (end <= 0 || end >= content.length) return end; + const endsAtHigh = isHighSurrogate(content.charCodeAt(end - 1)); + const followedByLow = isLowSurrogate(content.charCodeAt(end)); + return endsAtHigh && followedByLow ? end - 1 : end; +} + +function excerptWindow(content: string, requestedStart: number, excerptLen: number): string { + const boundedStart = Math.max(0, Math.min(requestedStart, content.length)); + const requestedEnd = Math.min(content.length, boundedStart + Math.max(0, excerptLen)); + const start = surrogateSafeWindowStart(content, boundedStart); + const end = Math.max(start, surrogateSafeWindowEnd(content, requestedEnd)); + return ensureWellFormed(content.slice(start, end)); +} + +/** Select the fixed-budget window containing the strongest unique query-term coverage. */ +function selectRelevantExcerpt( + content: string, + query: string, + excerptLen: number, + pageIdentity = '', +): string { + if (excerptLen <= 0) return ''; + if (content.length <= excerptLen) return ensureWellFormed(content); + + const uniqueTerms = Array.from(new Set( + excerptTokens(query) + .map(token => token.normalized) + .filter(term => term.length >= 2 && !EXCERPT_STOP_WORDS.has(term)), + )).map(normalized => ({ + normalized, + keys: excerptMatchKeys(normalized), + weight: Math.min(normalized.length, 12), + })); + if (uniqueTerms.length === 0) return excerptWindow(content, 0, excerptLen); + + const identityKeys = new Set( + excerptTokens(pageIdentity).flatMap(token => excerptMatchKeys(token.normalized)), + ); + const attributeTerms = uniqueTerms.filter( + term => !term.keys.some(key => identityKeys.has(key)), + ); + const terms = boundedExcerptTerms(attributeTerms.length > 0 ? attributeTerms : uniqueTerms); + const termByKey = new Map(); + for (const term of terms) { + for (const key of term.keys) { + if (!termByKey.has(key)) termByKey.set(key, term); + } + } + + const matches: MatchedExcerptToken[] = []; + for (const token of excerptTokens(content)) { + let term: ExcerptQueryTerm | undefined; + for (const key of excerptMatchKeys(token.normalized)) { + term = termByKey.get(key); + if (term) break; + } + if (term) matches.push({ ...token, term }); + } + if (matches.length === 0) return excerptWindow(content, 0, excerptLen); + + const termCounts = new Map(); + const maxStart = content.length - excerptLen; + let left = 0; + let currentScore = 0; + let bestScore = 0; + let bestStart = 0; + + for (let right = 0; right < matches.length; right++) { + const added = matches[right].term; + const addedCount = termCounts.get(added.normalized) ?? 0; + termCounts.set(added.normalized, addedCount + 1); + if (addedCount === 0) currentScore += added.weight; + + while ( + left <= right + && matches[right].end - matches[left].start > excerptLen + ) { + const removed = matches[left].term; + const remaining = (termCounts.get(removed.normalized) ?? 1) - 1; + if (remaining === 0) { + termCounts.delete(removed.normalized); + currentScore -= removed.weight; + } else { + termCounts.set(removed.normalized, remaining); + } + left++; + } + + while (left < right) { + const redundant = matches[left].term; + const count = termCounts.get(redundant.normalized) ?? 0; + if (count <= 1) break; + termCounts.set(redundant.normalized, count - 1); + left++; + } + + if (left > right) continue; + const earliestStart = Math.max(0, matches[right].end - excerptLen); + const contextualStart = Math.max( + earliestStart, + matches[left].start - Math.floor(excerptLen / 3), + ); + const candidateStart = surrogateSafeWindowStart( + content, + Math.min(contextualStart, maxStart), + ); + if ( + currentScore > bestScore + || (currentScore === bestScore && candidateStart < bestStart) + ) { + bestScore = currentScore; + bestStart = candidateStart; + } + } + + return excerptWindow(content, bestStart, excerptLen); +} + /** * Render gather results into the per-block strings the prompt builder uses. * Pages are rendered as `excerpt`; * takes are rendered via the renderTakesBlock helper from sanitize.ts. */ -export function renderPagesBlock(pages: SearchResult[], excerptLen = 600): string { +export function renderPagesBlock( + pages: SearchResult[], + excerptLen = 600, + query = '', +): string { return pages.map((p, idx) => { - const slug = String((p as unknown as { slug?: string }).slug ?? ''); - const excerpt = String( - (p as unknown as { compiled_truth?: string; chunk_text?: string; snippet?: string }).chunk_text - ?? (p as unknown as { compiled_truth?: string }).compiled_truth - ?? (p as unknown as { snippet?: string }).snippet - ?? '', - ).slice(0, excerptLen); + const page = p as unknown as { + slug?: string; + title?: string; + compiled_truth?: string; + chunk_text?: string; + snippet?: string; + }; + const slug = String(page.slug ?? ''); + const title = String(page.title ?? ''); + const slugIdentity = slug.split('/').pop()?.replace(/[-_]/g, ' ') ?? ''; + const content = String(page.chunk_text ?? page.compiled_truth ?? page.snippet ?? ''); + const excerpt = selectRelevantExcerpt( + content, + query, + excerptLen, + `${title} ${slugIdentity}`, + ); return `\n${excerpt}\n`; }).join('\n\n'); } diff --git a/src/core/think/index.ts b/src/core/think/index.ts index aaa370e96..05bdbd3b7 100644 --- a/src/core/think/index.ts +++ b/src/core/think/index.ts @@ -289,7 +289,7 @@ export async function runThink( }); // Render evidence blocks for the prompt - const pagesBlock = renderPagesBlock(gather.pages); + const pagesBlock = renderPagesBlock(gather.pages, 600, opts.question); const takesForPrompt = gather.takes.map(takesHitToTakeForPrompt); const { rendered: takesBlock, sanitizedCount } = renderTakesBlock(takesForPrompt); if (sanitizedCount > 0) { diff --git a/test/think-pages-block.test.ts b/test/think-pages-block.test.ts new file mode 100644 index 000000000..79dc65e3a --- /dev/null +++ b/test/think-pages-block.test.ts @@ -0,0 +1,212 @@ +import { describe, expect, test } from 'bun:test'; +import { renderPagesBlock } from '../src/core/think/gather.ts'; +import type { SearchResult } from '../src/core/types.ts'; + +function searchResult( + content: string, + title = 'Widget Co', + slug = 'companies/widget-co', +): SearchResult { + return { + slug, + title, + chunk_text: content, + } as SearchResult; +} + +function renderedExcerpt(rendered: string): string { + const match = rendered.match(/]*>\n([\s\S]*?)\n<\/page>/); + if (!match) throw new Error('rendered page block did not contain an excerpt'); + return match[1]; +} + +describe('renderPagesBlock', () => { + test('selects question-relevant facts beyond the leading 600 characters', () => { + const prefix = [ + '# Widget Co', + 'General company background and operating context. '.repeat(18), + ].join('\n'); + expect(prefix.length).toBeGreaterThan(600); + + const facts = [ + 'Enterprise pricing: the plan costs 125 credits per month.', + 'The annual option includes priority support.', + ].join(' '); + const content = `${prefix}\n${facts}\n${'Other context. '.repeat(80)}`; + const question = 'What is Widget Co enterprise pricing in credits per month?'; + + const rendered = renderPagesBlock([searchResult(content)], 600, question); + + expect(rendered).toContain('Enterprise pricing'); + expect(rendered).toContain('125 credits per month'); + expect(rendered).toContain('annual option includes priority support'); + }); + + test('matches query terms at token boundaries instead of inside unrelated words', () => { + const content = `${'partial context '.repeat(70)}\nExact art marker lives here.\n${'tail '.repeat(200)}`; + const rendered = renderPagesBlock([searchResult(content)], 120, 'Where is the art marker?'); + + expect(rendered).toContain('Exact art marker'); + }); + + test('keeps original offsets aligned when compatibility normalization expands earlier text', () => { + const content = `${'ffi '.repeat(260)}${'filler '.repeat(30)}\nTarget booking evidence.\n${'tail '.repeat(200)}`; + const rendered = renderPagesBlock( + [searchResult(content)], + 180, + 'Where is the target booking evidence?', + ); + + expect(rendered).toContain('Target booking evidence'); + }); + + test('never splits a surrogate pair at a selected window boundary', () => { + const content = `${'a'.repeat(599)}🚀 target evidence ${'z'.repeat(1000)}`; + const rendered = renderPagesBlock([searchResult(content)], 600, 'target evidence'); + const excerpt = renderedExcerpt(rendered); + + expect(rendered.isWellFormed()).toBe(true); + expect(excerpt).toContain('target evidence'); + expect(excerpt.length).toBeLessThanOrEqual(600); + }); + + test('keeps every scored term when a candidate starts inside a surrogate pair', () => { + const content = `🚀alpha${'.'.repeat(9)}omega${'.'.repeat(40)}`; + const rendered = renderPagesBlock([searchResult(content)], 20, 'alpha omega'); + const excerpt = renderedExcerpt(rendered); + + expect(excerpt.isWellFormed()).toBe(true); + expect(excerpt).toContain('alpha'); + expect(excerpt).toContain('omega'); + }); + + test('prefers the queried attribute over entity-title terms', () => { + const prefix = `# Widget Co\n${'General company background. '.repeat(40)}`; + const fact = '## Pricing\nThe plan costs 125 credits per month.'; + const content = `${prefix}\n${fact}\n${'Other notes. '.repeat(80)}`; + expect(content.indexOf(fact)).toBeGreaterThan(600); + + const rendered = renderPagesBlock( + [searchResult(content)], + 120, + "What is Widget Co's pricing?", + ); + + expect(renderedExcerpt(rendered)).toContain('125 credits per month'); + }); + + test('matches common inflections such as price and pricing', () => { + const prefix = 'General background without commercial details. '.repeat(25); + const fact = '## Pricing\nThe plan costs 125 credits per month.'; + const content = `${prefix}\n${fact}`; + expect(content.indexOf(fact)).toBeGreaterThan(600); + + const rendered = renderPagesBlock([searchResult(content)], 120, 'What is the price?'); + + expect(renderedExcerpt(rendered)).toContain('125 credits per month'); + }); + + test('considers windows that begin at a matched term', () => { + const content = `${'.'.repeat(100)}alpha${'.'.repeat(45)}omega${'.'.repeat(200)}`; + const rendered = renderPagesBlock([searchResult(content)], 60, 'alpha omega'); + const excerpt = renderedExcerpt(rendered); + + expect(excerpt).toContain('alpha'); + expect(excerpt).toContain('omega'); + }); + + test('retains late occurrences when query terms repeat early', () => { + const content = [ + 'alpha '.repeat(20), + 'x'.repeat(200), + 'omega '.repeat(20), + 'y'.repeat(700), + 'decisive alpha omega evidence', + ].join(''); + const rendered = renderPagesBlock([searchResult(content)], 80, 'alpha omega'); + + expect(renderedExcerpt(rendered)).toContain('decisive alpha omega evidence'); + }); + + test('finds a relevant middle occurrence between repeated edge terms', () => { + const content = [ + 'organization '.repeat(20), + 'x'.repeat(400), + 'decisive organization pricing evidence', + 'y'.repeat(400), + 'organization '.repeat(20), + ].join(''); + const rendered = renderPagesBlock( + [searchResult(content)], + 80, + 'organization pricing', + ); + + expect(renderedExcerpt(rendered)).toContain('decisive organization pricing evidence'); + }); + + test('keeps trailing fact context after dense repeated matches', () => { + const content = `${'organization '.repeat(100)}decisive organization pricing evidence`; + const rendered = renderPagesBlock( + [searchResult(content)], + 80, + 'organization pricing', + ); + + expect(renderedExcerpt(rendered)).toContain('decisive organization pricing evidence'); + }); + + test('retains terms from the end of long questions', () => { + const leadingTerms = Array.from({ length: 24 }, (_, i) => `term${i}`).join(' '); + const content = `${'Generic background. '.repeat(60)}\ntarget evidence lives here.`; + const rendered = renderPagesBlock( + [searchResult(content)], + 100, + `${leadingTerms} target evidence`, + ); + + expect(renderedExcerpt(rendered)).toContain('target evidence lives here'); + }); + + test('matches CJK query bigrams without whitespace token boundaries', () => { + const prefix = '一般背景。'.repeat(160); + const fact = '企业版价格:每月125积分。'; + const content = `${prefix}\n${fact}\n${'其他信息。'.repeat(80)}`; + expect(content.indexOf(fact)).toBeGreaterThan(600); + + const rendered = renderPagesBlock( + [searchResult(content, '示例公司', 'companies/example')], + 100, + '企业版价格是多少?', + ); + + expect(renderedExcerpt(rendered)).toContain('每月125积分'); + }); + + test('preserves the leading fixed-budget fallback when no query token matches', () => { + const content = '0123456789'.repeat(100); + const rendered = renderPagesBlock( + [searchResult(content)], + 600, + 'unmatched question', + ); + + expect(renderedExcerpt(rendered)).toBe(content.slice(0, 600)); + }); + + test('keeps a complete surrogate pair ending exactly at the fallback budget', () => { + const content = `${'a'.repeat(598)}🚀tail`; + const rendered = renderPagesBlock([searchResult(content)], 600); + const excerpt = renderedExcerpt(rendered); + + expect(excerpt).toBe(content.slice(0, 600)); + expect(excerpt.isWellFormed()).toBe(true); + }); + + test('preserves leading truncation for callers that omit the question', () => { + const content = 'abcdefghij'.repeat(100); + const rendered = renderPagesBlock([searchResult(content)], 600); + + expect(renderedExcerpt(rendered)).toBe(content.slice(0, 600)); + }); +}); diff --git a/test/think-pipeline.serial.test.ts b/test/think-pipeline.serial.test.ts index 8c85e6d0b..adec8360b 100644 --- a/test/think-pipeline.serial.test.ts +++ b/test/think-pipeline.serial.test.ts @@ -179,6 +179,64 @@ describe('runThink (with stub client)', () => { expect(result.warnings).not.toContain('LLM_OUTPUT_NOT_JSON'); }); + test('passes the question into page excerpt selection', async () => { + const prefix = [ + '# Widget Co', + 'General company background and operating context. '.repeat(18), + ].join('\n'); + const lateFact = 'Enterprise pricing: the plan costs 125 credits per month.'; + const content = `${prefix}\n${lateFact}\n${'Other context. '.repeat(80)}`; + let pageId: number | undefined; + let capturedUser = ''; + const stubClient: ThinkLLMClient = { + create: async (params) => { + const userMessage = params.messages[0]?.content; + capturedUser = typeof userMessage === 'string' + ? userMessage + : JSON.stringify(userMessage); + return { + id: 'msg_excerpt_wiring', + type: 'message', + role: 'assistant', + model: 'stub', + stop_reason: 'end_turn', + stop_sequence: null, + usage: { input_tokens: 1, output_tokens: 1, cache_creation_input_tokens: 0, cache_read_input_tokens: 0, server_tool_use: null, service_tier: null }, + content: [{ + type: 'text', + text: JSON.stringify({ answer: 'stubbed answer', citations: [], gaps: [] }), + }], + }; + }, + }; + + try { + const page = await engine.putPage('companies/widget-co', { + title: 'Widget Co', type: 'company', compiled_truth: content, + }); + pageId = page.id; + await engine.executeRaw('DELETE FROM content_chunks WHERE page_id = $1', [page.id]); + await engine.executeRaw( + `INSERT INTO content_chunks (page_id, chunk_index, chunk_text, chunk_source) + VALUES ($1, 0, $2, 'compiled_truth')`, + [page.id, content], + ); + + const result = await runThink(engine, { + question: 'What is Widget Co enterprise pricing in credits per month?', + client: stubClient, + withTrajectory: false, + }); + + expect(result.pagesGathered).toBeGreaterThan(0); + expect(capturedUser).toContain(lateFact); + } finally { + if (pageId !== undefined) { + await engine.executeRaw('DELETE FROM pages WHERE id = $1', [pageId]); + } + } + }); + test('handles malformed LLM output gracefully (regex citation fallback)', async () => { const stubClient: ThinkLLMClient = { create: async () => ({