From 4d6acc97e5c2311352b3676a47753a63b2be294c Mon Sep 17 00:00:00 2001 From: digua Date: Wed, 26 Aug 2026 13:35:11 +0800 Subject: [PATCH] fix(ai): harden global memory relevance ranking Use dependency-independent Chinese tokenization, preserve prompt capacity for the highest relevance match, and handle short symbolic and numeric terms without distorting bounded recall. --- .../__tests__/memory-optional-jieba.test.ts | 43 +++++++++++++ .../src/ai/__tests__/memory.test.ts | 63 +++++++++++++++++++ packages/node-runtime/src/ai/memory.ts | 48 +++++++++++--- 3 files changed, 147 insertions(+), 7 deletions(-) create mode 100644 packages/node-runtime/src/ai/__tests__/memory-optional-jieba.test.ts diff --git a/packages/node-runtime/src/ai/__tests__/memory-optional-jieba.test.ts b/packages/node-runtime/src/ai/__tests__/memory-optional-jieba.test.ts new file mode 100644 index 00000000..4b366ce4 --- /dev/null +++ b/packages/node-runtime/src/ai/__tests__/memory-optional-jieba.test.ts @@ -0,0 +1,43 @@ +import assert from 'node:assert/strict' +import test from 'node:test' +import type { AIMemoryEntry } from '@openchatlab/shared-types' + +test('ranks Chinese memories without loading the optional Jieba segmenter', async (t) => { + let segmentCalls = 0 + await t.mock.module('../../nlp/segmenter', { + exports: { + segment() { + segmentCalls++ + throw new Error('optional Jieba dependency is unavailable') + }, + }, + }) + const { rankMemoryEntries } = await import('../memory') + const entries: AIMemoryEntry[] = [ + createMemoryEntry('unrelated', '回答保持简短', 200), + createMemoryEntry('relevant', 'ChatLab CLI 的最近范围默认是 90 天', 100), + ] + + const result = rankMemoryEntries(entries, { + query: '查看 ChatLab CLI 最近联系人', + locale: 'zh-CN', + }) + + assert.equal(result.retrievalMode, 'relevance') + assert.equal(result.entries[0]?.id, 'relevant') + assert.equal(segmentCalls, 0) +}) + +function createMemoryEntry(id: string, content: string, updatedAt: number): AIMemoryEntry { + return { + id, + scopeType: 'global', + scopeId: null, + content, + sourceType: 'user', + sourceAIChatId: null, + sourceMessageId: null, + createdAt: updatedAt, + updatedAt, + } +} diff --git a/packages/node-runtime/src/ai/__tests__/memory.test.ts b/packages/node-runtime/src/ai/__tests__/memory.test.ts index b14d4f20..9bc1d10f 100644 --- a/packages/node-runtime/src/ai/__tests__/memory.test.ts +++ b/packages/node-runtime/src/ai/__tests__/memory.test.ts @@ -395,6 +395,22 @@ describe('memory relevance ranking', () => { assert.doesNotMatch(prompt, /ai-19/) }) + it('reserves prompt capacity for relevance before long baseline preferences', () => { + const prompt = buildGlobalMemoryPrompt( + [ + createMemoryEntry('recent-a', '甲'.repeat(2_000), 300), + createMemoryEntry('recent-b', '乙'.repeat(2_000), 200), + createMemoryEntry('relevant', '用户所说的“最近”默认指最近 90 天', 100, { sourceType: 'ai' }), + ], + 'zh-CN', + '最近按多少天' + ) + + assert.match(prompt, /relevant.*最近 90 天/) + assert.match(prompt, /recent-a/) + assert.doesNotMatch(prompt, /recent-b/) + }) + it('continues scanning for shorter memories when one entry exceeds the remaining character budget', () => { const prompt = buildGlobalMemoryPrompt([ createMemoryEntry('long-a', 'a'.repeat(1_894), 400, { sourceType: 'ai' }), @@ -425,6 +441,53 @@ describe('memory relevance ranking', () => { second.entries.map((entry) => entry.id) ) }) + + it('does not treat terse normalized content as a reverse substring match', () => { + const result = rankMemoryEntries( + [createMemoryEntry('cpp', 'C++', 200), createMemoryEntry('alice', 'Alice is my colleague', 100)], + { query: 'Who is Alice?', locale: 'en-US' } + ) + + assert.equal(result.retrievalMode, 'relevance') + assert.equal(result.matchedCount, 1) + assert.deepEqual( + result.entries.map((entry) => entry.id), + ['alice', 'cpp'] + ) + }) + + it('matches a symbolic short query without treating its normalized letter as a substring', () => { + const result = rankMemoryEntries( + [ + createMemoryEntry('unrelated', 'ChatLab configuration note', 200), + createMemoryEntry('cpp', 'I prefer C++', 100), + ], + { query: 'C++', locale: 'en-US' } + ) + + assert.equal(result.retrievalMode, 'relevance') + assert.equal(result.matchedCount, 1) + assert.deepEqual( + result.entries.map((entry) => entry.id), + ['cpp', 'unrelated'] + ) + }) + + it('retains numeric tokens that distinguish otherwise similar memories', () => { + const result = rankMemoryEntries( + [ + createMemoryEntry('port-4000', 'The API runs on port 4000', 200), + createMemoryEntry('port-3000', 'The API runs on port 3000', 100), + ], + { query: 'Should I use port 3000?', locale: 'en-US' } + ) + + assert.equal(result.retrievalMode, 'relevance') + assert.deepEqual( + result.entries.map((entry) => entry.id), + ['port-3000', 'port-4000'] + ) + }) }) describe('entity memory prompt', () => { diff --git a/packages/node-runtime/src/ai/memory.ts b/packages/node-runtime/src/ai/memory.ts index 4127d3d7..b018928a 100644 --- a/packages/node-runtime/src/ai/memory.ts +++ b/packages/node-runtime/src/ai/memory.ts @@ -18,6 +18,9 @@ const AI_MEMORY_SCHEMA_VERSION = 2 const AI_MEMORY_PROMPT_MAX_ENTRIES = 20 const AI_MEMORY_PROMPT_MAX_CONTENT_CHARS = 4_000 const AI_MEMORY_PROMPT_BASE_USER_ENTRIES = 5 +const AI_MEMORY_SUBSTRING_MIN_CHARS = 3 +const NUMERIC_MEMORY_TOKEN_REGEX = /\p{Number}+/gu +const SYMBOL_SUFFIX_MEMORY_TOKEN_REGEX = /[\p{Letter}\p{Number}]+(?:\+\+|#)/gu export type AIMemoryRetrievalMode = 'relevance' | 'recent_fallback' @@ -68,13 +71,20 @@ export function rankMemoryEntries(entries: AIMemoryEntry[], options: RankMemoryE } export function buildGlobalMemoryPrompt(entries: AIMemoryEntry[], locale = 'zh-CN', query = ''): string { - const rankedEntries = rankMemoryEntries(entries, { query, locale }).entries + const ranking = rankMemoryEntries(entries, { query, locale }) + const rankedEntries = ranking.entries const recentUserEntries = [...entries] .filter((entry) => entry.sourceType === 'user') .sort(compareRecentMemoryEntries) .slice(0, AI_MEMORY_PROMPT_BASE_USER_ENTRIES) - const baseUserIds = new Set(recentUserEntries.map((entry) => entry.id)) - const prioritizedEntries = [...recentUserEntries, ...rankedEntries.filter((entry) => !baseUserIds.has(entry.id))] + const prioritizedEntries: AIMemoryEntry[] = [] + const prioritizedIds = new Set() + const topRelevantEntries = ranking.retrievalMode === 'relevance' ? rankedEntries.slice(0, 1) : [] + for (const entry of [...topRelevantEntries, ...recentUserEntries, ...rankedEntries]) { + if (prioritizedIds.has(entry.id)) continue + prioritizedIds.add(entry.id) + prioritizedEntries.push(entry) + } const selected: AIMemoryEntry[] = [] let contentChars = 0 for (const entry of prioritizedEntries) { @@ -447,15 +457,35 @@ function normalizeSearchText(value: string): string { } function tokenizeMemoryText(value: string, locale: SupportedLocale): Set { - const tokens = new Set( - segment(value, locale, { posFilterMode: 'all', enableStopwords: true }).map(normalizeSearchText).filter(Boolean) - ) + const words = locale.startsWith('zh') + ? segmentChineseMemoryText(value, locale) + : segment(value, locale, { posFilterMode: 'all', enableStopwords: true }) + const tokens = new Set(words.map(normalizeSearchText).filter((token) => token.length >= 2)) + for (const token of extractNumericMemoryTokens(value)) tokens.add(token) + for (const token of extractSymbolSuffixMemoryTokens(value)) tokens.add(token) if (locale.startsWith('zh')) { for (const token of extractCjkBigrams(value)) tokens.add(token) } return tokens } +function extractNumericMemoryTokens(value: string): string[] { + return value.normalize('NFKC').match(NUMERIC_MEMORY_TOKEN_REGEX) ?? [] +} + +function extractSymbolSuffixMemoryTokens(value: string): string[] { + return value.normalize('NFKC').toLocaleLowerCase().match(SYMBOL_SUFFIX_MEMORY_TOKEN_REGEX) ?? [] +} + +function segmentChineseMemoryText(value: string, locale: SupportedLocale): string[] { + try { + const segmenter = new Intl.Segmenter(locale, { granularity: 'word' }) + return [...segmenter.segment(value)].filter((item) => item.isWordLike).map((item) => item.segment) + } catch { + return value.match(/[\p{Letter}\p{Number}]+/gu) ?? [] + } +} + function extractCjkBigrams(value: string): string[] { const runs = value.normalize('NFKC').match(/[\p{Script=Han}]+/gu) ?? [] return runs.flatMap((run) => @@ -471,7 +501,11 @@ function scoreMemoryEntry( ): number { const normalizedContent = normalizeSearchText(content) let score = 0 - if (normalizedContent.includes(normalizedQuery) || normalizedQuery.includes(normalizedContent)) score += 3 + const matchesFullQuery = + normalizedQuery.length >= AI_MEMORY_SUBSTRING_MIN_CHARS && normalizedContent.includes(normalizedQuery) + const isMeaningfulReverseMatch = + normalizedContent.length >= AI_MEMORY_SUBSTRING_MIN_CHARS && normalizedQuery.includes(normalizedContent) + if (matchesFullQuery || isMeaningfulReverseMatch) score += 3 if (queryTokens.size === 0) return score const contentTokens = tokenizeMemoryText(content, locale)