mirror of
https://github.com/hellodigua/ChatLab.git
synced 2026-10-25 17:54:34 +08:00
feat(ai): preload selected entity memories
This commit is contained in:
1 parent
04916041b4
commit
db582d852e
4 files changed
+231
-13
No files matched your search
@@ -41,7 +41,8 @@ describe('cross-chat agent prompt', () => {
|
||||
assert.match(prompt, /不构成永久锁定范围/)
|
||||
assert.match(prompt, /本轮最新选择/)
|
||||
assert.match(prompt, /不得复用历史消息中的旧实体/)
|
||||
assert.match(prompt, /都必须按本轮稳定 ID 重新调用 memory_read/)
|
||||
assert.match(prompt, /当前稳定 ID.*自动注入/)
|
||||
assert.match(prompt, /需要更多.*memory_read/)
|
||||
assert.match(prompt, /交集、并集/)
|
||||
assert.match(prompt, /不要仅因为消息中出现了 @联系人就机械调用/)
|
||||
assert.match(prompt, /roster_only/)
|
||||
@@ -104,7 +105,8 @@ describe('cross-chat agent prompt', () => {
|
||||
assert.match(prompt, /strict intersection containing every participant/)
|
||||
assert.match(prompt, /latest selection for this turn/)
|
||||
assert.match(prompt, /never stale entities from earlier messages/)
|
||||
assert.match(prompt, /must call memory_read again with the current stable IDs/)
|
||||
assert.match(prompt, /current stable IDs.*automatically injected/)
|
||||
assert.match(prompt, /need more.*memory_read/i)
|
||||
assert.match(prompt, /private-chat frequency ranking/)
|
||||
assert.match(prompt, /never substitute keyword hits/i)
|
||||
assert.match(prompt, /coverage\.complete=true/)
|
||||
@@ -156,6 +158,18 @@ describe('cross-chat agent prompt', () => {
|
||||
assert.match(injected, /最近默认表示 90 天/)
|
||||
assert.match(injected, /ai.*重新查询.*聊天证据/)
|
||||
})
|
||||
|
||||
it('keeps current entity memories in a separate prompt section', () => {
|
||||
const prompt = buildCrossChatSystemPrompt(
|
||||
'zh-CN',
|
||||
new Date('2026-08-23T12:00:00Z'),
|
||||
'- 全局偏好',
|
||||
'- [scope=contact] 当前联系人背景'
|
||||
)
|
||||
|
||||
assert.match(prompt, /当前已注入的全局用户偏好:\n- 全局偏好/)
|
||||
assert.match(prompt, /当前已注入的联系人和群聊记忆:\n- \[scope=contact\] 当前联系人背景/)
|
||||
})
|
||||
})
|
||||
|
||||
describe('cross-chat agent lifecycle', () => {
|
||||
@@ -176,10 +190,17 @@ describe('cross-chat agent lifecycle', () => {
|
||||
api: 'openai-completions',
|
||||
provider: 'test',
|
||||
} as unknown as PiModel<PiApi>
|
||||
const memoryService = { list: () => [] } as unknown as AIMemoryService
|
||||
const memoryScopes: Array<{ scopeType: string; scopeId: string | null }> = []
|
||||
const memoryService = {
|
||||
list: (scope: { scopeType: string; scopeId: string | null }) => {
|
||||
memoryScopes.push(scope)
|
||||
return []
|
||||
},
|
||||
} as unknown as AIMemoryService
|
||||
|
||||
await runCrossChatAgent({
|
||||
userMessage: '分析一下',
|
||||
entityRefs: [{ type: 'contact', contactKey: 'contact-a', displayName: '联系人 A' }],
|
||||
aiChatId: 'global-chat-1',
|
||||
piModel,
|
||||
apiKey: 'test-key',
|
||||
@@ -191,6 +212,10 @@ describe('cross-chat agent lifecycle', () => {
|
||||
})
|
||||
|
||||
assert.equal(compressionReads, 0)
|
||||
assert.deepEqual(memoryScopes, [
|
||||
{ scopeType: 'global', scopeId: null },
|
||||
{ scopeType: 'contact', scopeId: 'contact-a' },
|
||||
])
|
||||
assert.deepEqual(
|
||||
events.filter((event) => event.type === 'status').map((event) => event.status?.phase),
|
||||
['aborted']
|
||||
|
||||
@@ -3,7 +3,8 @@ import { mkdtempSync, rmSync } from 'node:fs'
|
||||
import { tmpdir } from 'node:os'
|
||||
import { join } from 'node:path'
|
||||
import { describe, it } from 'node:test'
|
||||
import { AI_MEMORY_CONTENT_MAX_CHARS, AIMemoryService } from '../memory'
|
||||
import type { AIMemoryEntry, AIMemoryScope } from '@openchatlab/shared-types'
|
||||
import { AI_MEMORY_CONTENT_MAX_CHARS, AIMemoryService, buildEntityMemoryPrompt } from '../memory'
|
||||
|
||||
const sqliteNativeBinding = process.env.CHATLAB_TEST_SQLITE_NATIVE_BINDING
|
||||
|
||||
@@ -214,3 +215,104 @@ describe('AIMemoryService', () => {
|
||||
}
|
||||
})
|
||||
})
|
||||
|
||||
describe('entity memory prompt', () => {
|
||||
it('injects only current contacts and groups, shares the budget fairly, and ignores private sessions', () => {
|
||||
const contactEntries: AIMemoryEntry[] = Array.from({ length: 20 }, (_, index) => ({
|
||||
id: `contact-memory-${index}`,
|
||||
scopeType: 'contact',
|
||||
scopeId: 'contact-a',
|
||||
content: `联系人记忆 ${index}`,
|
||||
sourceType: index === 0 ? 'ai' : 'user',
|
||||
sourceAIChatId: null,
|
||||
sourceMessageId: null,
|
||||
createdAt: 100 - index,
|
||||
updatedAt: 100 - index,
|
||||
}))
|
||||
const groupEntry: AIMemoryEntry = {
|
||||
id: 'group-memory',
|
||||
scopeType: 'group',
|
||||
scopeId: 'group-a',
|
||||
content: '群聊长期背景',
|
||||
sourceType: 'user',
|
||||
sourceAIChatId: null,
|
||||
sourceMessageId: null,
|
||||
createdAt: 100,
|
||||
updatedAt: 100,
|
||||
}
|
||||
const requestedScopes: AIMemoryScope[] = []
|
||||
const loadEntries = (scope: AIMemoryScope): AIMemoryEntry[] => {
|
||||
requestedScopes.push(scope)
|
||||
if (scope.scopeType === 'contact' && scope.scopeId === 'contact-a') return contactEntries
|
||||
if (scope.scopeType === 'group' && scope.scopeId === 'group-a') return [groupEntry]
|
||||
return []
|
||||
}
|
||||
|
||||
const prompt = buildEntityMemoryPrompt(
|
||||
[
|
||||
{ type: 'contact', contactKey: 'contact-a', displayName: '小红' },
|
||||
{ type: 'contact', contactKey: 'contact-a', displayName: '重复选择' },
|
||||
{ type: 'session', sessionId: 'private-a', displayName: '私聊', sessionType: 'private' },
|
||||
{ type: 'session', sessionId: 'group-a', displayName: '项目群', sessionType: 'group' },
|
||||
],
|
||||
loadEntries,
|
||||
'zh-CN'
|
||||
)
|
||||
|
||||
assert.deepEqual(requestedScopes, [
|
||||
{ scopeType: 'contact', scopeId: 'contact-a' },
|
||||
{ scopeType: 'group', scopeId: 'group-a' },
|
||||
])
|
||||
assert.match(prompt, /contact-a.*小红.*contact-memory-0.*source=ai/)
|
||||
assert.match(prompt, /group-a.*项目群.*group-memory.*群聊长期背景/)
|
||||
assert.match(prompt, /部分当前实体记忆未注入.*memory_read/)
|
||||
assert.doesNotMatch(prompt, /private-a/)
|
||||
})
|
||||
|
||||
it('rebuilds the prompt from the latest entity refs without retaining the previous contact', () => {
|
||||
const entries = new Map<string, AIMemoryEntry[]>([
|
||||
[
|
||||
'contact:contact-a',
|
||||
[
|
||||
{
|
||||
id: 'memory-a',
|
||||
scopeType: 'contact',
|
||||
scopeId: 'contact-a',
|
||||
content: '只属于联系人 A',
|
||||
sourceType: 'user',
|
||||
sourceAIChatId: null,
|
||||
sourceMessageId: null,
|
||||
createdAt: 1,
|
||||
updatedAt: 1,
|
||||
},
|
||||
],
|
||||
],
|
||||
[
|
||||
'contact:contact-b',
|
||||
[
|
||||
{
|
||||
id: 'memory-b',
|
||||
scopeType: 'contact',
|
||||
scopeId: 'contact-b',
|
||||
content: '只属于联系人 B',
|
||||
sourceType: 'user',
|
||||
sourceAIChatId: null,
|
||||
sourceMessageId: null,
|
||||
createdAt: 1,
|
||||
updatedAt: 1,
|
||||
},
|
||||
],
|
||||
],
|
||||
])
|
||||
const loadEntries = (scope: AIMemoryScope) => entries.get(`${scope.scopeType}:${scope.scopeId}`) ?? []
|
||||
|
||||
const prompt = buildEntityMemoryPrompt(
|
||||
[{ type: 'contact', contactKey: 'contact-b', displayName: '联系人 B' }],
|
||||
loadEntries,
|
||||
'zh-CN'
|
||||
)
|
||||
|
||||
assert.match(prompt, /只属于联系人 B/)
|
||||
assert.doesNotMatch(prompt, /只属于联系人 A/)
|
||||
})
|
||||
})
|
||||
@@ -6,7 +6,7 @@ import { DEFAULT_CONTEXT_COMPRESSION_CONFIG, checkAndCompress, createCompression
|
||||
import { createAiTranslate } from './i18n'
|
||||
import { initTokenizer } from './tokenizer'
|
||||
import type { AIChatManager } from './chats'
|
||||
import { buildGlobalMemoryPrompt, type AIMemoryService } from './memory'
|
||||
import { buildEntityMemoryPrompt, buildGlobalMemoryPrompt, type AIMemoryService } from './memory'
|
||||
import { AgentEventHandler, type AgentStreamChunk } from './agent/event-handler'
|
||||
import { appendEntityRefsForModel } from './agent/history'
|
||||
import { DEFAULT_MAX_TOOL_ROUNDS } from './agent/constants'
|
||||
@@ -70,7 +70,8 @@ export async function runCrossChatAgent(options: RunCrossChatAgentOptions): Prom
|
||||
toolCount: tools.length,
|
||||
})
|
||||
const globalMemoryPrompt = buildGlobalMemoryPrompt(memoryService.list({ scopeType: 'global', scopeId: null }), locale)
|
||||
const systemPrompt = buildCrossChatSystemPrompt(locale, new Date(), globalMemoryPrompt)
|
||||
const entityMemoryPrompt = buildEntityMemoryPrompt(entityRefs, (scope) => memoryService.list(scope), locale)
|
||||
const systemPrompt = buildCrossChatSystemPrompt(locale, new Date(), globalMemoryPrompt, entityMemoryPrompt)
|
||||
const handler = new AgentEventHandler({ onChunk: onEvent, context: {}, systemPrompt })
|
||||
let cachedMessages: PiMessage[] = []
|
||||
|
||||
@@ -205,7 +206,12 @@ export async function runCrossChatAgent(options: RunCrossChatAgentOptions): Prom
|
||||
}
|
||||
}
|
||||
|
||||
export function buildCrossChatSystemPrompt(locale = 'zh-CN', now = new Date(), globalMemoryPrompt = ''): string {
|
||||
export function buildCrossChatSystemPrompt(
|
||||
locale = 'zh-CN',
|
||||
now = new Date(),
|
||||
globalMemoryPrompt = '',
|
||||
entityMemoryPrompt = ''
|
||||
): string {
|
||||
const dateLocale = locale.startsWith('zh') ? 'zh-CN' : 'en-US'
|
||||
const currentDate = now.toLocaleDateString(dateLocale, {
|
||||
year: 'numeric',
|
||||
@@ -254,12 +260,13 @@ export function buildCrossChatSystemPrompt(locale = 'zh-CN', now = new Date(), g
|
||||
|
||||
长期记忆规则:
|
||||
- 只保存跨 AI 对话仍有长期价值的用户偏好、身份背景、明确纠正,或已有聊天证据支持且会影响后续调查方向的稳定结论。不要保存本轮临时要求、普通统计、聊天原文、工具过程、普通回答或未经证据支持的猜测。
|
||||
- 本轮存在 <chatlab_entity_refs> 时,任何针对“当前对象/所选实体”的记忆读取都必须按本轮稳定 ID 重新调用 memory_read;即使历史中回答过相同问题,也不得直接复用旧实体的记忆或旧答案。
|
||||
- 联系人和群聊记忆必须先取得稳定 contactKey 或 group sessionId。对具体实体展开实质分析前先调用 memory_read;写入前也先读取同一作用域,已有相同含义时不要重复创建。
|
||||
- 本轮选择的联系人和群聊记忆已按当前稳定 ID 自动注入;优先使用这些记忆,不得复用历史消息中旧实体的记忆。当前实体没有已注入记忆、需要更多记忆或实体由文本临时解析得到时,再调用 memory_read。
|
||||
- 联系人和群聊记忆必须使用稳定 contactKey 或 group sessionId。写入前先调用 memory_read 读取同一作用域,已有相同含义时不要重复创建。
|
||||
- 用户明确提供、要求记住或纠正的内容使用 source_type=user;你根据聊天数据形成的结论使用 source_type=ai,并写清必要的观察时间和不确定性。
|
||||
- source_type=ai 的记忆只是调查线索,后续使用时必须重新查询当前聊天数据和原始聊天证据,不能把旧结论当作当前事实。
|
||||
- 用户纠正旧内容时按 memory ID 调用 memory_write 精确更新;确认失效或用户要求忘记时调用 memory_forget。没有真正长期价值时不要调用 memory_write。
|
||||
${globalMemoryPrompt ? `\n当前已注入的全局用户偏好:\n${globalMemoryPrompt}` : ''}
|
||||
${entityMemoryPrompt ? `\n当前已注入的联系人和群聊记忆:\n${entityMemoryPrompt}` : ''}
|
||||
|
||||
你可以使用工具。如果需要你没有的信息,请调用提供的函数。`
|
||||
}
|
||||
@@ -303,12 +310,13 @@ Tool and conclusion rules:
|
||||
|
||||
Long-term memory rules:
|
||||
- Save only preferences, identity/background facts, explicit corrections, or evidence-backed stable conclusions that remain useful across AI conversations. Never save temporary instructions, ordinary statistics, transcripts, tool traces, ordinary answers, or unsupported guesses.
|
||||
- When the current turn includes <chatlab_entity_refs>, every memory question about "the current selection" must call memory_read again with the current stable IDs. Never reuse an earlier entity's memory or answer merely because the same question appeared in history.
|
||||
- Obtain a stable contactKey or group sessionId before reading or writing entity memory. Call memory_read before substantive entity analysis and before writing that scope; do not create a duplicate with the same meaning.
|
||||
- Memories keyed by the current stable IDs are automatically injected for the selected contacts and groups. Prefer those memories and never reuse memories for stale entities from history. If the current entity has no injected memory, you need more entries, or the entity was resolved from plain text during this turn, call memory_read.
|
||||
- Entity memory must use a stable contactKey or group sessionId. Call memory_read for the same scope before writing, and do not create a duplicate with the same meaning.
|
||||
- Use source_type=user only for facts the user explicitly provides, asks to remember, or corrects. Use source_type=ai for conclusions derived from chat data, including the observation period and uncertainty.
|
||||
- Treat source_type=ai memory only as an investigation lead. Re-query current chat data and original chat evidence before using it as a fact.
|
||||
- Update a corrected memory by stable memory ID with memory_write. Call memory_forget when the user asks to forget it or confirms it is invalid. Do not call memory_write when nothing has genuine long-term value.
|
||||
${globalMemoryPrompt ? `\nInjected global user preferences:\n${globalMemoryPrompt}` : ''}
|
||||
${entityMemoryPrompt ? `\nInjected memories for the current contacts and groups:\n${entityMemoryPrompt}` : ''}
|
||||
|
||||
You have access to tools. If you need information you don't have, use the provided functions.`
|
||||
}
|
||||
@@ -2,17 +2,30 @@ import { randomUUID } from 'node:crypto'
|
||||
import * as fs from 'node:fs'
|
||||
import * as path from 'node:path'
|
||||
import Database from 'better-sqlite3'
|
||||
import type { AIMemoryEntry, AIMemoryScope, AIMemoryScopeType, AIMemorySourceType } from '@openchatlab/shared-types'
|
||||
import type {
|
||||
AIEntityRef,
|
||||
AIMemoryEntry,
|
||||
AIMemoryScope,
|
||||
AIMemoryScopeType,
|
||||
AIMemorySourceType,
|
||||
} from '@openchatlab/shared-types'
|
||||
import { appLogger } from '../logging/app-logger'
|
||||
|
||||
export const AI_MEMORY_CONTENT_MAX_CHARS = 2_000
|
||||
const AI_MEMORY_SCHEMA_VERSION = 1
|
||||
const AI_MEMORY_PROMPT_MAX_ENTRIES = 20
|
||||
const AI_MEMORY_PROMPT_MAX_CONTENT_CHARS = 4_000
|
||||
|
||||
export function buildGlobalMemoryPrompt(entries: AIMemoryEntry[], locale = 'zh-CN'): string {
|
||||
const selected: AIMemoryEntry[] = []
|
||||
let contentChars = 0
|
||||
for (const entry of entries) {
|
||||
if (selected.length >= 20 || contentChars + entry.content.length > 4_000) break
|
||||
if (
|
||||
selected.length >= AI_MEMORY_PROMPT_MAX_ENTRIES ||
|
||||
contentChars + entry.content.length > AI_MEMORY_PROMPT_MAX_CONTENT_CHARS
|
||||
) {
|
||||
break
|
||||
}
|
||||
selected.push(entry)
|
||||
contentChars += entry.content.length
|
||||
}
|
||||
@@ -29,6 +42,76 @@ export function buildGlobalMemoryPrompt(entries: AIMemoryEntry[], locale = 'zh-C
|
||||
return lines.join('\n')
|
||||
}
|
||||
|
||||
interface EntityMemoryBucket {
|
||||
scope: AIMemoryScope
|
||||
displayName: string
|
||||
entries: AIMemoryEntry[]
|
||||
}
|
||||
|
||||
interface SelectedEntityMemory {
|
||||
bucket: EntityMemoryBucket
|
||||
entry: AIMemoryEntry
|
||||
}
|
||||
|
||||
export function buildEntityMemoryPrompt(
|
||||
entityRefs: AIEntityRef[] | undefined,
|
||||
loadEntries: (scope: AIMemoryScope) => AIMemoryEntry[],
|
||||
locale = 'zh-CN'
|
||||
): string {
|
||||
const buckets: EntityMemoryBucket[] = []
|
||||
const seenScopes = new Set<string>()
|
||||
|
||||
for (const ref of entityRefs ?? []) {
|
||||
const scope: AIMemoryScope | null =
|
||||
ref.type === 'contact'
|
||||
? { scopeType: 'contact', scopeId: ref.contactKey }
|
||||
: ref.sessionType === 'group'
|
||||
? { scopeType: 'group', scopeId: ref.sessionId }
|
||||
: null
|
||||
if (!scope?.scopeId) continue
|
||||
|
||||
const scopeKey = `${scope.scopeType}:${scope.scopeId}`
|
||||
if (seenScopes.has(scopeKey)) continue
|
||||
seenScopes.add(scopeKey)
|
||||
buckets.push({
|
||||
scope,
|
||||
displayName: ref.displayName,
|
||||
entries: loadEntries(scope),
|
||||
})
|
||||
}
|
||||
|
||||
const candidates: SelectedEntityMemory[] = []
|
||||
const maxBucketSize = Math.max(0, ...buckets.map((bucket) => bucket.entries.length))
|
||||
for (let index = 0; index < maxBucketSize; index++) {
|
||||
for (const bucket of buckets) {
|
||||
const entry = bucket.entries[index]
|
||||
if (entry) candidates.push({ bucket, entry })
|
||||
}
|
||||
}
|
||||
|
||||
const selected: SelectedEntityMemory[] = []
|
||||
let contentChars = 0
|
||||
for (const candidate of candidates) {
|
||||
if (selected.length >= AI_MEMORY_PROMPT_MAX_ENTRIES) break
|
||||
if (contentChars + candidate.entry.content.length > AI_MEMORY_PROMPT_MAX_CONTENT_CHARS) continue
|
||||
selected.push(candidate)
|
||||
contentChars += candidate.entry.content.length
|
||||
}
|
||||
if (selected.length === 0) return ''
|
||||
|
||||
const lines = selected.map(({ bucket, entry }) => {
|
||||
return `- [scope=${bucket.scope.scopeType}; scope_id=${JSON.stringify(bucket.scope.scopeId)}; display_name=${JSON.stringify(bucket.displayName)}; id=${entry.id}; source=${entry.sourceType}] ${entry.content}`
|
||||
})
|
||||
if (selected.length < candidates.length) {
|
||||
lines.push(
|
||||
locale.startsWith('zh')
|
||||
? '- 部分当前实体记忆未注入;需要时调用 memory_read 读取。'
|
||||
: '- Some memories for the current entities were not injected; call memory_read when needed.'
|
||||
)
|
||||
}
|
||||
return lines.join('\n')
|
||||
}
|
||||
|
||||
export interface CreateAIMemoryInput extends AIMemoryScope {
|
||||
content: string
|
||||
sourceType: AIMemorySourceType
|
||||
|
||||
Reference in new issue
Block a user